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