wehub-resource-sync
140 lines
4.6 KiB
Python
140 lines
4.6 KiB
Python
import sys
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.actor import ActorHandle
|
|
from ray.exceptions import RayTaskError, TaskCancelledError
|
|
from ray.util.state import list_workers
|
|
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class EndpointActor:
|
|
def __init__(self, *, injected_executor_delay_s: float, tokens_per_request: int):
|
|
self._tokens_per_request = tokens_per_request
|
|
# In this test we simulate conditions leading to use-after-free conditions,
|
|
# by injecting delays into worker's thread-pool executor
|
|
self._inject_delay_in_core_worker_executor(
|
|
target_delay_s=injected_executor_delay_s,
|
|
max_workers=1,
|
|
)
|
|
|
|
async def aio_stream(self):
|
|
for i in range(self._tokens_per_request):
|
|
yield i
|
|
|
|
@classmethod
|
|
def _inject_delay_in_core_worker_executor(
|
|
cls, target_delay_s: float, max_workers: int
|
|
):
|
|
if target_delay_s > 0:
|
|
|
|
class DelayedThreadPoolExecutor(ThreadPoolExecutor):
|
|
def submit(self, fn, /, *args, **kwargs):
|
|
def __slowed_fn():
|
|
print(
|
|
f">>> [DelayedThreadPoolExecutor] Starting executing "
|
|
f"function with delay {target_delay_s}s"
|
|
)
|
|
|
|
time.sleep(target_delay_s)
|
|
fn(*args, **kwargs)
|
|
|
|
return super().submit(__slowed_fn)
|
|
|
|
executor = DelayedThreadPoolExecutor(max_workers=max_workers)
|
|
ray._private.worker.global_worker.core_worker.reset_event_loop_executor(
|
|
executor
|
|
)
|
|
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class CallerActor:
|
|
def __init__(
|
|
self,
|
|
downstream: ActorHandle,
|
|
):
|
|
self._h = downstream
|
|
|
|
async def run(self):
|
|
print(">>> [Caller] Starting consuming stream")
|
|
|
|
async_obj_ref_gen = self._h.aio_stream.options(num_returns="streaming").remote()
|
|
async for ref in async_obj_ref_gen:
|
|
r = await ref
|
|
if r == 1:
|
|
print(">>> [Caller] Cancelling generator")
|
|
ray.cancel(async_obj_ref_gen, recursive=False)
|
|
|
|
# NOTE: This delay is crucial to let already scheduled task to report
|
|
# generated item (report_streaming_generator_output) before we
|
|
# will tear down this stream
|
|
delay_after_cancellation_s = 2
|
|
|
|
print(f">>> [Caller] **Sleeping** {delay_after_cancellation_s}s")
|
|
time.sleep(delay_after_cancellation_s)
|
|
else:
|
|
print(f">>> [Caller] Received {r}")
|
|
|
|
print(">>> [Caller] Completed consuming stream")
|
|
|
|
|
|
@pytest.mark.parametrize("injected_executor_delay_s", [0, 2])
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster",
|
|
[
|
|
{
|
|
"num_nodes": 2,
|
|
"num_cpus": 1,
|
|
}
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_segfault_report_streaming_generator_output(
|
|
ray_start_cluster, injected_executor_delay_s: float
|
|
):
|
|
"""
|
|
This is a "smoke" test attempting to emulate condition, when using Ray's async
|
|
streaming generator, that leads to worker crashing with SIGSEGV.
|
|
|
|
For more details summarizing these conditions, please refer to
|
|
https://github.com/ray-project/ray/issues/43771#issuecomment-1982301654
|
|
"""
|
|
|
|
caller = CallerActor.remote(
|
|
EndpointActor.remote(
|
|
injected_executor_delay_s=injected_executor_delay_s,
|
|
tokens_per_request=100,
|
|
),
|
|
)
|
|
|
|
worker_state_before = [(a.worker_id, a.exit_type) for a in list_workers()]
|
|
print(">>> Workers state before: ", worker_state_before)
|
|
|
|
try:
|
|
ray.get(caller.run.remote())
|
|
except Exception as exc:
|
|
# There is a small chance that the task cancellation signal will arrive
|
|
# late at the executor, after the task has already finished. In that
|
|
# case, the task will complete normally, with no exception thrown.
|
|
# Thus, we wrap ray.get in a try-catch instead of asserting an
|
|
# exception.
|
|
assert isinstance(exc, RayTaskError)
|
|
assert isinstance(exc.cause, TaskCancelledError)
|
|
|
|
worker_state_after = [(a.worker_id, a.exit_type) for a in list_workers()]
|
|
print(">>> Workers state after: ", worker_state_after)
|
|
|
|
worker_ids, worker_exit_types = zip(*worker_state_after)
|
|
# Make sure no workers crashed
|
|
assert (
|
|
"SYSTEM_ERROR" not in worker_exit_types
|
|
), f"Unexpected crashed worker(s) in {worker_ids}"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|