Skip to content

[Data] Parallelize ray.put of read tasks during input generation - #158

Open
rubinfergersoniv wants to merge 1 commit into
pinterest/main-2.52.1from
rfergerson/parallelize-read-task-put
Open

[Data] Parallelize ray.put of read tasks during input generation#158
rubinfergersoniv wants to merge 1 commit into
pinterest/main-2.52.1from
rfergerson/parallelize-read-task-put

Conversation

@rubinfergersoniv

@rubinfergersoniv rubinfergersoniv commented Jul 27, 2026

Copy link
Copy Markdown

Summary

Problem (Why)

Previously plan_read_op serialized read tasks into the object store one at a time (ray.put per task) and issued a separate get_local_object_locations RPC per task inside _derive_metadata. With many read tasks this per-task IPC latency accumulates serially and slows down input generation.

Solution (What + How)

ray.put releases the GIL during the plasma write and raylet RPC, so the work is IPC-bound. This change overlaps that latency with a ThreadPoolExecutor (executor.map preserves input order) and fetches all serialized sizes in a single get_local_object_locations call.

The pool size defaults to CPython's I/O-bound heuristic (min(32, cpu_count + 4)) and is overridable via RAY_DATA_READ_TASK_PUT_MAX_WORKERS for benchmarking. Falls back to the sequential path when there is <=1 task or a single worker.

Test Plan

...

Previously plan_read_op serialized read tasks into the object store one
at a time (`ray.put` per task) and issued a separate
`get_local_object_locations` RPC per task inside `_derive_metadata`. With
many read tasks this per-task IPC latency accumulates serially and slows
down input generation.

`ray.put` releases the GIL during the plasma write and raylet RPC, so the
work is IPC-bound. This change overlaps that latency with a
`ThreadPoolExecutor` (`executor.map` preserves input order) and fetches
all serialized sizes in a single `get_local_object_locations` call.

The pool size defaults to CPython's I/O-bound heuristic
(`min(32, cpu_count + 4)`) and is overridable via
`RAY_DATA_READ_TASK_PUT_MAX_WORKERS` for benchmarking. Falls back to the
sequential path when there is <=1 task or a single worker.

Co-authored-by: Cursor <cursoragent@cursor.com>
@github-actions

Copy link
Copy Markdown

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

@github-actions github-actions Bot added the stale label Aug 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant