chore: import upstream snapshot with attribution

python/ray/tests/test_streaming_generator_regression.py

@@ -0,0 +1,139 @@
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+
+import ray
+from ray.actor import ActorHandle
+from ray.exceptions import RayTaskError, TaskCancelledError
+from ray.util.state import list_workers
+
+
+@ray.remote(num_cpus=1)
+class EndpointActor:
+    def __init__(self, *, injected_executor_delay_s: float, tokens_per_request: int):
+        self._tokens_per_request = tokens_per_request
+        # In this test we simulate conditions leading to use-after-free conditions,
+        # by injecting delays into worker's thread-pool executor
+        self._inject_delay_in_core_worker_executor(
+            target_delay_s=injected_executor_delay_s,
+            max_workers=1,
+        )
+
+    async def aio_stream(self):
+        for i in range(self._tokens_per_request):
+            yield i
+
+    @classmethod
+    def _inject_delay_in_core_worker_executor(
+        cls, target_delay_s: float, max_workers: int
+    ):
+        if target_delay_s > 0:
+
+            class DelayedThreadPoolExecutor(ThreadPoolExecutor):
+                def submit(self, fn, /, *args, **kwargs):
+                    def __slowed_fn():
+                        print(
+                            f">>> [DelayedThreadPoolExecutor] Starting executing "
+                            f"function with delay {target_delay_s}s"
+                        )
+
+                        time.sleep(target_delay_s)
+                        fn(*args, **kwargs)
+
+                    return super().submit(__slowed_fn)
+
+            executor = DelayedThreadPoolExecutor(max_workers=max_workers)
+            ray._private.worker.global_worker.core_worker.reset_event_loop_executor(
+                executor
+            )
+
+
+@ray.remote(num_cpus=1)
+class CallerActor:
+    def __init__(
+        self,
+        downstream: ActorHandle,
+    ):
+        self._h = downstream
+
+    async def run(self):
+        print(">>> [Caller] Starting consuming stream")
+
+        async_obj_ref_gen = self._h.aio_stream.options(num_returns="streaming").remote()
+        async for ref in async_obj_ref_gen:
+            r = await ref
+            if r == 1:
+                print(">>> [Caller] Cancelling generator")
+                ray.cancel(async_obj_ref_gen, recursive=False)
+
+                # NOTE: This delay is crucial to let already scheduled task to report
+                #       generated item (report_streaming_generator_output) before we
+                #       will tear down this stream
+                delay_after_cancellation_s = 2
+
+                print(f">>> [Caller] **Sleeping** {delay_after_cancellation_s}s")
+                time.sleep(delay_after_cancellation_s)
+            else:
+                print(f">>> [Caller] Received {r}")
+
+        print(">>> [Caller] Completed consuming stream")
+
+
+@pytest.mark.parametrize("injected_executor_delay_s", [0, 2])
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_nodes": 2,
+            "num_cpus": 1,
+        }
+    ],
+    indirect=True,
+)
+def test_segfault_report_streaming_generator_output(
+    ray_start_cluster, injected_executor_delay_s: float
+):
+    """
+    This is a "smoke" test attempting to emulate condition, when using Ray's async
+    streaming generator, that leads to worker crashing with SIGSEGV.
+
+    For more details summarizing these conditions, please refer to
+    https://github.com/ray-project/ray/issues/43771#issuecomment-1982301654
+    """
+
+    caller = CallerActor.remote(
+        EndpointActor.remote(
+            injected_executor_delay_s=injected_executor_delay_s,
+            tokens_per_request=100,
+        ),
+    )
+
+    worker_state_before = [(a.worker_id, a.exit_type) for a in list_workers()]
+    print(">>> Workers state before: ", worker_state_before)
+
+    try:
+        ray.get(caller.run.remote())
+    except Exception as exc:
+        # There is a small chance that the task cancellation signal will arrive
+        # late at the executor, after the task has already finished. In that
+        # case, the task will complete normally, with no exception thrown.
+        # Thus, we wrap ray.get in a try-catch instead of asserting an
+        # exception.
+        assert isinstance(exc, RayTaskError)
+        assert isinstance(exc.cause, TaskCancelledError)
+
+    worker_state_after = [(a.worker_id, a.exit_type) for a in list_workers()]
+    print(">>> Workers state after: ", worker_state_after)
+
+    worker_ids, worker_exit_types = zip(*worker_state_after)
+    # Make sure no workers crashed
+    assert (
+        "SYSTEM_ERROR" not in worker_exit_types
+    ), f"Unexpected crashed worker(s) in {worker_ids}"
+
+
+if __name__ == "__main__":
+
+    sys.exit(pytest.main(["-sv", __file__]))