[Data] RFC: build a map task's timing from a declared list of steps

[marwan116]

RFC / experiment, not a merge candidate. This answers @iamjustinhsu's review comment on #65807 asking for an approach that is explicit about what gets timed. Opened as a draft so the diff is reviewable; see Open questions for what is deliberately unfinished.

Why these changes are needed

The review comment:

I think with all the lazy iterators it's getting harder to interpret what is getting timed, and what is being wrapped, since iterators don't show upstream iterators well.

That is fair, and it is the root of the subtlety in #65807: timing is emergent from the pull-chain rather than declared. You have to simulate the chain in your head to know what a number means.

What this change does

A task's transform becomes a flat list of steps.

@dataclass(frozen=True)
class TimedStep:
    label: str          # "MapBatches(score) body"
    bucket: MapTransformPhase
    stage_idx: int
    apply: Callable[[Iterable], Iterable]

_pre_process and _post_process are already Iterable -> Iterable, so they are the step bodies verbatim; only the wrapped fn needs a closure to bind the task context. apply_transform then reduces to:

steps = self.get_timed_steps(ctx, report, decomposed=decomposed)
iter = input_blocks
for idx in range(len(steps)):
    iter = udf_time_scope.wrap(idx, iter)
return iter

A step measures only itself, with no coordination.

def __next__(self):
    start = time.perf_counter()
    try:
        if self._iter is None:
            self._iter = iter(self._apply(self._upstream))
        return next(self._iter)
    finally:
        self._totals[self._idx] += time.perf_counter() - start

That total is inclusive of everything upstream. The chain is linear -- step k is only ever pulled by step k + 1 -- so all of a step's time lies inside its consumer's windows, and subtracting neighbouring totals recovers each step's own work:

def _self_times(self):
    return [max(0.0, total - (self.inclusive[i - 1] if i else 0.0))
            for i, total in enumerate(self.inclusive)]

#65807 does the same arithmetic, but per item, threading a shared cursor through every __next__. Doing it once over a list at drain time is easier to follow and measurably cheaper: 580 ns/item against 1016 over a three-stage chain.

Deferring apply to the first pull removes the eager special case. A write performs its whole upload while its iterable is being built, which on #65807 needed a separate timing window (attribute_call). Here that work happens inside the same window as the pulls, so there is one timing point rather than two.

Two special cases dissolve

• Eager stage bodies. [Data] Break "UDF time" down into input prep, function body and output build #65807 added attribute_call for these. Deferring apply into the first pull means there is nothing left to special-case -- the reviewer's second side note, satisfied by deletion rather than by handling.
• accurate_map_phase_timing stops being a second code path and becomes a rewrite of the step list -- fold a stage's three steps into one, whose bucket is None because it spans all three. A phase no step carries is then absent from the grouping, so "not measured" falls out of the data instead of being tracked by a flag. This is the same shape as the pre/post-process elision suggested in the follow-up comment, which would be another rewrite of the same list. Policy becomes data.

UDFTimeScope is renamed to TransformClock, and the udf_time_scope= keyword to clock= (41 call sites across 10 files). The old name no longer described the object: it times every step in the chain, including the ones that aren't user code, which is the same mismatch raised here.

Accruing per step rather than per phase also unifies #65807 and #65810: per-phase and per-fused-stage are two groupings of one array. #65810's separate stage_totals and its parallel += would disappear -- along with the bug a rebase surfaced there, where the eager window credited the phase bucket but silently skipped the stage bucket. This PR doesn't carry that grouping, since #65810 isn't in its base.

Where the reviewer's sketch does not work, and why

The suggestion was an eager per-input-block loop. Two measurements say that cannot express what Ray Data already does.

Batches span input blocks. 4 input blocks of 1 row each with batch_size=4:

batch_len values: [4, 4, 4, 4]     # one batch of 4, not four of 1

A single fused operator can contain several independent batching points:

OPERATORS: ['MapBatches(f)->MapBatches(g)']
f sees batches of [1, 2]      g sees batches of [3]

g's batches of 3 are re-assembled from f's outputs of 2, inside one task. No unit -- not a block, not a batch -- survives the chain intact, so for input_block in blocks: cannot be the outer loop; it would silently turn batch_size=4 into four batches of 1. The only eager alternative materializes at every stage, trading streaming for memory.

So laziness is structural. This PR keeps it and makes the timing declared instead.

Measured equivalence with #65807

case	#65807	this PR
ReadRange->Project->MapBatches(map1)->MapBatches(map2), 0.60s of real UDF work	Function body 614.3ms	611.03ms
datasink sleeping 0.3s/block over 2 blocks	Built-in stages 606ms (via attribute_call)	602.9ms, no special case
standalone ReadRange (no UDF)	0.0	0
UDF time vs Remote wall time	99.6%	99.8%

Isolated unit check of the baton -- three steps, only the middle one sleeping, 4 items x 50ms:

wall           = 0.2156s
total_s        = 0.2156s
  input_prep   = 0.0000s
  udf_body     = 0.2155s
  output_build = 0.0000s

Exact attribution, no leakage into neighbouring buckets.

Per-item cost is unchanged: over a 3-stage chain, 1019 ns/item for the baton against 1016 ns/item for the subtraction cursor in #65807 (31 ns/item with no timing at all). The 33x over baseline is inherent to per-item timing, which is why accurate_map_phase_timing still has to exist for row transforms.

Testing

PYTHONPATH=python/ray/data/tests pytest python/ray/data/tests/test_stats.py
pytest -q python/ray/data/tests/test_map_transformer.py python/ray/data/tests/unit/test_auto_batch_size.py
PYTHONPATH=python/ray/data/tests pytest python/ray/data/tests/test_checkpoint.py -k "transform or checkpoint_map"

test_stats.py: 108 collected, 105 passed, 2 skipped, 1 failed — the failure being test_streaming_exec_schedule_percentiles_populated, which fails identically on #65807 because the dashboard is not built locally. That matches #65807's baseline exactly, with test_stats.py unmodified: every test written for the phase breakdown there passes against this implementation untouched, which is the main evidence that this is behaviour-preserving.

test_map_transformer.py + unit/test_auto_batch_size.py: 9 passed. test_checkpoint.py (the cases that call apply_transform directly): 2 passed.

One test change was needed, and it is an API change rather than a behavioural one. test_chained_transforms_dont_double_count_udf_time and test_chained_transforms_total_is_independent_of_distribution read scope.attributed_s, the running total this PR removes; they now read sum(scope.totals). Both are #65776's core invariant tests, and both pass with only the accessor changed.

One behaviour change worth reviewing

A transform's body used to run when the chain was built; it now runs on the first pull:

#65807                              this PR
  -> calling apply_transform          -> calling apply_transform
  BODY RAN                            <- apply_transform returned
  <- apply_transform returned         -> pulling first item
  -> pulling first item               BODY RAN

So a transform that fails while being set up rather than while yielding now surfaces from the first next() instead of from apply_transform. That is safe where it matters, because iterate_with_retry -- which _map_task wraps the transform in -- puts construction and iteration inside the same try:

try:
    iterable = iterable_factory()
    for item_index, item in enumerate(iterable):
        ...
except Exception as e:
    ...retry...

so retried_map_errors and max_map_retries behave identically either way.

It also removes a hazard rather than adding one. Today _pre_process runs eagerly at construction, so with batch_size="auto" a later stage peeks at a real block and pulls data through earlier stages before their timers exist -- the regression Cursor Bugbot caught on #65776, and the reason that PR must install each stage's timer before building the next. When nothing runs until the first pull, every timer is already in place and that class of ordering bug cannot occur.

Open questions

• RayTurbo's fused transform fns collapse two MapTransformFns into one (is_udf = self._is_udf or next._is_udf), so one step can correspond to two logical operators. TimedStep.label would need care there.
• Should Built-in stages remain its own bucket? The review suggests merging built-in stages into UDFs. Under this design that is a one-line change, since every step carries is_udf. I lean towards keeping the split -- "is this time my code or Ray Data's?" is the actionable question for the workload that motivated [Data] Fix double-counted UDF time in fused map operators #65776 -- but the architecture no longer forces the answer either way.
• TimedStep.label would let [Data] Report UDF time per fused stage, behind a flag #65810 report MapBatches(score) instead of UDF stage 0, fixing a limitation documented there.

Related issue number

Follow-up to #65776 (merged), #65807 and #65810. No open issue or PR covers this; gh pr list --search "map transformer timing" returns only #65807 and #65810.

Checks

• I've signed off every commit (git commit -s).
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've added any new APIs to the API Reference. (no new public API)
• I've made sure the tests are passing. (see Testing)
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

AI assistance

AI assistance (Claude Code) was used to prototype this approach, run the measurements above, and draft this description. Per AGENTS.md the author has reviewed every changed line. This is an RFC opened for design discussion rather than a merge request.