chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,346 @@
|
||||
# coding: utf-8
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray._common.test_utils import (
|
||||
run_string_as_driver,
|
||||
wait_for_condition,
|
||||
)
|
||||
from ray._private.test_utils import run_string_as_driver_nonblocking
|
||||
from ray.util.state import list_workers
|
||||
from ray.util.state.common import WorkerState
|
||||
|
||||
|
||||
def get_workers() -> List[WorkerState]:
|
||||
"""Return non-driver workers."""
|
||||
return list_workers(
|
||||
filters=[("worker_type", "=", "WORKER"), ("is_alive", "=", "True")],
|
||||
raise_on_missing_output=False,
|
||||
)
|
||||
|
||||
|
||||
# Test that when `redis_address` and `job_config` is not set in
|
||||
# `ray.init(...)`, Raylet will start `num_cpus` Python workers for the driver.
|
||||
def test_initial_workers(shutdown_only):
|
||||
ray.init(num_cpus=2)
|
||||
wait_for_condition(lambda: len(get_workers()) == 2)
|
||||
|
||||
|
||||
# This test case starts some driver processes. Each driver process submits
|
||||
# some tasks and collect the PIDs of the workers used by the driver. The
|
||||
# drivers output the PID list which will be read by the test case itself. The
|
||||
# test case will compare the PIDs used by different drivers and make sure that
|
||||
# all the PIDs don't overlap. If overlapped, it means that tasks owned by
|
||||
# different drivers were scheduled to the same worker process, that is, tasks
|
||||
# of different jobs were not correctly isolated during execution.
|
||||
def test_multi_drivers(shutdown_only):
|
||||
info = ray.init(num_cpus=10)
|
||||
|
||||
driver_code = """
|
||||
import os
|
||||
import sys
|
||||
import ray
|
||||
|
||||
|
||||
ray.init(address="{}")
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def get_pid(self):
|
||||
return os.getpid()
|
||||
|
||||
@ray.remote
|
||||
def get_pid():
|
||||
return os.getpid()
|
||||
|
||||
pid_objs = []
|
||||
# Submit some normal tasks and get the PIDs of workers which execute the tasks.
|
||||
pid_objs = pid_objs + [get_pid.remote() for _ in range(2)]
|
||||
# Create some actors and get the PIDs of actors.
|
||||
actors = [Actor.remote() for _ in range(2)]
|
||||
pid_objs = pid_objs + [actor.get_pid.remote() for actor in actors]
|
||||
|
||||
pids = set([ray.get(obj) for obj in pid_objs])
|
||||
# Write pids to stdout
|
||||
print("PID:" + str.join(",", [str(_) for _ in pids]))
|
||||
|
||||
ray.shutdown()
|
||||
""".format(
|
||||
info["address"]
|
||||
)
|
||||
|
||||
driver_count = 3
|
||||
processes = [
|
||||
run_string_as_driver_nonblocking(driver_code) for _ in range(driver_count)
|
||||
]
|
||||
outputs = []
|
||||
for p in processes:
|
||||
out = p.stdout.read().decode("ascii")
|
||||
err = p.stderr.read().decode("ascii")
|
||||
p.wait()
|
||||
if p.returncode != 0:
|
||||
print(
|
||||
"Driver with PID {} returned error code {}".format(p.pid, p.returncode)
|
||||
)
|
||||
print("STDOUT:\n{}".format(out))
|
||||
print("STDERR:\n{}".format(err))
|
||||
outputs.append((p, out))
|
||||
|
||||
all_worker_pids = set()
|
||||
for p, out in outputs:
|
||||
assert p.returncode == 0
|
||||
for line in out.splitlines():
|
||||
if line.startswith("PID:"):
|
||||
worker_pids = [int(_) for _ in line.split(":")[1].split(",")]
|
||||
assert len(worker_pids) > 0
|
||||
for worker_pid in worker_pids:
|
||||
assert worker_pid not in all_worker_pids, (
|
||||
"Worker process with PID {} is shared" + " by multiple drivers."
|
||||
).format(worker_pid)
|
||||
all_worker_pids.add(worker_pid)
|
||||
|
||||
|
||||
class SignalFile:
|
||||
def __init__(self):
|
||||
self._tmpdir = tempfile.TemporaryDirectory()
|
||||
self._tmppath = os.path.join(self._tmpdir.name, "signal.txt")
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self._tmpdir.cleanup()
|
||||
|
||||
def wait(self):
|
||||
while not os.path.exists(self._tmppath):
|
||||
time.sleep(0.1)
|
||||
|
||||
def send(self):
|
||||
with open(self._tmppath, "w") as f:
|
||||
f.write("go!")
|
||||
f.flush()
|
||||
f.close()
|
||||
|
||||
|
||||
def test_kill_idle_workers(shutdown_only):
|
||||
# Avoid starting initial workers by setting num_cpus to 0.
|
||||
ray.init(num_cpus=0)
|
||||
assert len(get_workers()) == 0
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
class Actor:
|
||||
pass
|
||||
|
||||
# Worker 1 should be alive running the actor.
|
||||
a = Actor.remote()
|
||||
ray.get(a.__ray_ready__.remote())
|
||||
assert len(get_workers()) == 1
|
||||
|
||||
# NOTE(edoakes): I tried writing this test using a SignalActor instead of a file
|
||||
# to coordinate the tasks, but it failed because the idle workers weren't killed.
|
||||
with SignalFile() as signal:
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
def foo():
|
||||
signal.wait()
|
||||
|
||||
# Worker 2 should be alive running foo.
|
||||
obj1 = foo.remote()
|
||||
wait_for_condition(lambda: len(get_workers()) == 2)
|
||||
|
||||
# Worker 3 should be alive running foo.
|
||||
obj2 = foo.remote()
|
||||
wait_for_condition(lambda: len(get_workers()) == 3)
|
||||
|
||||
# Signal the tasks to unblock and wait for them to complete.
|
||||
signal.send()
|
||||
ray.get([obj1, obj2])
|
||||
|
||||
# Worker 2 and 3 now become idle and should be killed.
|
||||
wait_for_condition(lambda: len(get_workers()) == 1)
|
||||
|
||||
# Worker 1 should also be killed when the actor exits.
|
||||
del a
|
||||
wait_for_condition(lambda: len(get_workers()) == 0)
|
||||
|
||||
|
||||
def test_worker_capping_run_many_small_tasks(shutdown_only):
|
||||
ray.init(num_cpus=2)
|
||||
|
||||
with SignalFile() as signal:
|
||||
|
||||
@ray.remote(num_cpus=0.5)
|
||||
def foo():
|
||||
signal.wait()
|
||||
|
||||
# Run more tasks than `num_cpus`, but the CPU resource requirement is
|
||||
# still within `num_cpus`.
|
||||
obj_refs = [foo.remote() for _ in range(4)]
|
||||
wait_for_condition(lambda: len(get_workers()) == 4)
|
||||
|
||||
# Unblock the tasks.
|
||||
signal.send()
|
||||
ray.get(obj_refs)
|
||||
|
||||
# After the tasks finish, some workers are killed to keep the total
|
||||
# number of workers <= num_cpus.
|
||||
wait_for_condition(lambda: len(get_workers()) == 2)
|
||||
|
||||
# The two remaining workers stay alive forever.
|
||||
for _ in range(10):
|
||||
assert len(get_workers()) == 2
|
||||
|
||||
|
||||
def test_worker_capping_run_chained_tasks(shutdown_only):
|
||||
ray.init(num_cpus=2)
|
||||
|
||||
with SignalFile() as signal:
|
||||
|
||||
@ray.remote(num_cpus=0.5)
|
||||
def foo(x):
|
||||
if x > 1:
|
||||
return ray.get(foo.remote(x - 1)) + x
|
||||
else:
|
||||
signal.wait()
|
||||
return x
|
||||
|
||||
# Run a chain of tasks which exceed `num_cpus` in amount, but the CPU
|
||||
# resource requirement is still within `num_cpus`.
|
||||
obj = foo.remote(4)
|
||||
wait_for_condition(lambda: len(get_workers()) == 4)
|
||||
|
||||
# Unblock the tasks.
|
||||
signal.send()
|
||||
ray.get(obj)
|
||||
|
||||
# After finished the tasks, some workers are killed to keep the total
|
||||
# number of workers <= num_cpus.
|
||||
wait_for_condition(lambda: len(get_workers()) == 2)
|
||||
|
||||
# The two remaining workers stay alive forever.
|
||||
for _ in range(10):
|
||||
assert len(get_workers()) == 2
|
||||
|
||||
|
||||
def test_worker_registration_failure_after_driver_exit(shutdown_only):
|
||||
info = ray.init(num_cpus=2)
|
||||
wait_for_condition(lambda: len(get_workers()) == 2)
|
||||
|
||||
driver_code = """
|
||||
import os
|
||||
import ray
|
||||
import time
|
||||
|
||||
|
||||
ray.init(address="{}")
|
||||
|
||||
@ray.remote
|
||||
def foo():
|
||||
pass
|
||||
|
||||
obj_refs = [foo.remote() for _ in range(1000)]
|
||||
ray.get(obj_refs[0])
|
||||
os._exit(0)
|
||||
""".format(
|
||||
info["address"]
|
||||
)
|
||||
|
||||
# Run a driver that spawns many tasks and blocks until the first result is ready,
|
||||
# so at least one worker should have registered.
|
||||
try:
|
||||
run_string_as_driver(driver_code)
|
||||
except subprocess.CalledProcessError:
|
||||
# The driver exits with non-zero status Windows due to ungraceful os._exit.
|
||||
pass
|
||||
|
||||
# Verify that the workers spawned by the old driver go away.
|
||||
wait_for_condition(lambda: len(get_workers()) <= 2)
|
||||
|
||||
|
||||
def test_not_killing_workers_that_own_objects(shutdown_only):
|
||||
idle_worker_kill_interval_ms = 10
|
||||
|
||||
# Set the small interval for worker capping
|
||||
# so that we can easily trigger it.
|
||||
ray.init(
|
||||
num_cpus=0,
|
||||
_system_config={
|
||||
"kill_idle_workers_interval_ms": idle_worker_kill_interval_ms,
|
||||
},
|
||||
)
|
||||
|
||||
# Create a nested tasks to start 4 workers each of which owns an object.
|
||||
with SignalFile() as signal:
|
||||
expected_num_workers = 4
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
def nested(i):
|
||||
# Each of these tasks owns an object so it shouldn't be killed.
|
||||
if i >= expected_num_workers - 1:
|
||||
signal.wait()
|
||||
return [ray.put(np.ones(1 * 1024 * 1024, dtype=np.uint8))]
|
||||
else:
|
||||
return [ray.put(np.ones(1 * 1024 * 1024, dtype=np.uint8))] + ray.get(
|
||||
nested.remote(i + 1)
|
||||
)
|
||||
|
||||
# Wait for all the workers to start up.
|
||||
outer_ref = nested.remote(0)
|
||||
wait_for_condition(lambda: len(get_workers()) == expected_num_workers)
|
||||
|
||||
# Unblock the tasks.
|
||||
signal.send()
|
||||
inner_ref = ray.get(outer_ref)
|
||||
|
||||
# Sleep for 10x the idle worker kill interval and verify that those workers
|
||||
# aren't killed because they own objects that are in scope.
|
||||
time.sleep((10 * idle_worker_kill_interval_ms) / 1000.0)
|
||||
assert len(get_workers()) == expected_num_workers
|
||||
del inner_ref
|
||||
|
||||
|
||||
def test_kill_idle_workers_that_are_behind_owned_workers(shutdown_only):
|
||||
# When the first N idle workers own objects, and if we have N+N
|
||||
# total idle workers, we should make sure other N workers are killed.
|
||||
# It is because the idle workers are killed in the FIFO order.
|
||||
N = 4
|
||||
ray.init(
|
||||
num_cpus=1,
|
||||
_system_config={
|
||||
"kill_idle_workers_interval_ms": 10,
|
||||
"worker_lease_timeout_milliseconds": 0,
|
||||
},
|
||||
)
|
||||
|
||||
@ray.remote
|
||||
def nested(i):
|
||||
if i >= (N * 2) - 1:
|
||||
return [ray.put(np.ones(1 * 1024 * 1024, dtype=np.uint8))]
|
||||
elif i >= N:
|
||||
return [ray.put(np.ones(1 * 1024 * 1024, dtype=np.uint8))] + ray.get(
|
||||
nested.remote(i + 1)
|
||||
)
|
||||
else:
|
||||
return [1] + ray.get(nested.remote(i + 1))
|
||||
|
||||
# The first N workers don't own objects
|
||||
# and the later N workers do.
|
||||
ref = ray.get(nested.remote(0))
|
||||
assert len(ref) == N * 2
|
||||
num_workers = len(get_workers())
|
||||
assert num_workers == N * 2
|
||||
|
||||
# Make sure there are only N workers left after worker capping.
|
||||
wait_for_condition(lambda: len(get_workers()) == N)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
Reference in New Issue
Block a user