Files
2026-07-13 13:17:40 +08:00

610 lines
20 KiB
Python

import os
import sys
import time
import numpy as np
import pytest
import ray
import ray._private.ray_constants as ray_constants
import ray.exceptions
from ray._common.test_utils import Semaphore, SignalActor, wait_for_condition
from ray._private.internal_api import memory_summary
from ray.util.state import list_tasks
# Task status.
WAITING_FOR_DEPENDENCIES = "PENDING_ARGS_AVAIL"
FINISHED = "FINISHED"
WAITING_FOR_EXECUTION = "SUBMITTED_TO_WORKER"
@pytest.fixture
def config(request):
config = {
"health_check_initial_delay_ms": 5000,
"health_check_period_ms": 100,
"health_check_failure_threshold": 20,
"object_timeout_milliseconds": 200,
}
yield config
@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
@pytest.mark.parametrize("reconstruction_enabled", [False, True])
def test_nondeterministic_output(config, ray_start_cluster, reconstruction_enabled):
config["max_direct_call_object_size"] = 100
config["task_retry_delay_ms"] = 100
config["object_timeout_milliseconds"] = 200
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, resources={"node1": 1}, object_store_memory=10**8
)
cluster.add_node(num_cpus=1, object_store_memory=10**8)
cluster.wait_for_nodes()
@ray.remote
def nondeterministic_object():
if np.random.rand() < 0.5:
return np.zeros(10**5, dtype=np.uint8)
else:
return 0
@ray.remote
def dependent_task(x):
return
for _ in range(10):
obj = nondeterministic_object.options(resources={"node1": 1}).remote()
for _ in range(3):
ray.get(dependent_task.remote(obj))
x = dependent_task.remote(obj)
cluster.remove_node(node_to_kill, allow_graceful=False)
node_to_kill = cluster.add_node(
num_cpus=1, resources={"node1": 1}, object_store_memory=10**8
)
ray.get(x)
@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
def test_reconstruction_hangs(config, ray_start_cluster):
config["max_direct_call_object_size"] = 100
config["task_retry_delay_ms"] = 100
config["object_timeout_milliseconds"] = 200
config["fetch_warn_timeout_milliseconds"] = 1000
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, resources={"node1": 1}, object_store_memory=10**8
)
cluster.add_node(num_cpus=1, object_store_memory=10**8)
cluster.wait_for_nodes()
@ray.remote
def sleep():
# Task takes longer than the reconstruction timeout.
time.sleep(3)
return np.zeros(10**5, dtype=np.uint8)
@ray.remote
def dependent_task(x):
return
obj = sleep.options(resources={"node1": 1}).remote()
for _ in range(3):
ray.get(dependent_task.remote(obj))
x = dependent_task.remote(obj)
cluster.remove_node(node_to_kill, allow_graceful=False)
node_to_kill = cluster.add_node(
num_cpus=1, resources={"node1": 1}, object_store_memory=10**8
)
ray.get(x)
def test_lineage_evicted(config, ray_start_cluster):
config["max_lineage_bytes"] = 10_000
cluster = ray_start_cluster
# Head node with no resources.
cluster.add_node(
num_cpus=0,
_system_config=config,
object_store_memory=10**8,
enable_object_reconstruction=True,
)
ray.init(address=cluster.address)
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
cluster.wait_for_nodes()
@ray.remote
def large_object():
return np.zeros(10**7, dtype=np.uint8)
@ray.remote
def chain(x):
return x
@ray.remote
def dependent_task(x):
return x
obj = large_object.remote()
for _ in range(5):
obj = chain.remote(obj)
ray.get(dependent_task.remote(obj))
cluster.remove_node(node_to_kill, allow_graceful=False)
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
ray.get(dependent_task.remote(obj))
# Lineage now exceeds the eviction factor.
for _ in range(100):
obj = chain.remote(obj)
ray.get(dependent_task.remote(obj))
cluster.remove_node(node_to_kill, allow_graceful=False)
cluster.add_node(num_cpus=1, object_store_memory=10**8)
try:
ray.get(dependent_task.remote(obj))
assert False
except ray.exceptions.RayTaskError as e:
assert "ObjectReconstructionFailedError" in str(e)
assert "LINEAGE_EVICTED" in str(e)
@pytest.mark.parametrize("reconstruction_enabled", [False, True])
def test_multiple_returns(config, ray_start_cluster, reconstruction_enabled):
# Workaround to reset the config to the default value.
if not reconstruction_enabled:
config["lineage_pinning_enabled"] = False
cluster = ray_start_cluster
# Head node with no resources.
cluster.add_node(
num_cpus=0,
_system_config=config,
enable_object_reconstruction=reconstruction_enabled,
)
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=2)
def two_large_objects():
return (np.zeros(10**7, dtype=np.uint8), np.zeros(10**7, dtype=np.uint8))
@ray.remote
def dependent_task(x):
return
obj1, obj2 = two_large_objects.remote()
ray.get(dependent_task.remote(obj1))
cluster.add_node(num_cpus=1, resources={"node": 1}, object_store_memory=10**8)
ray.get(dependent_task.options(resources={"node": 1}).remote(obj1))
cluster.remove_node(node_to_kill, allow_graceful=False)
wait_for_condition(
lambda: not all(node["Alive"] for node in ray.nodes()), timeout=10
)
if reconstruction_enabled:
ray.get(dependent_task.remote(obj1))
ray.get(dependent_task.remote(obj2))
else:
with pytest.raises(ray.exceptions.RayTaskError):
ray.get(dependent_task.remote(obj1))
ray.get(dependent_task.remote(obj2))
with pytest.raises(ray.exceptions.ObjectReconstructionFailedError):
ray.get(obj2)
@pytest.mark.parametrize("reconstruction_enabled", [False, True])
def test_nested(config, ray_start_cluster, reconstruction_enabled):
config["fetch_fail_timeout_milliseconds"] = 10_000
# Workaround to reset the config to the default value.
if not reconstruction_enabled:
config["lineage_pinning_enabled"] = False
cluster = ray_start_cluster
# Head node with no resources.
cluster.add_node(
num_cpus=0,
_system_config=config,
enable_object_reconstruction=reconstruction_enabled,
)
ray.init(address=cluster.address)
done_signal = SignalActor.remote()
exit_signal = SignalActor.remote()
ray.get(done_signal.wait.remote(should_wait=False))
ray.get(exit_signal.wait.remote(should_wait=False))
# 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
def dependent_task(x):
return
@ray.remote
def large_object():
return np.zeros(10**7, dtype=np.uint8)
@ray.remote
def nested(done_signal, exit_signal):
ref = ray.put(np.zeros(10**7, dtype=np.uint8))
# Flush object store.
for _ in range(20):
ray.put(np.zeros(10**7, dtype=np.uint8))
dep = dependent_task.options(resources={"node": 1}).remote(ref)
ray.get(done_signal.send.remote(clear=True))
ray.get(dep)
return ray.get(ref)
ref = nested.remote(done_signal, exit_signal)
# Wait for task to get scheduled on the node to kill.
ray.get(done_signal.wait.remote())
# Wait for ray.put object to get transferred to the other node.
cluster.add_node(num_cpus=2, resources={"node": 10}, object_store_memory=10**8)
ray.get(dependent_task.remote(ref))
# Destroy the task's output.
cluster.remove_node(node_to_kill, allow_graceful=False)
wait_for_condition(
lambda: not all(node["Alive"] for node in ray.nodes()), timeout=10
)
if reconstruction_enabled:
ray.get(ref, timeout=60)
else:
with pytest.raises(ray.exceptions.ObjectReconstructionFailedError):
ray.get(ref, timeout=60)
@pytest.mark.parametrize("reconstruction_enabled", [False, True])
def test_spilled(config, ray_start_cluster, reconstruction_enabled):
# Workaround to reset the config to the default value.
if not reconstruction_enabled:
config["lineage_pinning_enabled"] = False
cluster = ray_start_cluster
# Head node with no resources.
cluster.add_node(
num_cpus=0,
_system_config=config,
enable_object_reconstruction=reconstruction_enabled,
)
ray.init(address=cluster.address)
# Node to place the initial object.
node_to_kill = cluster.add_node(
num_cpus=1, resources={"node1": 1}, object_store_memory=10**8
)
cluster.wait_for_nodes()
@ray.remote(max_retries=1 if reconstruction_enabled else 0)
def large_object():
return np.zeros(10**7, dtype=np.uint8)
@ray.remote
def dependent_task(x):
return
obj = large_object.options(resources={"node1": 1}).remote()
ray.get(dependent_task.options(resources={"node1": 1}).remote(obj))
# Force spilling.
objs = [large_object.options(resources={"node1": 1}).remote() for _ in range(20)]
for o in objs:
ray.get(o)
cluster.remove_node(node_to_kill, allow_graceful=False)
node_to_kill = cluster.add_node(
num_cpus=1, resources={"node1": 1}, object_store_memory=10**8
)
if reconstruction_enabled:
ray.get(dependent_task.remote(obj), timeout=60)
else:
with pytest.raises(ray.exceptions.RayTaskError):
ray.get(dependent_task.remote(obj), timeout=60)
with pytest.raises(ray.exceptions.ObjectReconstructionFailedError):
ray.get(obj, timeout=60)
def test_memory_util(config, ray_start_cluster):
cluster = ray_start_cluster
# Head node with no resources.
cluster.add_node(
num_cpus=0,
resources={"head": 1},
_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, resources={"node1": 1}, object_store_memory=10**8
)
cluster.wait_for_nodes()
@ray.remote
def large_object(sema=None):
if sema is not None:
ray.get(sema.acquire.remote())
return np.zeros(10**7, dtype=np.uint8)
@ray.remote
def dependent_task(x, sema):
ray.get(sema.acquire.remote())
return x
def stats():
info = memory_summary(cluster.address, line_wrap=False)
print(info)
info = info.split("\n")
reconstructing_waiting = [
fields
for fields in [[part.strip() for part in line.split("|")] for line in info]
if len(fields) == 9
and fields[4] == WAITING_FOR_DEPENDENCIES
and fields[5] == "2"
]
reconstructing_scheduled = [
fields
for fields in [[part.strip() for part in line.split("|")] for line in info]
if len(fields) == 9
and fields[4] == WAITING_FOR_EXECUTION
and fields[5] == "2"
]
reconstructing_finished = [
fields
for fields in [[part.strip() for part in line.split("|")] for line in info]
if len(fields) == 9 and fields[4] == FINISHED and fields[5] == "2"
]
return (
len(reconstructing_waiting),
len(reconstructing_scheduled),
len(reconstructing_finished),
)
sema = Semaphore.options(resources={"head": 1}).remote(value=0)
obj = large_object.options(resources={"node1": 1}).remote(sema)
x = dependent_task.options(resources={"node1": 1}).remote(obj, sema)
ref = dependent_task.options(resources={"node1": 1}).remote(x, sema)
ray.get(sema.release.remote())
ray.get(sema.release.remote())
ray.get(sema.release.remote())
ray.get(ref)
wait_for_condition(lambda: stats() == (0, 0, 0))
del ref
cluster.remove_node(node_to_kill, allow_graceful=False)
node_to_kill = cluster.add_node(
num_cpus=1, resources={"node1": 1}, object_store_memory=10**8
)
ref = dependent_task.remote(x, sema)
wait_for_condition(lambda: stats() == (1, 1, 0))
ray.get(sema.release.remote())
wait_for_condition(lambda: stats() == (0, 1, 1))
ray.get(sema.release.remote())
ray.get(sema.release.remote())
ray.get(ref)
wait_for_condition(lambda: stats() == (0, 0, 2))
@pytest.mark.parametrize("override_max_retries", [False, True])
def test_override_max_retries(ray_start_cluster, override_max_retries):
cluster = ray_start_cluster
cluster.add_node(num_cpus=1)
max_retries = ray_constants.DEFAULT_TASK_MAX_RETRIES
runtime_env = {}
if override_max_retries:
max_retries = 1
runtime_env["env_vars"] = {"RAY_TASK_MAX_RETRIES": str(max_retries)}
os.environ["RAY_TASK_MAX_RETRIES"] = str(max_retries)
# Since we're setting the OS environment variable after the driver process
# is already started, we need to set it a second time for the workers with
# runtime_env.
ray.init(cluster.address, runtime_env=runtime_env)
try:
@ray.remote
class ExecutionCounter:
def __init__(self):
self.count = 0
def inc(self):
self.count += 1
def pop(self):
count = self.count
self.count = 0
return count
@ray.remote
def f(counter):
ray.get(counter.inc.remote())
sys.exit(-1)
counter = ExecutionCounter.remote()
with pytest.raises(ray.exceptions.WorkerCrashedError):
ray.get(f.remote(counter))
assert ray.get(counter.pop.remote()) == max_retries + 1
# Check max_retries override still works.
with pytest.raises(ray.exceptions.WorkerCrashedError):
ray.get(f.options(max_retries=0).remote(counter))
assert ray.get(counter.pop.remote()) == 1
@ray.remote
def nested(counter):
ray.get(f.remote(counter))
# Check override works through nested tasks.
with pytest.raises(ray.exceptions.RayTaskError):
ray.get(nested.remote(counter))
assert ray.get(counter.pop.remote()) == max_retries + 1
finally:
if override_max_retries:
del os.environ["RAY_TASK_MAX_RETRIES"]
@pytest.mark.parametrize("reconstruction_enabled", [False, True])
def test_reconstruct_freed_object(config, ray_start_cluster, reconstruction_enabled):
# Workaround to reset the config to the default value.
if not reconstruction_enabled:
config["lineage_pinning_enabled"] = False
cluster = ray_start_cluster
# Head node with no resources.
cluster.add_node(
num_cpus=0,
_system_config=config,
enable_object_reconstruction=reconstruction_enabled,
)
ray.init(address=cluster.address)
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
cluster.wait_for_nodes()
@ray.remote
def large_object():
return np.zeros(10**7, dtype=np.uint8)
@ray.remote
def dependent_task(x):
return np.zeros(10**7, dtype=np.uint8)
obj = large_object.remote()
x = dependent_task.remote(obj)
ray.get(dependent_task.remote(x))
ray.internal.free(obj)
cluster.remove_node(node_to_kill, allow_graceful=False)
cluster.add_node(num_cpus=1, object_store_memory=10**8)
if reconstruction_enabled:
ray.get(x)
else:
with pytest.raises(ray.exceptions.ObjectReconstructionFailedError):
ray.get(x)
with pytest.raises(ray.exceptions.ObjectFreedError):
ray.get(obj)
def test_object_reconstruction_dead_actor(config, ray_start_cluster):
# Test to make sure that if object reconstruction fails
# due to dead actor, pending_creation is set back to false.
# https://github.com/ray-project/ray/issues/47606
cluster = ray_start_cluster
cluster.add_node(num_cpus=0, _system_config=config)
ray.init(address=cluster.address)
node1 = cluster.add_node(resources={"node1": 1})
node2 = cluster.add_node(resources={"node2": 1})
@ray.remote(max_restarts=0, max_task_retries=-1, resources={"node1": 0.1})
class Worker:
def func_in(self):
return np.random.rand(1024000)
@ray.remote(max_retries=-1, resources={"node2": 0.1})
def func_out(data):
return np.random.rand(1024000)
worker = Worker.remote()
ref_in = worker.func_in.remote()
ref_out = func_out.remote(ref_in)
ray.wait([ref_in, ref_out], num_returns=2, timeout=None, fetch_local=False)
def func_out_resubmitted():
tasks = list_tasks(filters=[("name", "=", "func_out")])
assert len(tasks) == 2
assert (
tasks[0]["state"] == "PENDING_NODE_ASSIGNMENT"
or tasks[1]["state"] == "PENDING_NODE_ASSIGNMENT"
)
return True
cluster.remove_node(node2, allow_graceful=False)
# ref_out will reconstruct, wait for the lease request to reach raylet.
wait_for_condition(func_out_resubmitted)
cluster.remove_node(node1, allow_graceful=False)
# ref_in is lost and the reconstruction will
# fail with ActorDiedError
node1 = cluster.add_node(resources={"node1": 1})
node2 = cluster.add_node(resources={"node2": 1})
with pytest.raises(ray.exceptions.RayTaskError) as exc_info:
ray.get(ref_out)
assert "input arguments for this task could not be computed" in str(exc_info.value)
def test_object_reconstruction_pending_creation(config, ray_start_cluster):
# Test to make sure that an object being reconstructured
# has pending_creation set to true.
config["fetch_fail_timeout_milliseconds"] = (
5000 if sys.platform == "linux" else 9000
)
cluster = ray_start_cluster
cluster.add_node(num_cpus=0, resources={"head": 1}, _system_config=config)
ray.init(address=cluster.address)
@ray.remote(num_cpus=0, resources={"head": 0.1})
class Counter:
def __init__(self):
self.count = 0
def inc(self):
self.count = self.count + 1
return self.count
counter = Counter.remote()
@ray.remote(num_cpus=1, max_retries=-1)
def generator(counter):
if ray.get(counter.inc.remote()) == 1:
# first attempt
yield np.zeros(10**6, dtype=np.uint8)
time.sleep(10000000)
yield np.zeros(10**6, dtype=np.uint8)
else:
time.sleep(10000000)
yield np.zeros(10**6, dtype=np.uint8)
time.sleep(10000000)
yield np.zeros(10**6, dtype=np.uint8)
worker = cluster.add_node(num_cpus=8)
gen = generator.remote(counter)
obj = next(gen)
cluster.remove_node(worker, allow_graceful=False)
# After removing the node, the generator task will be retried
# and the obj will be reconstructured and has pending_creation set to true.
cluster.add_node(num_cpus=8)
# This should raise GetTimeoutError instead of ObjectFetchTimedOutError
with pytest.raises(ray.exceptions.GetTimeoutError):
ray.get(obj, timeout=10)
if __name__ == "__main__":
sys.exit(pytest.main(["-sv", __file__]))