feat(dataset_tools): prototype mutable (resizable) steps (#1483)

[NJManganelli]

🤖 AI text below 🤖

Note

Targets scikit-hep:master; builds on #1579, the union-form PR, and the metadata PR. Until those merge, this PR's diff also includes their commits; it collapses to just the mutable-steps changes once they land. Review after them.

What this adds

New coffea/dataset_tools/mutable_steps.py: a send-channel generator protocol for step iteration where the consumer renegotiates the step size mid-stream. This is the shape a resource-monitoring scheduler needs to shrink chunks under worker memory pressure, and it matches the Computable.gen_steps send-channel in the #1470 compute protocol, so consumers written against it translate directly.

• resizable_steps(start, stop, size) yields [begin, end] pairs; a sent size applies from the next step, is at most the requested size (a resize triggered by resource exhaustion can never grow a step), and re-tiles the remainder evenly (n = ceil(remaining/size), actual = ceil(remaining/n)) — the semantics @btovar proposed in Dynamic chunking in coffea.compute #1483.
• iter_file_steps / iter_dataset_steps tile a file spec's covered regions (merging adjacent stored steps) and chain across a dataset's files, with resize requests carrying across region and file boundaries.
• remaining_regions / completed_spec are interval-level resumption helpers for completed ranges that need not align with the stored steps (resized steps generally don't); completed_spec results accumulate via the ordinary step arithmetic.
• WallTimeStepPolicy + run_adaptive_steps are a toy single-consumer driver that measures each step's wall time (injectable clock, so the control loop is testable without real waiting) and resizes toward a target seconds-per-step, with damped growth and min/max clamps.

Prototype module — not re-exported from coffea.dataset_tools; no executor is wired to it. The intent is to prototype the protocol shape so a resource-aware coffea.compute / TaskVine backend can adopt it.

Try it (against the bundled test files)

import shutil, tempfile, os
from coffea.dataset_tools import DataGroupSpec, preprocess
from coffea.dataset_tools.mutable_steps import (
    iter_dataset_steps, run_adaptive_steps, WallTimeStepPolicy,
)

# two files (a repo sample + a copy) so a resize can carry across a file boundary
tmp = tempfile.mkdtemp()
copy = os.path.join(tmp, "nano_dy_copy.root")
shutil.copy("tests/samples/nano_dy.root", copy)
fileset = DataGroupSpec({"dy": {"files": {"tests/samples/nano_dy.root": "Events", copy: "Events"}}})
available, _ = preprocess(fileset, step_size=10, save_form=False, backend="iterative")
dataset = available["dy"]        # each file: 40 entries

# iterate steps; send(new_size) renegotiates the size from the next step onward
gen = iter_dataset_steps(dataset, step_size=10)
fname, step = next(gen)
print("first step:", step)                       # [0, 10]
rest = []
try:
    item = gen.send(4)                            # shrink to <= 4 entries/step
    while True:
        rest.append(item[1]); item = next(gen)
except StopIteration:
    pass
print("after send(4):", rest[:6])                 # [[10,14],[14,18],[18,22],[22,26],[26,30],[30,34]]

# toy adaptive driver: resize toward a target seconds-per-step from measured wall time
class Clock:
    def __init__(self): self.t = 0.0
    def __call__(self): return self.t
clock = Clock()
def work(fname, step):
    clock.t += (step[1] - step[0]) * 0.01         # pretend 0.01 s/entry
run = run_adaptive_steps(dataset, work, step_size=40,
                         policy=WallTimeStepPolicy(target_seconds=0.1), clock=clock)
print("adaptive step sizes:", run.step_sizes)     # [40, 10, 10, 10, 10] -- converges to ~10

Tests

tests/test_dataset_tools_mutable_steps.py covers even tiling, shrink/grow-at-most-requested semantics, invalid sizes, region merging and disjoint-region iteration, resize carrying across files, interval subtraction / completed-spec accumulation, and the adaptive driver (convergence to target, damped growth with clamps).

---

Developed with Claude Fable 5. Prototype — API may change; aligns with the #1470 gen_steps send-channel.