wehub-resource-sync
413 lines
13 KiB
Python
413 lines
13 KiB
Python
import asyncio
|
|
import gc
|
|
import json
|
|
import os
|
|
import random
|
|
import signal
|
|
import sys
|
|
import time
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from pydantic import BaseModel
|
|
|
|
import ray
|
|
from ray._common.test_utils import SignalActor
|
|
|
|
RECONSTRUCTION_CONFIG = {
|
|
"health_check_failure_threshold": 10,
|
|
"health_check_period_ms": 100,
|
|
"health_check_timeout_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,
|
|
}
|
|
|
|
|
|
def assert_no_leak():
|
|
gc.collect()
|
|
core_worker = ray._private.worker.global_worker.core_worker
|
|
ref_counts = core_worker.get_all_reference_counts()
|
|
print(ref_counts)
|
|
for rc in ref_counts.values():
|
|
assert rc["local"] == 0
|
|
assert rc["submitted"] == 0
|
|
assert core_worker.get_memory_store_size() == 0
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="SIGKILL is not available on Windows"
|
|
)
|
|
def test_caller_death(monkeypatch, shutdown_only):
|
|
"""
|
|
Test the case where caller of a streaming generator actor task dies
|
|
while the streaming generator task is executing. The streaming
|
|
generator task should still finish and won't block other actor tasks.
|
|
This means that `ReportGeneratorItemReturns` RPC should fail and it shouldn't
|
|
be retried indefinitely.
|
|
"""
|
|
ray.init()
|
|
|
|
@ray.remote
|
|
class Callee:
|
|
def gen(self, caller_pid):
|
|
os.kill(caller_pid, signal.SIGKILL)
|
|
yield [1] * 1024 * 1024
|
|
|
|
def ping(self):
|
|
pass
|
|
|
|
@ray.remote
|
|
def caller(callee):
|
|
ray.get(callee.gen.remote(os.getpid()))
|
|
|
|
callee = Callee.remote()
|
|
o = caller.remote(callee)
|
|
ray.wait([o])
|
|
# Make sure gen will finish and ping can run.
|
|
ray.get(callee.ping.remote())
|
|
|
|
|
|
def test_intermediate_generator_object_recovery_while_generator_running(
|
|
ray_start_cluster,
|
|
):
|
|
"""
|
|
1. Streaming producer starts on worker1.
|
|
2. consumer consumes value 1 from producer on worker2 and finishes.
|
|
3. Run an extra consumer on worker2 to track when reconstruction is triggered.
|
|
4. Add worker3.
|
|
5. worker2 dies.
|
|
6. Try to get consumer output.
|
|
7. Therefore Ray tries to reconstruct value 1 from producer.
|
|
8. Get the reconstructed extra_consumer_ref (assures 7 happened).
|
|
9. Streaming producer should be cancelled and resubmitted.
|
|
10. Retry for consumer should complete.
|
|
"""
|
|
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=0) # head
|
|
ray.init(address=cluster.address)
|
|
cluster.add_node(num_cpus=1, resources={"producer": 1}) # worker1
|
|
worker2 = cluster.add_node(num_cpus=1, resources={"consumer": 1})
|
|
|
|
@ray.remote(num_cpus=1, resources={"producer": 1})
|
|
def producer():
|
|
for _ in range(3):
|
|
yield np.zeros(10 * 1024 * 1024, dtype=np.uint8)
|
|
|
|
@ray.remote(num_cpus=1, resources={"consumer": 1})
|
|
def consumer(np_arr):
|
|
return np_arr
|
|
|
|
streaming_ref = producer.options(_generator_backpressure_num_objects=1).remote()
|
|
consumer_ref = consumer.remote(next(streaming_ref))
|
|
extra_consumer_ref = consumer.remote(np.zeros(10 * 1024 * 1024, dtype=np.uint8))
|
|
|
|
ray.wait([consumer_ref, extra_consumer_ref], num_returns=2, fetch_local=False)
|
|
|
|
cluster.add_node(num_cpus=1, resources={"consumer": 1}) # worker3
|
|
cluster.remove_node(worker2, allow_graceful=True)
|
|
|
|
# Make sure reconstruction was triggered.
|
|
assert ray.get(extra_consumer_ref).size == (10 * 1024 * 1024)
|
|
# Allow first streaming generator attempt to finish
|
|
ray.get([next(streaming_ref), next(streaming_ref)])
|
|
|
|
assert ray.get(consumer_ref).size == (10 * 1024 * 1024)
|
|
|
|
|
|
def test_actor_intermediate_generator_object_recovery_while_generator_running(
|
|
ray_start_cluster,
|
|
):
|
|
"""
|
|
1. Producer actor and its generator producer task start on worker1.
|
|
2. consumer consumes value 1 from producer on worker2 and finishes.
|
|
3. Run an extra consumer on worker2 to track when reconstruction is triggered.
|
|
4. Add worker3.
|
|
5. worker2 dies.
|
|
6. Ray tries to reconstruct value 1 from producer.
|
|
7. Get the reconstructed extra_consumer_ref (assures 6 happened).
|
|
8. Ray tries and fails to cancel the producer task.
|
|
9. Get the next two values to relieve backpressure and allow producer to finish.
|
|
10. Ray resubmits the producer generator task.
|
|
11. Retry for consumer should complete.
|
|
"""
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=0) # head
|
|
ray.init(address=cluster.address)
|
|
cluster.add_node(num_cpus=1, resources={"producer": 1}) # worker 1
|
|
worker2 = cluster.add_node(num_cpus=1, resources={"consumer": 1})
|
|
|
|
@ray.remote(num_cpus=1, resources={"producer": 1}, max_task_retries=-1)
|
|
class Producer:
|
|
def producer(self):
|
|
for _ in range(3):
|
|
yield np.zeros(10 * 1024 * 1024, dtype=np.uint8)
|
|
|
|
@ray.remote(num_cpus=1, resources={"consumer": 1})
|
|
def consumer(np_arr):
|
|
return np_arr
|
|
|
|
producer_actor = Producer.remote()
|
|
streaming_ref = producer_actor.producer.options(
|
|
_generator_backpressure_num_objects=1
|
|
).remote()
|
|
consumer_ref = consumer.remote(next(streaming_ref))
|
|
extra_consumer_ref = consumer.remote(np.zeros(10 * 1024 * 1024, dtype=np.uint8))
|
|
|
|
ray.wait([consumer_ref, extra_consumer_ref], num_returns=2, fetch_local=False)
|
|
|
|
cluster.add_node(num_cpus=1, resources={"consumer": 1}) # worker 3
|
|
cluster.remove_node(worker2, allow_graceful=True)
|
|
|
|
# Make sure reconstruction was triggered.
|
|
ray.get(extra_consumer_ref)
|
|
# Allow first streaming generator attempt to finish
|
|
ray.get([next(streaming_ref), next(streaming_ref)])
|
|
|
|
assert ray.get(consumer_ref).size == (10 * 1024 * 1024)
|
|
|
|
|
|
@pytest.mark.parametrize("backpressure", [False, True])
|
|
@pytest.mark.parametrize("delay_latency", [0.1, 1])
|
|
@pytest.mark.parametrize("threshold", [1, 3])
|
|
def test_many_tasks_lineage_reconstruction_mini_stress_test(
|
|
monkeypatch, ray_start_cluster, backpressure, delay_latency, threshold
|
|
):
|
|
"""Test a workload that spawns many tasks and relies on lineage reconstruction."""
|
|
if not backpressure:
|
|
if delay_latency == 0.1 and threshold == 1:
|
|
return
|
|
elif delay_latency == 1:
|
|
return
|
|
|
|
with monkeypatch.context() as m:
|
|
m.setenv(
|
|
"RAY_testing_asio_delay_us",
|
|
"CoreWorkerService.grpc_server.ReportGeneratorItemReturns=10000:1000000",
|
|
)
|
|
m.setenv(
|
|
"RAY_testing_rpc_failure",
|
|
json.dumps(
|
|
{
|
|
"CoreWorkerService.grpc_client.ReportGeneratorItemReturns": {
|
|
"num_failures": 5,
|
|
"req_failure_prob": 25,
|
|
"resp_failure_prob": 25,
|
|
"in_flight_failure_prob": 25,
|
|
}
|
|
}
|
|
),
|
|
)
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(
|
|
num_cpus=1,
|
|
resources={"head": 1},
|
|
_system_config=RECONSTRUCTION_CONFIG,
|
|
enable_object_reconstruction=True,
|
|
)
|
|
ray.init(address=cluster.address)
|
|
|
|
if backpressure:
|
|
threshold = 1
|
|
else:
|
|
threshold = -1
|
|
|
|
@ray.remote(
|
|
max_retries=-1,
|
|
_generator_backpressure_num_objects=threshold,
|
|
)
|
|
def dynamic_generator(num_returns):
|
|
for i in range(num_returns):
|
|
time.sleep(0.1)
|
|
yield np.ones(1_000_000, dtype=np.int8) * i
|
|
|
|
@ray.remote(num_cpus=0, resources={"head": 1})
|
|
def driver():
|
|
unready = [dynamic_generator.remote(10) for _ in range(5)]
|
|
ready = []
|
|
while unready:
|
|
for a in unready:
|
|
print(a._generator_ref)
|
|
ready, unready = ray.wait(
|
|
unready, num_returns=len(unready), timeout=0.1
|
|
)
|
|
for r in ready:
|
|
try:
|
|
ref = next(r)
|
|
print(ref)
|
|
ray.get(ref)
|
|
except StopIteration:
|
|
pass
|
|
else:
|
|
unready.append(r)
|
|
return None
|
|
|
|
ref = driver.remote()
|
|
|
|
nodes = []
|
|
for _ in range(4):
|
|
nodes.append(cluster.add_node(num_cpus=1, object_store_memory=10**8))
|
|
cluster.wait_for_nodes()
|
|
|
|
for _ in range(10):
|
|
time.sleep(0.1)
|
|
node_to_kill = random.choices(nodes)[0]
|
|
nodes.remove(node_to_kill)
|
|
cluster.remove_node(node_to_kill, allow_graceful=False)
|
|
nodes.append(cluster.add_node(num_cpus=1, object_store_memory=10**8))
|
|
|
|
ray.get(ref)
|
|
del ref
|
|
|
|
assert_no_leak()
|
|
|
|
|
|
def test_local_gc_not_hang(shutdown_only, monkeypatch):
|
|
"""Verify the generator doesn't deadlock when a local GC is triggered."""
|
|
with monkeypatch.context() as m:
|
|
m.setenv("RAY_local_gc_interval_s", 1)
|
|
|
|
ray.init()
|
|
|
|
@ray.remote(_generator_backpressure_num_objects=1)
|
|
def f():
|
|
for _ in range(5):
|
|
yield 1
|
|
|
|
gen = f.remote()
|
|
time.sleep(5)
|
|
|
|
# It should not hang.
|
|
for ref in gen:
|
|
ray.get(gen)
|
|
|
|
|
|
def test_sync_async_mix_regression_test(shutdown_only):
|
|
"""Verify when sync and async tasks are mixed up
|
|
it doesn't raise a segfault
|
|
|
|
https://github.com/ray-project/ray/issues/41346
|
|
"""
|
|
|
|
class PayloadPydantic(BaseModel):
|
|
class Error(BaseModel):
|
|
msg: str
|
|
code: int
|
|
type: str
|
|
|
|
text: Optional[str] = None
|
|
ts: Optional[float] = None
|
|
reason: Optional[str] = None
|
|
error: Optional[Error] = None
|
|
|
|
ray.init()
|
|
|
|
@ray.remote
|
|
class B:
|
|
def __init__(self, a):
|
|
self.a = a
|
|
|
|
async def stream(self):
|
|
async for ref in self.a.stream.remote(1):
|
|
print("stream")
|
|
await ref
|
|
|
|
async def start(self):
|
|
await asyncio.gather(*[self.stream() for _ in range(2)])
|
|
|
|
@ray.remote
|
|
class A:
|
|
def stream(self, i):
|
|
payload = PayloadPydantic(
|
|
text="Test output",
|
|
ts=time.time(),
|
|
reason="Success!",
|
|
)
|
|
|
|
for _ in range(10):
|
|
yield payload
|
|
|
|
async def aio_stream(self):
|
|
for _ in range(10):
|
|
yield 1
|
|
|
|
a = A.remote()
|
|
b = B.remote(a)
|
|
ray.get(b.start.remote())
|
|
|
|
|
|
@pytest.mark.parametrize("use_asyncio", [False, True])
|
|
def test_cancel(shutdown_only, use_asyncio):
|
|
"""Test concurrent task cancellation with generator task.
|
|
|
|
Once the caller receives an ack that the executor has cancelled the task
|
|
execution, the caller should receive a TaskCancelledError for the next
|
|
ObjectRef that it tries to read from the generator. This should happen even
|
|
if the caller has already received values for the next object indices in
|
|
the stream. Also, we should not apply the usual logic that reorders
|
|
out-of-order reports if the task was cancelled; waiting for the
|
|
intermediate indices to appear would hang the caller."""
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def ready(self):
|
|
return
|
|
|
|
def stream(self, signal):
|
|
cancelled_ref = signal.wait.remote()
|
|
|
|
i = 0
|
|
done_at = time.time() + 1
|
|
while time.time() < done_at:
|
|
yield i
|
|
i += 1
|
|
|
|
ready, _ = ray.wait([cancelled_ref], timeout=0)
|
|
if not ready:
|
|
# Continue executing for one second after the driver
|
|
# cancels. This is to make sure that we receive the cancel
|
|
# signal while the task is still running.
|
|
done_at = time.time() + 1
|
|
|
|
async def async_stream(self, signal):
|
|
cancelled_ref = signal.wait.remote()
|
|
|
|
i = 0
|
|
done_at = time.time() + 1
|
|
while time.time() < done_at:
|
|
yield i
|
|
i += 1
|
|
|
|
ready, _ = ray.wait([cancelled_ref], timeout=0)
|
|
if not ready:
|
|
# Continue executing for one second after the driver
|
|
# cancels. This is to make sure that we receive the cancel
|
|
# signal while the task is still running.
|
|
done_at = time.time() + 1
|
|
|
|
signal = SignalActor.remote()
|
|
a = Actor.remote()
|
|
ray.get(a.ready.remote())
|
|
if use_asyncio:
|
|
gen = a.async_stream.remote(signal)
|
|
else:
|
|
gen = a.stream.remote(signal)
|
|
|
|
try:
|
|
for i, ref in enumerate(gen):
|
|
assert i == ray.get(ref)
|
|
print(i)
|
|
if i == 0:
|
|
ray.cancel(gen)
|
|
signal.send.remote()
|
|
except ray.exceptions.TaskCancelledError:
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|