[2/N] Streaming Generator. Support core worker APIs + cython generato…

…r interface. (#35324) (#35682)

This is the second PR to support streaming generator.

The detailed design and API proposal can be found from https://docs.google.com/document/d/1hAASLe2sCoay23raqxqwJdSDiJWNMcNhlTwWJXsJOU4/edit#heading=h.w91y1fgnpu0m.
The Execution plan can be found from https://docs.google.com/document/d/1hAASLe2sCoay23raqxqwJdSDiJWNMcNhlTwWJXsJOU4/edit#heading=h.kxktymq5ihf7.
There will be 4 PRs to enable streaming generator for Ray Serve (phase 1).

 This PR -> introduce cpp interfaces to handle intermediate task return [1/N] Streaming Generator. Cpp interfaces and implementation #35291
 Support core worker APIs + cython generator interface. [2/N] Streaming Generator. Support core worker APIs + cython generator interface. #35324 < --- this PR
 E2e integration [3/N] Streaming Generator. E2e integration #35325 (review)
 Support async actors
This PR implements the Cython generator interface that users can use to obtain a next available object reference.
---------

Signed-off-by: SangBin Cho <rkooo567@gmail.com>

python/ray/_raylet.pyx

@@ -198,6 +198,149 @@ class ObjectRefGenerator:
         return len(self._refs)
 
 
+class ObjectRefStreamEoFError(RayError):
+    pass
+
+
+class StreamingObjectRefGenerator:
+    def __init__(self, generator_ref: ObjectRef, worker: "Worker"):
+        # The reference to a generator task.
+        self._generator_ref = generator_ref
+        # The last time generator task has completed.
+        self._generator_task_completed_time = None
+        # The exception raised from a generator task.
+        self._generator_task_exception = None
+        # Ray's worker class. ray._private.worker.global_worker
+        self.worker = worker
+        assert hasattr(worker, "core_worker")
+        self.worker.core_worker.create_object_ref_stream(self._generator_ref)
+
+    def __iter__(self) -> "StreamingObjectRefGenerator":
+        return self
+
+    def __next__(self) -> ObjectRef:
+        """Waits until a next ref is available and returns the object ref.
+
+        Raises StopIteration if there's no more objects
+        to generate.
+
+        The object ref will contain an exception if the task fails.
+        When the generator task returns N objects, it can return
+        up to N + 1 objects (if there's a system failure, the
+        last object will contain a system level exception).
+        """
+        return self._next()
+
+    def _next(
+            self,
+            timeout_s: float = -1,
+            sleep_interval_s: float = 0.0001,
+            unexpected_network_failure_timeout_s: float = 30) -> ObjectRef:
+        """Waits for timeout_s and returns the object ref if available.
+
+        If an object is not available within the given timeout, it
+        returns a nil object reference.
+
+        If -1 timeout is provided, it means it waits infinitely.
+
+        Waiting is implemented as busy waiting. You can control
+        the busy waiting interval via sleep_interval_s.
+
+        Raises StopIteration if there's no more objects
+        to generate.
+
+        The object ref will contain an exception if the task fails.
+        When the generator task returns N objects, it can return
+        up to N + 1 objects (if there's a system failure, the
+        last object will contain a system level exception).
+
+        Args:
+            timeout_s: If the next object is not ready within
+                this timeout, it returns the nil object ref.
+            sleep_interval_s: busy waiting interval.
+            unexpected_network_failure_timeout_s: If the
+                task is finished, but the next ref is not
+                available within this time, it will hard fail
+                the generator.
+        """
+        obj = self._handle_next()
+        last_time = time.time()
+
+        # The generator ref will be None if the task succeeds.
+        # It will contain an exception if the task fails by
+        # a system error.
+        while obj.is_nil():
+            if self._generator_task_exception:
+                # The generator task has failed already.
+                # We raise StopIteration
+                # to conform the next interface in Python.
+                raise StopIteration from None
+            else:
+                # Otherwise, we should ray.get on the generator
+                # ref to find if the task has a system failure.
+                # Return the generator ref that contains the system
+                # error as soon as possible.
+                r, _ = ray.wait([self._generator_ref], timeout=0)
+                if len(r) > 0:
+                    try:
+                        ray.get(r)
+                    except Exception as e:
+                        # If it has failed, return the generator task ref
+                        # so that the ref will raise an exception.
+                        self._generator_task_exception = e
+                        return self._generator_ref
+                    finally:
+                        if self._generator_task_completed_time is None:
+                            self._generator_task_completed_time = time.time()
+
+            # Currently, since the ordering of intermediate result report
+            # is not guaranteed, it is possible that althoug the task
+            # has succeeded, all of the object references are not reported
+            # (e.g., when there are network failures).
+            # If all the object refs are not reported to the generator
+            # within 30 seconds, we consider is as an unreconverable error.
+            if self._generator_task_completed_time:
+                if (time.time() - self._generator_task_completed_time
+                        > unexpected_network_failure_timeout_s):
+                    # It means the next wasn't reported although the task
+                    # has been terminated 30 seconds ago.
+                    self._generator_task_exception = AssertionError
+                    assert False, "Unexpected network failure occured."
+
+            if timeout_s != -1 and time.time() - last_time > timeout_s:
+                return ObjectRef.nil()
+
+            # 100us busy waiting
+            time.sleep(sleep_interval_s)
+            obj = self._handle_next()
+        return obj
+
+    def _handle_next(self) -> ObjectRef:
+        try:
+            if hasattr(self.worker, "core_worker"):
+                obj = self.worker.core_worker.try_read_next_object_ref_stream(
+                    self._generator_ref)
+                return obj
+            else:
+                raise ValueError(
+                    "Cannot access the core worker. "
+                    "Did you already shutdown Ray via ray.shutdown()?")
+        except ObjectRefStreamEoFError:
+            raise StopIteration from None
+
+    def __del__(self):
+        if hasattr(self.worker, "core_worker"):
+            # NOTE: This can be called multiple times
+            # because python doesn't guarantee __del__ is called
+            # only once.
+            self.worker.core_worker.delete_object_ref_stream(self._generator_ref)
+
+    def __getstate__(self):
+        raise TypeError(
+            "You cannot return or pass a generator to other task. "
+            "Serializing a StreamingObjectRefGenerator is not allowed.")
+
+
 cdef int check_status(const CRayStatus& status) nogil except -1:
     if status.ok():
         return 0
@@ -209,6 +352,8 @@ cdef int check_status(const CRayStatus& status) nogil except -1:
         raise ObjectStoreFullError(message)
     elif status.IsOutOfDisk():
         raise OutOfDiskError(message)
+    elif status.IsObjectRefStreamEoF():
+        raise ObjectRefStreamEoFError(message)
     elif status.IsInterrupted():
         raise KeyboardInterrupt()
     elif status.IsTimedOut():
@@ -3139,6 +3284,33 @@ cdef class CoreWorker:
         CCoreWorkerProcess.GetCoreWorker() \
             .RecordTaskLogEnd(out_end_offset, err_end_offset)
 
+    def create_object_ref_stream(self, ObjectRef generator_id):
+        cdef:
+            CObjectID c_generator_id = generator_id.native()
+
+        CCoreWorkerProcess.GetCoreWorker().CreateObjectRefStream(c_generator_id)
+
+    def delete_object_ref_stream(self, ObjectRef generator_id):
+        cdef:
+            CObjectID c_generator_id = generator_id.native()
+
+        CCoreWorkerProcess.GetCoreWorker().DelObjectRefStream(c_generator_id)
+
+    def try_read_next_object_ref_stream(self, ObjectRef generator_id):
+        cdef:
+            CObjectID c_generator_id = generator_id.native()
+            CObjectReference c_object_ref
+
+        check_status(
+            CCoreWorkerProcess.GetCoreWorker().TryReadObjectRefStream(
+                c_generator_id, &c_object_ref))
+        return ObjectRef(
+            c_object_ref.object_id(),
+            c_object_ref.owner_address().SerializeAsString(),
+            "",
+            # Already added when the ref is updated.
+            skip_adding_local_ref=True)
+
 cdef void async_callback(shared_ptr[CRayObject] obj,
                          CObjectID object_ref,
                          void *user_callback) with gil:

python/ray/includes/common.pxd

@@ -99,6 +99,9 @@ cdef extern from "ray/common/status.h" namespace "ray" nogil:
         @staticmethod
         CRayStatus NotFound()
 
+        @staticmethod
+        CRayStatus ObjectRefStreamEoF()
+
         c_bool ok()
         c_bool IsOutOfMemory()
         c_bool IsKeyError()
@@ -118,6 +121,7 @@ cdef extern from "ray/common/status.h" namespace "ray" nogil:
         c_bool IsObjectUnknownOwner()
         c_bool IsRpcError()
         c_bool IsOutOfResource()
+        c_bool IsObjectRefStreamEoF()
 
         c_string ToString()
         c_string CodeAsString()

python/ray/includes/libcoreworker.pxd

@@ -147,6 +147,11 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
             const CObjectID& return_id,
             shared_ptr[CRayObject] *return_object,
             const CObjectID& generator_id)
+        void DelObjectRefStream(const CObjectID &generator_id)
+        void CreateObjectRefStream(const CObjectID &generator_id)
+        CRayStatus TryReadObjectRefStream(
+            const CObjectID &generator_id,
+            CObjectReference *object_ref_out)
         CObjectID AllocateDynamicReturnId()
 
         CJobID GetCurrentJobId()

python/ray/tests/BUILD

@@ -46,6 +46,7 @@ py_test_module_list(
     "test_gcs_fault_tolerance.py",
     "test_gcs_utils.py",
     "test_generators.py",
+    "test_streaming_generator.py",
     "test_metrics_agent.py",
     "test_metrics_head.py",
     "test_component_failures_2.py",

python/ray/tests/test_streaming_generator.py

@@ -0,0 +1,143 @@
+import pytest
+import sys
+import time
+
+from unittest.mock import patch, Mock
+
+import ray
+from ray._raylet import StreamingObjectRefGenerator, ObjectRefStreamEoFError
+from ray.cloudpickle import dumps
+from ray.exceptions import WorkerCrashedError
+
+
+class MockedWorker:
+    def __init__(self, mocked_core_worker):
+        self.core_worker = mocked_core_worker
+
+    def reset_core_worker(self):
+        """Emulate the case ray.shutdown is called
+        and the core_worker instance is GC'ed.
+        """
+        self.core_worker = None
+
+
+@pytest.fixture
+def mocked_worker():
+    mocked_core_worker = Mock()
+    mocked_core_worker.try_read_next_object_ref_stream.return_value = None
+    mocked_core_worker.delete_object_ref_stream.return_value = None
+    mocked_core_worker.create_object_ref_stream.return_value = None
+    worker = MockedWorker(mocked_core_worker)
+    yield worker
+
+
+def test_streaming_object_ref_generator_basic_unit(mocked_worker):
+    """
+    Verify the basic case:
+    create a generator -> read values -> nothing more to read -> delete.
+    """
+    with patch("ray.wait") as mocked_ray_wait:
+        c = mocked_worker.core_worker
+        generator_ref = ray.ObjectRef.from_random()
+        generator = StreamingObjectRefGenerator(generator_ref, mocked_worker)
+        c.try_read_next_object_ref_stream.return_value = ray.ObjectRef.nil()
+        c.create_object_ref_stream.assert_called()
+
+        # Test when there's no new ref, it returns a nil.
+        mocked_ray_wait.return_value = [], [generator_ref]
+        ref = generator._next(timeout_s=0)
+        assert ref.is_nil()
+
+        # When the new ref is available, next should return it.
+        for _ in range(3):
+            new_ref = ray.ObjectRef.from_random()
+            c.try_read_next_object_ref_stream.return_value = new_ref
+            ref = generator._next(timeout_s=0)
+            assert new_ref == ref
+
+        # When try_read_next_object_ref_stream raises a
+        # ObjectRefStreamEoFError, it should raise a stop iteration.
+        c.try_read_next_object_ref_stream.side_effect = ObjectRefStreamEoFError(
+            ""
+        )  # noqa
+        with pytest.raises(StopIteration):
+            ref = generator._next(timeout_s=0)
+
+        # Make sure we cannot serialize the generator.
+        with pytest.raises(TypeError):
+            dumps(generator)
+
+        del generator
+        c.delete_object_ref_stream.assert_called()
+
+
+def test_streaming_object_ref_generator_task_failed_unit(mocked_worker):
+    """
+    Verify when a task is failed by a system error,
+    the generator ref is returned.
+    """
+    with patch("ray.get") as mocked_ray_get:
+        with patch("ray.wait") as mocked_ray_wait:
+            c = mocked_worker.core_worker
+            generator_ref = ray.ObjectRef.from_random()
+            generator = StreamingObjectRefGenerator(generator_ref, mocked_worker)
+
+            # Simulate the worker failure happens.
+            mocked_ray_wait.return_value = [generator_ref], []
+            mocked_ray_get.side_effect = WorkerCrashedError()
+
+            c.try_read_next_object_ref_stream.return_value = ray.ObjectRef.nil()
+            ref = generator._next(timeout_s=0)
+            # If the generator task fails by a systsem error,
+            # meaning the ref will raise an exception
+            # it should be returned.
+            print(ref)
+            print(generator_ref)
+            assert ref == generator_ref
+
+            # Once exception is raised, it should always
+            # raise stopIteration regardless of what
+            # the ref contains now.
+            with pytest.raises(StopIteration):
+                ref = generator._next(timeout_s=0)
+
+
+def test_streaming_object_ref_generator_network_failed_unit(mocked_worker):
+    """
+    Verify when a task is finished, but if the next ref is not available
+    on time, it raises an assertion error.
+
+    TODO(sang): Once we move the task subimssion path to use pubsub
+    to guarantee the ordering, we don't need this test anymore.
+    """
+    with patch("ray.get") as mocked_ray_get:
+        with patch("ray.wait") as mocked_ray_wait:
+            c = mocked_worker.core_worker
+            generator_ref = ray.ObjectRef.from_random()
+            generator = StreamingObjectRefGenerator(generator_ref, mocked_worker)
+
+            # Simulate the task has finished.
+            mocked_ray_wait.return_value = [generator_ref], []
+            mocked_ray_get.return_value = None
+
+            # If StopIteration is not raised within
+            # unexpected_network_failure_timeout_s second,
+            # it should fail.
+            c.try_read_next_object_ref_stream.return_value = ray.ObjectRef.nil()
+            ref = generator._next(timeout_s=0, unexpected_network_failure_timeout_s=1)
+            assert ref == ray.ObjectRef.nil()
+            time.sleep(1)
+            with pytest.raises(AssertionError):
+                generator._next(timeout_s=0, unexpected_network_failure_timeout_s=1)
+            # After that StopIteration should be raised.
+            with pytest.raises(StopIteration):
+                generator._next(timeout_s=0, unexpected_network_failure_timeout_s=1)
+
+
+if __name__ == "__main__":
+    import os
+
+    if os.environ.get("PARALLEL_CI"):
+        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
+    else:
+        sys.exit(pytest.main(["-sv", __file__]))