Refactor ID Serial 1: Separate ObjectID and TaskID from UniqueID (#4776)

* Enable BaseId.

* Change TaskID and make python test pass

* Remove unnecessary functions and fix test failure and change TaskID to
16 bytes.

* Java code change draft

* Refine

* Lint

* Update java/api/src/main/java/org/ray/api/id/TaskId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Update java/api/src/main/java/org/ray/api/id/BaseId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Update java/api/src/main/java/org/ray/api/id/BaseId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Update java/api/src/main/java/org/ray/api/id/ObjectId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Address comment

* Lint

* Fix SINGLE_PROCESS

* Fix comments

* Refine code

* Refine test

* Resolve conflict

python/ray/_raylet.pyx

@@ -88,11 +88,11 @@ def compute_put_id(TaskID task_id, int64_t put_index):
     if put_index < 1 or put_index > kMaxTaskPuts:
         raise ValueError("The range of 'put_index' should be [1, %d]"
                          % kMaxTaskPuts)
-    return ObjectID(ComputePutId(task_id.native(), put_index).binary())
+    return ObjectID(CObjectID.for_put(task_id.native(), put_index).binary())
 
 
 def compute_task_id(ObjectID object_id):
-    return TaskID(ComputeTaskId(object_id.native()).binary())
+    return TaskID(object_id.native().task_id().binary())
 
 
 cdef c_bool is_simple_value(value, int *num_elements_contained):

python/ray/actor.py

@@ -17,7 +17,6 @@ from ray.function_manager import FunctionDescriptor
 import ray.ray_constants as ray_constants
 import ray.signature as signature
 import ray.worker
-from ray.utils import _random_string
 from ray import (ObjectID, ActorID, ActorHandleID, ActorClassID, TaskID,
                  DriverID)
 
@@ -308,7 +307,7 @@ class ActorClass(object):
             raise Exception("Actors cannot be created before ray.init() "
                             "has been called.")
 
-        actor_id = ActorID(_random_string())
+        actor_id = ActorID.from_random()
         # The actor cursor is a dummy object representing the most recent
         # actor method invocation. For each subsequent method invocation,
         # the current cursor should be added as a dependency, and then
@@ -670,7 +669,7 @@ class ActorHandle(object):
             # to release, since it could be unpickled and submit another
             # dependent task at any time. Therefore, we notify the backend of a
             # random handle ID that will never actually be used.
-            new_actor_handle_id = ActorHandleID(_random_string())
+            new_actor_handle_id = ActorHandleID.from_random()
         # Notify the backend to expect this new actor handle. The backend will
         # not release the cursor for any new handles until the first task for
         # each of the new handles is submitted.
@@ -780,7 +779,7 @@ def make_actor(cls, num_cpus, num_gpus, resources, max_reconstructions):
     Class.__module__ = cls.__module__
     Class.__name__ = cls.__name__
 
-    class_id = ActorClassID(_random_string())
+    class_id = ActorClassID.from_random()
 
     return ActorClass(Class, class_id, max_reconstructions, num_cpus, num_gpus,
                       resources)

python/ray/includes/common.pxd

@@ -81,15 +81,9 @@ cdef extern from "ray/status.h" namespace "ray::StatusCode" nogil:
 
 
 cdef extern from "ray/id.h" namespace "ray" nogil:
-    const CTaskID FinishTaskId(const CTaskID &task_id)
-    const CObjectID ComputeReturnId(const CTaskID &task_id,
-                                    int64_t return_index)
-    const CObjectID ComputePutId(const CTaskID &task_id, int64_t put_index)
-    const CTaskID ComputeTaskId(const CObjectID &object_id)
     const CTaskID GenerateTaskId(const CDriverID &driver_id,
                                  const CTaskID &parent_task_id,
                                  int parent_task_counter)
-    int64_t ComputeObjectIndex(const CObjectID &object_id)
 
 
 cdef extern from "ray/gcs/format/gcs_generated.h" nogil:

python/ray/includes/unique_ids.pxd

@@ -1,12 +1,35 @@
 from libcpp cimport bool as c_bool
 from libcpp.string cimport string as c_string
-from libc.stdint cimport uint8_t
+from libc.stdint cimport uint8_t, int64_t
 
 cdef extern from "ray/id.h" namespace "ray" nogil:
-    cdef cppclass CUniqueID "ray::UniqueID":
+    cdef cppclass CBaseID[T]:
+        @staticmethod
+        T from_random()
+
+        @staticmethod
+        T from_binary(const c_string &binary)
+
+        @staticmethod
+        const T nil()
+
+        @staticmethod
+        size_t size()
+
+        size_t hash() const
+        c_bool is_nil() const
+        c_bool operator==(const CBaseID &rhs) const
+        c_bool operator!=(const CBaseID &rhs) const
+        const uint8_t *data() const;
+
+        c_string binary() const;
+        c_string hex() const;
+
+    cdef cppclass CUniqueID "ray::UniqueID"(CBaseID):
         CUniqueID()
-        CUniqueID(const c_string &binary)
-        CUniqueID(const CUniqueID &from_id)
+
+        @staticmethod
+        size_t size()
 
         @staticmethod
         CUniqueID from_random()
@@ -17,15 +40,8 @@ cdef extern from "ray/id.h" namespace "ray" nogil:
         @staticmethod
         const CUniqueID nil()
 
-        size_t hash() const
-        c_bool is_nil() const
-        c_bool operator==(const CUniqueID& rhs) const
-        c_bool operator!=(const CUniqueID& rhs) const
-        const uint8_t *data() const
-        uint8_t *mutable_data()
-        size_t size() const
-        c_string binary() const
-        c_string hex() const
+        @staticmethod
+        size_t size()
 
     cdef cppclass CActorCheckpointID "ray::ActorCheckpointID"(CUniqueID):
 
@@ -67,16 +83,40 @@ cdef extern from "ray/id.h" namespace "ray" nogil:
         @staticmethod
         CDriverID from_binary(const c_string &binary)
 
-    cdef cppclass CTaskID "ray::TaskID"(CUniqueID):
+    cdef cppclass CTaskID "ray::TaskID"(CBaseID[CTaskID]):
 
         @staticmethod
         CTaskID from_binary(const c_string &binary)
 
-    cdef cppclass CObjectID" ray::ObjectID"(CUniqueID):
+        @staticmethod
+        const CTaskID nil()
+
+        @staticmethod
+        size_t size()
+
+    cdef cppclass CObjectID" ray::ObjectID"(CBaseID[CObjectID]):
 
         @staticmethod
         CObjectID from_binary(const c_string &binary)
 
+        @staticmethod
+        const CObjectID nil()
+
+        @staticmethod
+        CObjectID for_put(const CTaskID &task_id, int64_t index);
+
+        @staticmethod
+        CObjectID for_task_return(const CTaskID &task_id, int64_t index);
+
+        @staticmethod
+        size_t size()
+
+        c_bool is_put()
+
+        int64_t object_index() const 
+
+        CTaskID task_id() const
+
     cdef cppclass CWorkerID "ray::WorkerID"(CUniqueID):
 
         @staticmethod

python/ray/includes/unique_ids.pxi

@@ -6,10 +6,8 @@ See https://github.com/ray-project/ray/issues/3721.
 
 # WARNING: Any additional ID types defined in this file must be added to the
 # _ID_TYPES list at the bottom of this file.
-from ray.includes.common cimport (
-    ComputePutId,
-    ComputeTaskId,
-)
+import os
+
 from ray.includes.unique_ids cimport (
     CActorCheckpointID,
     CActorClassID,
@@ -28,12 +26,12 @@ from ray.includes.unique_ids cimport (
 from ray.utils import decode
 
 
-def check_id(b):
+def check_id(b, size=kUniqueIDSize):
     if not isinstance(b, bytes):
         raise TypeError("Unsupported type: " + str(type(b)))
-    if len(b) != kUniqueIDSize:
+    if len(b) != size:
         raise ValueError("ID string needs to have length " +
-                         str(kUniqueIDSize))
+                         str(size))
 
 
 cdef extern from "ray/constants.h" nogil:
@@ -41,28 +39,27 @@ cdef extern from "ray/constants.h" nogil:
     cdef int64_t kMaxTaskPuts
 
 
-cdef class UniqueID:
-    cdef CUniqueID data
+cdef class BaseID:
 
-    def __init__(self, id):
-        check_id(id)
-        self.data = CUniqueID.from_binary(id)
+    # To avoid the error of "Python int too large to convert to C ssize_t",
+    # here `cdef size_t` is required.
+    cdef size_t hash(self):
+        pass
 
-    @classmethod
-    def from_binary(cls, id_bytes):
-        if not isinstance(id_bytes, bytes):
-            raise TypeError("Expect bytes, got " + str(type(id_bytes)))
-        return cls(id_bytes)
+    def binary(self):
+        pass
 
-    @classmethod
-    def nil(cls):
-        return cls(CUniqueID.nil().binary())
+    def size(self):
+        pass
 
-    def __hash__(self):
-        return self.data.hash()
+    def hex(self):
+        pass
 
     def is_nil(self):
-        return self.data.is_nil()
+        pass
+
+    def __hash__(self):
+        return self.hash()
 
     def __eq__(self, other):
         return type(self) == type(other) and self.binary() == other.binary()
@@ -70,18 +67,9 @@ cdef class UniqueID:
     def __ne__(self, other):
         return self.binary() != other.binary()
 
-    def size(self):
-        return self.data.size()
-
-    def binary(self):
-        return self.data.binary()
-
     def __bytes__(self):
         return self.binary()
 
-    def hex(self):
-        return decode(self.data.hex())
-
     def __hex__(self):
         return self.hex()
 
@@ -98,11 +86,52 @@ cdef class UniqueID:
         # NOTE: The hash function used here must match the one in
         # GetRedisContext in src/ray/gcs/tables.h. Changes to the
         # hash function should only be made through std::hash in
-        # src/common/common.h
+        # src/common/common.h.
+        # Do not use __hash__ that returns signed uint64_t, which
+        # is different from std::hash in c++ code.
+        return self.hash()
+
+
+cdef class UniqueID(BaseID):
+    cdef CUniqueID data
+
+    def __init__(self, id):
+        check_id(id)
+        self.data = CUniqueID.from_binary(id)
+
+    @classmethod
+    def from_binary(cls, id_bytes):
+        if not isinstance(id_bytes, bytes):
+            raise TypeError("Expect bytes, got " + str(type(id_bytes)))
+        return cls(id_bytes)
+
+    @classmethod
+    def nil(cls):
+        return cls(CUniqueID.nil().binary())
+
+    
+    @classmethod
+    def from_random(cls):
+        return cls(os.urandom(CUniqueID.size()))
+
+    def size(self):
+        return CUniqueID.size()
+
+    def binary(self):
+        return self.data.binary()
+
+    def hex(self):
+        return decode(self.data.hex())
+    
+    def is_nil(self):
+        return self.data.is_nil()
+
+    cdef size_t hash(self):
         return self.data.hash()
 
 
-cdef class ObjectID(UniqueID):
+cdef class ObjectID(BaseID):
+    cdef CObjectID data
 
     def __init__(self, id):
         check_id(id)
@@ -111,16 +140,67 @@ cdef class ObjectID(UniqueID):
     cdef CObjectID native(self):
         return <CObjectID>self.data
 
+    def size(self):
+        return CObjectID.size()
 
-cdef class TaskID(UniqueID):
+    def binary(self):
+        return self.data.binary()
+
+    def hex(self):
+        return decode(self.data.hex())
+    
+    def is_nil(self):
+        return self.data.is_nil()
+
+    cdef size_t hash(self):
+        return self.data.hash()
+
+    @classmethod
+    def nil(cls):
+        return cls(CObjectID.nil().binary())
+
+    @classmethod
+    def from_random(cls):
+        return cls(os.urandom(CObjectID.size()))
+
+
+cdef class TaskID(BaseID):
+    cdef CTaskID data
 
     def __init__(self, id):
-        check_id(id)
+        check_id(id, CTaskID.size())
         self.data = CTaskID.from_binary(<c_string>id)
 
     cdef CTaskID native(self):
         return <CTaskID>self.data
 
+    def size(self):
+        return CTaskID.size()
+
+    def binary(self):
+        return self.data.binary()
+
+    def hex(self):
+        return decode(self.data.hex())
+
+    def is_nil(self):
+        return self.data.is_nil()
+
+    cdef size_t hash(self):
+        return self.data.hash()
+
+    @classmethod
+    def nil(cls):
+        return cls(CTaskID.nil().binary())
+
+    @classmethod
+    def size(cla):
+        return CTaskID.size()
+
+    @classmethod
+    def from_random(cls):
+        return cls(os.urandom(CTaskID.size()))
+
 
 cdef class ClientID(UniqueID):
 

python/ray/monitor.py

@@ -16,8 +16,8 @@ import ray.cloudpickle as pickle
 import ray.gcs_utils
 import ray.utils
 import ray.ray_constants as ray_constants
-from ray.utils import (binary_to_hex, binary_to_object_id, hex_to_binary,
-                       setup_logger)
+from ray.utils import (binary_to_hex, binary_to_object_id, binary_to_task_id,
+                       hex_to_binary, setup_logger)
 
 logger = logging.getLogger(__name__)
 
@@ -169,8 +169,12 @@ class Monitor(object):
                 driver_object_id_bins.add(object_id.binary())
 
         def to_shard_index(id_bin):
-            return binary_to_object_id(id_bin).redis_shard_hash() % len(
-                self.state.redis_clients)
+            if len(id_bin) == ray.TaskID.size():
+                return binary_to_task_id(id_bin).redis_shard_hash() % len(
+                    self.state.redis_clients)
+            else:
+                return binary_to_object_id(id_bin).redis_shard_hash() % len(
+                    self.state.redis_clients)
 
         # Form the redis keys to delete.
         sharded_keys = [[] for _ in range(len(self.state.redis_clients))]

python/ray/tests/test_basic.py

@@ -7,6 +7,7 @@ import collections
 from concurrent.futures import ThreadPoolExecutor
 import json
 import logging
+from multiprocessing import Process
 import os
 import random
 import re
@@ -28,7 +29,6 @@ import pytest
 import ray
 import ray.tests.cluster_utils
 import ray.tests.utils
-from ray.utils import _random_string
 
 logger = logging.getLogger(__name__)
 
@@ -2630,12 +2630,33 @@ def test_object_id_properties():
         ray.ObjectID(id_bytes + b"1234")
     with pytest.raises(ValueError, match=r".*needs to have length 20.*"):
         ray.ObjectID(b"0123456789")
-    object_id = ray.ObjectID(_random_string())
+    object_id = ray.ObjectID.from_random()
     assert not object_id.is_nil()
     assert object_id.binary() != id_bytes
     id_dumps = pickle.dumps(object_id)
     id_from_dumps = pickle.loads(id_dumps)
     assert id_from_dumps == object_id
+    file_prefix = "test_object_id_properties"
+
+    # Make sure the ids are fork safe.
+    def write(index):
+        str = ray.ObjectID.from_random().hex()
+        with open("{}{}".format(file_prefix, index), "w") as fo:
+            fo.write(str)
+
+    def read(index):
+        with open("{}{}".format(file_prefix, index), "r") as fi:
+            for line in fi:
+                return line
+
+    processes = [Process(target=write, args=(_, )) for _ in range(4)]
+    for process in processes:
+        process.start()
+    for process in processes:
+        process.join()
+    hexes = {read(i) for i in range(4)}
+    [os.remove("{}{}".format(file_prefix, i)) for i in range(4)]
+    assert len(hexes) == 4
 
 
 @pytest.fixture
@@ -2768,7 +2789,7 @@ def test_pandas_parquet_serialization():
 
 
 def test_socket_dir_not_existing(shutdown_only):
-    random_name = ray.ObjectID(_random_string()).hex()
+    random_name = ray.ObjectID.from_random().hex()
     temp_raylet_socket_dir = "/tmp/ray/tests/{}".format(random_name)
     temp_raylet_socket_name = os.path.join(temp_raylet_socket_dir,
                                            "raylet_socket")

python/ray/tests/test_failure.py

@@ -15,7 +15,6 @@ import redis
 
 import ray
 import ray.ray_constants as ray_constants
-from ray.utils import _random_string
 from ray.tests.cluster_utils import Cluster
 from ray.tests.utils import (
     relevant_errors,
@@ -667,7 +666,7 @@ def test_warning_for_dead_node(ray_start_cluster_2_nodes):
 
 
 def test_raylet_crash_when_get(ray_start_regular):
-    nonexistent_id = ray.ObjectID(_random_string())
+    nonexistent_id = ray.ObjectID.from_random()
 
     def sleep_to_kill_raylet():
         # Don't kill raylet before default workers get connected.

python/ray/utils.py

@@ -216,6 +216,10 @@ def binary_to_object_id(binary_object_id):
     return ray.ObjectID(binary_object_id)
 
 
+def binary_to_task_id(binary_task_id):
+    return ray.TaskID(binary_task_id)
+
+
 def binary_to_hex(identifier):
     hex_identifier = binascii.hexlify(identifier)
     if sys.version_info >= (3, 0):

python/ray/worker.py

@@ -198,7 +198,7 @@ class Worker(object):
                 # to the current task ID may not be correct. Generate a
                 # random task ID so that the backend can differentiate
                 # between different threads.
-                self._task_context.current_task_id = TaskID(_random_string())
+                self._task_context.current_task_id = TaskID.from_random()
                 if getattr(self, "_multithreading_warned", False) is not True:
                     logger.warning(
                         "Calling ray.get or ray.wait in a separate thread "
@@ -1725,7 +1725,7 @@ def connect(node,
     else:
         # This is the code path of driver mode.
         if driver_id is None:
-            driver_id = DriverID(_random_string())
+            driver_id = DriverID.from_random()
 
         if not isinstance(driver_id, DriverID):
             raise TypeError("The type of given driver id must be DriverID.")
@@ -1834,6 +1834,7 @@ def connect(node,
     # Create an object store client.
     worker.plasma_client = thread_safe_client(
         plasma.connect(node.plasma_store_socket_name, None, 0, 300))
+    driver_id_str = _random_string()
 
     # If this is a driver, set the current task ID, the task driver ID, and set
     # the task index to 0.
@@ -1865,7 +1866,7 @@ def connect(node,
             function_descriptor.get_function_descriptor_list(),
             [],  # arguments.
             0,  # num_returns.
-            TaskID(_random_string()),  # parent_task_id.
+            TaskID(driver_id_str[:TaskID.size()]),  # parent_task_id.
             0,  # parent_counter.
             ActorID.nil(),  # actor_creation_id.
             ObjectID.nil(),  # actor_creation_dummy_object_id.
@@ -1894,7 +1895,7 @@ def connect(node,
         node.raylet_socket_name,
         ClientID(worker.worker_id),
         (mode == WORKER_MODE),
-        DriverID(worker.current_task_id.binary()),
+        DriverID(driver_id_str),
     )
 
     # Start the import thread