[direct task] For serialized object IDs, check with owner before declaring object unreconstructable (#6286)

* Track borrowed vs owned objects

* Serialize owner address with object ID

* serialize owner task id

* Deserialize object IDs

* Pass direct task ID instead of plasma ID

* it works

* Fix ref count test

* Add unit test

* update warning

* we own ray.put objects

* missing file

* doc

* Fix unit test

* comments

* Fix py2

* lint

* update

python/ray/_raylet.pyx

@@ -41,6 +41,7 @@ from libcpp.vector cimport vector as c_vector
 from cython.operator import dereference, postincrement
 
 from ray.includes.common cimport (
+    CAddress,
     CLanguage,
     CRayObject,
     CRayStatus,
@@ -1052,11 +1053,29 @@ cdef class CoreWorker:
         # Note: faster to not release GIL for short-running op.
         self.core_worker.get().RemoveObjectIDReference(c_object_id)
 
-    def promote_object_to_plasma(self, ObjectID object_id):
+    def serialize_and_promote_object_id(self, ObjectID object_id):
         cdef:
             CObjectID c_object_id = object_id.native()
-        self.core_worker.get().PromoteObjectToPlasma(c_object_id)
-        return object_id.with_plasma_transport_type()
+            CTaskID c_owner_id = CTaskID.Nil()
+            CAddress c_owner_address = CAddress()
+        self.core_worker.get().PromoteToPlasmaAndGetOwnershipInfo(
+                c_object_id, &c_owner_id, &c_owner_address)
+        return (object_id,
+                TaskID(c_owner_id.Binary()),
+                c_owner_address.SerializeAsString())
+
+    def deserialize_and_register_object_id(
+            self, const c_string &object_id_binary, const c_string
+            &owner_id_binary, const c_string &serialized_owner_address):
+        cdef:
+            CObjectID c_object_id = CObjectID.FromBinary(object_id_binary)
+            CTaskID c_owner_id = CTaskID.FromBinary(owner_id_binary)
+            CAddress c_owner_address = CAddress()
+        c_owner_address.ParseFromString(serialized_owner_address)
+        self.core_worker.get().RegisterOwnershipInfoAndResolveFuture(
+                c_object_id,
+                c_owner_id,
+                c_owner_address)
 
     # TODO: handle noreturn better
     cdef store_task_outputs(

python/ray/includes/common.pxd

@@ -132,6 +132,10 @@ cdef extern from "ray/protobuf/common.pb.h" nogil:
         pass
     cdef cppclass CTaskType "ray::TaskType":
         pass
+    cdef cppclass CAddress "ray::rpc::Address":
+        CAddress()
+        const c_string &SerializeAsString()
+        void ParseFromString(const c_string &serialized)
 
 
 # This is a workaround for C++ enum class since Cython has no corresponding

python/ray/includes/libcoreworker.pxd

@@ -17,6 +17,7 @@ from ray.includes.unique_ids cimport (
     CObjectID,
 )
 from ray.includes.common cimport (
+    CAddress,
     CActorCreationOptions,
     CBuffer,
     CRayFunction,
@@ -116,6 +117,12 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
         void AddObjectIDReference(const CObjectID &object_id)
         void RemoveObjectIDReference(const CObjectID &object_id)
         void PromoteObjectToPlasma(const CObjectID &object_id)
+        void PromoteToPlasmaAndGetOwnershipInfo(const CObjectID &object_id,
+                                                CTaskID *owner_id,
+                                                CAddress *owner_address)
+        void RegisterOwnershipInfoAndResolveFuture(
+                const CObjectID &object_id, const CTaskID &owner_id, const
+                CAddress &owner_address)
 
         CRayStatus SetClientOptions(c_string client_name, int64_t limit)
         CRayStatus Put(const CRayObject &object, CObjectID *object_id)

python/ray/includes/task.pxd

@@ -27,7 +27,6 @@ cdef extern from "ray/protobuf/common.pb.h" nogil:
     cdef cppclass RpcTask "ray::rpc::Task":
         RpcTaskSpec *mutable_task_spec()
 
-
 cdef extern from "ray/protobuf/gcs.pb.h" nogil:
     cdef cppclass TaskTableData "ray::rpc::TaskTableData":
         RpcTask *mutable_task()

python/ray/serialization.py

@@ -157,19 +157,52 @@ class SerializationContext(object):
                 serialization_context)
 
             def id_serializer(obj):
-                if isinstance(obj, ray.ObjectID) and obj.is_direct_call_type():
-                    obj = self.worker.core_worker.promote_object_to_plasma(obj)
                 return pickle.dumps(obj)
 
             def id_deserializer(serialized_obj):
                 return pickle.loads(serialized_obj)
 
+            def object_id_serializer(obj):
+                owner_id = ""
+                owner_address = ""
+                if obj.is_direct_call_type():
+                    worker = ray.worker.get_global_worker()
+                    worker.check_connected()
+                    obj, owner_id, owner_address = (
+                        worker.core_worker.serialize_and_promote_object_id(obj)
+                    )
+                obj = obj.__reduce__()
+                owner_id = owner_id.__reduce__() if owner_id else owner_id
+                return pickle.dumps((obj, owner_id, owner_address))
+
+            def object_id_deserializer(serialized_obj):
+                obj_id, owner_id, owner_address = pickle.loads(serialized_obj)
+                # Must deserialize the object in the core worker before we
+                # create the ObjectID to ensure that the reference is added
+                # before we increment its count to 1.
+                if owner_id:
+                    worker = ray.worker.get_global_worker()
+                    worker.check_connected()
+                    # UniqueIDs are serialized as
+                    # (class name, (unique bytes,)).
+                    worker.core_worker.deserialize_and_register_object_id(
+                        obj_id[1][0], owner_id[1][0], owner_address)
+                obj_id = obj_id[0](obj_id[1][0])
+                return obj_id
+
             for id_type in ray._raylet._ID_TYPES:
-                serialization_context.register_type(
-                    id_type,
-                    "{}.{}".format(id_type.__module__, id_type.__name__),
-                    custom_serializer=id_serializer,
-                    custom_deserializer=id_deserializer)
+                if id_type == ray._raylet.ObjectID:
+                    serialization_context.register_type(
+                        id_type,
+                        "{}.{}".format(id_type.__module__, id_type.__name__),
+                        custom_serializer=object_id_serializer,
+                        custom_deserializer=object_id_deserializer)
+                else:
+                    serialization_context.register_type(
+                        id_type,
+                        "{}.{}".format(id_type.__module__, id_type.__name__),
+                        custom_serializer=id_serializer,
+                        custom_deserializer=id_deserializer)
 
             # We register this serializer on each worker instead of calling
             # _register_custom_serializer from the driver so that isinstance
@@ -188,16 +221,46 @@ class SerializationContext(object):
                 custom_deserializer=actor_handle_deserializer)
 
             def id_serializer(obj):
-                if isinstance(obj, ray.ObjectID) and obj.is_direct_call_type():
-                    obj = self.worker.core_worker.promote_object_to_plasma(obj)
                 return obj.__reduce__()
 
             def id_deserializer(serialized_obj):
                 return serialized_obj[0](*serialized_obj[1])
 
+            def object_id_serializer(obj):
+                owner_id = ""
+                owner_address = ""
+                if obj.is_direct_call_type():
+                    worker = ray.worker.get_global_worker()
+                    worker.check_connected()
+                    obj, owner_id, owner_address = (
+                        worker.core_worker.serialize_and_promote_object_id(obj)
+                    )
+                obj = id_serializer(obj)
+                owner_id = id_serializer(owner_id) if owner_id else owner_id
+                return (obj, owner_id, owner_address)
+
+            def object_id_deserializer(serialized_obj):
+                obj_id, owner_id, owner_address = serialized_obj
+                # Must deserialize the object in the core worker before we
+                # create the ObjectID to ensure that the reference is added
+                # before we increment its count to 1.
+                if owner_id:
+                    worker = ray.worker.get_global_worker()
+                    worker.check_connected()
+                    # UniqueIDs are serialized as
+                    # (class name, (unique bytes,)).
+                    worker.core_worker.deserialize_and_register_object_id(
+                        obj_id[1][0], owner_id[1][0], owner_address)
+                obj_id = id_deserializer(obj_id)
+                return obj_id
+
             for id_type in ray._raylet._ID_TYPES:
-                self._register_cloudpickle_serializer(id_type, id_serializer,
-                                                      id_deserializer)
+                if id_type == ray._raylet.ObjectID:
+                    self._register_cloudpickle_serializer(
+                        id_type, object_id_serializer, object_id_deserializer)
+                else:
+                    self._register_cloudpickle_serializer(
+                        id_type, id_serializer, id_deserializer)
 
     def initialize(self):
         """ Register custom serializers """

python/ray/tests/test_failure.py

@@ -951,15 +951,13 @@ def test_direct_call_serialized_id_eviction(ray_start_cluster):
     ray.get(get.remote([obj]))
 
 
-@pytest.mark.skip(
-    "Uncomment once eviction errors for serialized IDs are implemented")
 @pytest.mark.parametrize(
     "ray_start_cluster", [{
         "num_nodes": 2,
-        "num_cpus": 10,
+        "num_cpus": 1,
     }, {
         "num_nodes": 1,
-        "num_cpus": 20,
+        "num_cpus": 2,
     }],
     indirect=True)
 def test_direct_call_serialized_id(ray_start_cluster):
@@ -971,16 +969,33 @@ def test_direct_call_serialized_id(ray_start_cluster):
         return 1
 
     @ray.remote
-    def get(obj_ids):
+    def dependent_task(x):
+        return x
+
+    @ray.remote
+    def get(obj_ids, test_dependent_task):
         print("get", obj_ids)
         obj_id = obj_ids[0]
-        assert ray.get(obj_id) == 1
+        if test_dependent_task:
+            assert ray.get(dependent_task.remote(obj_id)) == 1
+        else:
+            assert ray.get(obj_id) == 1
 
     small_object = small_object.options(is_direct_call=True)
+    dependent_task = dependent_task.options(is_direct_call=True)
     get = get.options(is_direct_call=True)
 
     obj = small_object.remote()
-    ray.get(get.remote([obj]))
+    ray.get(get.remote([obj], False))
+
+    obj = small_object.remote()
+    ray.get(get.remote([obj], True))
+
+    obj = ray.put(1)
+    ray.get(get.remote([obj], False))
+
+    obj = ray.put(1)
+    ray.get(get.remote([obj], True))
 
 
 if __name__ == "__main__":