[Core] Move request id creation to worker to address plasma get perf regression (#58390)

## Description
This PR address the performance regression introduced in the [PR to make
ray.get thread safe](#57911).
Specifically, the previous PR requires the worker to block and wait for
AsyncGet to return with a reply of the request id needed for correctly
cleaning up get requests. This additional synchronous step causes the
plasma store Get to regress in performance.

This PR moves the request id generation step to the plasma store,
removing the blocking step to fix the perf regression.

## Related issues
- [PR which introduced perf
regression](#57911)
- [PR which observed the
regression](#58175)

## Additional information
New performance of the change measured by `ray microbenchmark`.
<img width="485" height="17" alt="image"
src="https://github.com/user-attachments/assets/b96b9676-3735-4e94-9ade-aaeb7514f4d0"
/>

Original performance prior to the change. Here we focus on the
regressing `single client get calls (Plasma Store)` metric, where our
new performance returns us back to the original 10k per second range
compared to the existing sub 5k per second.
<img width="811" height="355" alt="image"
src="https://github.com/user-attachments/assets/d1fecf82-708e-48c4-9879-34c59a5e056c"
/>

---------

Signed-off-by: davik <davik@anyscale.com>
Co-authored-by: davik <davik@anyscale.com>

Author: Kunchd

src/ray/core_worker/store_provider/plasma_store_provider.cc

@@ -75,7 +75,8 @@ CoreWorkerPlasmaStoreProvider::CoreWorkerPlasmaStoreProvider(
       store_client_(std::move(store_client)),
       reference_counter_(reference_counter),
       check_signals_(std::move(check_signals)),
-      fetch_batch_size_(fetch_batch_size) {
+      fetch_batch_size_(fetch_batch_size),
+      get_request_counter_(0) {
   if (get_current_call_site != nullptr) {
     get_current_call_site_ = get_current_call_site;
   } else {
@@ -275,8 +276,8 @@ Status CoreWorkerPlasmaStoreProvider::Get(
     // 1. Make the request to pull all objects into local plasma if not local already.
     std::vector<rpc::Address> owner_addresses =
         reference_counter_.GetOwnerAddresses(batch_ids);
-    StatusOr<ipc::ScopedResponse> status_or_cleanup =
-        raylet_ipc_client_->AsyncGetObjects(batch_ids, owner_addresses);
+    StatusOr<ipc::ScopedResponse> status_or_cleanup = raylet_ipc_client_->AsyncGetObjects(
+        batch_ids, owner_addresses, get_request_counter_.fetch_add(1));
     RAY_RETURN_NOT_OK(status_or_cleanup.status());
     get_request_cleanup_handlers.emplace_back(std::move(status_or_cleanup.value()));
 

src/ray/core_worker/store_provider/plasma_store_provider.h

@@ -261,6 +261,7 @@ class CoreWorkerPlasmaStoreProvider {
   // Pointer to the shared buffer tracker.
   std::shared_ptr<BufferTracker> buffer_tracker_;
   int64_t fetch_batch_size_ = 0;
+  std::atomic<int64_t> get_request_counter_;
 };
 
 }  // namespace core

src/ray/flatbuffers/node_manager.fbs

@@ -40,8 +40,6 @@ enum MessageType:int {
   DisconnectClientReply,
   // Request the Raylet to pull a set of objects to the local node.
   AsyncGetObjectsRequest,
-  // Reply contains the request id that will be used to clean up the request.
-  AsyncGetObjectsReply,
   // Cleanup a given get request on the raylet.
   CancelGetRequest,
   // Notify the current worker is blocked for objects to become available. The raylet
@@ -134,10 +132,7 @@ table AsyncGetObjectsRequest {
   // Object IDs that we want the Raylet to pull locally.
   object_ids: [string];
   owner_addresses: [Address];
-}
-
-table AsyncGetObjectsReply {
-  request_id: long;
+  get_request_id: long;
 }
 
 table CancelGetRequest {

src/ray/raylet/lease_dependency_manager.cc

@@ -115,8 +115,10 @@ void LeaseDependencyManager::CancelWaitRequest(const WorkerID &worker_id) {
   wait_requests_.erase(req_iter);
 }
 
-GetRequestId LeaseDependencyManager::StartGetRequest(
-    const WorkerID &worker_id, std::vector<rpc::ObjectReference> &&required_objects) {
+void LeaseDependencyManager::StartGetRequest(
+    const WorkerID &worker_id,
+    std::vector<rpc::ObjectReference> &&required_objects,
+    int64_t get_request_id) {
   std::vector<ObjectID> object_ids;
   object_ids.reserve(required_objects.size());
 
@@ -130,16 +132,12 @@ GetRequestId LeaseDependencyManager::StartGetRequest(
   uint64_t new_pull_request_id = object_manager_.Pull(
       std::move(required_objects), BundlePriority::GET_REQUEST, {"", false});
 
-  const GetRequestId get_request_id = get_request_counter_++;
-
   const std::pair<WorkerID, GetRequestId> worker_and_request_ids =
       std::make_pair(worker_id, get_request_id);
 
   get_requests_.emplace(std::move(worker_and_request_ids),
                         std::make_pair(std::move(object_ids), new_pull_request_id));
   worker_to_requests_[worker_id].emplace(get_request_id);
-
-  return get_request_id;
 }
 
 void LeaseDependencyManager::CancelGetRequest(const WorkerID &worker_id,

src/ray/raylet/lease_dependency_manager.h

@@ -135,9 +135,11 @@ class LeaseDependencyManager : public LeaseDependencyManagerInterface {
 
   /// \param worker_id The ID of the worker that called `ray.get`.
   /// \param required_objects The objects required by the worker.
-  /// \return the request id which will be used for cleanup.
-  GetRequestId StartGetRequest(const WorkerID &worker_id,
-                               std::vector<rpc::ObjectReference> &&required_objects);
+  /// \param get_request_id The ID of the get request. It is used by the worker to clean
+  /// up a GetRequest.
+  void StartGetRequest(const WorkerID &worker_id,
+                       std::vector<rpc::ObjectReference> &&required_objects,
+                       int64_t get_request_id);
 
   /// Cleans up either an inflight or finished get request. Cancels the underlying
   /// pull if necessary.
@@ -302,9 +304,6 @@ class LeaseDependencyManager : public LeaseDependencyManagerInterface {
   /// dependencies are all local or not.
   absl::flat_hash_map<LeaseID, std::unique_ptr<LeaseDependencies>> queued_lease_requests_;
 
-  /// Used to generate monotonically increasing get request ids.
-  GetRequestId get_request_counter_;
-
   // Maps a GetRequest to the PullRequest Id and the set of ObjectIDs.
   // Used to cleanup a finished or cancel an inflight get request.
   // TODO(57911): This can be slimmed down. We do not need to track the ObjectIDs.

src/ray/raylet/node_manager.cc

@@ -1624,21 +1624,7 @@ void NodeManager::HandleAsyncGetObjectsRequest(
   auto request = flatbuffers::GetRoot<protocol::AsyncGetObjectsRequest>(message_data);
   std::vector<rpc::ObjectReference> refs =
       FlatbufferToObjectReferences(*request->object_ids(), *request->owner_addresses());
-  int64_t request_id = AsyncGet(client, refs);
-  flatbuffers::FlatBufferBuilder fbb;
-  auto get_reply = protocol::CreateAsyncGetObjectsReply(fbb, request_id);
-  fbb.Finish(get_reply);
-  Status status = client->WriteMessage(
-      static_cast<int64_t>(protocol::MessageType::AsyncGetObjectsReply),
-      fbb.GetSize(),
-      fbb.GetBufferPointer());
-  if (!status.ok()) {
-    DisconnectClient(client,
-                     /*graceful=*/false,
-                     rpc::WorkerExitType::SYSTEM_ERROR,
-                     absl::StrFormat("Could not send AsyncGetObjectsReply because of %s",
-                                     status.ToString()));
-  }
+  AsyncGet(client, refs, request->get_request_id());
 }
 
 void NodeManager::ProcessWaitRequestMessage(
@@ -2373,15 +2359,16 @@ void NodeManager::HandleNotifyWorkerUnblocked(
   }
 }
 
-int64_t NodeManager::AsyncGet(const std::shared_ptr<ClientConnection> &client,
-                              std::vector<rpc::ObjectReference> &object_refs) {
+void NodeManager::AsyncGet(const std::shared_ptr<ClientConnection> &client,
+                           std::vector<rpc::ObjectReference> &object_refs,
+                           int64_t get_request_id) {
   std::shared_ptr<WorkerInterface> worker = worker_pool_.GetRegisteredWorker(client);
   if (!worker) {
     worker = worker_pool_.GetRegisteredDriver(client);
   }
   RAY_CHECK(worker);
-  return lease_dependency_manager_.StartGetRequest(worker->WorkerId(),
-                                                   std::move(object_refs));
+  lease_dependency_manager_.StartGetRequest(
+      worker->WorkerId(), std::move(object_refs), get_request_id);
 }
 
 void NodeManager::AsyncWait(const std::shared_ptr<ClientConnection> &client,

src/ray/raylet/node_manager.h

@@ -429,9 +429,11 @@ class NodeManager : public rpc::NodeManagerServiceHandler,
   /// \param client The client that is requesting the objects.
   /// \param object_refs The objects that are requested.
   ///
-  /// \return the request_id that will be used to cancel the get request.
-  int64_t AsyncGet(const std::shared_ptr<ClientConnection> &client,
-                   std::vector<rpc::ObjectReference> &object_refs);
+  /// \param get_request_id The ID of the get request. It is used by the worker to clean
+  /// up a GetRequest.
+  void AsyncGet(const std::shared_ptr<ClientConnection> &client,
+                std::vector<rpc::ObjectReference> &object_refs,
+                int64_t get_request_id);
 
   /// Cancel all ongoing get requests from the client.
   ///

src/ray/raylet/tests/lease_dependency_manager_test.cc

@@ -243,15 +243,14 @@ TEST_F(LeaseDependencyManagerTest, TestLeaseArgEviction) {
 TEST_F(LeaseDependencyManagerTest, TestCancelingSingleGetRequestForWorker) {
   WorkerID worker_id = WorkerID::FromRandom();
   int num_requests = 5;
-  std::vector<GetRequestId> requests;
-  for (int i = 0; i < num_requests; i++) {
+  for (int64_t i = 0; i < num_requests; i++) {
     ObjectID argument_id = ObjectID::FromRandom();
-    requests.emplace_back(lease_dependency_manager_.StartGetRequest(
-        worker_id, ObjectIdsToRefs({argument_id})));
+    lease_dependency_manager_.StartGetRequest(
+        worker_id, ObjectIdsToRefs({argument_id}), i);
   }
   ASSERT_EQ(object_manager_mock_.active_get_requests.size(), num_requests);
-  for (int i = 0; i < num_requests; i++) {
-    lease_dependency_manager_.CancelGetRequest(worker_id, requests[i]);
+  for (int64_t i = 0; i < num_requests; i++) {
+    lease_dependency_manager_.CancelGetRequest(worker_id, i);
     ASSERT_EQ(object_manager_mock_.active_get_requests.size(), num_requests - (i + 1));
   }
   AssertNoLeaks();
@@ -260,11 +259,10 @@ TEST_F(LeaseDependencyManagerTest, TestCancelingSingleGetRequestForWorker) {
 TEST_F(LeaseDependencyManagerTest, TestCancelingAllGetRequestsForWorker) {
   WorkerID worker_id = WorkerID::FromRandom();
   int num_requests = 5;
-  std::vector<GetRequestId> requests;
-  for (int i = 0; i < num_requests; i++) {
+  for (int64_t i = 0; i < num_requests; i++) {
     ObjectID argument_id = ObjectID::FromRandom();
-    requests.emplace_back(lease_dependency_manager_.StartGetRequest(
-        worker_id, ObjectIdsToRefs({argument_id})));
+    lease_dependency_manager_.StartGetRequest(
+        worker_id, ObjectIdsToRefs({argument_id}), i);
   }
   ASSERT_EQ(object_manager_mock_.active_get_requests.size(), num_requests);
   lease_dependency_manager_.CancelGetRequest(worker_id);

src/ray/raylet_ipc_client/fake_raylet_ipc_client.h

@@ -55,7 +55,8 @@ class FakeRayletIpcClient : public RayletIpcClientInterface {
 
   StatusOr<ScopedResponse> AsyncGetObjects(
       const std::vector<ObjectID> &object_ids,
-      const std::vector<rpc::Address> &owner_addresses) override {
+      const std::vector<rpc::Address> &owner_addresses,
+      int64_t get_request_id) override {
     return ScopedResponse();
   }
 

src/ray/raylet_ipc_client/raylet_ipc_client.cc

@@ -195,23 +195,22 @@ Status RayletIpcClient::ActorCreationTaskDone() {
 
 StatusOr<ScopedResponse> RayletIpcClient::AsyncGetObjects(
     const std::vector<ObjectID> &object_ids,
-    const std::vector<rpc::Address> &owner_addresses) {
+    const std::vector<rpc::Address> &owner_addresses,
+    int64_t get_request_id) {
   RAY_CHECK(object_ids.size() == owner_addresses.size());
   flatbuffers::FlatBufferBuilder fbb;
   auto object_ids_message = flatbuf::to_flatbuf(fbb, object_ids);
-  auto message = protocol::CreateAsyncGetObjectsRequest(
-      fbb, object_ids_message, AddressesToFlatbuffer(fbb, owner_addresses));
+  auto message =
+      protocol::CreateAsyncGetObjectsRequest(fbb,
+                                             object_ids_message,
+                                             AddressesToFlatbuffer(fbb, owner_addresses),
+                                             get_request_id);
   fbb.Finish(message);
   std::vector<uint8_t> reply;
   // TODO(57923): This should be FATAL. Local sockets are reliable. If a worker is unable
   // to communicate with the raylet, there's no way to recover.
-  RAY_RETURN_NOT_OK(AtomicRequestReply(MessageType::AsyncGetObjectsRequest,
-                                       MessageType::AsyncGetObjectsReply,
-                                       &reply,
-                                       &fbb));
-  auto reply_message = flatbuffers::GetRoot<protocol::AsyncGetObjectsReply>(reply.data());
-  int64_t request_id = reply_message->request_id();
-  return ScopedResponse([this, request_id_to_cleanup = request_id]() {
+  RAY_RETURN_NOT_OK(WriteMessage(MessageType::AsyncGetObjectsRequest, &fbb));
+  return ScopedResponse([this, request_id_to_cleanup = get_request_id]() {
     return CancelGetRequest(request_id_to_cleanup);
   });
 }