[core] Fix race condition b/w object eviction & repinning for recovery. (#53934)

This PR has two components:

### Fixing the Race Condition b/w Object Eviction and Repining for
Object Recovery

The sequence of events documents the life of an object called A.

1. Owner sends object eviction notifications to all locations as part of
Object A going out of scope. (done through worker<->worker pubsub)
2. Workers that have a copy of Object A mark it for asynchronous
deletion.
3. Owner decides to recover the object and repin a copy of object A that
has already been marked for asynchronous deletion.
4. Owner submits a task that has Object A as an argument.
5. The task executor tries to fetch Object A and the pull manager hangs
forever.

The fix is to check to see if an object is pending deletion before
pinning it. If it is, the owner cannot pin the object and has to either
try a different location or trigger reconstruction by resubmitting the
task that created the object.

### (The Joys of) Writing a Unit Test in Node Manager

There are a few problems with writing unit tests in the NodeManager

1. Despite heroics from @rueian and @dayshah recently, the NodeManager
still had a few dependencies that needed to be put behind interfaces so
fakes could be injected.
2. The gRPC handlers are all private to the class so they cannot be
called directly. Instead the unit tests need to create a gRPC client.
3. The NodeManager has a lot of dependencies and the dependencies have
interdependencies so we currently dedicate 200 lines to constructing a
test fixture.
4. There is a lack of usable fakes e.g. for the PlasmaClient for writing
unit tests.

### Follow Ups
1. Create a shim between the NodeManager and it's gRPC handlers so they
can be unit tested directly without creating gRPC clients.
2. Create a utility function that lets you construct a NodeManager for
testing optionally overriding interesting dependencies (otherwise they
will default to a Noop fake).
3. Move the fakes created in this PR into their own (more complete
implementations).
4. Create a Fake for SubscriberInterface.

---------

Signed-off-by: irabbani <[email protected]>
Signed-off-by: Ibrahim Rabbani <[email protected]>
Signed-off-by: elliot-barn <[email protected]>

Author: israbbani

BUILD.bazel

@@ -567,11 +567,22 @@ ray_cc_library(
     ],
 )
 
+ray_cc_library(
+    name ="local_object_manager_interface",
+    hdrs = ["src/ray/raylet/local_object_manager_interface.h"],
+    deps = [
+        "//src/ray/common:id",
+        "//src/ray/common:ray_object",
+        "//src/ray/protobuf:node_manager_cc_proto",
+    ],
+)
+
 ray_cc_library(
     name = "local_object_manager",
     srcs = ["src/ray/raylet/local_object_manager.cc"],
     hdrs = ["src/ray/raylet/local_object_manager.h"],
     deps = [
+        ":local_object_manager_interface",
         ":worker_pool",
         ":worker_rpc",
         "//src/ray/common:id",
@@ -615,6 +626,7 @@ ray_cc_library(
         ],
     }),
     deps = [
+        ":local_object_manager_interface",
         ":local_object_manager",
         ":node_manager_fbs",
         ":node_manager_rpc",
@@ -765,6 +777,8 @@ ray_cc_test(
     tags = ["team:core"],
     deps = [
         ":ray_mock",
+        ":local_object_manager_interface",
+        "//src/ray/util:macros",
         ":raylet_lib",
         "@com_google_googletest//:gtest_main",
     ],

src/ray/raylet/local_object_manager.cc

@@ -139,24 +139,33 @@ void LocalObjectManager::ReleaseFreedObject(const ObjectID &object_id) {
 
   // Try to evict all copies of the object from the cluster.
   if (free_objects_period_ms_ >= 0) {
-    objects_to_free_.push_back(object_id);
+    objects_pending_deletion_.emplace(object_id);
   }
-  if (objects_to_free_.size() == free_objects_batch_size_ ||
+  if (objects_pending_deletion_.size() == free_objects_batch_size_ ||
       free_objects_period_ms_ == 0) {
     FlushFreeObjects();
   }
 }
 
 void LocalObjectManager::FlushFreeObjects() {
-  if (!objects_to_free_.empty()) {
-    RAY_LOG(DEBUG) << "Freeing " << objects_to_free_.size() << " out-of-scope objects";
-    on_objects_freed_(objects_to_free_);
-    objects_to_free_.clear();
+  if (!objects_pending_deletion_.empty()) {
+    RAY_LOG(DEBUG) << "Freeing " << objects_pending_deletion_.size()
+                   << " out-of-scope objects";
+    // TODO(irabbani): CORE-1640 will modify as much as the plasma API as is
+    // reasonable to remove usage of vectors in favor of sets.
+    std::vector<ObjectID> objects_to_delete(objects_pending_deletion_.begin(),
+                                            objects_pending_deletion_.end());
+    on_objects_freed_(objects_to_delete);
+    objects_pending_deletion_.clear();
   }
   ProcessSpilledObjectsDeleteQueue(free_objects_batch_size_);
   last_free_objects_at_ms_ = current_time_ms();
 }
 
+bool LocalObjectManager::ObjectPendingDeletion(const ObjectID &object_id) {
+  return objects_pending_deletion_.find(object_id) != objects_pending_deletion_.end();
+}
+
 void LocalObjectManager::SpillObjectUptoMaxThroughput() {
   if (RayConfig::instance().object_spilling_config().empty()) {
     return;

src/ray/raylet/local_object_manager.h

@@ -18,6 +18,7 @@
 #include <memory>
 #include <queue>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "ray/common/id.h"
@@ -26,9 +27,9 @@
 #include "ray/object_manager/common.h"
 #include "ray/object_manager/object_directory.h"
 #include "ray/pubsub/subscriber.h"
+#include "ray/raylet/local_object_manager_interface.h"
 #include "ray/raylet/worker_pool.h"
 #include "ray/rpc/worker/core_worker_client_pool.h"
-#include "src/ray/protobuf/node_manager.pb.h"
 
 namespace ray {
 
@@ -39,7 +40,7 @@ inline constexpr int64_t kDefaultSpilledObjectDeleteRetries = 3;
 
 /// This class implements memory management for primary objects, objects that
 /// have been freed, and objects that have been spilled.
-class LocalObjectManager {
+class LocalObjectManager : public LocalObjectManagerInterface {
  public:
   LocalObjectManager(
       const NodeID &node_id,
@@ -58,19 +59,19 @@ class LocalObjectManager {
       pubsub::SubscriberInterface *core_worker_subscriber,
       IObjectDirectory *object_directory)
       : self_node_id_(node_id),
-        self_node_address_(self_node_address),
+        self_node_address_(std::move(self_node_address)),
         self_node_port_(self_node_port),
         io_service_(io_service),
         free_objects_period_ms_(free_objects_period_ms),
         free_objects_batch_size_(free_objects_batch_size),
         io_worker_pool_(io_worker_pool),
         owner_client_pool_(owner_client_pool),
-        on_objects_freed_(on_objects_freed),
+        on_objects_freed_(std::move(on_objects_freed)),
         last_free_objects_at_ms_(current_time_ms()),
         min_spilling_size_(RayConfig::instance().min_spilling_size()),
         num_active_workers_(0),
         max_active_workers_(max_io_workers),
-        is_plasma_object_spillable_(is_plasma_object_spillable),
+        is_plasma_object_spillable_(std::move(is_plasma_object_spillable)),
         is_external_storage_type_fs_(is_external_storage_type_fs),
         max_fused_object_count_(max_fused_object_count),
         next_spill_error_log_bytes_(RayConfig::instance().verbose_spill_logs()),
@@ -95,12 +96,12 @@ class LocalObjectManager {
   void PinObjectsAndWaitForFree(const std::vector<ObjectID> &object_ids,
                                 std::vector<std::unique_ptr<RayObject>> &&objects,
                                 const rpc::Address &owner_address,
-                                const ObjectID &generator_id = ObjectID::Nil());
+                                const ObjectID &generator_id = ObjectID::Nil()) override;
 
   /// Spill objects as much as possible as fast as possible up to the max throughput.
   ///
   /// \return True if spilling is in progress.
-  void SpillObjectUptoMaxThroughput();
+  void SpillObjectUptoMaxThroughput() override;
 
   /// TODO(dayshah): This function is only used for testing, we should remove and just
   /// keep SpillObjectsInternal.
@@ -110,7 +111,7 @@ class LocalObjectManager {
   /// \param callback A callback to call once the objects have been spilled, or
   /// there is an error.
   void SpillObjects(const std::vector<ObjectID> &objects_ids,
-                    std::function<void(const ray::Status &)> callback);
+                    std::function<void(const ray::Status &)> callback) override;
 
   /// Restore a spilled object from external storage back into local memory.
   /// Note: This is no-op if the same restoration request is in flight or the requested
@@ -121,14 +122,19 @@ class LocalObjectManager {
   /// \param object_url The URL where the object is spilled.
   /// \param callback A callback to call when the restoration is done.
   /// Status will contain the error during restoration, if any.
-  void AsyncRestoreSpilledObject(const ObjectID &object_id,
-                                 int64_t object_size,
-                                 const std::string &object_url,
-                                 std::function<void(const ray::Status &)> callback);
+  void AsyncRestoreSpilledObject(
+      const ObjectID &object_id,
+      int64_t object_size,
+      const std::string &object_url,
+      std::function<void(const ray::Status &)> callback) override;
 
   /// Clear any freed objects. This will trigger the callback for freed
   /// objects.
-  void FlushFreeObjects();
+  void FlushFreeObjects() override;
+
+  /// Returns true if the object has been marked for deletion through the
+  /// eviction notification.
+  bool ObjectPendingDeletion(const ObjectID &object_id) override;
 
   /// Judge if objects are deletable from pending_delete_queue and delete them if
   /// necessary.
@@ -138,7 +144,7 @@ class LocalObjectManager {
   ///
   /// \param max_batch_size Maximum number of objects that can be deleted by one
   /// invocation.
-  void ProcessSpilledObjectsDeleteQueue(uint32_t max_batch_size);
+  void ProcessSpilledObjectsDeleteQueue(uint32_t max_batch_size) override;
 
   /// Return True if spilling is in progress.
   /// This is a narrow interface that is accessed by plasma store.
@@ -147,31 +153,31 @@ class LocalObjectManager {
   /// which is against the general raylet design.
   ///
   /// \return True if spilling is still in progress. False otherwise.
-  bool IsSpillingInProgress();
+  bool IsSpillingInProgress() override;
 
   /// Populate object store stats.
   ///
   /// \param reply Output parameter.
-  void FillObjectStoreStats(rpc::GetNodeStatsReply *reply) const;
+  void FillObjectStoreStats(rpc::GetNodeStatsReply *reply) const override;
 
   /// Record object spilling stats to metrics.
-  void RecordMetrics() const;
+  void RecordMetrics() const override;
 
   /// Return the spilled object URL if the object is spilled locally,
   /// or the empty string otherwise.
   /// If the external storage is cloud, this will always return an empty string.
   /// In that case, the URL is supposed to be obtained by the object directory.
-  std::string GetLocalSpilledObjectURL(const ObjectID &object_id);
+  std::string GetLocalSpilledObjectURL(const ObjectID &object_id) override;
 
   /// Get the current bytes used by primary object copies. This number includes
   /// bytes used by objects currently being spilled.
-  int64_t GetPrimaryBytes() const;
+  int64_t GetPrimaryBytes() const override;
 
   /// Returns true if we have objects spilled to the local
   /// filesystem.
-  bool HasLocallySpilledObjects() const;
+  bool HasLocallySpilledObjects() const override;
 
-  std::string DebugString() const;
+  std::string DebugString() const override;
 
  private:
   struct LocalObjectInfo {
@@ -284,7 +290,7 @@ class LocalObjectManager {
   /// from plasma. The cache is flushed when it reaches the
   /// free_objects_batch_size, or if objects have been in the cache for longer
   /// than the config's free_objects_period, whichever occurs first.
-  std::vector<ObjectID> objects_to_free_;
+  absl::flat_hash_set<ObjectID> objects_pending_deletion_;
 
   /// The total size of the objects that are currently being
   /// spilled from this node, in bytes.

src/ray/raylet/local_object_manager_interface.h

@@ -0,0 +1,74 @@
+// Copyright 2025 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "ray/common/id.h"
+#include "ray/common/ray_object.h"
+#include "src/ray/protobuf/node_manager.pb.h"
+
+namespace ray {
+
+namespace raylet {
+
+class LocalObjectManagerInterface {
+ public:
+  virtual ~LocalObjectManagerInterface() = default;
+
+  virtual void PinObjectsAndWaitForFree(const std::vector<ObjectID> &,
+                                        std::vector<std::unique_ptr<RayObject>> &&,
+                                        const rpc::Address &,
+                                        const ObjectID & = ObjectID::Nil()) = 0;
+
+  virtual void SpillObjectUptoMaxThroughput() = 0;
+
+  /// TODO(dayshah): This function is only used for testing, we should remove and just
+  /// keep SpillObjectsInternal.
+  virtual void SpillObjects(const std::vector<ObjectID> &,
+                            std::function<void(const ray::Status &)>) = 0;
+
+  virtual void AsyncRestoreSpilledObject(const ObjectID &,
+                                         int64_t,
+                                         const std::string &,
+                                         std::function<void(const ray::Status &)>) = 0;
+
+  virtual void FlushFreeObjects() = 0;
+
+  virtual bool ObjectPendingDeletion(const ObjectID &) = 0;
+
+  virtual void ProcessSpilledObjectsDeleteQueue(uint32_t) = 0;
+
+  virtual bool IsSpillingInProgress() = 0;
+
+  virtual void FillObjectStoreStats(rpc::GetNodeStatsReply *) const = 0;
+
+  virtual void RecordMetrics() const = 0;
+
+  virtual std::string GetLocalSpilledObjectURL(const ObjectID &) = 0;
+
+  virtual int64_t GetPrimaryBytes() const = 0;
+
+  virtual bool HasLocallySpilledObjects() const = 0;
+
+  virtual std::string DebugString() const = 0;
+};
+
+};  // namespace raylet
+
+};  // namespace ray

src/ray/raylet/main.cc

@@ -31,6 +31,8 @@
 #include "ray/common/task/task_common.h"
 #include "ray/gcs/gcs_client/gcs_client.h"
 #include "ray/object_manager/ownership_object_directory.h"
+#include "ray/raylet/local_object_manager.h"
+#include "ray/raylet/local_object_manager_interface.h"
 #include "ray/raylet/raylet.h"
 #include "ray/stats/stats.h"
 #include "ray/util/cmd_line_utils.h"
@@ -40,7 +42,6 @@
 #include "ray/util/stream_redirection_options.h"
 #include "ray/util/subreaper.h"
 #include "scheduling/cluster_task_manager.h"
-#include "src/ray/protobuf/gcs.pb.h"
 
 using json = nlohmann::json;
 
@@ -262,7 +263,7 @@ int main(int argc, char *argv[]) {
   std::unique_ptr<ray::raylet::WorkerPoolInterface> worker_pool;
   /// Manages all local objects that are pinned (primary
   /// copies), freed, and/or spilled.
-  std::unique_ptr<ray::raylet::LocalObjectManager> local_object_manager;
+  std::unique_ptr<ray::raylet::LocalObjectManagerInterface> local_object_manager;
   /// These classes make up the new scheduler. ClusterResourceScheduler is
   /// responsible for maintaining a view of the cluster state w.r.t resource
   /// usage. ClusterTaskManager is responsible for queuing, spilling back, and

src/ray/raylet/node_manager.cc

@@ -42,6 +42,7 @@
 #include "ray/gcs/pb_util.h"
 #include "ray/object_manager/ownership_object_directory.h"
 #include "ray/raylet/format/node_manager_generated.h"
+#include "ray/raylet/local_object_manager_interface.h"
 #include "ray/raylet/scheduling/cluster_task_manager.h"
 #include "ray/raylet/worker_killing_policy.h"
 #include "ray/raylet/worker_pool.h"
@@ -121,7 +122,7 @@ NodeManager::NodeManager(
     ClusterTaskManagerInterface &cluster_task_manager,
     IObjectDirectory &object_directory,
     ObjectManagerInterface &object_manager,
-    LocalObjectManager &local_object_manager,
+    LocalObjectManagerInterface &local_object_manager,
     DependencyManager &dependency_manager,
     WorkerPoolInterface &worker_pool,
     absl::flat_hash_map<WorkerID, std::shared_ptr<WorkerInterface>> &leased_workers,
@@ -2490,11 +2491,16 @@ void NodeManager::HandlePinObjectIDs(rpc::PinObjectIDsRequest request,
     auto object_id_it = object_ids.begin();
     auto result_it = results.begin();
     while (object_id_it != object_ids.end()) {
-      if (*result_it == nullptr) {
+      // Note: It is safe to call ObjectPendingDeletion here because the asynchronous
+      // deletion can only happen on the same thread as the call to HandlePinObjectIDs.
+      // Therefore, a new object cannot be marked for deletion while this function is
+      // executing.
+      if (*result_it == nullptr ||
+          local_object_manager_.ObjectPendingDeletion(*object_id_it)) {
         RAY_LOG(DEBUG).WithField(*object_id_it)
             << "Failed to get object in the object store. This should only happen when "
-               "the owner tries to pin a "
-            << "secondary copy and it's evicted in the meantime";
+               "the owner tries to pin an object and it's already been deleted or is "
+               "marked for deletion.";
         object_id_it = object_ids.erase(object_id_it);
         result_it = results.erase(result_it);
         reply->add_successes(false);