Send active object IDs to the raylet (#5803)

* Send active object IDs to the raylet

* comment

* comments

* dedup

* signed int in config

* comments

* Remove object ID from monitor

* Fix test

* re-add check

* fix cast

* check if core worker

* Add comment

* Reservoir sampling

* Fix lint

* Pointer return

* tmp

* Fix merge

* Initialize object ids properly

* Fix lint

python/ray/_raylet.pyx

@@ -725,3 +725,15 @@ cdef class CoreWorker:
         check_status(self.core_worker.get().SerializeActorHandle(
             c_actor_id, &output))
         return output
+
+    def add_active_object_id(self, ObjectID object_id):
+        cdef:
+            CObjectID c_object_id = object_id.native()
+        with nogil:
+            self.core_worker.get().AddActiveObjectID(c_object_id)
+
+    def remove_active_object_id(self, ObjectID object_id):
+        cdef:
+            CObjectID c_object_id = object_id.native()
+        with nogil:
+            self.core_worker.get().RemoveActiveObjectID(c_object_id)

python/ray/includes/libcoreworker.pxd

@@ -90,3 +90,5 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
         CActorID DeserializeAndRegisterActorHandle(const c_string &bytes)
         CRayStatus SerializeActorHandle(const CActorID &actor_id, c_string
                                         *bytes)
+        void AddActiveObjectID(const CObjectID &object_id)
+        void RemoveActiveObjectID(const CObjectID &object_id)

python/ray/includes/ray_config.pxd

@@ -12,7 +12,7 @@ cdef extern from "ray/common/ray_config.h" nogil:
 
         int64_t handler_warning_timeout_ms() const
 
-        int64_t heartbeat_timeout_milliseconds() const
+        int64_t raylet_heartbeat_timeout_milliseconds() const
 
         int64_t debug_dump_period_milliseconds() const
 

python/ray/includes/ray_config.pxi

@@ -10,8 +10,8 @@ cdef class Config:
         return RayConfig.instance().handler_warning_timeout_ms()
 
     @staticmethod
-    def heartbeat_timeout_milliseconds():
-        return RayConfig.instance().heartbeat_timeout_milliseconds()
+    def raylet_heartbeat_timeout_milliseconds():
+        return RayConfig.instance().raylet_heartbeat_timeout_milliseconds()
 
     @staticmethod
     def debug_dump_period_milliseconds():

python/ray/includes/unique_ids.pxi

@@ -22,6 +22,7 @@ from ray.includes.unique_ids cimport (
     CWorkerID,
 )
 
+import ray
 from ray.utils import decode
 
 
@@ -128,13 +129,32 @@ cdef class UniqueID(BaseID):
 
 
 cdef class ObjectID(BaseID):
-    cdef CObjectID data
-    cdef object buffer_ref
+    cdef:
+        CObjectID data
+        object buffer_ref
+        # Flag indicating whether or not this object ID was added to the set
+        # of active IDs in the core worker so we know whether we should clean
+        # it up.
+        c_bool in_core_worker
 
     def __init__(self, id):
         check_id(id)
         self.data = CObjectID.FromBinary(<c_string>id)
 
+        worker = ray.worker.global_worker
+        # TODO(edoakes): there are dummy object IDs being created in
+        # includes/task.pxi before the core worker is initialized.
+        if hasattr(worker, "core_worker"):
+            worker.core_worker.add_active_object_id(self)
+            self.in_core_worker = True
+        else:
+            self.in_core_worker = False
+
+    def __dealloc__(self):
+        worker = ray.worker.global_worker
+        if self.in_core_worker and hasattr(worker, "core_worker"):
+            worker.core_worker.remove_active_object_id(self)
+
     cdef CObjectID native(self):
         return <CObjectID>self.data
 

python/ray/monitor.py

@@ -354,7 +354,8 @@ class Monitor(object):
 
             # Wait for a heartbeat interval before processing the next round of
             # messages.
-            time.sleep(ray._config.heartbeat_timeout_milliseconds() * 1e-3)
+            time.sleep(
+                ray._config.raylet_heartbeat_timeout_milliseconds() * 1e-3)
 
     def run(self):
         try:

python/ray/state.py

@@ -30,7 +30,7 @@ def _parse_client_table(redis_client):
     Returns:
         A list of information about the nodes in the cluster.
     """
-    NIL_CLIENT_ID = ray.ObjectID.nil().binary()
+    NIL_CLIENT_ID = ray.ClientID.nil().binary()
     message = redis_client.execute_command(
         "RAY.TABLE_LOOKUP", gcs_utils.TablePrefix.Value("CLIENT"), "",
         NIL_CLIENT_ID)