[karticam]

This is PR 6 of a series of PRs to implement move semantics for plasma objects in core.

The complete design doc is here: https://docs.google.com/document/d/1ZJhgRmwVyKfNCR21CWf6xMcc0B6HJwSPQsfebTRvnnE/edit?tab=t.0#heading=h.7wcod1f9afpg

After #67093, a moved object that the consumer's raylet pulls is pinned there as the new primary copy. This PR covers the cases where the bytes are already on the consumer's node. This can happen if:

1. the owner is on the same node as the consumer and does a ray.get(), or
2. the consumer task pulled the object, but was retried due to worker failure midway and was later rescheduled on the same node.
   This PR handles the different edge cases and races in those scenarios. Every moved arg of an actor task now ends up pinned on the consumer's node, whichever way it got there, unless that node already held the primary copy (this is the case when producer and consumer are already on the same node)

Background

When an actor task is scheduled on a worker, the worker asks its raylet to make the task's args local (WaitForActorCallArgs). For each arg that is not local, the raylet starts a wait pull through the lease dependency manager, and it replies to the worker once every arg is local. The worker then reads the args from plasma itself. If an arg has gone missing by then, the worker's core worker starts its own get pull for it.

Two facts about this path drive the changes below:

• The raylet learns that an object is local or gone from the store's "object added" and "object deleted" notifications, which it handles on its main thread. So the raylet's view can lag the store by one queued notification: an object the raylet believes is local may have been evicted a moment ago.
• A plain received copy is unprotected between its seal and the moment the pull manager takes a reference on it in HandleObjectAdded. The store may evict it in that gap. [core][move semantics][5/n] Raylet pins a pulled consume_once object as the primary copy and reports it to the owner #67093 closed that gap for moves by keeping the create reference until after the pin, but a copy that was sealed as a plain copy, before its move was recorded, has no such protection.

#67041 records a moved arg when its wait pull starts, and #67093 pins the copy as the new primary when the store reports it local. Both assume the consumer's raylet is the one pulling the object. This PR covers the cases where it is not.

Changes

1. The raylet handles each moved arg once, when the task's args are requested. The new NodeManager::HandleMovedActorCallArgs runs before the wait that fetches the args, and does one of three things per arg flagged as a move:

   • If this node already holds the primary copy, nothing. There is nothing to move. "Holds the primary copy" means pinned in memory, pinned and being spilled, or spilled from this node; the new LocalObjectManager::HoldsPrimaryCopy checks all three. For the spilled case, a spilled primary is not in plasma, so without it the raylet would record a move, the pull would restore the object from its own disk, and the restored copy would be pinned a second time and reported as a move to the node it already lives on.
   • If a plain copy is local, for example left by the owner's earlier ray.get on this node, pin it as the new primary right now with PinAndReportMovedObject. No transfer happens.
   • Otherwise, record the move. The wait right after starts the pull, and when the store reports the object local, HandleObjectLocal pins it as the new primary and reports it to the owner, as in [core][move semantics][5/n] Raylet pins a pulled consume_once object as the primary copy and reports it to the owner #67093.

   The lease dependency manager goes back to its shape before [core][move semantics][4/n] Consumer raylet records which pulls are moves #67041. The per-arg decision needs the local object manager, which the lease dependency manager cannot see, so the marking moves up into the request handler, and the is_move parameter that [core][move semantics][4/n] Consumer raylet records which pulls are moves #67041 added to StartOrUpdateWaitRequest goes away.

2. A move record names the worker that asked for it, and dies with that worker. ObjectManager now keeps object to worker id. An object is moved by exactly one task, so a later mark for the same object replaces the worker; the only way that happens is a retry of the task on a restarted actor on the same node, and it is logged.

   Why the worker matters: normally a record belongs to a pull, and that pull's cancel erases it. After a failed pin (change 3) the record is written with no pull owning it, waiting for the worker's own fetch of its args to start one. If the worker dies before that fetch, no pull ever comes, and nothing would erase the record. A stale record is harmful: the next plain pull of that object on this node, say the owner's ray.get, would be treated as a move, pinned here and reported to the owner. So DisconnectClient now calls the new EraseMovePullsRequestedBy, which erases the dead worker's records that no pull needs. A record that a live pull still needs is left alone: that pull's cancel erases it, and until then the copy must still arrive as a move. PullManager::HasObjectPullRequest answers whether any pull, active or queued, still names the object.

3. HandleObjectLocal recovers from a pin miss. The pin at the top of HandleObjectLocal from [core][move semantics][5/n] Raylet pins a pulled consume_once object as the primary copy and reports it to the owner #67093 already runs for any object with a record, so it now also covers a copy that arrives through someone else's pull, typically the owner's ray.get on the consumer's node that was in flight when the move was requested. This PR adds what happens when that pin finds nothing in the store. There are two ways to get there:

   • At request time (change 1, second branch): the raylet believes a plain copy is local, but the store evicted it a moment ago and the "object deleted" notification is still queued. The pin misses. The handler records the move anyway. The lease dependency manager also still believes the object is local, so it starts no pull, and the wait completes at once. The worker then finds the object missing, starts its own get pull, and the record makes that pull a move.
   • At arrival time: the copy was sealed as a plain copy, so nothing held it, and the store evicted it in the gap before the pull manager took its reference. The move was recorded in between, by a request that saw the pull in flight. HandleObjectLocal finds the record, tries to pin, and misses. It then calls the lease dependency manager as usual, which cancels the wait pull and, if that was the last pull, erases the record. So after that call the record is written back for the same worker. As above, the worker's own fetch pulls the object again, and that pull is a move.

   In both cases the record has no pull between the miss and the worker's fetch; change 2 is what cleans it up if the worker dies in that window. PinAndReportMovedObject also returns early when this node already holds the primary copy, since it now has two callers that can see the same object. A pin that fails because the owner is dead or the object is being freed writes no record, since no later pull will bring the object here.

Scenarios

P is the producer's node, C the consumer's node, and O the owner's node. In the "Owner placement" column, "O = C" means the owner runs on the consumer's node, "O = P" means the owner runs on the producer's node, and "O on its own node" means the owner runs on a third node that is neither. "O = head" means the owner is the driver on the head node, which is also neither P nor C. "P = C" in row 6 means the producer and the consumer run on the same node. Each row ends with C holding a pinned primary copy of the object and P's copy untouched (the owner frees it in a later PR), unless the row says otherwise. Only the owner may read a consume_once object without moving it, so every "copy on C before the move" comes from the owner's own read on C. Tests are in python/ray/tests/test_consume_once_cross_node.py unless a C++ test is named.

#	Scenario	Owner placement	What happens	Test
1	Owner reads the object, then it is moved	O = C; O = P; O on its own node	On C: pinned in place, no transfer (zero-copy and pickled). On P: the owner reads the primary in place, C pulls as a move. On its own node: a plain copy stays there and the move ignores it	test_consume_of_local_copy_pins_it_without_transfer, test_topology_owner_read_then_consume
2	Owner's copy on C was evicted before the move	O = C	If the raylet already saw the eviction: ordinary record and pull. If not (pin miss at request time): record with no pull, the worker's own fetch pulls it as a move	test_consume_of_evicted_local_copy_pulls_it_as_move; node_manager_test.TestHandleMovedActorCallArgs
3	Owner's fetch of the objects (ray.wait(fetch_local=True) or a blocking ray.get) is in flight when they are moved	O = C; O = P; O on its own node	The move shares the owner's pull on C; each object is created once on C and pinned in HandleObjectLocal	test_consume_during_owner_fetch_pulls_once, test_consume_during_owner_get_pulls_once, test_topology_owner_fetch_in_flight_then_consume
4	Copy sealed through the owner's pull and evicted before the pull manager held it (pin miss at arrival time)	O = C	The record is written back after the wait pull's cancel; the worker's own fetch pulls it as a move	node_manager_test.TestHandleObjectLocalKeepsRecordAfterPinMiss
5	Consumer worker dies while its moves are in flight	O = head	Objects stay readable and P's copies untouched. A record with no pull is swept at disconnect; a record the owner's pull still needs is kept and the arriving copy is pinned, which is harmless since no one else may move it. C serves later moves	test_consumer_death_mid_move_keeps_object; object_manager_test.MovePullRecordFollowsItsRequesterWithoutAPull
6	Producer and consumer on the same node, primaries in memory or spilled to disk	P = C	The node already holds the primary copies: nothing is recorded, a spilled one is restored, nothing new is pinned	test_same_node_consume_keeps_single_copy, test_same_node_consume_of_spilled_primary_is_not_a_move; local_object_manager_test.TestHoldsPrimaryCopy
7	The #67093 cases rerun with every owner placement: plain move with a later owner read, c.remote(r, r) and two args, producer node death after the move, producer copies spilled	O = head; O = P; O = C; O on its own node	Unchanged behaviour	test_topology_move_pins_consumer_copy, test_topology_repeated_and_two_args, test_topology_move_survives_producer_node_death, and the #67093 tests

Record lifetime on its own: a record follows its last pull, and without a pull its requesting worker (object_manager_test); HasObjectPullRequest is true from a pull's start to its cancel, active or not (pull_manager_test).

Ran all of the above plus object_buffer_pool_test, local_lease_manager_test, lease_dependency_manager_test and the targets of the previous PRs, the full //src/ray/... build, test_consume_once.py (32 passed), and, as a check that a node that never sees a move behaves as before, test_object_manager (13), test_basic_2 (24), test_actor (89), test_object_spilling (26), test_reconstruction (21), all passed.

No new benchmarks: on a node that never receives a move the added work is one reader-locked lookup per HandleObjectLocal and one scan of the flags per actor task request.

[gemini-code-assist]

Code Review

This pull request refactors the tracking of "move pull" records in Ray by associating them with the specific requesting worker ID. This allows unowned move records to be cleaned up when a worker disconnects, preventing stale records from incorrectly turning subsequent plain pulls into moves. The feedback recommends addressing a potential denial-of-service vulnerability by replacing a crashing RAY_CHECK_EQ with safe validation on RPC input sizes, avoiding test flakiness by replacing time.sleep(1) with wait_for_condition, and adopting the safer C++11 it = map.erase(it) pattern for erasing elements during map iteration.

[gemini-code-assist]

Using RAY_CHECK_EQ to validate input sizes from RPC messages (which can be sent by worker processes) is a robustness and security risk. If a worker sends a malformed message with mismatched is_move and refs sizes, it will crash the shared raylet daemon. It is safer to log a warning and return early instead of crashing.

  if (is_move.size() != refs.size()) {\n    RAY_LOG(WARNING) << \"Received HandleMovedActorCallArgs with mismatched is_move size: \"\n                     << is_move.size() << \" vs \" << refs.size();\n    return;\n  }

[gemini-code-assist]

Using time.sleep(1) in tests can lead to flakiness and slows down the test suite. Since we want to wait for the primary-moved report to reach the owner, we can use wait_for_condition to check that the owner's object directory has been updated with the new location (node_b.node_id).

Suggested change

	time.sleep(1)
	wait_for_condition(lambda: node_b.node_id in ray.get(owner.locations.remote())[0])

[gemini-code-assist]

Using it = move_pulls_.erase(it) is the standard C++11 way to erase elements from a map while iterating. It is cleaner, safer, and more idiomatic than the post-increment trick (move_pulls_.erase(it++)), especially for flat hash maps where iterator invalidation rules can be subtle.

    if (it->second == worker_id && !pull_manager_->HasObjectPullRequest(it->first)) {\n      it = move_pulls_.erase(it);\n    } else {\n      ++it;\n    }

[gemini-code-assist]

Using it = move_pulls.erase(it) is the standard C++11 way to erase elements from a map while iterating. It is cleaner, safer, and more idiomatic than the post-increment trick (move_pulls.erase(it++)).

    for (auto it = move_pulls.begin(); it != move_pulls.end();) {\n      if (it->second == worker_id) {\n        it = move_pulls.erase(it);\n      } else {\n        ++it;\n      }\n    }

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes 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 39c23b7. Configure here.}

[cursor]

Missing Doxygen on new APIs

Low Severity

⚠️ Document functions and classes with Doxygen /** ... */ block comments using @ tags (@brief, and @param/@return where applicable).

This violates the Doxygen documentation rule for new C++ APIs. HoldsPrimaryCopy has no comment on the class declaration, HasObjectPullRequest omits @param/@return, and the new ObjectManagerInterface methods (MarkMovePull, GetMovePullRequester, EraseMovePullsRequestedBy) have no comments after the previous ones were removed.

Additional Locations (2)

• src/ray/object_manager/pull_manager.h#L151-L153
• src/ray/object_manager/object_manager_interface.h#L35-L40

 

^{Triggered by project rule: Bugbot Rules}

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