[Data] Break "UDF time" down into input prep, function body and output build

[marwan116]

Why are these changes needed?

ds.stats() reports one figure per operator for a map stage's transform and calls it "UDF time". It isn't the user's function: it spans turning input blocks into batches or rows, calling the function, and assembling the output back into blocks. So a slow operator tells you nothing about whether to optimise your code or your data layout.

That bites hardest when rows carry Python objects, because both ends of the window get expensive. A UDF that never reads an object column is still charged for unpickling it, and nothing in the output says so.

What this change does

Renames the figure to say what it measures, and breaks it down. ds.stats() prints * Block transform time: where it printed * UDF time:, and under verbose_stats_logs splits it into three lines that sum to it:

* Block transform time: 4.02s min, 4.11s max, 4.06s mean, 16.2s total
	* Input prep: ...           forming batches or rows, where object unpickling lands
	* Function body: ...        the stage bodies, yours and Ray Data's alike
	* Output block build: ...   assembling output back into blocks

The total and all three phases are metric_fields on OpRuntimeMetrics, so they reach Prometheus per operator tagged dataset and operator. The three phases stay None rather than zero when a chain is timed only as a whole, so a dashboard can tell "not measured" from "took no time".

Default output is unchanged: the breakdown renders only under verbose_stats_logs, the same treatment extra_metrics gets, and the figures are always on get_stats_summary() regardless. Row transforms (map, filter) report only a total by default, since a timer per row costs roughly 0.6µs across a three-stage chain; DataContext.accurate_map_phase_timing opts them in.

The docs gain a worked example that runs a fused ReadRange->Project->MapBatches(map1)->MapBatches(map2) and reads the three figures off its output, because "which stages land in which bucket?" was the first thing review asked.

Under the hood

map_transformer.py is rewritten by #65855, merged in here: a task's transform is now a flat list of TimedSteps you can print and assert on, each timing itself with no coordination. Subtracting neighbouring totals once at drain recovers each step's own work, at 580 ns/item against the 1016 the per-item cursor cost.

Timing is no longer gated on the chain containing a UDF, so reads and writes are measured too and MapTransformFn.is_udf is gone, along with the separate Built-in stages bucket. Removing the gate surfaced an ordering bug in the checkpoint 2PC chain, fixed in c89db16.

What needs review

1. The name. ray_data_block_transform_time_s is introduced here, so it is free to change now and expensive after a release. @iamjustinhsu suggested total_task_per_block_s.
2. The checkpoint fix wants a checkpoint owner. It belongs in this PR, since the bug is created by this PR's deferred apply and doesn't exist on master, but nobody else has looked at it.
3. Folding Built-in stages into Function body costs the "my code or Ray Data's?" split until per-stage attribution lands in [Data] Split block transform time per fused stage, behind a flag #65810. @iamjustinhsu asked for it; worth confirming the interim is fine.
4. Reads and writes now report a figure where they reported nothing, which changes ds.stats() output for every read operator.

Related issue number

None. Closed #55052 raised adjacent complaints about operator-level metric granularity and was addressed by ds.explain(); it did not cover this.

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. (DataContext.accurate_map_phase_timing is documented in the DataContext docstring; no new public API)
• I've made sure the tests are passing.
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

Testing

Eight tests added:

test	asserts
test_block_transform_time_phases_sum_to_the_total	the components add up, which is what makes this a decomposition rather than three loose numbers
test_block_transform_time_phases_separate_object_serde	input prep dominates the body for a UDF that only counts rows beside an object column
test_row_transform_phases_are_opt_in	row transforms report only a total by default, and the flag adds detail without moving it
test_eagerly_consuming_stage_body_is_timed	a datasink's I/O reaches Function body instead of being reported as microseconds
test_operator_without_a_udf_is_still_timed	a standalone read reports its phases, and they stay inside the task's wall time
test_map_transformer.py::test_every_output_block_is_timed	every output block of a task is measured, not just the first
test_op_runtime_metrics.py::test_phase_metrics_stay_none_when_unmeasured	a chain timed only as a whole exports no phases, rather than zeros beside a non-zero total
test_op_runtime_metrics.py::test_phase_metrics_accumulate_when_measured	the None default doesn't swallow a measured phase

Existing tests that assert on the rendered ds.stats() string keep their non-verbose expectations unchanged, which is what gating the breakdown buys; the verbose ones gained three lines. The new OpRuntimeMetrics keys are canonicalized to a placeholder in those expectations, as obj_store_mem_used already is, because a trivial UDF's phases round to zero or not depending on the run.

Run locally against a source build, with PYTHONPATH=python/ray/data/tests:

file	result
test_stats.py	105 passed, 2 skipped, 1 failed
test_map.py	252 passed, 1 failed
test_checkpoint.py	60 passed, 13 skipped, 4 errors
test_consumption.py	28 passed, 2 skipped, 3 failed
test_map_transformer.py, unit/test_auto_batch_size.py, test_op_runtime_metrics.py, test_context.py	all passed
datasource/{test_datasink,test_file_datasink,test_parquet}.py	failures identical to the merge base

Every failure above was A/B'd against master and traces to something absent from my machine rather than to this change: the dashboard API server isn't running, pandas 3.0 breaks an untouched conversion helper, and the rest are s3 fixtures. pre-commit run --all-files passes.

AI assistance

AI assistance (Claude Code) was used to investigate the attribution gap, write the change and the tests, and draft this description. Per AGENTS.md, the author has reviewed every changed line and the tests reported above were run locally.

Duplicate check

gh pr list --repo ray-project/ray --state open was searched for udf time phase breakdown input prep and map_transformer phase timing; neither returns anything. The open PRs touching nearby code — #61379 (caching the serialized map transformer), #64090 (per-actor placement groups), #65663 (exposing backpressure policy in diagnostics) — do not touch this timing.

[gemini-code-assist]

Code Review

This pull request decomposes the overall 'UDF time' metric in Ray Data into four distinct phases: input preparation, UDF body execution, output block building, and other fused stages. It introduces these granular metrics across execution stats, operator summaries, and block metadata, with an opt-in configuration (accurate_map_phase_timing) for row-based transforms to avoid overhead. The review feedback highlights a critical missing update to BlockExecStatsBuilder.build that would cause a runtime TypeError when constructing block stats, and suggests a formatting improvement to ensure core UDF phases are consistently displayed in verbose logs even when their sum is zero.

[iamjustinhsu]

Nice, left some questions and some higher-level feedback since I think the code is getting a bit complicated

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit c3e49a8. Configure here.}

[marwan116]

@iamjustinhsu this has moved a lot since your review — worth re-reading the description rather than the old threads.

• [Data] RFC: build a map task's timing from a declared list of steps #65855 is merged in, so map_transformer.py builds on the prototype you suggested.
• Both your requests are done:
  • udf_time is now block_transform_time_s
  • Built-in stages is now folded into Function body.
• Removing the UDF gate surfaced a 2PC ordering bug in the checkpoint write chain. Fixed here, and it's the part I'd most like another pair of eyes on.

Four things need your call; they're listed under What needs review in the description.

[dstrodtman]

Docs-team style pass from Douglas Strodtman (Anyscale docs). Claude Code assisted; I read every comment below and stand behind each one.

Scope: prose style, grammar, and docs conventions only. I'm deliberately staying out of the naming (UDF time → Block transform time) and the phase-model questions in @iamjustinhsu's open threads — those are the Data team's call, and they look like they're still in motion. The new prose reads clearly; the notes below are all mechanical.

Four small fixes inline, all serial-comma and filler-word nits against the Google developer style guide (Ray's stated fallback for anything the Ray writing-style guide doesn't cover). Nothing here blocks — take or leave each individually. No technical review implied, and no approval: a Data maintainer is already reviewing.

[iamjustinhsu]

Nice, I think once we add the data_dashboard panels it should be good, or a follow up would be fine too

[iamjustinhsu]

I think we should also update the data_dashboard_panels.py so that we can get accurate metrics as well. The metric name appends the ray_data_ prefix, so in this case you have created ray_data_block_transform_time_s, etc... which can be queried with promQL

[marwan116]

Added in 175de2f. Two panels in the Outputs row next to Block Generation Time: Block Transform Time, and Block Transform Time by Phase stacking the three phases. Both divide by num_task_outputs_generated like the neighbouring panel, since these are per-output-block sums.

One thing to flag: the phase panel is empty for a row transform that did not opt into accurate_map_phase_timing. Those gauges stay None and Gauge.set skips None, so nothing is recorded rather than three misleading zeros. The panel description says so.

[iamjustinhsu]

Can you elaborate more on why this is being eagerly pulled? I don't quite understand the 2nd sentence

[marwan116]

Rewritten in ff53e7e:

# Each stage runs on the first pull from it, so nothing upstream has
# run yet. `prepare_checkpoint` is what leaves the pending checkpoints
# on `ctx`. Drain first, or there is nothing here to commit.

Additionally, TransformClock.chain doc string now concretely explains how the chain is built and executed.

[iamjustinhsu]

I think this docstring is slightly misleading: this function does time, but the timing is delegated to the TimeStep classes

[marwan116]

Fixed by renaming rather than rewording, in cbf57bb. A Step is inert — a body plus the metadata the summary groups by — so timed_steps() was the misleading part. It is Step / MapTransformFn.steps / MapTransformer.get_steps now, and _TimedStep stays as the wrapper that does the timing.

[iamjustinhsu]

I think it makes more sense to avoid _TimeStep accepting a mutable reference to inclusive. We can probably just init to 0s in _TimedStep, remove this line, and then in _self_times, we can call some _TimedStep.drain() to get the inclusive times.

[marwan116]

Done in 2ca8f2e. _TimedStep holds its own total and hands it over through drain(); TransformClock collects them in order before subtracting neighbours, which it did anyway. The shared list and the in-place reset are both gone, along with the comment that was guarding them.