Revert "[core] Fix bugs in data locality (#24698)" (#25035)

This reverts commit eaec96d.

python/ray/tests/test_scheduling_2.py

@@ -647,50 +647,6 @@ def check_backlog_info():
     cluster.shutdown()
 
 
-def test_data_locality_spilled_objects(
-    ray_start_cluster_enabled, fs_only_object_spilling_config
-):
-    cluster = ray_start_cluster_enabled
-    object_spilling_config, _ = fs_only_object_spilling_config
-    cluster.add_node(
-        num_cpus=1,
-        object_store_memory=100 * 1024 * 1024,
-        _system_config={
-            "min_spilling_size": 1,
-            "object_spilling_config": object_spilling_config,
-        },
-    )
-    ray.init(cluster.address)
-    cluster.add_node(
-        num_cpus=1, object_store_memory=100 * 1024 * 1024, resources={"remote": 1}
-    )
-
-    @ray.remote(resources={"remote": 1})
-    def f():
-        return (
-            np.zeros(50 * 1024 * 1024, dtype=np.uint8),
-            ray.runtime_context.get_runtime_context().node_id,
-        )
-
-    @ray.remote
-    def check_locality(x):
-        _, node_id = x
-        assert node_id == ray.runtime_context.get_runtime_context().node_id
-
-    # Check locality works when dependent task is already submitted by the time
-    # the upstream task finishes.
-    for _ in range(5):
-        ray.get(check_locality.remote(f.remote()))
-
-    # Check locality works when some objects were spilled.
-    xs = [f.remote() for _ in range(5)]
-    ray.wait(xs, num_returns=len(xs), fetch_local=False)
-    for i, x in enumerate(xs):
-        task = check_locality.remote(x)
-        print(i, x, task)
-        ray.get(task)
-
-
 if __name__ == "__main__":
     import pytest
 

src/mock/ray/core_worker/lease_policy.h

@@ -30,7 +30,7 @@ class MockLocalityDataProviderInterface : public LocalityDataProviderInterface {
   MOCK_METHOD(absl::optional<LocalityData>,
               GetLocalityData,
               (const ObjectID &object_id),
-              (const override));
+              (override));
 };
 
 }  // namespace core

src/ray/core_worker/lease_policy.h

@@ -32,8 +32,7 @@ struct LocalityData {
 /// Interface for providers of locality data to the lease policy.
 class LocalityDataProviderInterface {
  public:
-  virtual absl::optional<LocalityData> GetLocalityData(
-      const ObjectID &object_id) const = 0;
+  virtual absl::optional<LocalityData> GetLocalityData(const ObjectID &object_id) = 0;
 
   virtual ~LocalityDataProviderInterface() {}
 };

src/ray/core_worker/reference_count.cc

@@ -1254,7 +1254,7 @@ bool ReferenceCounter::HandleObjectSpilled(const ObjectID &object_id,
 }
 
 absl::optional<LocalityData> ReferenceCounter::GetLocalityData(
-    const ObjectID &object_id) const {
+    const ObjectID &object_id) {
   absl::MutexLock lock(&mutex_);
   // Uses the reference table to return locality data for an object.
   auto it = object_id_refs_.find(object_id);
@@ -1281,10 +1281,7 @@ absl::optional<LocalityData> ReferenceCounter::GetLocalityData(
   //   locations.
   // - If we don't own this object, this will contain a snapshot of the object locations
   //   at future resolution time.
-  auto node_ids = it->second.locations;
-  if (!it->second.spilled_node_id.IsNil()) {
-    node_ids.emplace(it->second.spilled_node_id);
-  }
+  const auto &node_ids = it->second.locations;
 
   // We should only reach here if we have valid locality data to return.
   absl::optional<LocalityData> locality_data(

src/ray/core_worker/reference_count.h

@@ -448,7 +448,7 @@ class ReferenceCounter : public ReferenceCounterInterface,
   ///
   /// \param[in] object_id Object whose locality data we want.
   /// \return Locality data.
-  absl::optional<LocalityData> GetLocalityData(const ObjectID &object_id) const;
+  absl::optional<LocalityData> GetLocalityData(const ObjectID &object_id);
 
   /// Report locality data for object. This is used by the FutureResolver to report
   /// locality data for borrowed refs.

src/ray/core_worker/task_manager.cc

@@ -252,14 +252,11 @@ void TaskManager::CompletePendingTask(const TaskID &task_id,
     const auto nested_refs =
         VectorFromProtobuf<rpc::ObjectReference>(return_object.nested_inlined_refs());
     if (return_object.in_plasma()) {
-      // NOTE(swang): We need to add the location of the object before marking
-      // it as local in the in-memory store so that the data locality policy
-      // will choose the right raylet for any queued dependent tasks.
-      const auto pinned_at_raylet_id = NodeID::FromBinary(worker_addr.raylet_id());
-      reference_counter_->UpdateObjectPinnedAtRaylet(object_id, pinned_at_raylet_id);
       // Mark it as in plasma with a dummy object.
       RAY_CHECK(
           in_memory_store_->Put(RayObject(rpc::ErrorType::OBJECT_IN_PLASMA), object_id));
+      const auto pinned_at_raylet_id = NodeID::FromBinary(worker_addr.raylet_id());
+      reference_counter_->UpdateObjectPinnedAtRaylet(object_id, pinned_at_raylet_id);
     } else {
       // NOTE(swang): If a direct object was promoted to plasma, then we do not
       // record the node ID that it was pinned at, which means that we will not

src/ray/core_worker/test/lease_policy_test.cc

@@ -40,15 +40,15 @@ class MockLocalityDataProvider : public LocalityDataProviderInterface {
   MockLocalityDataProvider(absl::flat_hash_map<ObjectID, LocalityData> locality_data)
       : locality_data_(locality_data) {}
 
-  absl::optional<LocalityData> GetLocalityData(const ObjectID &object_id) const {
+  absl::optional<LocalityData> GetLocalityData(const ObjectID &object_id) {
     num_locality_data_fetches++;
     return locality_data_[object_id];
   };
 
   ~MockLocalityDataProvider() {}
 
-  mutable int num_locality_data_fetches = 0;
-  mutable absl::flat_hash_map<ObjectID, LocalityData> locality_data_;
+  int num_locality_data_fetches = 0;
+  absl::flat_hash_map<ObjectID, LocalityData> locality_data_;
 };
 
 absl::optional<rpc::Address> MockNodeAddrFactory(const NodeID &node_id) {

src/ray/core_worker/test/reference_count_test.cc

@@ -696,15 +696,6 @@ TEST_F(ReferenceCountTest, TestGetLocalityData) {
   ASSERT_EQ(locality_data_obj1->nodes_containing_object,
             absl::flat_hash_set<NodeID>({node1}));
 
-  // Include spilled locations in locality data.
-  rc->RemoveObjectLocation(obj1, node1);
-  locality_data_obj1 = rc->GetLocalityData(obj1);
-  ASSERT_EQ(locality_data_obj1->nodes_containing_object, absl::flat_hash_set<NodeID>({}));
-  rc->HandleObjectSpilled(obj1, "spill_loc", node1);
-  locality_data_obj1 = rc->GetLocalityData(obj1);
-  ASSERT_EQ(locality_data_obj1->nodes_containing_object,
-            absl::flat_hash_set<NodeID>({node1}));
-
   // Borrowed object with defined object size and at least one node location should
   // return valid locality data.
   rc->AddLocalReference(obj2, "file.py:43");

src/ray/core_worker/test/task_manager_test.cc

@@ -346,31 +346,6 @@ TEST_F(TaskManagerTest, TestLineageEvicted) {
   ASSERT_FALSE(reference_counter_->HasReference(return_id));
 }
 
-TEST_F(TaskManagerTest, TestLocalityDataAdded) {
-  auto spec = CreateTaskHelper(1, {});
-  auto return_id = spec.ReturnId(0);
-  auto node_id = NodeID::FromRandom();
-  int object_size = 100;
-  store_->GetAsync(return_id, [&](std::shared_ptr<RayObject> obj) {
-    // By the time the return object is available to get, we should be able
-    // to get the locality data too.
-    auto locality_data = reference_counter_->GetLocalityData(return_id);
-    ASSERT_TRUE(locality_data.has_value());
-    ASSERT_EQ(locality_data->object_size, object_size);
-    ASSERT_TRUE(locality_data->nodes_containing_object.contains(node_id));
-  });
-
-  rpc::PushTaskReply reply;
-  auto return_object = reply.add_return_objects();
-  return_object->set_object_id(return_id.Binary());
-  return_object->set_in_plasma(true);
-  return_object->set_size(object_size);
-  rpc::Address worker_addr;
-  worker_addr.set_raylet_id(node_id.Binary());
-  manager_.AddPendingTask(rpc::Address(), spec, "", 0);
-  manager_.CompletePendingTask(spec.TaskId(), reply, worker_addr);
-}
-
 // Test to make sure that the task spec and dependencies for an object are
 // pinned when lineage pinning is enabled in the ReferenceCounter.
 TEST_F(TaskManagerLineageTest, TestLineagePinned) {