[Data] Added tracking of block serialization time by alexeykudinkin · Pull Request #60574 · ray-project/ray

alexeykudinkin · 2026-01-29T01:07:24Z

Changes

Modified Ray Core's generator handling sequence to inject back object creation & serialization durations
Updated task_completion_time_excl_backpressure_s to track both UDF block generation time AND block serialization overhead

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Updated `task_completion_time_excl_backpressure_s` to include block serde time Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

…des serde time; Updated tests; Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

gemini-code-assist

Code Review

This pull request introduces tracking for block serialization time, which is a valuable addition for performance monitoring in Ray Data. The core change involves modifying the generator execution flow in _raylet.pyx to feed back serialization duration to the caller using gen.send(). This new metric is then integrated throughout the data stack, including BlockExecStats, OpRuntimeMetrics, and relevant operators like MapOperator and HashShuffleOperator. The tests have also been updated to cover these new metrics.

The implementation is clean and effective. The use of yield expressions to pass data back into generators is a good pattern for this use case. I have one minor suggestion to improve code clarity.

gemini-code-assist · 2026-01-29T01:09:52Z

python/ray/data/_internal/execution/operators/map_transformer.py

+    def udf_time_s(self, reset: bool) -> float:
+        cur_time_s = self._udf_time_s
+        self._udf_time_s = 0
+        return cur_time_s


The reset parameter in udf_time_s is not used in the function body; the timer self._udf_time_s is always reset to 0. Since the only call site in map_operator.py passes reset=True, the behavior is correct for now. However, to improve clarity and prevent potential misuse in the future, I suggest removing the reset parameter from the method signature and updating the call in map_operator.py:772 to map_transformer.udf_time_s().

Suggested change

def udf_time_s(self, reset: bool) -> float:

cur_time_s = self._udf_time_s

self._udf_time_s = 0

return cur_time_s

def udf_time_s(self) -> float:

cur_time_s = self._udf_time_s

self._udf_time_s = 0

return cur_time_s

python/ray/data/_internal/execution/operators/map_transformer.py

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

python/ray/data/block.py

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

edoakes

From offline discussion, core change looks ok to me. Let's consider it a private API for now. I will follow up to add a disclaimer & context in code comments.

Defer to others to review the data changes. You may want to consider an alternative name to "serialization time" since it is not purely serialization, but also includes memory allocation & copy time. Perhaps "block_write_time_s" or "block_output_time_s"

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

cursor

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

python/ray/data/_internal/execution/operators/hash_shuffle.py

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

goutamvenkat-anyscale · 2026-01-31T00:16:15Z

python/ray/_raylet.pyx

+
+@dataclass(frozen=True)
+class StreamingGeneratorStats:
+    object_creation_dur_s: float


Can we rename this to object_serialization_time_s

See the conversation above. This is not just serialization time, it includes object serialization time

goutamvenkat-anyscale · 2026-01-31T00:24:03Z

python/ray/data/_internal/execution/operators/map_transformer.py

        self._init_fn = init_fn if init_fn is not None else lambda: None
        self._output_block_size_option_override = output_block_size_option_override
-        self._udf_time = 0
+        self._udf_time_s = 0


Nit: Maybe a type var to denote secs?

That's what _s prefix is for

goutamvenkat-anyscale · 2026-01-31T00:32:49Z

python/ray/data/_internal/execution/operators/map_operator.py

+            # Yield block and retrieve its Ray object serialization timing
+            stats: StreamingGeneratorStats = yield block
+            if stats:
+                exec_stats.block_ser_time_s = stats.object_creation_dur_s


when operators are fused, can we assert in that test this value is 0?

That won't be 0

Changes --- 1. Modified Ray Core's generator handling sequence to inject back object creation & serialization durations 2. Updated `task_completion_time_excl_backpressure_s` to track both UDF block generation time AND block serialization overhead ## Related issues > Link related issues: "Fixes ray-project#1234", "Closes ray-project#1234", or "Related to ray-project#1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: 400Ping <jiekaichang@apache.org>

Changes --- 1. Modified Ray Core's generator handling sequence to inject back object creation & serialization durations 2. Updated `task_completion_time_excl_backpressure_s` to track both UDF block generation time AND block serialization overhead ## Related issues > Link related issues: "Fixes ray-project#1234", "Closes ray-project#1234", or "Related to ray-project#1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Sirui Huang <ray.huang@anyscale.com>

Changes --- 1. Modified Ray Core's generator handling sequence to inject back object creation & serialization durations 2. Updated `task_completion_time_excl_backpressure_s` to track both UDF block generation time AND block serialization overhead ## Related issues > Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

Changes --- 1. Modified Ray Core's generator handling sequence to inject back object creation & serialization durations 2. Updated `task_completion_time_excl_backpressure_s` to track both UDF block generation time AND block serialization overhead ## Related issues > Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Changes --- 1. Modified Ray Core's generator handling sequence to inject back object creation & serialization durations 2. Updated `task_completion_time_excl_backpressure_s` to track both UDF block generation time AND block serialization overhead ## Related issues > Link related issues: "Fixes ray-project#1234", "Closes ray-project#1234", or "Related to ray-project#1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Adel Nour <ans9868@nyu.edu>

Changes --- 1. Modified Ray Core's generator handling sequence to inject back object creation & serialization durations 2. Updated `task_completion_time_excl_backpressure_s` to track both UDF block generation time AND block serialization overhead ## Related issues > Link related issues: "Fixes ray-project#1234", "Closes ray-project#1234", or "Related to ray-project#1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

alexeykudinkin added 12 commits January 28, 2026 10:47

Wired in serialization duration tracking into Ray Core generator

d93bc3d

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added wiring of serialization duration into async gen unrolling

89942ca

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed udf_time to be properly reset per block

e392681

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Properly track block serialization time in MapOperator

6fa647e

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Properly track block serialization time in HashShuffle Op base class

34ef64f

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added block_serialization_time_s;

fababc1

Updated `task_completion_time_excl_backpressure_s` to include block serde time Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Updated tests

5cb8233

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

d5bd5f7

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added text to verify task_completion_time_excl_backpressure_s inclu…

4d34d0e

…des serde time; Updated tests; Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Updating tests

b0119af

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed tests

e1a196d

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

35aa1c8

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin requested review from a team as code owners January 29, 2026 01:07

gemini-code-assist bot reviewed Jan 29, 2026

View reviewed changes

cursor bot reviewed Jan 29, 2026

View reviewed changes

python/ray/data/_internal/execution/operators/map_transformer.py Show resolved Hide resolved

ray-gardener bot added data Ray Data-related issues observability Issues related to the Ray Dashboard, Logging, Metrics, Tracing, and/or Profiling labels Jan 29, 2026

alexeykudinkin added 2 commits January 28, 2026 18:53

Fixed reset param handling

fa5eb26

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Made obj_store_mem_pending_task_outputs metric property

db9956b

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin added the go add ONLY when ready to merge, run all tests label Jan 29, 2026

cursor bot reviewed Jan 29, 2026

View reviewed changes

python/ray/data/block.py Show resolved Hide resolved

Fixed tests

ef96658

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

edoakes approved these changes Jan 30, 2026

View reviewed changes

alexeykudinkin added 3 commits January 30, 2026 13:39

Added _StreamingGeneratorStats

fcd78b6

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Updated usages

8cf9a5b

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

9d1aa20

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

cursor bot reviewed Jan 30, 2026

View reviewed changes

python/ray/data/_internal/execution/operators/hash_shuffle.py Outdated Show resolved Hide resolved

Fixed potential NPEs

deb1036

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin enabled auto-merge (squash) January 30, 2026 23:12

goutamvenkat-anyscale reviewed Jan 31, 2026

View reviewed changes

goutamvenkat-anyscale approved these changes Jan 31, 2026

View reviewed changes

alexeykudinkin merged commit 2d85efa into master Jan 31, 2026
7 checks passed

alexeykudinkin deleted the ak/serde-mtrcs-add branch January 31, 2026 01:55

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Data] Added tracking of block serialization time #60574

[Data] Added tracking of block serialization time #60574
alexeykudinkin merged 19 commits intomasterfrom
ak/serde-mtrcs-add

alexeykudinkin commented Jan 29, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jan 29, 2026

Uh oh!

Uh oh!

Uh oh!

edoakes left a comment

Uh oh!

cursor bot left a comment

Uh oh!

Uh oh!

goutamvenkat-anyscale Jan 31, 2026

Uh oh!

alexeykudinkin Jan 31, 2026

Uh oh!

goutamvenkat-anyscale Jan 31, 2026

Uh oh!

alexeykudinkin Jan 31, 2026

Uh oh!

goutamvenkat-anyscale Jan 31, 2026

Uh oh!

alexeykudinkin Jan 31, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

alexeykudinkin commented Jan 29, 2026

Changes

Related issues

Additional information

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

edoakes left a comment

Choose a reason for hiding this comment

Uh oh!

cursor bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

goutamvenkat-anyscale Jan 31, 2026

Choose a reason for hiding this comment

Uh oh!

alexeykudinkin Jan 31, 2026

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Jan 31, 2026

Choose a reason for hiding this comment

Uh oh!

alexeykudinkin Jan 31, 2026

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Jan 31, 2026

Choose a reason for hiding this comment

Uh oh!

alexeykudinkin Jan 31, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants