chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,771 @@
|
||||
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__]))
|
||||
Reference in New Issue
Block a user