ray-project--ray/python/ray/tests/test_streaming_generator.py

import asyncio
import gc
import sys
import threading
import time
from unittest.mock import Mock, patch

import numpy as np
import pytest

import ray
from ray._common.test_utils import wait_for_condition
from ray._raylet import ObjectRefGenerator, ObjectRefStreamEndOfStreamError
from ray.cloudpickle import dumps
from ray.exceptions import WorkerCrashedError
from ray.experimental.state.api import list_objects


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

    def check_connected(self):
        return True


@pytest.fixture
def mocked_worker():
    mocked_core_worker = Mock()
    mocked_core_worker.try_read_next_object_ref_stream.return_value = None
    mocked_core_worker.try_read_next_object_ref_stream_n.return_value = None
    mocked_core_worker.async_delete_object_ref_stream.return_value = None
    mocked_core_worker.create_object_ref_stream.return_value = None
    mocked_core_worker.peek_object_ref_stream.return_value = [], []
    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:
        with patch("ray.get") as mocked_ray_get:
            c = mocked_worker.core_worker
            generator_ref = ray.ObjectRef.from_random()
            generator = ObjectRefGenerator(generator_ref, mocked_worker)

            # Make sure we cannot serialize the generator.
            with pytest.raises(TypeError):
                dumps(generator)

            # Test when there's no new ref, it returns a nil.
            new_ref = ray.ObjectRef.from_random()
            c.peek_object_ref_stream.return_value = (new_ref, False)
            mocked_ray_wait.return_value = [], [new_ref]
            ref = generator._next_sync(timeout_s=0)
            assert ref.is_nil()

            # When the new ref is available, next should return it.
            # When peek_object_ref_stream returns is_ready = True,
            # it shouldn't call ray.wait.
            new_ref = ray.ObjectRef.from_random()
            c.peek_object_ref_stream.return_value = (new_ref, True)
            c.try_read_next_object_ref_stream.return_value = new_ref
            ref = generator._next_sync(timeout_s=0)
            assert new_ref == ref

            # When the new ref is available, next should return it.
            # When peek_object_ref_stream returns is_ready = False,
            # it should wait until ray.wait returns.
            for _ in range(3):
                new_ref = ray.ObjectRef.from_random()
                c.peek_object_ref_stream.return_value = (new_ref, False)
                mocked_ray_wait.return_value = [new_ref], []
                c.try_read_next_object_ref_stream.return_value = new_ref
                ref = generator._next_sync(timeout_s=0)
                assert new_ref == ref

            # When try_read_next_object_ref_stream raises a
            # ObjectRefStreamEndOfStreamError, it should raise a stop iteration.
            new_ref = ray.ObjectRef.from_random()
            c.peek_object_ref_stream.return_value = (new_ref, True)
            c.try_read_next_object_ref_stream.side_effect = (
                ObjectRefStreamEndOfStreamError("")
            )  # noqa
            mocked_ray_get.return_value = None

            with pytest.raises(StopIteration):
                generator._next_sync(timeout_s=0)


def test_streaming_object_ref_generator_consume_bulk_unit(mocked_worker):
    c = mocked_worker.core_worker
    generator_ref = ray.ObjectRef.from_random()
    generator = ObjectRefGenerator(generator_ref, mocked_worker)

    generator._consume_next_ref_n(2)

    c.try_read_next_object_ref_stream_n.assert_called_once_with(generator_ref, 2)


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 = ObjectRefGenerator(generator_ref, mocked_worker)

            # Simulate the worker failure happens.
            next_ref = ray.ObjectRef.from_random()
            c.peek_object_ref_stream.return_value = (next_ref, False)
            mocked_ray_wait.return_value = [next_ref], []
            mocked_ray_get.side_effect = WorkerCrashedError()

            c.try_read_next_object_ref_stream.side_effect = (
                ObjectRefStreamEndOfStreamError("")
            )  # noqa
            ref = generator._next_sync(timeout_s=0)
            # If the generator task fails by a systsem error,
            # meaning the ref will raise an exception
            # it should be returned.
            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_sync(timeout_s=0)


def test_generator_basic(shutdown_only):
    ray.init(num_cpus=1)

    """Basic cases"""
    print("Test basic case")

    @ray.remote
    def f():
        for i in range(5):
            yield i

    gen = f.remote()
    i = 0
    for ref in gen:
        print(ray.get(ref))
        assert i == ray.get(ref)
        del ref
        i += 1

    """Exceptions"""
    print("Test exceptions")

    @ray.remote
    def f():
        for i in range(5):
            if i == 2:
                raise ValueError
            yield i

    gen = f.remote()
    print(ray.get(next(gen)))
    print(ray.get(next(gen)))
    with pytest.raises(ray.exceptions.RayTaskError) as e:
        print(ray.get(next(gen)))
    with pytest.raises(StopIteration):
        ray.get(next(gen))
    with pytest.raises(StopIteration):
        ray.get(next(gen))

    """Generator Task failure"""
    print("Test task failures")

    @ray.remote
    class A:
        def getpid(self):
            import os

            return os.getpid()

        def f(self):
            for i in range(5):
                time.sleep(1)
                yield i

    a = A.remote()
    gen = a.f.remote()
    i = 0
    for ref in gen:
        if i == 2:
            ray.kill(a)
        if i == 3:
            with pytest.raises(ray.exceptions.RayActorError) as e:
                ray.get(ref)
            assert "The actor is dead because it was killed by `ray.kill`" in str(
                e.value
            )
            break
        assert i == ray.get(ref)
        del ref
        i += 1
    for _ in range(10):
        with pytest.raises(StopIteration):
            next(gen)

    """Retry exceptions"""
    print("Test retry exceptions")

    @ray.remote
    class Actor:
        def __init__(self):
            self.should_kill = True

        def should_kill(self):
            return self.should_kill

        async def set(self, wait_s):
            await asyncio.sleep(wait_s)
            self.should_kill = False

    @ray.remote(retry_exceptions=[ValueError], max_retries=10)
    def f(a):
        for i in range(5):
            should_kill = ray.get(a.should_kill.remote())
            if i == 3 and should_kill:
                raise ValueError
            yield i

    a = Actor.remote()
    gen = f.remote(a)
    assert ray.get(next(gen)) == 0
    assert ray.get(next(gen)) == 1
    assert ray.get(next(gen)) == 2
    a.set.remote(3)
    assert ray.get(next(gen)) == 3
    assert ray.get(next(gen)) == 4
    with pytest.raises(StopIteration):
        ray.get(next(gen))

    """Cancel"""
    print("Test cancel")

    @ray.remote
    def f():
        for i in range(5):
            time.sleep(5)
            yield i

    gen = f.remote()
    assert ray.get(next(gen)) == 0
    ray.cancel(gen)
    with pytest.raises(ray.exceptions.RayTaskError) as e:
        assert ray.get(next(gen)) == 1
    assert "was cancelled" in str(e.value)
    with pytest.raises(StopIteration):
        next(gen)


def test_streaming_generator_bad_exception_not_failing(shutdown_only, capsys):
    """This test verifies when a return value cannot be stored
        e.g., because it holds a lock) if it handles failures gracefully.

    Previously, when it happens, there was a check failure. This verifies
    the check failure doesn't happen anymore.
    """
    ray.init()

    class UnserializableException(Exception):
        def __init__(self):
            self.lock = threading.Lock()

    @ray.remote
    def f():
        raise UnserializableException
        yield 1  # noqa

    for ref in f.remote():
        with pytest.raises(ray.exceptions.RayTaskError):
            ray.get(ref)
    captured = capsys.readouterr()
    lines = captured.err.strip().split("\n")

    # Verify check failure doesn't happen because we handle the error
    # properly.
    for line in lines:
        assert "Check failed:" not in line


@pytest.mark.parametrize("crash_type", ["exception", "worker_crash"])
def test_generator_streaming_no_leak_upon_failures(
    monkeypatch, shutdown_only, crash_type
):
    with monkeypatch.context() as m:
        m.setenv(
            "RAY_testing_asio_delay_us",
            "CoreWorkerService.grpc_server.ReportGeneratorItemReturns=100000:1000000",
        )
        ray.init(num_cpus=1)

        @ray.remote
        def g():
            try:
                gen = f.remote()
                for ref in gen:
                    print(ref)
                    ray.get(ref)
            except Exception:
                print("exception!")
                del ref

            del gen
            gc.collect()

            # Only the ref g is alive.
            def verify():
                print(list_objects())
                return len(list_objects()) == 1

            wait_for_condition(verify)
            return True

        @ray.remote
        def f():
            for i in range(10):
                time.sleep(0.2)
                if i == 4:
                    if crash_type == "exception":
                        raise ValueError
                    else:
                        sys.exit(9)
                yield 2

        for _ in range(5):
            ray.get(g.remote())


@pytest.mark.parametrize("use_actors", [False, True])
@pytest.mark.parametrize("store_in_plasma", [False, True])
def test_generator_streaming(shutdown_only, use_actors, store_in_plasma):
    """Verify the generator is working in a streaming fashion."""
    ray.init()
    remote_generator_fn = None
    if use_actors:

        @ray.remote
        class Generator:
            def __init__(self):
                pass

            def generator(self, num_returns, store_in_plasma):
                for i in range(num_returns):
                    if store_in_plasma:
                        yield np.ones(1_000_000, dtype=np.int8) * i
                    else:
                        yield [i]

        g = Generator.remote()
        remote_generator_fn = g.generator
    else:

        @ray.remote(max_retries=0)
        def generator(num_returns, store_in_plasma):
            for i in range(num_returns):
                if store_in_plasma:
                    yield np.ones(1_000_000, dtype=np.int8) * i
                else:
                    yield [i]

        remote_generator_fn = generator

    """Verify num_returns="streaming" is streaming"""
    gen = remote_generator_fn.remote(3, store_in_plasma)
    i = 0
    for ref in gen:
        id = ref.hex()
        if store_in_plasma:
            expected = np.ones(1_000_000, dtype=np.int8) * i
            assert np.array_equal(ray.get(ref), expected)
        else:
            expected = [i]
            assert ray.get(ref) == expected

        del ref

        wait_for_condition(
            lambda id=id: len(list_objects(filters=[("object_id", "=", id)])) == 0
        )
        i += 1


def test_generator_dist_chain(ray_start_cluster):
    """E2E test to verify chain of generator works properly."""
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=0, object_store_memory=1 * 1024 * 1024 * 1024)
    ray.init()
    cluster.add_node(num_cpus=1)
    cluster.add_node(num_cpus=1)
    cluster.add_node(num_cpus=1)
    cluster.add_node(num_cpus=1)

    @ray.remote
    class ChainActor:
        def __init__(self, child=None):
            self.child = child

        def get_data(self):
            if not self.child:
                for _ in range(10):
                    time.sleep(0.1)
                    yield np.ones(5 * 1024 * 1024)
            else:
                for data in self.child.get_data.remote():
                    yield ray.get(data)

    chain_actor = ChainActor.remote()
    chain_actor_2 = ChainActor.remote(chain_actor)
    chain_actor_3 = ChainActor.remote(chain_actor_2)
    chain_actor_4 = ChainActor.remote(chain_actor_3)

    for ref in chain_actor_4.get_data.remote():
        assert np.array_equal(np.ones(5 * 1024 * 1024), ray.get(ref))
        print("getting the next data")
        del ref


def test_generator_slow_pinning_requests(monkeypatch, shutdown_only):
    """
    Verify when the Object pinning request from the raylet
    is reported slowly, there's no refernece leak.
    """
    with monkeypatch.context() as m:
        m.setenv(
            "RAY_testing_asio_delay_us",
            "CoreWorkerService.grpc_server.PubsubLongPolling=1000000:1000000",
        )

        @ray.remote
        def f():
            yield np.ones(5 * 1024 * 1024)

        for ref in f.remote():
            del ref

        print(list_objects())


@pytest.mark.parametrize("store_in_plasma", [False, True])
def test_actor_streaming_generator(shutdown_only, store_in_plasma):
    """Test actor/async actor with sync/async generator interfaces."""
    ray.init()

    @ray.remote
    class Actor:
        def f(self, ref):
            for i in range(3):
                yield i

        async def async_f(self, ref):
            for i in range(3):
                await asyncio.sleep(0.1)
                yield i

        def g(self):
            return 3

    a = Actor.remote()
    if store_in_plasma:
        arr = np.random.rand(5 * 1024 * 1024)
    else:
        arr = 3

    def verify_sync_task_executor():
        generator = a.f.remote(ray.put(arr))
        # Verify it works with next.
        assert isinstance(generator, ObjectRefGenerator)
        assert ray.get(next(generator)) == 0
        assert ray.get(next(generator)) == 1
        assert ray.get(next(generator)) == 2
        with pytest.raises(StopIteration):
            ray.get(next(generator))

        # Verify it works with for.
        generator = a.f.remote(ray.put(3))
        for index, ref in enumerate(generator):
            assert index == ray.get(ref)

    def verify_async_task_executor():
        # Verify it works with next.
        generator = a.async_f.remote(ray.put(arr))
        assert isinstance(generator, ObjectRefGenerator)
        assert ray.get(next(generator)) == 0
        assert ray.get(next(generator)) == 1
        assert ray.get(next(generator)) == 2

        # Verify it works with for.
        generator = a.f.remote(ray.put(3))
        for index, ref in enumerate(generator):
            assert index == ray.get(ref)

    async def verify_sync_task_async_generator():
        # Verify anext
        async_generator = a.f.remote(ray.put(arr))
        assert isinstance(async_generator, ObjectRefGenerator)
        for expected in range(3):
            ref = await async_generator.__anext__()
            assert await ref == expected
        with pytest.raises(StopAsyncIteration):
            await async_generator.__anext__()

        # Verify async for.
        async_generator = a.f.remote(ray.put(arr))
        expected = 0
        async for ref in async_generator:
            value = await ref
            assert expected == value
            expected += 1

    async def verify_async_task_async_generator():
        async_generator = a.async_f.remote(ray.put(arr))
        assert isinstance(async_generator, ObjectRefGenerator)
        for expected in range(3):
            ref = await async_generator.__anext__()
            assert await ref == expected
        with pytest.raises(StopAsyncIteration):
            await async_generator.__anext__()

        # Verify async for.
        async_generator = a.async_f.remote(ray.put(arr))
        expected = 0
        async for ref in async_generator:
            value = await ref
            assert expected == value
            expected += 1

    verify_sync_task_executor()
    verify_async_task_executor()
    asyncio.run(verify_sync_task_async_generator())
    asyncio.run(verify_async_task_async_generator())


def test_streaming_generator_num_objects_per_yield(shutdown_only):
    ray.init()

    @ray.remote(_num_objects_per_yield=2)
    def generator():
        for i in range(3):
            stats = yield i, f"metadata-{i}"
            assert stats is None or stats.object_creation_dur_s >= 0

    gen = generator.remote()
    for i in range(3):
        assert ray.get(next(gen)) == i
        assert ray.get(next(gen)) == f"metadata-{i}"

    with pytest.raises(StopIteration):
        next(gen)

    @ray.remote
    def per_call():
        yield 1, 2

    with pytest.raises(ValueError, match="_num_objects_per_yield"):
        per_call.options(_num_objects_per_yield=2).remote()


def test_actor_streaming_generator_num_objects_per_yield(shutdown_only):
    ray.init()

    @ray.remote
    class Actor:
        @ray.method(_num_objects_per_yield=2)
        def decorated(self):
            yield "block", "metadata"

        def per_call(self):
            yield 1, 2

    actor = Actor.remote()

    gen = actor.decorated.remote()
    assert ray.get(next(gen)) == "block"
    assert ray.get(next(gen)) == "metadata"
    with pytest.raises(StopIteration):
        next(gen)

    with pytest.raises(ValueError, match="_num_objects_per_yield"):
        actor.per_call.options(_num_objects_per_yield=2).remote()


def test_streaming_generator_num_objects_per_yield_invalid_yield(shutdown_only):
    ray.init()

    @ray.remote(_num_objects_per_yield=2)
    def generator():
        yield (1,)

    gen = generator.remote()
    with pytest.raises(ValueError, match="_num_objects_per_yield=2"):
        ray.get(next(gen))

    with pytest.raises(StopIteration):
        next(gen)


def test_streaming_generator_num_objects_per_yield_serialization_failure(shutdown_only):
    ray.init()

    @ray.remote(_num_objects_per_yield=2)
    def generator():
        yield threading.Lock(), 1

    gen = generator.remote()
    with pytest.raises(ray.exceptions.RayTaskError):
        ray.get(next(gen))
    with pytest.raises(ray.exceptions.RayTaskError):
        ray.get(next(gen))

    with pytest.raises(StopIteration):
        next(gen)


def test_streaming_generator_num_objects_per_yield_partial_store_failure(
    shutdown_only,
):
    ray.init()

    @ray.remote(_num_objects_per_yield=2)
    def generator():
        # If a later object fails after an earlier object has been stored, the
        # caller should still receive a ref for every object in the grouped yield.
        yield 1, threading.Lock()

    gen = generator.remote()
    assert ray.get(next(gen)) == 1
    with pytest.raises(ray.exceptions.RayTaskError):
        ray.get(next(gen))

    with pytest.raises(StopIteration):
        next(gen)


def test_streaming_generator_num_objects_per_yield_failure_not_retried(
    shutdown_only,
):
    ray.init()

    @ray.remote(_num_objects_per_yield=2, retry_exceptions=True, max_retries=1)
    def generator():
        # Once grouped-yield IDs are allocated, the whole group must be
        # reported so those temporary refs are cleared instead of retrying.
        yield 1, threading.Lock()

    gen = generator.remote()
    assert ray.get(next(gen)) == 1
    with pytest.raises(ray.exceptions.RayTaskError):
        ray.get(next(gen))

    with pytest.raises(StopIteration):
        next(gen)


def test_streaming_generator_exception(shutdown_only):
    # Verify the exceptions are correctly raised.
    # Also verify the followup next will raise StopIteration.
    ray.init()

    @ray.remote
    class Actor:
        def f(self):
            raise ValueError
            yield 1  # noqa

        async def async_f(self):
            raise ValueError
            yield 1  # noqa

    a = Actor.remote()
    g = a.f.remote()
    with pytest.raises(ValueError):
        ray.get(next(g))

    with pytest.raises(StopIteration):
        ray.get(next(g))

    with pytest.raises(StopIteration):
        ray.get(next(g))

    g = a.async_f.remote()
    with pytest.raises(ValueError):
        ray.get(next(g))

    with pytest.raises(StopIteration):
        ray.get(next(g))

    with pytest.raises(StopIteration):
        ray.get(next(g))


def test_next_sync_timeout_when_generator_ref_unavailable(
    monkeypatch, ray_start_cluster
):
    """_next_sync(timeout_s) must not block in its end-of-stream handling.

    After all yielded refs are consumed, ``_next_sync`` calls
    ``ray.get(generator_ref)`` to distinguish a normal end of the stream
    from a task failure. If that object is unavailable — e.g. it lived in
    the plasma store of a node that died — the get must be bounded by the
    caller's timeout (reporting "not ready yet" with a nil ref) instead of
    blocking the caller until the object is reconstructed. A blocked caller
    can deadlock: reconstruction needs a CPU, and the blocked caller may be
    what releases one (see ray-project/ray#63701).
    """
    # Force the generator task's return object into plasma (instead of
    # being inlined with the owner) so that it is lost when its node dies.
    monkeypatch.setenv("RAY_max_direct_call_object_size", "0")
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=0)
    ray.init(address=cluster.address)
    worker_node = cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()

    @ray.remote(num_returns="streaming")
    def gen():
        yield 1

    g = gen.remote()
    # Consume the single yielded ref so that the stream is exhausted.
    first = g._next_sync(timeout_s=30)
    assert not first.is_nil()

    # Lose the generator's return object together with its node.
    cluster.remove_node(worker_node)

    # End-of-stream handling must honor the timeout: a nil ref, not a block.
    start = time.monotonic()
    assert g._next_sync(timeout_s=0).is_nil()
    assert time.monotonic() - start < 10

    # The async counterpart must honor the timeout the same way.
    start = time.monotonic()
    assert asyncio.run(g._next_async(timeout_s=0)).is_nil()
    assert time.monotonic() - start < 10

    # Once a node is back, lineage reconstruction restores the return object
    # and the stream terminates normally.
    cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()
    deadline = time.monotonic() + 60
    while time.monotonic() < deadline:
        try:
            ref = g._next_sync(timeout_s=1)
        except StopIteration:
            break
        assert ref.is_nil()
    else:
        pytest.fail("Generator did not finish after the node was restored.")


if __name__ == "__main__":
    sys.exit(pytest.main(["-sv", __file__]))