Skip to content

feat(model): Expose chunks_task_count_override on the v1 iterator - #344

Open
leongdl wants to merge 5 commits into
OpenJobDescription:mainlinefrom
leongdl:feat/expose-chunks-task-count-override
Open

feat(model): Expose chunks_task_count_override on the v1 iterator#344
leongdl wants to merge 5 commits into
OpenJobDescription:mainlinefrom
leongdl:feat/expose-chunks-task-count-override

Conversation

@leongdl

@leongdl leongdl commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

What

Adds chunks_task_count_override to StepParameterSpaceIterator.__new__ on the Rust-backed v1 iterator, wired to openjd-model's existing StepParameterSpaceIterator::new_with_chunk_override.

Why

A statically chunked space could not be walked at any granularity but the template's own. The only override the binding exposed is the chunks_default_task_count setter, which accepts adaptive spaces only:

>>> it = StepParameterSpaceIterator(space=static_chunked_space)
>>> it.chunks_default_task_count = 1
ValueError: The parameter space does not use adaptive chunking, so cannot modify chunks_default_task_count.

The capability already existed everywhere else:

Layer Static-chunk override
openjd-rs (openjd-model) new_with_chunk_override(space, Some(n)) — used internally by create_job for task counting and by openjd-cli's step_total_tasks
Pure-Python reference (openjd.model) chunks_task_count_override constructor kwarg
PyO3 binding (this repo) absent

So this is a missing wrapper rather than a missing feature, and it closes a v0/v1 parity gap.

The consumer case is storing one task per chunk value. Given 1-20 chunked five at a time, iteration yields 1-5, 6-10, 11-15, 16-20, but a caller that needs individual tasks wants 1-1, 2-2, …. On the pure-Python path that is chunks_task_count_override=1; on v1 there was no equivalent, so callers had to refuse such templates.

Behaviour

Mirrors the pure-Python reference: the override replaces defaultTaskCount and turns adaptive chunking off, and is ignored when the space has no chunked parameter.

it = StepParameterSpaceIterator(step=step, chunks_task_count_override=1)
[p["Frame"].value for p in it]   # 1-1, 2-2, ... 10-10  (was: 1-5, 6-10)
len(it)                          # 10  (was: 2)

Two details worth a reviewer's attention.

__getitem__ carries the override. It builds a fresh iterator, so without threading the override through, random access would report the template's chunks while iteration reported the overridden ones. Covered by a test using NONCONTIGUOUS, since a contiguous chunked space is always sequential and declines get (see below).

chunks_task_count_override=0 raises ValueError. openjd-model applies .max(1) to the override, so 0 would silently mean 1, and the chunks_default_task_count setter already rejects 0. This is a deliberate, small divergence from the pure-Python reference, which does not validate the argument. Happy to drop the check if you would rather match the reference exactly.

Two pre-existing gaps found along the way

Neither is caused by this change — both reproduce without the new argument — so they are recorded as strict=True xfails in test/openjd/model_v1/test_known_gaps.py rather than fixed in passing. Happy to split either into its own PR.

  1. Chunk metadata is missing for static spaces. chunks_parameter_name and chunks_default_task_count are both derived from adaptive detection (chunks_param_name and adaptive_chunk_size are built from adaptive_info in step_param_space.rs), so both return None for a statically chunked space. v0 returns "Frame" and 5. Neither value is unknowable — both are in the template.
  2. Contiguous chunked spaces refuse indexing. needs_sequential = adaptive || has_contiguous_chunks(space), and get() returns None whenever sequential, so it[0] raises IndexError. v0 answers 1-5 for the same space. len() works, so the count is known; only get declines.

Testing

  • hatch run test — 5528 passed, 24 skipped, 5 xfailed
  • hatch run lint, hatch run typing — clean
  • cargo fmt --check, cargo build --all-targets, cargo clippy --all-targets -- -D warnings, cargo test — clean

New tests in test/openjd/model_v1/test_step_param_space_iter.py::TestChunksTaskCountOverride build the space through decode_job_template + create_job, so they exercise the path a consumer actually hits. They cover: the un-overridden baseline, override=1 yielding individual tasks, an intermediate size regrouping the space, len() reflecting the override, indexing agreeing with iteration, adaptive being turned off (and len() becoming answerable as a result), the override being ignored for an unchunked space, 0 being rejected, the setter still refusing static spaces, and containment round-tripping.

specs/python-model-interface.md and src/openjd/_openjd_rs.pyi are updated. The stub was hand-edited: scripts/generate_stubs.sh needs the patched pyo3-stub-gen, which does not compile in my environment (8 errors in pyo3-stub-gen itself). Worth regenerating on a machine with the patched tool to confirm it matches.

Signed-off-by: David Leong <leongdl@amazon.com>
@leongdl
leongdl requested a review from a team as a code owner September 1, 2026 02:23
Comment thread test/openjd/model_v1/test_step_param_space_iter.py Fixed
Signed-off-by: David Leong <leongdl@amazon.com>
Comment thread rust-bindings/src/model/step_param_space.rs
Comment thread rust-bindings/src/model/step_param_space.rs Outdated
Comment thread specs/python-model-interface.md Outdated
Signed-off-by: David Leong <leongdl@amazon.com>
Comment thread rust-bindings/src/model/step_param_space.rs
Comment thread specs/python-model-interface.md Outdated
…case

Signed-off-by: David Leong <leongdl@amazon.com>
Comment thread rust-bindings/src/model/step_param_space.rs
Comment thread test/openjd/model_v1/test_step_param_space_iter.py
Comment thread specs/python-model-interface.md
Signed-off-by: David Leong <leongdl@amazon.com>
Comment thread src/openjd/_openjd_rs.pyi
Comment thread specs/python-model-interface.md
Comment thread specs/python-model-interface.md
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants