[core] (FreeObjects 2/n) Remove original FreeObjects RPC and original ObjectEviction pubsub #63181

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

Sparks0219 merged 48 commits into ray-project:master from aaronscalene:aaron/gut-object-eviction-pubsub

Jun 30, 2026

src/mock/ray/object_manager/object_manager.h

-Original file line number
+Diff line change
@@ Expand Up / @@ -41,10 +41,7 @@ class MockObjectManager : public ObjectManagerInterface { @@
                   (const TaskMetricsKey &task_key),
                   (const, override));
       MOCK_METHOD(int, GetServerPort, (), (const, override));
-      MOCK_METHOD(void,
-                  FreeObjects,
-                  (const std::vector<ObjectID> &object_ids, bool local_only),
-                  (override));
+      MOCK_METHOD(void, FreeObjects, (const std::vector<ObjectID> &object_ids), (override));
       MOCK_METHOD(bool, IsPlasmaObjectSpillable, (const ObjectID &object_id), (override));
       MOCK_METHOD(int64_t, GetUsedMemory, (), (const, override));
       MOCK_METHOD(bool, PullManagerHasPullsQueued, (), (const, override));
@@ Expand Down @@

src/ray/core_worker/core_worker.cc

-Original file line number
+Diff line change
@@ Expand Up / @@ -3802,60 +3802,6 @@ void CoreWorker::HandleWaitForActorRefDeleted( @@
       }
     }
-    void CoreWorker::ProcessSubscribeForObjectEviction(
-        const rpc::WorkerObjectEvictionSubMessage &message) {
-      // Send a response to trigger unpinning the object when it is no longer in scope.
-      auto unpin_object = [this](const ObjectID &object_id) {
-        RAY_LOG(DEBUG).WithField(object_id) << "Object is deleted. Unpinning the object.";
-        rpc::PubMessage pub_message;
-        pub_message.set_key_id(object_id.Binary());
-        pub_message.set_channel_type(rpc::ChannelType::WORKER_OBJECT_EVICTION);
-        pub_message.mutable_worker_object_eviction_message()->set_object_id(
-            object_id.Binary());
-        object_info_publisher_->Publish(std::move(pub_message));
-      };
-      const auto object_id = ObjectID::FromBinary(message.object_id());
-      const auto intended_worker_id = WorkerID::FromBinary(message.intended_worker_id());
-      if (intended_worker_id != worker_context_->GetWorkerID()) {
-        RAY_LOG(INFO).WithField(object_id)
-            << "The SubscribeForObjectEviction message for object is for worker "
-            << intended_worker_id << ", but the current worker is "
-            << worker_context_->GetWorkerID() << ". The RPC will be no-op.";
-        unpin_object(object_id);
-        return;
-      }
-      if (message.has_generator_id()) {
-        // For dynamically generated return values, the raylet may subscribe to
-        // eviction events before we know about the object. This can happen when we
-        // receive the subscription request before the reply from the task that
-        // created the object. Add the dynamically created object to our ref
-        // counter so that we know that it exists.
-        const auto generator_id = ObjectID::FromBinary(message.generator_id());
-        RAY_CHECK(!generator_id.IsNil());
-        if (task_manager_->ObjectRefStreamExists(generator_id)) {
-          // ObjectRefStreamExists is used to distinguigsh num_returns="dynamic" vs
-          // "streaming".
-          task_manager_->TemporarilyOwnGeneratorReturnRefIfNeeded(object_id, generator_id);
-        } else {
-          reference_counter_->AddDynamicReturn(object_id, generator_id);
-        }
-      }
-      // Returns true if the object was present and the callback was added. It might have
-      // already been evicted by the time we get this request, in which case we should
-      // respond immediately so the raylet unpins the object.
-      if (!reference_counter_->AddObjectOutOfScopeOrFreedCallback(object_id, unpin_object)) {
-        // If the object is already evicted (callback cannot be set), unregister the
-        // subscription & publish the message so that the subscriber knows it.
-        unpin_object(object_id);
-        RAY_LOG(DEBUG).WithField(object_id) << "Reference for object has already been freed.";
-      }
-    }
     StatusSet<StatusT::InvalidArgument> CoreWorker::ProcessSubscribeMessage(
         const rpc::SubMessage &sub_message,
         rpc::ChannelType channel_type,
@@ Expand All @@
         return result;
       }
-      if (!sub_message.has_worker_object_eviction_message() &&
-          !sub_message.has_worker_ref_removed_message() &&
+      if (!sub_message.has_worker_ref_removed_message() &&
           !sub_message.has_worker_object_locations_message()) {
         return StatusT::InvalidArgument(
             absl::StrFormat("Unexpected subscribe command has been received: %s"
-                            "Expected worker_object_eviction, worker_ref_removed, or "
+                            "Expected worker_ref_removed or "
                             "worker_object_locations message",
                             sub_message.DebugString()));
       }
-      if (sub_message.has_worker_object_eviction_message()) {
-        ProcessSubscribeForObjectEviction(sub_message.worker_object_eviction_message());
-      } else if (sub_message.has_worker_ref_removed_message()) {
+      if (sub_message.has_worker_ref_removed_message()) {
         ProcessSubscribeForRefRemoved(sub_message.worker_ref_removed_message());
       } else {  // worker_object_locations_message case
         ProcessSubscribeObjectLocations(sub_message.worker_object_locations_message());
@@ Expand Down @@

src/ray/core_worker/core_worker.h

-Original file line number
+Diff line change
@@ Expand Up @@
                                       std::vector<rpc::ObjectReference> *arg_refs,
                                       std::vector<ObjectID> *pinned_ids);
-      /// Process a subscribe message for wait for object eviction.
-      /// The object eviction message will be published once the object
-      /// needs to be evicted.
-      void ProcessSubscribeForObjectEviction(
-          const rpc::WorkerObjectEvictionSubMessage &message);
       /// Process a subscribe message for wait for ref removed.
       /// It is used for the ref counting protocol. When the borrower
       /// stops using the reference, the message will be published to the owner.
@@ Expand Down @@

src/ray/core_worker/core_worker_process.cc

-Original file line number
+Diff line change
@@ Expand Up @@
           std::make_shared<pubsub::Publisher>(
               /*channels=*/
               std::vector<rpc::ChannelType>{
-                  rpc::ChannelType::WORKER_OBJECT_EVICTION,
                   rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL,
                   rpc::ChannelType::WORKER_OBJECT_LOCATIONS_CHANNEL},
               /*periodical_runner=*/*periodical_runner,
@@ Expand All @@
       auto object_info_subscriber = std::make_unique<pubsub::Subscriber>(
           /*subscriber_id=*/worker_context->GetWorkerID(),
           /*channels=*/
-          std::vector<rpc::ChannelType>{rpc::ChannelType::WORKER_OBJECT_EVICTION,
-                                        rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL,
+          std::vector<rpc::ChannelType>{rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL,
                                         rpc::ChannelType::WORKER_OBJECT_LOCATIONS_CHANNEL},
           /*max_command_batch_size*/ RayConfig::instance().max_command_batch_size(),
           /*get_client=*/
@@ Expand Down @@

src/ray/core_worker/tests/core_worker_test.cc

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -138,8 +138,7 @@ class CoreWorkerTest : public ::testing::Test {
  
        auto object_info_publisher = std::make_unique<pubsub::Publisher>(

            /*channels=*/

            std::vector<rpc::ChannelType>{rpc::ChannelType::WORKER_OBJECT_EVICTION,

                                          rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL,

            std::vector<rpc::ChannelType>{rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL,

                                          rpc::ChannelType::WORKER_OBJECT_LOCATIONS_CHANNEL},

            /*periodical_runner=*/*fake_periodical_runner_,

            /*clock=*/clock_,

    @@ -798,119 +797,9 @@ TEST(CoreWorkerPlasmaStoreProviderFastPath, SendsOnlyRemoteIdsToRayletOnMixed) {
  
      EXPECT_EQ(pulled, (absl::flat_hash_set<ObjectID>{ids[1], ids[3]}));

    }

    class CoreWorkerPubsubWorkerObjectEvictionChannelTest

        : public CoreWorkerTest,

          public ::testing::WithParamInterface<bool> {};

    TEST_P(CoreWorkerPubsubWorkerObjectEvictionChannelTest, HandlePubsubCommandBatchRetries) {

      // should_free_object: determines whether the object is freed from plasma. This is used

      // to trigger AddObjectOutOfScopeOrFreedCallback in HandlePubsubCommandBatch which

      // stores the unpin_object callback that publishes the message to the

      // WORKER_OBJECT_EVICTION channel

      // should_free_object == true: the object is freed from plasma and we expect the message

      // to the WORKER_OBJECT_EVICTION channel to be published.

      // should_free_object == false: the object is not freed and we expect the message to the

      // WORKER_OBJECT_EVICTION channel to not be published.

      bool should_free_object = GetParam();

      auto subscriber_id = NodeID::FromRandom();

      auto object_id = ObjectID::FromRandom();

      rpc::Address owner_address;

      owner_address.set_worker_id(core_worker_->GetWorkerID().Binary());

      reference_counter_->AddOwnedObject(object_id,

                                         {},

                                         owner_address,

                                         "",

                                         0,

                                         LineageReconstructionEligibility::INELIGIBLE_PUT,

                                         true);

      rpc::PubsubCommandBatchRequest command_batch_request;

      command_batch_request.set_subscriber_id(subscriber_id.Binary());

      auto *command = command_batch_request.add_commands();

      command->set_channel_type(rpc::ChannelType::WORKER_OBJECT_EVICTION);

      command->set_key_id(object_id.Binary());

      auto *sub_message = command->mutable_subscribe_message();

      auto *real_sub_message = sub_message->mutable_worker_object_eviction_message();

      real_sub_message->set_intended_worker_id(core_worker_->GetWorkerID().Binary());

      real_sub_message->set_object_id(object_id.Binary());

      *real_sub_message->mutable_subscriber_address() = rpc_address_;

      rpc::PubsubCommandBatchReply command_reply1;

      rpc::PubsubCommandBatchReply command_reply2;

      // Each call to HandlePubsubCommandBatch causes the reference counter to store the

      // unpin_object callback that publishes the WORKER_OBJECT_EVICTION message

      core_worker_->HandlePubsubCommandBatch(

          command_batch_request,

          &command_reply1,

          [](const Status &status, std::function<void()>, std::function<void()>) {

            ASSERT_TRUE(status.ok());

          });

      core_worker_->HandlePubsubCommandBatch(

          command_batch_request,

          &command_reply2,

          [](const Status &status, std::function<void()>, std::function<void()>) {

            ASSERT_TRUE(status.ok());

          });

      if (should_free_object) {

        // Triggers the unpin_object callbacks that publish the message to the

        // WORKER_OBJECT_EVICTION channel

        reference_counter_->FreePlasmaObjects({object_id});

      }

      rpc::PubsubLongPollingRequest request;

      request.set_subscriber_id(subscriber_id.Binary());

      request.set_max_processed_sequence_id(0);

      request.set_publisher_id("");

      rpc::PubsubLongPollingReply reply;

      // should_free_object == true: Each call to HandlePubsubCommandBatch adds an

      // unpin_object callback that is triggered via FreePlasmaObjects which publishes the

      // message to the WORKER_OBJECT_EVICTION channel, hence we have 1 publish per callback

      // so 2 in total. The long poll connection is closed

      // should_free_object == false: Since FreePlasmaObjects is not called, the unpin_object

      // callbacks are not triggered and we have 0 publishes. NOTE: The long poll connection

      // is not closed when should_free_object == false since there was no publish.

      core_worker_->HandlePubsubLongPolling(

          request,

          &reply,

          [](Status s, std::function<void()> success, std::function<void()> failure) {

            ASSERT_TRUE(s.ok());

          });

      int expected_messages = should_free_object ? 2 : 0;

      EXPECT_EQ(reply.pub_messages_size(), expected_messages);

      for (int i = 0; i < expected_messages; i++) {

        const auto &msg = reply.pub_messages(i);

        EXPECT_EQ(msg.channel_type(), rpc::ChannelType::WORKER_OBJECT_EVICTION);

        EXPECT_EQ(msg.key_id(), object_id.Binary());

        EXPECT_EQ(msg.sequence_id(), i + 1);

        EXPECT_EQ(msg.worker_object_eviction_message().object_id(), object_id.Binary());

      }

      if (!should_free_object) {

        // Since the long poll connection is not closed, we need to flush it. Otherwise this

        // can trigger undefined behavior since unlike in prod where grpc arena allocates the

        // reply, here we allocate the reply on the stack. Hence the normal order of

        // destruction is: reply goes out of scope -> publisher is destructed -> flushes the

        // reply which access freed memory

        clock_.AdvanceTime(absl::Milliseconds(RayConfig::instance().subscriber_timeout_ms()));

        object_info_publisher_->CheckDeadSubscribers();

      }

    }

    INSTANTIATE_TEST_SUITE_P(WorkerObjectEvictionChannel,

                             CoreWorkerPubsubWorkerObjectEvictionChannelTest,

                             ::testing::Values(true, false));

    TEST_F(CoreWorkerTest, HandlePubsubCommandBatchInvalidChannelType) {

      // Test that HandlePubsubCommandBatch returns InvalidArgument for an invalid channel

      // type. The publisher was created with only:

      // - WORKER_OBJECT_EVICTION

      // - WORKER_REF_REMOVED_CHANNEL

      // - WORKER_OBJECT_LOCATIONS_CHANNEL

      // Using a channel type that was not registered should return InvalidArgument.

    @@ -953,7 +842,7 @@ TEST_F(CoreWorkerTest,
  
      rpc::PubsubCommandBatchRequest command_batch_request;

      command_batch_request.set_subscriber_id(subscriber_id.Binary());

      auto *command = command_batch_request.add_commands();

      command->set_channel_type(rpc::ChannelType::WORKER_OBJECT_EVICTION);

      command->set_channel_type(rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL);

      command->set_key_id(object_id.Binary());

      rpc::PubsubCommandBatchReply command_reply;

    @@ -977,7 +866,7 @@ TEST_F(CoreWorkerTest,
  
      rpc::PubsubCommandBatchRequest command_batch_request;

      command_batch_request.set_subscriber_id(subscriber_id.Binary());

      auto *command = command_batch_request.add_commands();

      command->set_channel_type(rpc::ChannelType::WORKER_OBJECT_EVICTION);

      command->set_channel_type(rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL);

      command->set_key_id(object_id.Binary());

      command->mutable_subscribe_message();

src/ray/gcs/tests/pubsub_handler_test.cc

-Original file line number
+Diff line change
@@ Expand Up @@
       rpc::GcsSubscriberCommandBatchRequest request;
       request.set_subscriber_id(subscriber_id.Binary());
       auto *command = request.add_commands();
-      command->set_channel_type(rpc::ChannelType::WORKER_OBJECT_EVICTION);
+      command->set_channel_type(rpc::ChannelType::WORKER_REF_REMOVED_CHANNEL);
       command->set_key_id(key_id.Binary());
       command->mutable_subscribe_message();
@@ Expand Down @@

src/ray/object_manager/object_manager.cc

-Original file line number
+Diff line change
@@ Expand Up / @@ -24,12 +24,10 @@ @@
     #include <vector>
     #include "absl/time/clock.h"
-    #include "ray/asio/asio_util.h"
     #include "ray/common/filter_local_objects_util.h"
     #include "ray/common/protobuf_utils.h"
     #include "ray/object_manager/plasma/store_runner.h"
     #include "ray/object_manager/spilled_object_reader.h"
-    #include "ray/util/exponential_backoff.h"
     #include "ray/util/network_util.h"
     #include "ray/util/time.h"
@@ Expand Down Expand Up / @@ -670,93 +668,8 @@ void ObjectManager::HandlePull(rpc::PullRequest request, @@
       send_reply_callback(Status::OK(), nullptr, nullptr);
     }
-    void ObjectManager::HandleFreeObjects(rpc::FreeObjectsRequest request,
-                                          rpc::FreeObjectsReply *reply,
-                                          rpc::SendReplyCallback send_reply_callback) {
-      std::vector<ObjectID> object_ids;
-      for (const auto &e : request.object_ids()) {
-        object_ids.emplace_back(ObjectID::FromBinary(e));
-      }
-      FreeObjects(object_ids, /* local_only */ true);
-      send_reply_callback(Status::OK(), nullptr, nullptr);
-    }
-    // TODO(#63213) will delete local_only=false and related dead code
-    void ObjectManager::FreeObjects(const std::vector<ObjectID> &object_ids,
-                                    bool local_only) {
+    void ObjectManager::FreeObjects(const std::vector<ObjectID> &object_ids) {
       buffer_pool_.FreeObjects(object_ids);
-      if (!local_only) {
-        std::vector<std::pair<NodeID, std::shared_ptr<rpc::ObjectManagerClientInterface>>>
-            rpc_clients;
-        // TODO(#56414): optimize this so we don't have to send a free objects request for
-        // every object to every node
-        const auto &node_info_map = gcs_client_.Nodes().GetAllNodeAddressAndLiveness();
-        for (const auto &[node_id, _] : node_info_map) {
-          if (node_id == self_node_id_) {
-            continue;
-          }
-          auto rpc_client = GetRpcClient(node_id);
-          if (rpc_client != nullptr) {
-            rpc_clients.emplace_back(node_id, std::move(rpc_client));
-          }
-        }
-        rpc_service_.post(
-            [this, object_ids, rpc_clients = std::move(rpc_clients)]() {
-              SpreadFreeObjectsRequest(object_ids, rpc_clients);
-            },
-            "ObjectManager.FreeObjects");
-      }
-    }
-    void ObjectManager::SpreadFreeObjectsRequest(
-        const std::vector<ObjectID> &object_ids,
-        const std::vector<
-            std::pair<NodeID, std::shared_ptr<rpc::ObjectManagerClientInterface>>>
-            &rpc_clients) {
-      // This code path should be called from node manager.
-      rpc::FreeObjectsRequest free_objects_request;
-      for (const auto &e : object_ids) {
-        free_objects_request.add_object_ids(e.Binary());
-      }
-      for (const auto &entry : rpc_clients) {
-        // NOTE: The callback for FreeObjects is posted back onto the main_service_ since
-        // RetryFreeObjects accesses remote_object_manager_clients_ which is not thread safe.
-        entry.second->FreeObjects(
-            free_objects_request,
-            [this, node_id = entry.first, free_objects_request](
-                const Status &status, const rpc::FreeObjectsReply &reply) {
-              if (!status.ok()) {
-                RetryFreeObjects(node_id, 0, free_objects_request);
-              }
-            });
-      }
-    }
-    void ObjectManager::RetryFreeObjects(
-        const NodeID &node_id,
-        uint32_t attempt_number,
-        const rpc::FreeObjectsRequest &free_objects_request) {
-      if (!remote_object_manager_clients_.contains(node_id)) {
-        return;
-      }
-      auto delay_ms = ExponentialBackoff::GetBackoffMs(attempt_number, 1000);
-      execute_after(
-          *main_service_,
-          [this, node_id, attempt_number, free_objects_request] {
-            auto it = remote_object_manager_clients_.find(node_id);
-            if (it == remote_object_manager_clients_.end()) {
-              return;
-            }
-            it->second->FreeObjects(
-                free_objects_request,
-                [this, node_id, attempt_number, free_objects_request](
-                    const Status &status, const rpc::FreeObjectsReply &reply) {
-                  if (!status.ok()) {
-                    RetryFreeObjects(node_id, attempt_number + 1, free_objects_request);
-                  }
-                });
-          },
-          std::chrono::milliseconds(delay_ms));
     }
     std::shared_ptr<rpc::ObjectManagerClientInterface> ObjectManager::GetRpcClient(
@@ Expand Down @@

Navigation Menu

Uh oh!

[core] (FreeObjects 2/n) Remove original FreeObjects RPC and original ObjectEviction pubsub #63181

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!