chore: import upstream snapshot with attribution

python/ray/includes/object_ref.pxi

@@ -0,0 +1,182 @@
+from ray.includes.unique_ids cimport CObjectID
+from ray.includes.optional cimport (
+    optional,
+    nullopt,
+)
+
+import asyncio
+import concurrent.futures
+import functools
+import logging
+import threading
+from typing import Callable, Any, Union, Optional
+from _collections_abc import GenericAlias
+from builtins import StopAsyncIteration
+
+import ray
+import cython
+
+logger = logging.getLogger(__name__)
+
+
+def _set_future_helper(
+    result: Any,
+    *,
+    py_future: Union[asyncio.Future, concurrent.futures.Future],
+):
+    # Issue #11030, #8841
+    # If this future has result set already, we just need to
+    # skip the set result/exception procedure.
+    if py_future.done():
+        return
+
+    if isinstance(result, RayTaskError):
+        exc = result.as_instanceof_cause()
+        # Convert StopIteration to RuntimeError to prevent segfaults due to
+        # Cpython's behavior w.r.t. PEP479
+        if isinstance(exc, StopIteration) or isinstance(exc, StopAsyncIteration):
+            runtime_error = RuntimeError(f"generator raised {type(exc).__name__}")
+            runtime_error.__cause__ = exc
+            py_future.set_exception(runtime_error)
+        else:
+            py_future.set_exception(exc)
+    elif isinstance(result, RayError):
+        # Directly raise exception for RayActorError
+        py_future.set_exception(result)
+    else:
+        py_future.set_result(result)
+
+
+cdef class ObjectRef(BaseID):
+    __class_getitem__ = classmethod(GenericAlias) # should match how typing.Generic works
+
+    def __cinit__(self):
+        self.in_core_worker = False
+
+    def __init__(
+            self, id, owner_addr="", call_site_data="",
+            skip_adding_local_ref=False, tensor_transport: Optional[str] = None):
+        self._set_id(id)
+        self.owner_addr = owner_addr
+        self.in_core_worker = False
+        self.call_site_data = call_site_data
+        self._tensor_transport = tensor_transport
+        worker = ray._private.worker.global_worker
+        # TODO(edoakes): We should be able to remove the in_core_worker flag.
+        # But there are still some dummy object refs being created outside the
+        # context of a core worker.
+        if hasattr(worker, "core_worker"):
+            if not skip_adding_local_ref:
+                worker.core_worker.add_object_ref_reference(self)
+            self.in_core_worker = True
+
+    def __dealloc__(self):
+        if self.in_core_worker:
+            try:
+                worker = ray._private.worker.global_worker
+                worker.core_worker.remove_object_ref_reference(self)
+            except Exception as e:
+                # There is a strange error in rllib that causes the above to
+                # fail. Somehow the global 'ray' variable corresponding to the
+                # imported package is None when this gets called. Unfortunately
+                # this is hard to debug because __dealloc__ is called during
+                # garbage collection so we can't get a good stack trace. In any
+                # case, there's not much we can do besides ignore it
+                # (re-importing ray won't help).
+                pass
+
+    cdef CObjectID native(self):
+        return <CObjectID>self.data
+
+    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()
+
+    def task_id(self):
+        return TaskID(self.data.TaskId().Binary())
+
+    def job_id(self):
+        return self.task_id().job_id()
+
+    def owner_address(self):
+        return self.owner_addr
+
+    def call_site(self):
+        return decode(self.call_site_data)
+
+    @classmethod
+    def size(cls):
+        return CObjectID.Size()
+
+    def _set_id(self, id):
+        check_id(id)
+        self.data = CObjectID.FromBinary(<c_string>id)
+
+    @classmethod
+    def nil(cls):
+        return cls(CObjectID.Nil().Binary())
+
+    @classmethod
+    def from_random(cls):
+        return cls(CObjectID.FromRandom().Binary())
+
+    def future(self) -> concurrent.futures.Future:
+        """Wrap ObjectRef with a concurrent.futures.Future
+
+        Note that the future cancellation will not cancel the correspoding
+        task when the ObjectRef representing return object of a task.
+        Additionally, future.running() will always be ``False`` even if the
+        underlying task is running.
+        """
+        py_future = concurrent.futures.Future()
+
+        self._on_completed(
+            functools.partial(_set_future_helper, py_future=py_future))
+
+        # A hack to keep a reference to the object ref for ref counting.
+        py_future.object_ref = self
+        return py_future
+
+    def __await__(self):
+        return self.as_future(_internal=True).__await__()
+
+    def as_future(self, _internal=False) -> asyncio.Future:
+        """Wrap ObjectRef with an asyncio.Future.
+
+        Note that the future cancellation will not cancel the correspoding
+        task when the ObjectRef representing return object of a task.
+        """
+        if not _internal:
+            logger.warning("ref.as_future() is deprecated in favor of "
+                           "asyncio.wrap_future(ref.future()).")
+        return asyncio.wrap_future(self.future())
+
+    def _on_completed(self, py_callback: Callable[[Any], None]):
+        """Register a callback that will be called after Object is ready.
+        If the ObjectRef is already ready, the callback will be called soon.
+        The callback should take the result as the only argument. The result
+        can be an exception object in case of task error.
+        """
+        core_worker = ray._private.worker.global_worker.core_worker
+        core_worker.set_get_async_callback(self, py_callback)
+        return self
+
+    def tensor_transport(self):
+        return self._tensor_transport
+
+    cdef optional[c_string] c_tensor_transport(self):
+        cdef:
+            optional[c_string] c_tensor_transport = nullopt
+            c_string c_tensor_transport_str
+        if self._tensor_transport is not None:
+            c_tensor_transport_str = self._tensor_transport.encode("utf-8")
+            c_tensor_transport.emplace(move(c_tensor_transport_str))
+        return c_tensor_transport