chore: import upstream snapshot with attribution

python/ray/tests/test_streaming_generator_4.py

@@ -0,0 +1,412 @@
+import asyncio
+import gc
+import json
+import os
+import random
+import signal
+import sys
+import time
+from typing import Optional
+
+import numpy as np
+import pytest
+from pydantic import BaseModel
+
+import ray
+from ray._common.test_utils import SignalActor
+
+RECONSTRUCTION_CONFIG = {
+    "health_check_failure_threshold": 10,
+    "health_check_period_ms": 100,
+    "health_check_timeout_ms": 100,
+    "health_check_initial_delay_ms": 0,
+    "max_direct_call_object_size": 100,
+    "task_retry_delay_ms": 100,
+    "object_timeout_milliseconds": 200,
+    "fetch_warn_timeout_milliseconds": 1000,
+}
+
+
+def assert_no_leak():
+    gc.collect()
+    core_worker = ray._private.worker.global_worker.core_worker
+    ref_counts = core_worker.get_all_reference_counts()
+    print(ref_counts)
+    for rc in ref_counts.values():
+        assert rc["local"] == 0
+        assert rc["submitted"] == 0
+    assert core_worker.get_memory_store_size() == 0
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="SIGKILL is not available on Windows"
+)
+def test_caller_death(monkeypatch, shutdown_only):
+    """
+    Test the case where caller of a streaming generator actor task dies
+    while the streaming generator task is executing. The streaming
+    generator task should still finish and won't block other actor tasks.
+    This means that `ReportGeneratorItemReturns` RPC should fail and it shouldn't
+    be retried indefinitely.
+    """
+    ray.init()
+
+    @ray.remote
+    class Callee:
+        def gen(self, caller_pid):
+            os.kill(caller_pid, signal.SIGKILL)
+            yield [1] * 1024 * 1024
+
+        def ping(self):
+            pass
+
+    @ray.remote
+    def caller(callee):
+        ray.get(callee.gen.remote(os.getpid()))
+
+    callee = Callee.remote()
+    o = caller.remote(callee)
+    ray.wait([o])
+    # Make sure gen will finish and ping can run.
+    ray.get(callee.ping.remote())
+
+
+def test_intermediate_generator_object_recovery_while_generator_running(
+    ray_start_cluster,
+):
+    """
+    1. Streaming producer starts on worker1.
+    2. consumer consumes value 1 from producer on worker2 and finishes.
+    3. Run an extra consumer on worker2 to track when reconstruction is triggered.
+    4. Add worker3.
+    5. worker2 dies.
+    6. Try to get consumer output.
+    7. Therefore Ray tries to reconstruct value 1 from producer.
+    8. Get the reconstructed extra_consumer_ref (assures 7 happened).
+    9. Streaming producer should be cancelled and resubmitted.
+    10. Retry for consumer should complete.
+    """
+
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=0)  # head
+    ray.init(address=cluster.address)
+    cluster.add_node(num_cpus=1, resources={"producer": 1})  # worker1
+    worker2 = cluster.add_node(num_cpus=1, resources={"consumer": 1})
+
+    @ray.remote(num_cpus=1, resources={"producer": 1})
+    def producer():
+        for _ in range(3):
+            yield np.zeros(10 * 1024 * 1024, dtype=np.uint8)
+
+    @ray.remote(num_cpus=1, resources={"consumer": 1})
+    def consumer(np_arr):
+        return np_arr
+
+    streaming_ref = producer.options(_generator_backpressure_num_objects=1).remote()
+    consumer_ref = consumer.remote(next(streaming_ref))
+    extra_consumer_ref = consumer.remote(np.zeros(10 * 1024 * 1024, dtype=np.uint8))
+
+    ray.wait([consumer_ref, extra_consumer_ref], num_returns=2, fetch_local=False)
+
+    cluster.add_node(num_cpus=1, resources={"consumer": 1})  # worker3
+    cluster.remove_node(worker2, allow_graceful=True)
+
+    # Make sure reconstruction was triggered.
+    assert ray.get(extra_consumer_ref).size == (10 * 1024 * 1024)
+    # Allow first streaming generator attempt to finish
+    ray.get([next(streaming_ref), next(streaming_ref)])
+
+    assert ray.get(consumer_ref).size == (10 * 1024 * 1024)
+
+
+def test_actor_intermediate_generator_object_recovery_while_generator_running(
+    ray_start_cluster,
+):
+    """
+    1. Producer actor and its generator producer task start on worker1.
+    2. consumer consumes value 1 from producer on worker2 and finishes.
+    3. Run an extra consumer on worker2 to track when reconstruction is triggered.
+    4. Add worker3.
+    5. worker2 dies.
+    6. Ray tries to reconstruct value 1 from producer.
+    7. Get the reconstructed extra_consumer_ref (assures 6 happened).
+    8. Ray tries and fails to cancel the producer task.
+    9. Get the next two values to relieve backpressure and allow producer to finish.
+    10. Ray resubmits the producer generator task.
+    11. Retry for consumer should complete.
+    """
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=0)  # head
+    ray.init(address=cluster.address)
+    cluster.add_node(num_cpus=1, resources={"producer": 1})  # worker 1
+    worker2 = cluster.add_node(num_cpus=1, resources={"consumer": 1})
+
+    @ray.remote(num_cpus=1, resources={"producer": 1}, max_task_retries=-1)
+    class Producer:
+        def producer(self):
+            for _ in range(3):
+                yield np.zeros(10 * 1024 * 1024, dtype=np.uint8)
+
+    @ray.remote(num_cpus=1, resources={"consumer": 1})
+    def consumer(np_arr):
+        return np_arr
+
+    producer_actor = Producer.remote()
+    streaming_ref = producer_actor.producer.options(
+        _generator_backpressure_num_objects=1
+    ).remote()
+    consumer_ref = consumer.remote(next(streaming_ref))
+    extra_consumer_ref = consumer.remote(np.zeros(10 * 1024 * 1024, dtype=np.uint8))
+
+    ray.wait([consumer_ref, extra_consumer_ref], num_returns=2, fetch_local=False)
+
+    cluster.add_node(num_cpus=1, resources={"consumer": 1})  # worker 3
+    cluster.remove_node(worker2, allow_graceful=True)
+
+    # Make sure reconstruction was triggered.
+    ray.get(extra_consumer_ref)
+    # Allow first streaming generator attempt to finish
+    ray.get([next(streaming_ref), next(streaming_ref)])
+
+    assert ray.get(consumer_ref).size == (10 * 1024 * 1024)
+
+
+@pytest.mark.parametrize("backpressure", [False, True])
+@pytest.mark.parametrize("delay_latency", [0.1, 1])
+@pytest.mark.parametrize("threshold", [1, 3])
+def test_many_tasks_lineage_reconstruction_mini_stress_test(
+    monkeypatch, ray_start_cluster, backpressure, delay_latency, threshold
+):
+    """Test a workload that spawns many tasks and relies on lineage reconstruction."""
+    if not backpressure:
+        if delay_latency == 0.1 and threshold == 1:
+            return
+        elif delay_latency == 1:
+            return
+
+    with monkeypatch.context() as m:
+        m.setenv(
+            "RAY_testing_asio_delay_us",
+            "CoreWorkerService.grpc_server.ReportGeneratorItemReturns=10000:1000000",
+        )
+        m.setenv(
+            "RAY_testing_rpc_failure",
+            json.dumps(
+                {
+                    "CoreWorkerService.grpc_client.ReportGeneratorItemReturns": {
+                        "num_failures": 5,
+                        "req_failure_prob": 25,
+                        "resp_failure_prob": 25,
+                        "in_flight_failure_prob": 25,
+                    }
+                }
+            ),
+        )
+        cluster = ray_start_cluster
+        cluster.add_node(
+            num_cpus=1,
+            resources={"head": 1},
+            _system_config=RECONSTRUCTION_CONFIG,
+            enable_object_reconstruction=True,
+        )
+        ray.init(address=cluster.address)
+
+        if backpressure:
+            threshold = 1
+        else:
+            threshold = -1
+
+        @ray.remote(
+            max_retries=-1,
+            _generator_backpressure_num_objects=threshold,
+        )
+        def dynamic_generator(num_returns):
+            for i in range(num_returns):
+                time.sleep(0.1)
+                yield np.ones(1_000_000, dtype=np.int8) * i
+
+        @ray.remote(num_cpus=0, resources={"head": 1})
+        def driver():
+            unready = [dynamic_generator.remote(10) for _ in range(5)]
+            ready = []
+            while unready:
+                for a in unready:
+                    print(a._generator_ref)
+                ready, unready = ray.wait(
+                    unready, num_returns=len(unready), timeout=0.1
+                )
+                for r in ready:
+                    try:
+                        ref = next(r)
+                        print(ref)
+                        ray.get(ref)
+                    except StopIteration:
+                        pass
+                    else:
+                        unready.append(r)
+            return None
+
+        ref = driver.remote()
+
+        nodes = []
+        for _ in range(4):
+            nodes.append(cluster.add_node(num_cpus=1, object_store_memory=10**8))
+        cluster.wait_for_nodes()
+
+        for _ in range(10):
+            time.sleep(0.1)
+            node_to_kill = random.choices(nodes)[0]
+            nodes.remove(node_to_kill)
+            cluster.remove_node(node_to_kill, allow_graceful=False)
+            nodes.append(cluster.add_node(num_cpus=1, object_store_memory=10**8))
+
+        ray.get(ref)
+        del ref
+
+        assert_no_leak()
+
+
+def test_local_gc_not_hang(shutdown_only, monkeypatch):
+    """Verify the generator doesn't deadlock when a local GC is triggered."""
+    with monkeypatch.context() as m:
+        m.setenv("RAY_local_gc_interval_s", 1)
+
+        ray.init()
+
+        @ray.remote(_generator_backpressure_num_objects=1)
+        def f():
+            for _ in range(5):
+                yield 1
+
+        gen = f.remote()
+        time.sleep(5)
+
+        # It should not hang.
+        for ref in gen:
+            ray.get(gen)
+
+
+def test_sync_async_mix_regression_test(shutdown_only):
+    """Verify when sync and async tasks are mixed up
+    it doesn't raise a segfault
+
+    https://github.com/ray-project/ray/issues/41346
+    """
+
+    class PayloadPydantic(BaseModel):
+        class Error(BaseModel):
+            msg: str
+            code: int
+            type: str
+
+        text: Optional[str] = None
+        ts: Optional[float] = None
+        reason: Optional[str] = None
+        error: Optional[Error] = None
+
+    ray.init()
+
+    @ray.remote
+    class B:
+        def __init__(self, a):
+            self.a = a
+
+        async def stream(self):
+            async for ref in self.a.stream.remote(1):
+                print("stream")
+                await ref
+
+        async def start(self):
+            await asyncio.gather(*[self.stream() for _ in range(2)])
+
+    @ray.remote
+    class A:
+        def stream(self, i):
+            payload = PayloadPydantic(
+                text="Test output",
+                ts=time.time(),
+                reason="Success!",
+            )
+
+            for _ in range(10):
+                yield payload
+
+        async def aio_stream(self):
+            for _ in range(10):
+                yield 1
+
+    a = A.remote()
+    b = B.remote(a)
+    ray.get(b.start.remote())
+
+
+@pytest.mark.parametrize("use_asyncio", [False, True])
+def test_cancel(shutdown_only, use_asyncio):
+    """Test concurrent task cancellation with generator task.
+
+    Once the caller receives an ack that the executor has cancelled the task
+    execution, the caller should receive a TaskCancelledError for the next
+    ObjectRef that it tries to read from the generator. This should happen even
+    if the caller has already received values for the next object indices in
+    the stream. Also, we should not apply the usual logic that reorders
+    out-of-order reports if the task was cancelled; waiting for the
+    intermediate indices to appear would hang the caller."""
+
+    @ray.remote
+    class Actor:
+        def ready(self):
+            return
+
+        def stream(self, signal):
+            cancelled_ref = signal.wait.remote()
+
+            i = 0
+            done_at = time.time() + 1
+            while time.time() < done_at:
+                yield i
+                i += 1
+
+                ready, _ = ray.wait([cancelled_ref], timeout=0)
+                if not ready:
+                    # Continue executing for one second after the driver
+                    # cancels. This is to make sure that we receive the cancel
+                    # signal while the task is still running.
+                    done_at = time.time() + 1
+
+        async def async_stream(self, signal):
+            cancelled_ref = signal.wait.remote()
+
+            i = 0
+            done_at = time.time() + 1
+            while time.time() < done_at:
+                yield i
+                i += 1
+
+                ready, _ = ray.wait([cancelled_ref], timeout=0)
+                if not ready:
+                    # Continue executing for one second after the driver
+                    # cancels. This is to make sure that we receive the cancel
+                    # signal while the task is still running.
+                    done_at = time.time() + 1
+
+    signal = SignalActor.remote()
+    a = Actor.remote()
+    ray.get(a.ready.remote())
+    if use_asyncio:
+        gen = a.async_stream.remote(signal)
+    else:
+        gen = a.stream.remote(signal)
+
+    try:
+        for i, ref in enumerate(gen):
+            assert i == ray.get(ref)
+            print(i)
+            if i == 0:
+                ray.cancel(gen)
+                signal.send.remote()
+    except ray.exceptions.TaskCancelledError:
+        pass
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))