wehub-resource-sync
772 lines
23 KiB
Python
772 lines
23 KiB
Python
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__]))
|