wehub-resource-sync
817 lines
26 KiB
Python
817 lines
26 KiB
Python
import gc
|
|
import sys
|
|
import time
|
|
from unittest.mock import Mock
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import (
|
|
wait_for_condition,
|
|
)
|
|
from ray._private.client_mode_hook import enable_client_mode
|
|
from ray.tests.conftest import call_ray_start_context
|
|
from ray.util.client.ray_client_helpers import (
|
|
ray_start_client_server_for_address,
|
|
)
|
|
|
|
|
|
def assert_no_leak():
|
|
def check():
|
|
gc.collect()
|
|
core_worker = ray._private.worker.global_worker.core_worker
|
|
ref_counts = core_worker.get_all_reference_counts()
|
|
for k, rc in ref_counts.items():
|
|
if rc["local"] != 0:
|
|
return False
|
|
if rc["submitted"] != 0:
|
|
return False
|
|
return True
|
|
|
|
wait_for_condition(check)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "linux2",
|
|
reason="This test requires Linux.",
|
|
)
|
|
# This test can spill many GiB to disk (the normal-return task may not OOM and
|
|
# instead materializes all returns), so it needs a longer timeout.
|
|
@pytest.mark.timeout(600)
|
|
def test_generator_oom(ray_start_regular_shared):
|
|
num_returns = 100
|
|
|
|
@ray.remote(max_retries=0)
|
|
def large_values(num_returns):
|
|
return [
|
|
np.random.randint(
|
|
np.iinfo(np.int8).max, size=(100_000_000, 1), dtype=np.int8
|
|
)
|
|
for _ in range(num_returns)
|
|
]
|
|
|
|
@ray.remote(max_retries=0)
|
|
def large_values_generator(num_returns):
|
|
for _ in range(num_returns):
|
|
yield np.random.randint(
|
|
np.iinfo(np.int8).max, size=(100_000_000, 1), dtype=np.int8
|
|
)
|
|
|
|
try:
|
|
# Worker may OOM using normal returns.
|
|
ray.get(large_values.options(num_returns=num_returns).remote(num_returns)[0])
|
|
except ray.exceptions.WorkerCrashedError:
|
|
pass
|
|
|
|
# Using a generator will allow the worker to finish.
|
|
ray.get(
|
|
large_values_generator.options(num_returns=num_returns).remote(num_returns)[0]
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("use_actors", [False, True])
|
|
@pytest.mark.parametrize("store_in_plasma", [False, True])
|
|
def test_generator_returns(ray_start_regular_shared, use_actors, store_in_plasma):
|
|
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
|
|
|
|
# Check cases when num_returns does not match the number of values returned
|
|
# by the generator.
|
|
num_returns = 3
|
|
|
|
try:
|
|
ray.get(
|
|
remote_generator_fn.options(num_returns=num_returns).remote(
|
|
num_returns - 1, store_in_plasma
|
|
)
|
|
)
|
|
assert False
|
|
except ray.exceptions.RayTaskError as e:
|
|
assert isinstance(e.as_instanceof_cause(), ValueError)
|
|
|
|
# TODO(swang): When generators return more values than expected, we log an
|
|
# error but the exception is not thrown to the application.
|
|
# https://github.com/ray-project/ray/issues/28689.
|
|
ray.get(
|
|
remote_generator_fn.options(num_returns=num_returns).remote(
|
|
num_returns + 1, store_in_plasma
|
|
)
|
|
)
|
|
|
|
# Check return values.
|
|
[
|
|
x[0]
|
|
for x in ray.get(
|
|
remote_generator_fn.options(num_returns=num_returns).remote(
|
|
num_returns, store_in_plasma
|
|
)
|
|
)
|
|
] == list(range(num_returns))
|
|
# Works for num_returns=1 if generator returns a single value.
|
|
assert (
|
|
ray.get(remote_generator_fn.options(num_returns=1).remote(1, store_in_plasma))[
|
|
0
|
|
]
|
|
== 0
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("use_actors", [False, True])
|
|
@pytest.mark.parametrize("store_in_plasma", [False, True])
|
|
@pytest.mark.parametrize("num_returns_type", ["dynamic", None])
|
|
def test_generator_errors(
|
|
ray_start_regular_shared, use_actors, store_in_plasma, num_returns_type
|
|
):
|
|
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 - 2):
|
|
if store_in_plasma:
|
|
yield np.ones(1_000_000, dtype=np.int8) * i
|
|
else:
|
|
yield [i]
|
|
raise Exception("error")
|
|
|
|
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 - 2):
|
|
if store_in_plasma:
|
|
yield np.ones(1_000_000, dtype=np.int8) * i
|
|
else:
|
|
yield [i]
|
|
raise Exception("error")
|
|
|
|
remote_generator_fn = generator
|
|
|
|
ref1, ref2, ref3 = remote_generator_fn.options(num_returns=3).remote(
|
|
3, store_in_plasma
|
|
)
|
|
ray.get(ref1)
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
ray.get(ref2)
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
ray.get(ref3)
|
|
|
|
dynamic_ref = remote_generator_fn.options(num_returns=num_returns_type).remote(
|
|
3, store_in_plasma
|
|
)
|
|
ref1, ref2 = ray.get(dynamic_ref)
|
|
ray.get(ref1)
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
ray.get(ref2)
|
|
|
|
|
|
@pytest.mark.parametrize("store_in_plasma", [False, True])
|
|
@pytest.mark.parametrize("num_returns_type", ["dynamic", None])
|
|
def test_dynamic_generator_retry_exception(
|
|
ray_start_regular_shared, store_in_plasma, num_returns_type
|
|
):
|
|
class CustomException(Exception):
|
|
pass
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class ExecutionCounter:
|
|
def __init__(self):
|
|
self.count = 0
|
|
|
|
def inc(self):
|
|
self.count += 1
|
|
return self.count
|
|
|
|
def get_count(self):
|
|
return self.count
|
|
|
|
def reset(self):
|
|
self.count = 0
|
|
|
|
@ray.remote(max_retries=1)
|
|
def generator(num_returns, store_in_plasma, counter):
|
|
for i in range(num_returns):
|
|
if store_in_plasma:
|
|
yield np.ones(1_000_000, dtype=np.int8) * i
|
|
else:
|
|
yield [i]
|
|
|
|
# Fail on first execution, succeed on next.
|
|
if ray.get(counter.inc.remote()) == 1:
|
|
raise CustomException("error")
|
|
|
|
counter = ExecutionCounter.remote()
|
|
dynamic_ref = generator.options(num_returns=num_returns_type).remote(
|
|
3, store_in_plasma, counter
|
|
)
|
|
ref1, ref2 = ray.get(dynamic_ref)
|
|
ray.get(ref1)
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
ray.get(ref2)
|
|
|
|
ray.get(counter.reset.remote())
|
|
dynamic_ref = generator.options(
|
|
num_returns=num_returns_type, retry_exceptions=[CustomException]
|
|
).remote(3, store_in_plasma, counter)
|
|
for i, ref in enumerate(ray.get(dynamic_ref)):
|
|
assert ray.get(ref)[0] == i
|
|
|
|
|
|
@pytest.mark.parametrize("use_actors", [False, True])
|
|
@pytest.mark.parametrize("store_in_plasma", [False, True])
|
|
@pytest.mark.parametrize("num_returns_type", ["dynamic", None])
|
|
def test_dynamic_generator(
|
|
ray_start_regular_shared, use_actors, store_in_plasma, num_returns_type
|
|
):
|
|
if not use_actors:
|
|
|
|
@ray.remote(num_returns=num_returns_type)
|
|
def dynamic_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 = dynamic_generator
|
|
else:
|
|
|
|
@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
|
|
|
|
@ray.remote
|
|
def read(gen):
|
|
for i, ref in enumerate(gen):
|
|
if ray.get(ref)[0] != i:
|
|
return False
|
|
return True
|
|
|
|
gen = ray.get(
|
|
remote_generator_fn.options(num_returns=num_returns_type).remote(
|
|
10, store_in_plasma
|
|
)
|
|
)
|
|
for i, ref in enumerate(gen):
|
|
assert ray.get(ref)[0] == i
|
|
|
|
# Test empty generator.
|
|
gen = ray.get(
|
|
remote_generator_fn.options(num_returns=num_returns_type).remote(
|
|
0, store_in_plasma
|
|
)
|
|
)
|
|
assert len(list(gen)) == 0
|
|
|
|
# Check that passing as task arg.
|
|
if num_returns_type == "dynamic":
|
|
gen = remote_generator_fn.options(num_returns=num_returns_type).remote(
|
|
10, store_in_plasma
|
|
)
|
|
assert ray.get(read.remote(gen))
|
|
assert ray.get(read.remote(ray.get(gen)))
|
|
else:
|
|
with pytest.raises(TypeError):
|
|
gen = remote_generator_fn.options(num_returns=num_returns_type).remote(
|
|
10, store_in_plasma
|
|
)
|
|
assert ray.get(read.remote(gen))
|
|
|
|
# Also works if we override num_returns with a static value.
|
|
ray.get(
|
|
read.remote(
|
|
remote_generator_fn.options(num_returns=10).remote(10, store_in_plasma)
|
|
)
|
|
)
|
|
|
|
if num_returns_type == "dynamic":
|
|
# Normal remote functions don't work with num_returns="dynamic".
|
|
# This should fail at decoration time, not at runtime.
|
|
with pytest.raises(ValueError, match="can only be used with generator"):
|
|
|
|
@ray.remote(num_returns=num_returns_type)
|
|
def static(num_returns):
|
|
return list(range(num_returns))
|
|
|
|
|
|
def test_dynamic_generator_gc_each_yield(ray_start_cluster):
|
|
# Need to shutdown when going from ray_start_regular_shared to ray_start_cluster
|
|
ray.shutdown()
|
|
|
|
num_returns = 5
|
|
|
|
@ray.remote(num_returns="dynamic")
|
|
def generator():
|
|
for i in range(num_returns):
|
|
yield np.ones((1000, 1000), dtype=np.uint8)
|
|
|
|
def check_ref_counts(expected):
|
|
ref_counts = (
|
|
ray._private.worker.global_worker.core_worker.get_all_reference_counts()
|
|
)
|
|
return len(ref_counts) == expected
|
|
|
|
dynamic_ref = ray.get(generator.remote())
|
|
|
|
for i, ref in enumerate(dynamic_ref):
|
|
gc.collect()
|
|
# assert references are released after each yield
|
|
wait_for_condition(lambda: check_ref_counts(num_returns - i))
|
|
ray.get(ref)
|
|
|
|
|
|
@pytest.mark.parametrize("num_returns_type", ["dynamic", None])
|
|
def test_dynamic_generator_distributed(ray_start_cluster, num_returns_type):
|
|
cluster = ray_start_cluster
|
|
# Head node with no resources.
|
|
cluster.add_node(num_cpus=0)
|
|
ray.init(address=cluster.address)
|
|
cluster.add_node(num_cpus=1)
|
|
cluster.wait_for_nodes()
|
|
|
|
@ray.remote(num_returns=num_returns_type)
|
|
def dynamic_generator(num_returns):
|
|
for i in range(num_returns):
|
|
yield np.ones(1_000_000, dtype=np.int8) * i
|
|
time.sleep(0.1)
|
|
|
|
gen = ray.get(dynamic_generator.remote(3))
|
|
for i, ref in enumerate(gen):
|
|
# Check that we can fetch the values from a different node.
|
|
assert ray.get(ref)[0] == i
|
|
|
|
|
|
@pytest.mark.parametrize("num_returns_type", ["dynamic", None])
|
|
def test_dynamic_generator_reconstruction(ray_start_cluster, num_returns_type):
|
|
config = {
|
|
"health_check_failure_threshold": 10,
|
|
"health_check_period_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,
|
|
"local_gc_min_interval_s": 1,
|
|
}
|
|
cluster = ray_start_cluster
|
|
# Head node with no resources.
|
|
cluster.add_node(
|
|
num_cpus=0, _system_config=config, enable_object_reconstruction=True
|
|
)
|
|
ray.init(address=cluster.address)
|
|
# Node to place the initial object.
|
|
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
cluster.wait_for_nodes()
|
|
|
|
@ray.remote(num_returns=num_returns_type)
|
|
def dynamic_generator(num_returns):
|
|
for i in range(num_returns):
|
|
# Random ray.put to make sure it's okay to interleave these with
|
|
# the dynamic returns.
|
|
if np.random.randint(2) == 1:
|
|
ray.put(np.ones(1_000_000, dtype=np.int8) * np.random.randint(100))
|
|
yield np.ones(1_000_000, dtype=np.int8) * i
|
|
|
|
@ray.remote
|
|
def fetch(x):
|
|
return x[0]
|
|
|
|
# Test recovery of all dynamic objects through re-execution.
|
|
gen = ray.get(dynamic_generator.remote(10))
|
|
cluster.remove_node(node_to_kill, allow_graceful=False)
|
|
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
refs = list(gen)
|
|
|
|
for i, ref in enumerate(refs):
|
|
print("fetching ", i)
|
|
assert ray.get(fetch.remote(ref)) == i
|
|
|
|
cluster.add_node(num_cpus=1, resources={"node2": 1}, object_store_memory=10**8)
|
|
|
|
# Fetch one of the ObjectRefs to another node. We should try to reuse this
|
|
# copy during recovery.
|
|
ray.get(fetch.options(resources={"node2": 1}).remote(refs[-1]))
|
|
cluster.remove_node(node_to_kill, allow_graceful=False)
|
|
for i, ref in enumerate(refs):
|
|
assert ray.get(fetch.remote(ref)) == i
|
|
del ref
|
|
|
|
del refs
|
|
del gen
|
|
assert_no_leak()
|
|
|
|
|
|
@pytest.mark.parametrize("too_many_returns", [False, True])
|
|
@pytest.mark.parametrize("num_returns_type", ["dynamic", None])
|
|
def test_dynamic_generator_reconstruction_nondeterministic(
|
|
ray_start_cluster, too_many_returns, num_returns_type
|
|
):
|
|
# The num_returns_type=None variants used to hang under the RocksDB GCS
|
|
# backend: RocksDB's per-write WAL fsync delayed the actor-death
|
|
# notification enough to expose a pre-existing reconstruction race, so the
|
|
# driver hung in list(gen). Fixed by making the death-notification tables
|
|
# (NODE, ACTOR) soft-durable, which skips the fsync on those tables, so
|
|
# these variants now pass and are no longer skipped. See the
|
|
# SoftDurableTables() comment in rocksdb_store_client.cc for detail.
|
|
config = {
|
|
"health_check_failure_threshold": 10,
|
|
"health_check_period_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,
|
|
"local_gc_min_interval_s": 1,
|
|
}
|
|
cluster = ray_start_cluster
|
|
# Head node with no resources.
|
|
cluster.add_node(
|
|
num_cpus=1,
|
|
_system_config=config,
|
|
enable_object_reconstruction=True,
|
|
resources={"head": 1},
|
|
)
|
|
ray.init(address=cluster.address)
|
|
# Node to place the initial object.
|
|
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
cluster.wait_for_nodes()
|
|
|
|
@ray.remote(num_cpus=1, resources={"head": 1})
|
|
class FailureSignal:
|
|
def __init__(self):
|
|
return
|
|
|
|
def ping(self):
|
|
return
|
|
|
|
@ray.remote(num_returns=num_returns_type)
|
|
def dynamic_generator(failure_signal):
|
|
num_returns = 10
|
|
try:
|
|
ray.get(failure_signal.ping.remote())
|
|
except ray.exceptions.RayActorError:
|
|
if too_many_returns:
|
|
num_returns += 1
|
|
else:
|
|
num_returns -= 1
|
|
for i in range(num_returns):
|
|
yield np.ones(1_000_000, dtype=np.int8) * i
|
|
|
|
@ray.remote
|
|
def fetch(x):
|
|
return
|
|
|
|
failure_signal = FailureSignal.remote()
|
|
gen = ray.get(dynamic_generator.remote(failure_signal))
|
|
cluster.remove_node(node_to_kill, allow_graceful=False)
|
|
ray.kill(failure_signal)
|
|
refs = list(gen)
|
|
if too_many_returns:
|
|
for i, ref in enumerate(refs):
|
|
assert np.array_equal(np.ones(1_000_000, dtype=np.int8) * i, ray.get(ref))
|
|
del ref
|
|
else:
|
|
if num_returns_type == "dynamic":
|
|
# If dynamic is specified, when the num_returns
|
|
# is different, all previous refs are failed.
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
for ref in refs:
|
|
ray.get(ref)
|
|
del ref
|
|
else:
|
|
# Otherwise, we can reconstruct the refs again.
|
|
# We allow it because the refs could have already obtained
|
|
# by the generator.
|
|
for i, ref in enumerate(refs):
|
|
assert np.array_equal(
|
|
np.ones(1_000_000, dtype=np.int8) * i, ray.get(ref)
|
|
)
|
|
del ref
|
|
# TODO(swang): If the re-executed task returns a different number of
|
|
# objects, we should throw an error for every return value.
|
|
# for ref in refs:
|
|
# with pytest.raises(ray.exceptions.RayTaskError):
|
|
# ray.get(ref)
|
|
del gen
|
|
del refs
|
|
if num_returns_type is None:
|
|
# TODO(sang): For some reasons, it fails when "dynamic"
|
|
# is used. We don't fix the issue because we will
|
|
# remove this flag soon anyway.
|
|
assert_no_leak()
|
|
|
|
|
|
@pytest.mark.parametrize("num_returns_type", ["dynamic", None])
|
|
def test_dynamic_generator_reconstruction_fails(ray_start_cluster, num_returns_type):
|
|
config = {
|
|
"health_check_failure_threshold": 10,
|
|
"health_check_period_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,
|
|
"local_gc_min_interval_s": 1,
|
|
}
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(
|
|
num_cpus=1,
|
|
_system_config=config,
|
|
enable_object_reconstruction=True,
|
|
resources={"head": 1},
|
|
)
|
|
ray.init(address=cluster.address)
|
|
# Node to place the initial object.
|
|
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
cluster.wait_for_nodes()
|
|
|
|
@ray.remote(num_cpus=1, resources={"head": 1})
|
|
class FailureSignal:
|
|
def __init__(self):
|
|
return
|
|
|
|
def ping(self):
|
|
return
|
|
|
|
@ray.remote(num_returns=num_returns_type)
|
|
def dynamic_generator(failure_signal):
|
|
num_returns = 10
|
|
for i in range(num_returns):
|
|
yield np.ones(1_000_000, dtype=np.int8) * i
|
|
if i == num_returns // 2:
|
|
# If this is the re-execution, fail the worker after partial yield.
|
|
try:
|
|
ray.get(failure_signal.ping.remote())
|
|
except ray.exceptions.RayActorError:
|
|
sys.exit(-1)
|
|
|
|
@ray.remote
|
|
def fetch(*refs):
|
|
pass
|
|
|
|
failure_signal = FailureSignal.remote()
|
|
gen = ray.get(dynamic_generator.remote(failure_signal))
|
|
refs = list(gen)
|
|
ray.get(fetch.remote(*refs))
|
|
cluster.remove_node(node_to_kill, allow_graceful=False)
|
|
done = fetch.remote(*refs)
|
|
ray.kill(failure_signal)
|
|
# Make sure we can get the error.
|
|
with pytest.raises(ray.exceptions.WorkerCrashedError):
|
|
for ref in refs:
|
|
ray.get(ref)
|
|
|
|
# Make sure other tasks can also get the error.
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
ray.get(done)
|
|
|
|
del ref, gen, refs, done, failure_signal
|
|
gc.collect()
|
|
assert_no_leak()
|
|
|
|
|
|
@pytest.mark.parametrize("num_returns_type", ["dynamic", None])
|
|
def test_dynamic_empty_generator_reconstruction_nondeterministic(
|
|
ray_start_cluster, num_returns_type
|
|
):
|
|
config = {
|
|
"health_check_failure_threshold": 10,
|
|
"health_check_period_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,
|
|
"local_gc_min_interval_s": 1,
|
|
}
|
|
cluster = ray_start_cluster
|
|
# Head node with no resources.
|
|
cluster.add_node(
|
|
num_cpus=0,
|
|
_system_config=config,
|
|
enable_object_reconstruction=True,
|
|
resources={"head": 1},
|
|
)
|
|
ray.init(address=cluster.address)
|
|
# Node to place the initial object.
|
|
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
cluster.wait_for_nodes()
|
|
|
|
@ray.remote(num_cpus=0, resources={"head": 1})
|
|
class ExecutionCounter:
|
|
def __init__(self):
|
|
self.count = 0
|
|
|
|
def inc(self):
|
|
self.count += 1
|
|
return self.count
|
|
|
|
def get_count(self):
|
|
return self.count
|
|
|
|
@ray.remote(num_returns=num_returns_type)
|
|
def maybe_empty_generator(exec_counter):
|
|
if ray.get(exec_counter.inc.remote()) > 1:
|
|
for i in range(3):
|
|
yield np.ones(1_000_000, dtype=np.int8) * i
|
|
|
|
@ray.remote
|
|
def check(empty_generator):
|
|
return len(list(empty_generator)) == 0
|
|
|
|
exec_counter = ExecutionCounter.remote()
|
|
gen = maybe_empty_generator.remote(exec_counter)
|
|
gen = ray.get(gen)
|
|
refs = list(gen)
|
|
assert ray.get(check.remote(refs))
|
|
cluster.remove_node(node_to_kill, allow_graceful=False)
|
|
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
assert ray.get(check.remote(refs))
|
|
|
|
# We should never reconstruct an empty generator.
|
|
assert ray.get(exec_counter.get_count.remote()) == 1
|
|
|
|
del gen, refs, exec_counter
|
|
assert_no_leak()
|
|
|
|
|
|
def test_yield_exception(ray_start_cluster):
|
|
@ray.remote
|
|
def f():
|
|
yield 1
|
|
yield 2
|
|
yield Exception("value")
|
|
yield 3
|
|
raise Exception("raise")
|
|
yield 5
|
|
|
|
gen = f.remote()
|
|
assert ray.get(next(gen)) == 1
|
|
assert ray.get(next(gen)) == 2
|
|
yield_exc = ray.get(next(gen))
|
|
assert isinstance(yield_exc, Exception)
|
|
assert str(yield_exc) == "value"
|
|
assert ray.get(next(gen)) == 3
|
|
with pytest.raises(Exception, match="raise"):
|
|
ray.get(next(gen))
|
|
with pytest.raises(StopIteration):
|
|
ray.get(next(gen))
|
|
|
|
|
|
def test_actor_yield_exception(ray_start_cluster):
|
|
@ray.remote
|
|
class A:
|
|
def f(self):
|
|
yield 1
|
|
yield 2
|
|
yield Exception("value")
|
|
yield 3
|
|
raise Exception("raise")
|
|
yield 5
|
|
|
|
a = A.remote()
|
|
gen = a.f.remote()
|
|
assert ray.get(next(gen)) == 1
|
|
assert ray.get(next(gen)) == 2
|
|
yield_exc = ray.get(next(gen))
|
|
assert isinstance(yield_exc, Exception)
|
|
assert str(yield_exc) == "value"
|
|
assert ray.get(next(gen)) == 3
|
|
with pytest.raises(Exception, match="raise"):
|
|
ray.get(next(gen))
|
|
with pytest.raises(StopIteration):
|
|
ray.get(next(gen))
|
|
|
|
|
|
def test_async_actor_yield_exception(ray_start_cluster):
|
|
@ray.remote
|
|
class A:
|
|
async def f(self):
|
|
yield 1
|
|
yield 2
|
|
yield Exception("value")
|
|
yield 3
|
|
raise Exception("raise")
|
|
yield 5
|
|
|
|
a = A.remote()
|
|
gen = a.f.remote()
|
|
assert ray.get(next(gen)) == 1
|
|
assert ray.get(next(gen)) == 2
|
|
yield_exc = ray.get(next(gen))
|
|
assert isinstance(yield_exc, Exception)
|
|
assert str(yield_exc) == "value"
|
|
assert ray.get(next(gen)) == 3
|
|
with pytest.raises(Exception, match="raise"):
|
|
ray.get(next(gen))
|
|
with pytest.raises(StopIteration):
|
|
ray.get(next(gen))
|
|
|
|
|
|
# Client server port of the shared Ray instance
|
|
SHARED_CLIENT_SERVER_PORT = 25555
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def call_ray_start_shared(request):
|
|
request = Mock()
|
|
request.param = (
|
|
"ray start --head --min-worker-port=0 --max-worker-port=0 --port 0 "
|
|
f"--ray-client-server-port={SHARED_CLIENT_SERVER_PORT}"
|
|
)
|
|
with call_ray_start_context(request) as address:
|
|
yield address
|
|
|
|
|
|
@pytest.mark.parametrize("store_in_plasma", [False, True])
|
|
def test_ray_client(call_ray_start_shared, store_in_plasma):
|
|
with ray_start_client_server_for_address(call_ray_start_shared):
|
|
enable_client_mode()
|
|
|
|
@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]
|
|
|
|
# TODO(swang): When generators return more values than expected, we log an
|
|
# error but the exception is not thrown to the application.
|
|
# https://github.com/ray-project/ray/issues/28689.
|
|
num_returns = 3
|
|
ray.get(
|
|
generator.options(num_returns=num_returns).remote(
|
|
num_returns + 1, store_in_plasma
|
|
)
|
|
)
|
|
|
|
# Check return values.
|
|
[
|
|
x[0]
|
|
for x in ray.get(
|
|
generator.options(num_returns=num_returns).remote(
|
|
num_returns, store_in_plasma
|
|
)
|
|
)
|
|
] == list(range(num_returns))
|
|
# Works for num_returns=1 if generator returns a single value.
|
|
assert (
|
|
ray.get(generator.options(num_returns=1).remote(1, store_in_plasma))[0] == 0
|
|
)
|
|
|
|
gen = ray.get(
|
|
generator.options(num_returns="dynamic").remote(3, store_in_plasma)
|
|
)
|
|
for i, ref in enumerate(gen):
|
|
assert ray.get(ref)[0] == i
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|