feat(model): Expose chunks_task_count_override on the v1 iterator - #344
Conversation
Signed-off-by: David Leong <leongdl@amazon.com>
Signed-off-by: David Leong <leongdl@amazon.com>
Signed-off-by: David Leong <leongdl@amazon.com>
…case Signed-off-by: David Leong <leongdl@amazon.com>
Signed-off-by: David Leong <leongdl@amazon.com>
f0c59de to
99fdecf
Compare
| // same chunk override, or indexing would report chunks that | ||
| // iteration never yields. | ||
| let iter = | ||
| StepParameterSpaceIterator::new_with_chunk_override(&self.space, self.chunk_override) |
There was a problem hiding this comment.
The invariant this change codifies — "indexing must not report chunks that iteration never yields" — is still violated on the setter path, which this PR leaves untouched.
set_chunks_default_task_count (line 320) mutates only the persistent self.iter (via its internal adaptive Arc<AtomicUsize>). __getitem__ rebuilds from self.space + self.chunk_override, and chunk_override is None for an adaptive space that took no constructor override. So for an adaptive CHUNK[INT] space:
it = StepParameterSpaceIterator(step=step) # template defaultTaskCount = 10
it.chunks_default_task_count = 5
next(it) # a 5-task chunk, e.g. "1-5"
it[0] # rebuilt fresh at defaultTaskCount=10 -> "1-10"self.len (line 112) has the same staleness: it is captured at construction, so the negative-index adjustment at line 218 uses the pre-mutation count. (__len__ itself raises for adaptive spaces, so the stale value is only observable through negative indexing.)
Two options that would close it without much code: have __getitem__ read the current chunk size off the live iter (iter.chunks_default_task_count()) and pass that as the override when the space is adaptive, or have the setter update a stored effective-chunk-size field that __getitem__ and the negative-index math both consult.
Not introduced here, but the comment added on these lines now asserts the property, and it does not hold for the mutation route.
| `CHUNK[INT]` parameter and turns adaptive chunking off, so a chunked space | ||
| can be walked at a caller-chosen granularity. Pass `1` to iterate individual | ||
| tasks. Ignored when the space has no chunked parameter, matching the | ||
| pure-Python reference -- though a non-positive value is still rejected in |
There was a problem hiding this comment.
This docstring looks hand-edited rather than regenerated, so the checked-in stub will not round-trip through scripts/generate_stubs.sh.
The Rust source (step_param_space.rs:141-143) writes an em dash:
pure-Python reference — though a non-positive value is still rejected
but here it is --, and the trailing clause "since validating an argument is cheaper to reason about than silently discarding a bad one" present in the Rust doc comment is dropped. pyo3-stub-gen copies /// comments verbatim and generate_stubs.sh has no em-dash or reflow post-processing step (its only sed rewrites are r#type, r#let, __next__, and the noqa line) — em dashes survive elsewhere in this same file, e.g. line 2405 in the __iter__ docstring just below.
So the next person who runs the documented regeneration flow (AGENTS.md:233-240) will get an unrelated diff on this block. Either regenerate the stub so it matches the macro output, or make the Rust doc comment read as -- / drop the trailing clause there too.
What
Adds
chunks_task_count_overridetoStepParameterSpaceIterator.__new__on the Rust-backed v1 iterator, wired toopenjd-model's existingStepParameterSpaceIterator::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_countsetter, which accepts adaptive spaces only:The capability already existed everywhere else:
openjd-rs(openjd-model)new_with_chunk_override(space, Some(n))— used internally bycreate_jobfor task counting and byopenjd-cli'sstep_total_tasksopenjd.model)chunks_task_count_overrideconstructor kwargSo 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-20chunked five at a time, iteration yields1-5,6-10,11-15,16-20, but a caller that needs individual tasks wants1-1,2-2, …. On the pure-Python path that ischunks_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
defaultTaskCountand turns adaptive chunking off, and is ignored when the space has no chunked parameter.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 usingNONCONTIGUOUS, since a contiguous chunked space is always sequential and declinesget(see below).chunks_task_count_override=0raisesValueError.openjd-modelapplies.max(1)to the override, so 0 would silently mean 1, and thechunks_default_task_countsetter 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=Truexfails intest/openjd/model_v1/test_known_gaps.pyrather than fixed in passing. Happy to split either into its own PR.chunks_parameter_nameandchunks_default_task_countare both derived from adaptive detection (chunks_param_nameandadaptive_chunk_sizeare built fromadaptive_infoinstep_param_space.rs), so both returnNonefor a statically chunked space. v0 returns"Frame"and5. Neither value is unknowable — both are in the template.needs_sequential = adaptive || has_contiguous_chunks(space), andget()returnsNonewheneversequential, soit[0]raisesIndexError. v0 answers1-5for the same space.len()works, so the count is known; onlygetdeclines.Testing
hatch run test— 5528 passed, 24 skipped, 5 xfailedhatch run lint,hatch run typing— cleancargo fmt --check,cargo build --all-targets,cargo clippy --all-targets -- -D warnings,cargo test— cleanNew tests in
test/openjd/model_v1/test_step_param_space_iter.py::TestChunksTaskCountOverridebuild the space throughdecode_job_template+create_job, so they exercise the path a consumer actually hits. They cover: the un-overridden baseline,override=1yielding individual tasks, an intermediate size regrouping the space,len()reflecting the override, indexing agreeing with iteration, adaptive being turned off (andlen()becoming answerable as a result), the override being ignored for an unchunked space,0being rejected, the setter still refusing static spaces, and containment round-tripping.specs/python-model-interface.mdandsrc/openjd/_openjd_rs.pyiare updated. The stub was hand-edited:scripts/generate_stubs.shneeds the patchedpyo3-stub-gen, which does not compile in my environment (8 errors inpyo3-stub-genitself). Worth regenerating on a machine with the patched tool to confirm it matches.