chore: import upstream snapshot with attribution

python/ray/includes/unique_ids.pxi

@@ -0,0 +1,436 @@
+"""This is a module for unique IDs in Ray.
+We define different types for different IDs for type safety.
+
+See https://github.com/ray-project/ray/issues/3721.
+"""
+
+import logging
+import os
+
+from ray.includes.unique_ids cimport (
+    CActorClassID,
+    CActorID,
+    CNodeID,
+    CConfigID,
+    CJobID,
+    CFunctionID,
+    CObjectID,
+    CTaskID,
+    CUniqueID,
+    CWorkerID,
+    CPlacementGroupID,
+    CClusterID,
+)
+
+
+import ray
+from ray._common.utils import decode
+
+logger = logging.getLogger(__name__)
+
+
+def check_id(b, size=kUniqueIDSize):
+    if not isinstance(b, bytes):
+        raise TypeError(f"Unsupported type: {type(b)}, expected bytes")
+    if len(b) != size:
+        raise ValueError("ID string needs to have length " +
+                         str(size) + ", got " + str(len(b)))
+
+
+def check_nil(id_obj):
+    if id_obj.is_nil():
+        raise ValueError(f"{id_obj.__class__.__name__} is nil")
+
+
+cdef extern from "ray/common/constants.h" nogil:
+    cdef int64_t kUniqueIDSize
+
+
+cdef class BaseID:
+
+    @classmethod
+    def from_binary(cls, id_bytes):
+        return cls(id_bytes)
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        raise NotImplementedError
+
+    cdef size_t hash(self):
+        raise NotImplementedError
+
+    def binary(self):
+        raise NotImplementedError
+
+    @classmethod
+    def size(cls):
+        raise NotImplementedError
+
+    def hex(self):
+        raise NotImplementedError
+
+    def is_nil(self):
+        raise NotImplementedError
+
+    def __hash__(self):
+        return self.hash()
+
+    def __eq__(self, other):
+        return type(self) == type(other) and self.binary() == other.binary()
+
+    def __ne__(self, other):
+        return type(self) != type(other) or self.binary() != other.binary()
+
+    def __bytes__(self):
+        return self.binary()
+
+    def __hex__(self):
+        return self.hex()
+
+    def __repr__(self):
+        return self.__class__.__name__ + "(" + self.hex() + ")"
+
+    def __str__(self):
+        return self.__repr__()
+
+    def __reduce__(self):
+        return type(self), (self.binary(),)
+
+    def redis_shard_hash(self):
+        # 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.
+        # 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.FromBinary(id)
+
+    @classmethod
+    def nil(cls):
+        return cls(CUniqueID.Nil().Binary())
+
+    @classmethod
+    def from_random(cls):
+        return cls(CUniqueID.FromRandom().Binary())
+
+    @classmethod
+    def size(cls):
+        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.IsNil()
+
+    cdef size_t hash(self):
+        return self.data.Hash()
+
+
+cdef class TaskID(BaseID):
+    cdef CTaskID data
+
+    def __init__(self, id):
+        check_id(id, CTaskID.Size())
+        self.data = CTaskID.FromBinary(<c_string>id)
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        binary_id = CTaskID.FromHex(<c_string>hex_id).Binary()
+        return cls(binary_id)
+
+    cdef CTaskID native(self):
+        return <CTaskID>self.data
+
+    @classmethod
+    def size(cls):
+        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.IsNil()
+
+    def actor_id(self):
+        return ActorID(self.data.ActorId().Binary())
+
+    def job_id(self):
+        check_nil(self)
+        return JobID(self.data.JobId().Binary())
+
+    cdef size_t hash(self):
+        return self.data.Hash()
+
+    @classmethod
+    def nil(cls):
+        return cls(CTaskID.Nil().Binary())
+
+    @classmethod
+    def for_fake_task(cls, job_id):
+        return cls(CTaskID.FromRandom(
+            CJobID.FromBinary(job_id.binary())).Binary())
+
+    @classmethod
+    def for_driver_task(cls, job_id):
+        return cls(CTaskID.ForDriverTask(
+            CJobID.FromBinary(job_id.binary())).Binary())
+
+    @classmethod
+    def for_actor_creation_task(cls, actor_id):
+        assert isinstance(actor_id, ActorID)
+        return cls(CTaskID.ForActorCreationTask(
+            CActorID.FromBinary(actor_id.binary())).Binary())
+
+    @classmethod
+    def for_actor_task(cls, job_id, parent_task_id,
+                       parent_task_counter, actor_id):
+        assert isinstance(job_id, JobID)
+        assert isinstance(parent_task_id, TaskID)
+        assert isinstance(actor_id, ActorID)
+        return cls(CTaskID.ForActorTask(
+            CJobID.FromBinary(job_id.binary()),
+            CTaskID.FromBinary(parent_task_id.binary()),
+            parent_task_counter,
+            CActorID.FromBinary(actor_id.binary())).Binary())
+
+    @classmethod
+    def for_normal_task(cls, job_id, parent_task_id, parent_task_counter):
+        assert isinstance(job_id, JobID)
+        assert isinstance(parent_task_id, TaskID)
+        return cls(CTaskID.ForNormalTask(
+            CJobID.FromBinary(job_id.binary()),
+            CTaskID.FromBinary(parent_task_id.binary()),
+            parent_task_counter).Binary())
+
+cdef class NodeID(UniqueID):
+
+    def __init__(self, id):
+        check_id(id)
+        self.data = CUniqueID.FromBinary(<c_string>id)
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        binary_id = CNodeID.FromHex(<c_string>hex_id).Binary()
+        return cls(binary_id)
+
+    cdef CNodeID native(self):
+        return <CNodeID>self.data
+
+
+cdef class JobID(BaseID):
+    cdef CJobID data
+
+    def __init__(self, id):
+        check_id(id, CJobID.Size())
+        self.data = CJobID.FromBinary(<c_string>id)
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        binary_id = CJobID.FromHex(<c_string>hex_id).Binary()
+        return cls(binary_id)
+
+    cdef CJobID native(self):
+        return <CJobID>self.data
+
+    @classmethod
+    def from_int(cls, value):
+        assert value < 2**32, "Maximum JobID integer is 2**32 - 1."
+        return cls(CJobID.FromInt(value).Binary())
+
+    @classmethod
+    def nil(cls):
+        return cls(CJobID.Nil().Binary())
+
+    @classmethod
+    def size(cls):
+        return CJobID.Size()
+
+    def int(self):
+        return self.data.ToInt()
+
+    def binary(self):
+        return self.data.Binary()
+
+    def hex(self):
+        return decode(self.data.Hex())
+
+    def is_nil(self):
+        return self.data.IsNil()
+
+    cdef size_t hash(self):
+        return self.data.Hash()
+
+cdef class WorkerID(UniqueID):
+
+    def __init__(self, id):
+        check_id(id)
+        self.data = CUniqueID.FromBinary(<c_string>id)
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        binary_id = CWorkerID.FromHex(<c_string>hex_id).Binary()
+        return cls(binary_id)
+
+    cdef CWorkerID native(self):
+        return <CWorkerID>self.data
+
+cdef class ActorID(BaseID):
+
+    def __init__(self, id):
+        self._set_id(id)
+
+    @classmethod
+    def of(cls, job_id, parent_task_id, parent_task_counter):
+        assert isinstance(job_id, JobID)
+        assert isinstance(parent_task_id, TaskID)
+        return cls(CActorID.Of(CJobID.FromBinary(job_id.binary()),
+                               CTaskID.FromBinary(parent_task_id.binary()),
+                               parent_task_counter).Binary())
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        binary_id = CActorID.FromHex(<c_string>hex_id).Binary()
+        return cls(binary_id)
+
+    @classmethod
+    def nil(cls):
+        return cls(CActorID.Nil().Binary())
+
+    @classmethod
+    def from_random(cls):
+        return cls(os.urandom(CActorID.Size()))
+
+    @classmethod
+    def size(cls):
+        return CActorID.Size()
+
+    def _set_id(self, id):
+        check_id(id, CActorID.Size())
+        self.data = CActorID.FromBinary(<c_string>id)
+
+    @property
+    def job_id(self):
+        check_nil(self)
+        return JobID(self.data.JobId().Binary())
+
+    def binary(self):
+        return self.data.Binary()
+
+    def hex(self):
+        return decode(self.data.Hex())
+
+    def is_nil(self):
+        return self.data.IsNil()
+
+    cdef size_t hash(self):
+        return self.data.Hash()
+
+    cdef CActorID native(self):
+        return <CActorID>self.data
+
+
+cdef class FunctionID(UniqueID):
+
+    def __init__(self, id):
+        check_id(id)
+        self.data = CUniqueID.FromBinary(<c_string>id)
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        binary_id = CFunctionID.FromHex(<c_string>hex_id).Binary()
+        return cls(binary_id)
+
+    cdef CFunctionID native(self):
+        return <CFunctionID>self.data
+
+
+cdef class ActorClassID(UniqueID):
+
+    def __init__(self, id):
+        check_id(id)
+        self.data = CUniqueID.FromBinary(<c_string>id)
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        binary_id = CActorClassID.FromHex(<c_string>hex_id).Binary()
+        return cls(binary_id)
+
+    cdef CActorClassID native(self):
+        return <CActorClassID>self.data
+
+cdef class ClusterID(UniqueID):
+
+    def __init__(self, id):
+        check_id(id)
+        self.data = CUniqueID.FromBinary(<c_string>id)
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        binary_id = CClusterID.FromHex(<c_string>hex_id).Binary()
+        return cls(binary_id)
+
+    cdef CClusterID native(self):
+        return <CClusterID>self.data
+
+
+# This type alias is for backward compatibility.
+ObjectID = ObjectRef
+
+cdef class PlacementGroupID(BaseID):
+    cdef CPlacementGroupID data
+
+    def __init__(self, id):
+        check_id(id, CPlacementGroupID.Size())
+        self.data = CPlacementGroupID.FromBinary(<c_string>id)
+
+    cdef CPlacementGroupID native(self):
+        return <CPlacementGroupID>self.data
+
+    @classmethod
+    def from_hex(cls, hex_id):
+        binary_id = CPlacementGroupID.FromHex(<c_string>hex_id).Binary()
+        return cls(binary_id)
+
+    @classmethod
+    def from_random(cls):
+        return cls(os.urandom(CPlacementGroupID.Size()))
+
+    @classmethod
+    def of(cls, job_id):
+        assert isinstance(job_id, JobID)
+        return cls(CPlacementGroupID.Of(CJobID.FromBinary(job_id.binary())).Binary())
+
+    @classmethod
+    def nil(cls):
+        return cls(CPlacementGroupID.Nil().Binary())
+
+    @classmethod
+    def size(cls):
+        return CPlacementGroupID.Size()
+
+    def binary(self):
+        return self.data.Binary()
+
+    def hex(self):
+        return decode(self.data.Hex())
+
+    def is_nil(self):
+        return self.data.IsNil()
+
+    cdef size_t hash(self):
+        return self.data.Hash()