diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index b95f993aca2bd..33932c8367219 100644 --- a/python/ray/_raylet.pyx +++ b/python/ray/_raylet.pyx @@ -303,8 +303,8 @@ cdef void prepare_args( else: args_vector.push_back( CTaskArg.PassByReference( - (core_worker.put_serialized_object( - serialized_arg)).native())) + (CObjectID.FromBinary(core_worker.put_serialized_cobject( + serialized_arg))))) cdef deserialize_args( const c_vector[shared_ptr[CRayObject]] &c_args, @@ -689,6 +689,11 @@ cdef class CoreWorker: def put_serialized_object(self, serialized_object, ObjectID object_id=None, c_bool pin_object=True): + return ObjectID(self.put_serialized_cobject(serialized_object, object_id, pin_object)) + + def put_serialized_cobject(self, serialized_object, + ObjectID object_id=None, + c_bool pin_object=True): cdef: CObjectID c_object_id shared_ptr[CBuffer] data @@ -710,7 +715,7 @@ cdef class CoreWorker: self.core_worker.get().Seal( c_object_id, pin_object and object_id is None)) - return ObjectID(c_object_id.Binary()) + return c_object_id.Binary() def wait(self, object_ids, int num_returns, int64_t timeout_ms, TaskID current_task_id): diff --git a/python/ray/tests/test_reference_counting.py b/python/ray/tests/test_reference_counting.py index f1f9e2d767389..f900a1e3946e8 100644 --- a/python/ray/tests/test_reference_counting.py +++ b/python/ray/tests/test_reference_counting.py @@ -195,17 +195,13 @@ def test_pending_task_dependency_pinning(shutdown_only): def pending(input1, input2): return - @ray.remote - def slow(dep): - pass - # The object that is ray.put here will go out of scope immediately, so if # pending task dependencies aren't considered, it will be evicted before # the ray.get below due to the subsequent ray.puts that fill up the object # store. np_array = np.zeros(40 * 1024 * 1024, dtype=np.uint8) random_id = ray.ObjectID.from_random() - oid = pending.remote(np_array, slow.remote(random_id)) + oid = pending.remote(np_array, random_id) for _ in range(2): ray.put(np_array)