[core] Admission control for pulling objects to the local node

python/ray/tests/test_object_manager.py

@@ -296,6 +296,83 @@ def driver():
     ray.get(driver.remote())
 
 
+@pytest.mark.timeout(30)
+def test_pull_bundles_admission_control(shutdown_only):
+    cluster = Cluster()
+    object_size = int(1e7)
+    num_objects = 10
+    num_tasks = 10
+    # Head node can fit all of the objects at once.
+    cluster.add_node(
+        num_cpus=0,
+        object_store_memory=2 * num_tasks * num_objects * object_size)
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+
+    # Worker node can only fit 1 task at a time.
+    cluster.add_node(
+        num_cpus=1, object_store_memory=1.5 * num_objects * object_size)
+    cluster.wait_for_nodes()
+
+    @ray.remote
+    def foo(*args):
+        return
+
+    args = []
+    for _ in range(num_tasks):
+        task_args = [
+            ray.put(np.zeros(object_size, dtype=np.uint8))
+            for _ in range(num_objects)
+        ]
+        args.append(task_args)
+
+    tasks = [foo.remote(*task_args) for task_args in args]
+    ray.get(tasks)
+
+
+@pytest.mark.timeout(30)
+def test_pull_bundles_admission_control_dynamic(shutdown_only):
+    # This test is the same as test_pull_bundles_admission_control, except that
+    # the object store's capacity starts off higher and is later consumed
+    # dynamically by concurrent workers.
+    cluster = Cluster()
+    object_size = int(1e7)
+    num_objects = 10
+    num_tasks = 10
+    # Head node can fit all of the objects at once.
+    cluster.add_node(
+        num_cpus=0,
+        object_store_memory=2 * num_tasks * num_objects * object_size)
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+
+    # Worker node can fit 2 tasks at a time.
+    cluster.add_node(
+        num_cpus=1, object_store_memory=2.5 * num_objects * object_size)
+    cluster.wait_for_nodes()
+
+    @ray.remote
+    def foo(*args):
+        return
+
+    @ray.remote
+    def allocate(*args):
+        return np.zeros(object_size, dtype=np.uint8)
+
+    args = []
+    for _ in range(num_tasks):
+        task_args = [
+            ray.put(np.zeros(object_size, dtype=np.uint8))
+            for _ in range(num_objects)
+        ]
+        args.append(task_args)
+
+    tasks = [foo.remote(*task_args) for task_args in args]
+    allocated = [allocate.remote() for _ in range(num_objects)]
+    ray.get(tasks)
+    del allocated
+
+
 if __name__ == "__main__":
     import pytest
     import sys

python/ray/tests/test_object_spilling.py

@@ -661,5 +661,64 @@ def test_release_during_plasma_fetch(tmp_path, shutdown_only):
     do_test_release_resource(tmp_path, expect_released=True)
 
 
+@pytest.mark.skipif(
+    platform.system() == "Windows", reason="Failing on Windows.")
+@pytest.mark.timeout(30)
+def test_spill_objects_on_object_transfer(object_spilling_config,
+                                          ray_start_cluster):
+    # This test checks that objects get spilled to make room for transferred
+    # objects.
+    cluster = ray_start_cluster
+    object_size = int(1e7)
+    num_objects = 10
+    num_tasks = 10
+    # Head node can fit all of the objects at once.
+    cluster.add_node(
+        num_cpus=0,
+        object_store_memory=2 * num_tasks * num_objects * object_size,
+        _system_config={
+            "max_io_workers": 1,
+            "automatic_object_spilling_enabled": True,
+            "object_store_full_max_retries": 4,
+            "object_store_full_delay_ms": 100,
+            "object_spilling_config": object_spilling_config,
+            "min_spilling_size": 0
+        })
+    cluster.wait_for_nodes()
+    ray.init(address=cluster.address)
+
+    # Worker node can fit 1 tasks at a time.
+    cluster.add_node(
+        num_cpus=1, object_store_memory=1.5 * num_objects * object_size)
+    cluster.wait_for_nodes()
+
+    @ray.remote
+    def foo(*args):
+        return
+
+    @ray.remote
+    def allocate(*args):
+        return np.zeros(object_size, dtype=np.uint8)
+
+    # Allocate some objects that must be spilled to make room for foo's
+    # arguments.
+    allocated = [allocate.remote() for _ in range(num_objects)]
+    ray.get(allocated)
+    print("done allocating")
+
+    args = []
+    for _ in range(num_tasks):
+        task_args = [
+            ray.put(np.zeros(object_size, dtype=np.uint8))
+            for _ in range(num_objects)
+        ]
+        args.append(task_args)
+
+    # Check that tasks scheduled to the worker node have enough room after
+    # spilling.
+    tasks = [foo.remote(*task_args) for task_args in args]
+    ray.get(tasks)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-sv", __file__]))

src/ray/core_worker/core_worker.cc

@@ -2213,6 +2213,7 @@ void CoreWorker::HandleGetObjectLocationsOwner(
   } else {
     status = Status::ObjectNotFound("Object " + object_id.Hex() + " not found");
   }
+  reply->set_object_size(reference_counter_->GetObjectSize(object_id));
   send_reply_callback(status, nullptr, nullptr);
 }
 

src/ray/core_worker/reference_count.cc

@@ -948,6 +948,15 @@ absl::optional<absl::flat_hash_set<NodeID>> ReferenceCounter::GetObjectLocations
   return it->second.locations;
 }
 
+size_t ReferenceCounter::GetObjectSize(const ObjectID &object_id) const {
+  absl::MutexLock lock(&mutex_);
+  auto it = object_id_refs_.find(object_id);
+  if (it == object_id_refs_.end()) {
+    return 0;
+  }
+  return it->second.object_size;
+}
+
 void ReferenceCounter::HandleObjectSpilled(const ObjectID &object_id) {
   absl::MutexLock lock(&mutex_);
   auto it = object_id_refs_.find(object_id);

src/ray/core_worker/reference_count.h

@@ -397,6 +397,12 @@ class ReferenceCounter : public ReferenceCounterInterface,
   absl::optional<absl::flat_hash_set<NodeID>> GetObjectLocations(
       const ObjectID &object_id) LOCKS_EXCLUDED(mutex_);
 
+  /// Get an object's size. This will return 0 if the object is out of scope.
+  ///
+  /// \param[in] object_id The object whose size to get.
+  /// \return Object size, or 0 if the object is out of scope.
+  size_t GetObjectSize(const ObjectID &object_id) const;
+
   /// Handle an object has been spilled to external storage.
   ///
   /// This notifies the primary raylet that the object is safe to release and

src/ray/gcs/accessor.h

@@ -320,7 +320,7 @@ class ObjectInfoAccessor {
   /// \param callback Callback that will be called after object has been added to GCS.
   /// \return Status
   virtual Status AsyncAddLocation(const ObjectID &object_id, const NodeID &node_id,
-                                  const StatusCallback &callback) = 0;
+                                  size_t object_size, const StatusCallback &callback) = 0;
 
   /// Add spilled location of object to GCS asynchronously.
   ///

src/ray/gcs/gcs_client/service_based_accessor.cc

@@ -1128,13 +1128,15 @@ Status ServiceBasedObjectInfoAccessor::AsyncGetAll(
 
 Status ServiceBasedObjectInfoAccessor::AsyncAddLocation(const ObjectID &object_id,
                                                         const NodeID &node_id,
+                                                        size_t object_size,
                                                         const StatusCallback &callback) {
   RAY_LOG(DEBUG) << "Adding object location, object id = " << object_id
                  << ", node id = " << node_id
                  << ", job id = " << object_id.TaskId().JobId();
   rpc::AddObjectLocationRequest request;
   request.set_object_id(object_id.Binary());
   request.set_node_id(node_id.Binary());
+  request.set_size(object_size);
 
   auto operation = [this, request, object_id, node_id,
                     callback](const SequencerDoneCallback &done_callback) {
@@ -1229,11 +1231,13 @@ Status ServiceBasedObjectInfoAccessor::AsyncSubscribeToLocations(
           rpc::ObjectLocationChange update;
           update.set_is_add(true);
           update.set_node_id(loc.manager());
+          update.set_size(result->size());
           notification.push_back(update);
         }
         if (!result->spilled_url().empty()) {
           rpc::ObjectLocationChange update;
           update.set_spilled_url(result->spilled_url());
+          update.set_size(result->size());
           notification.push_back(update);
         }
         subscribe(object_id, notification);

src/ray/gcs/gcs_client/service_based_accessor.h

@@ -333,7 +333,7 @@ class ServiceBasedObjectInfoAccessor : public ObjectInfoAccessor {
   Status AsyncGetAll(const MultiItemCallback<rpc::ObjectLocationInfo> &callback) override;
 
   Status AsyncAddLocation(const ObjectID &object_id, const NodeID &node_id,
-                          const StatusCallback &callback) override;
+                          size_t object_size, const StatusCallback &callback) override;
 
   Status AsyncAddSpilledUrl(const ObjectID &object_id, const std::string &spilled_url,
                             const StatusCallback &callback) override;

src/ray/gcs/gcs_server/gcs_object_manager.cc

@@ -51,6 +51,7 @@ void GcsObjectManager::HandleGetAllObjectLocations(
       object_table_data.set_manager(node_id.Binary());
       object_location_info.add_locations()->CopyFrom(object_table_data);
     }
+    object_location_info.set_size(item.second.object_size);
     reply->add_object_location_info_list()->CopyFrom(object_location_info);
   }
   RAY_LOG(DEBUG) << "Finished getting all object locations.";
@@ -78,7 +79,10 @@ void GcsObjectManager::HandleAddObjectLocation(
     RAY_LOG(DEBUG) << "Adding object spilled location, object id = " << object_id;
   }
 
-  auto on_done = [this, object_id, node_id, spilled_url, reply,
+  size_t size = request.size();
+  object_to_locations_[object_id].object_size = size;
+
+  auto on_done = [this, object_id, node_id, spilled_url, size, reply,
                   send_reply_callback](const Status &status) {
     if (status.ok()) {
       rpc::ObjectLocationChange notification;
@@ -89,6 +93,7 @@ void GcsObjectManager::HandleAddObjectLocation(
       if (!spilled_url.empty()) {
         notification.set_spilled_url(spilled_url);
       }
+      notification.set_size(size);
       RAY_CHECK_OK(gcs_pub_sub_->Publish(OBJECT_CHANNEL, object_id.Hex(),
                                          notification.SerializeAsString(), nullptr));
       RAY_LOG(DEBUG) << "Finished adding object location, job id = "
@@ -287,6 +292,7 @@ const ObjectLocationInfo GcsObjectManager::GenObjectLocationInfo(
       object_data.add_locations()->set_manager(node_id.Binary());
     }
     object_data.set_spilled_url(it->second.spilled_url);
+    object_data.set_size(it->second.object_size);
   }
   return object_data;
 }

src/ray/gcs/gcs_server/gcs_object_manager.h

@@ -65,6 +65,7 @@ class GcsObjectManager : public rpc::ObjectInfoHandler {
   struct LocationSet {
     absl::flat_hash_set<NodeID> locations;
     std::string spilled_url = "";
+    size_t object_size = 0;
   };
 
   /// Add a location of objects.

src/ray/object_manager/object_directory.cc

@@ -31,13 +31,21 @@ using ray::rpc::ObjectTableData;
 /// object table entries up to but not including this notification.
 bool UpdateObjectLocations(const std::vector<rpc::ObjectLocationChange> &location_updates,
                            std::shared_ptr<gcs::GcsClient> gcs_client,
-                           std::unordered_set<NodeID> *node_ids,
-                           std::string *spilled_url) {
+                           std::unordered_set<NodeID> *node_ids, std::string *spilled_url,
+                           size_t *object_size) {
   // location_updates contains the updates of locations of the object.
   // with GcsChangeMode, we can determine whether the update mode is
   // addition or deletion.
   bool isUpdated = false;
   for (const auto &update : location_updates) {
+    // The size can be 0 if the update was a deletion. This assumes that an
+    // object's size is always greater than 0.
+    // TODO(swang): If that's not the case, we should use a flag to check
+    // whether the size is set instead.
+    if (update.size() > 0) {
+      *object_size = update.size();
+    }
+
     if (!update.node_id().empty()) {
       NodeID node_id = NodeID::FromBinary(update.node_id());
       if (update.is_add() && 0 == node_ids->count(node_id)) {
@@ -73,9 +81,10 @@ bool UpdateObjectLocations(const std::vector<rpc::ObjectLocationChange> &locatio
 ray::Status ObjectDirectory::ReportObjectAdded(
     const ObjectID &object_id, const NodeID &node_id,
     const object_manager::protocol::ObjectInfoT &object_info) {
-  RAY_LOG(DEBUG) << "Reporting object added to GCS " << object_id;
+  size_t size = object_info.data_size + object_info.metadata_size;
+  RAY_LOG(DEBUG) << "Reporting object added to GCS " << object_id << " size " << size;
   ray::Status status =
-      gcs_client_->Objects().AsyncAddLocation(object_id, node_id, nullptr);
+      gcs_client_->Objects().AsyncAddLocation(object_id, node_id, size, nullptr);
   return status;
 }
 
@@ -119,14 +128,14 @@ void ObjectDirectory::HandleNodeRemoved(const NodeID &node_id) {
       // If the subscribed object has the removed node as a location, update
       // its locations with an empty update so that the location will be removed.
       UpdateObjectLocations({}, gcs_client_, &listener.second.current_object_locations,
-                            &listener.second.spilled_url);
+                            &listener.second.spilled_url, &listener.second.object_size);
       // Re-call all the subscribed callbacks for the object, since its
       // locations have changed.
       for (const auto &callback_pair : listener.second.callbacks) {
         // It is safe to call the callback directly since this is already running
         // in the subscription callback stack.
         callback_pair.second(object_id, listener.second.current_object_locations,
-                             listener.second.spilled_url);
+                             listener.second.spilled_url, listener.second.object_size);
       }
     }
   }
@@ -157,7 +166,7 @@ ray::Status ObjectDirectory::SubscribeObjectLocations(const UniqueID &callback_i
           // Update entries for this object.
           if (!UpdateObjectLocations(object_notifications, gcs_client_,
                                      &it->second.current_object_locations,
-                                     &it->second.spilled_url)) {
+                                     &it->second.spilled_url, &it->second.object_size)) {
             return;
           }
           // Copy the callbacks so that the callbacks can unsubscribe without interrupting
@@ -171,7 +180,7 @@ ray::Status ObjectDirectory::SubscribeObjectLocations(const UniqueID &callback_i
             // It is safe to call the callback directly since this is already running
             // in the subscription callback stack.
             callback_pair.second(object_id, it->second.current_object_locations,
-                                 it->second.spilled_url);
+                                 it->second.spilled_url, it->second.object_size);
           }
         };
     status = gcs_client_->Objects().AsyncSubscribeToLocations(
@@ -189,8 +198,9 @@ ray::Status ObjectDirectory::SubscribeObjectLocations(const UniqueID &callback_i
   if (listener_state.subscribed) {
     auto &locations = listener_state.current_object_locations;
     auto &spilled_url = listener_state.spilled_url;
-    io_service_.post([callback, locations, spilled_url, object_id]() {
-      callback(object_id, locations, spilled_url);
+    auto object_size = it->second.object_size;
+    io_service_.post([callback, locations, spilled_url, object_size, object_id]() {
+      callback(object_id, locations, spilled_url, object_size);
     });
   }
   return status;
@@ -223,8 +233,9 @@ ray::Status ObjectDirectory::LookupLocations(const ObjectID &object_id,
     // cached locations.
     auto &locations = it->second.current_object_locations;
     auto &spilled_url = it->second.spilled_url;
-    io_service_.post([callback, object_id, spilled_url, locations]() {
-      callback(object_id, locations, spilled_url);
+    auto object_size = it->second.object_size;
+    io_service_.post([callback, object_id, spilled_url, locations, object_size]() {
+      callback(object_id, locations, spilled_url, object_size);
     });
   } else {
     // We do not have any locations cached due to a concurrent
@@ -252,10 +263,12 @@ ray::Status ObjectDirectory::LookupLocations(const ObjectID &object_id,
 
           std::unordered_set<NodeID> node_ids;
           std::string spilled_url;
-          UpdateObjectLocations(notification, gcs_client_, &node_ids, &spilled_url);
+          size_t object_size = 0;
+          UpdateObjectLocations(notification, gcs_client_, &node_ids, &spilled_url,
+                                &object_size);
           // It is safe to call the callback directly since this is already running
           // in the GCS client's lookup callback stack.
-          callback(object_id, node_ids, spilled_url);
+          callback(object_id, node_ids, spilled_url, object_size);
         });
   }
   return status;