wehub-resource-sync
548 lines
14 KiB
Python
548 lines
14 KiB
Python
import collections
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from typing import List, Optional
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import (
|
|
run_string_as_driver,
|
|
wait_for_condition,
|
|
)
|
|
from ray._private.runtime_env.context import RuntimeEnvContext
|
|
from ray._private.runtime_env.plugin import RuntimeEnvPlugin
|
|
from ray._private.test_utils import (
|
|
get_load_metrics_report,
|
|
get_resource_usage,
|
|
run_string_as_driver_nonblocking,
|
|
)
|
|
from ray._private.utils import get_num_cpus
|
|
from ray.util.state import list_objects, list_workers
|
|
|
|
|
|
# This tests the queue transitions for infeasible tasks. This has been an issue
|
|
# in the past, e.g., https://github.com/ray-project/ray/issues/3275.
|
|
def test_infeasible_tasks(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
|
|
@ray.remote
|
|
def f():
|
|
return
|
|
|
|
cluster.add_node(resources={str(0): 100})
|
|
ray.init(address=cluster.address)
|
|
|
|
# Submit an infeasible task.
|
|
x_id = f._remote(args=[], kwargs={}, resources={str(1): 1})
|
|
|
|
# Add a node that makes the task feasible and make sure we can get the
|
|
# result.
|
|
cluster.add_node(resources={str(1): 100})
|
|
ray.get(x_id)
|
|
|
|
# Start a driver that submits an infeasible task and then let it exit.
|
|
driver_script = """
|
|
import ray
|
|
|
|
ray.init(address="{}")
|
|
|
|
@ray.remote(resources={})
|
|
def f():
|
|
{}pass # This is a weird hack to insert some blank space.
|
|
|
|
f.remote()
|
|
""".format(
|
|
cluster.address, "{str(2): 1}", " "
|
|
)
|
|
|
|
run_string_as_driver(driver_script)
|
|
|
|
# Now add a new node that makes the task feasible.
|
|
cluster.add_node(resources={str(2): 100})
|
|
|
|
# Make sure we can still run tasks on all nodes.
|
|
ray.get([f._remote(args=[], kwargs={}, resources={str(i): 1}) for i in range(3)])
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
["""ray start --head"""],
|
|
indirect=True,
|
|
)
|
|
def test_kill_driver_clears_backlog(call_ray_start):
|
|
driver = """
|
|
import ray
|
|
|
|
@ray.remote
|
|
def f():
|
|
import time
|
|
time.sleep(300)
|
|
|
|
refs = [f.remote() for _ in range(10000)]
|
|
|
|
ray.get(refs)
|
|
"""
|
|
proc = run_string_as_driver_nonblocking(driver)
|
|
ctx = ray.init(address=call_ray_start)
|
|
|
|
def get_backlog_and_pending():
|
|
resources_batch = get_resource_usage(
|
|
gcs_address=ctx.address_info["gcs_address"]
|
|
)
|
|
backlog = (
|
|
resources_batch.resource_load_by_shape.resource_demands[0].backlog_size
|
|
if resources_batch.resource_load_by_shape.resource_demands
|
|
else 0
|
|
)
|
|
|
|
pending = 0
|
|
demands = get_load_metrics_report(webui_url=ctx.address_info["webui_url"])[
|
|
"resourceDemand"
|
|
]
|
|
for demand in demands:
|
|
resource_dict, amount = demand
|
|
if "CPU" in resource_dict:
|
|
pending = amount
|
|
|
|
return pending, backlog
|
|
|
|
def check_backlog(expect_backlog) -> bool:
|
|
pending, backlog = get_backlog_and_pending()
|
|
if expect_backlog:
|
|
return pending > 0 and backlog > 0
|
|
else:
|
|
return pending == 0 and backlog == 0
|
|
|
|
wait_for_condition(
|
|
check_backlog, timeout=10, retry_interval_ms=1000, expect_backlog=True
|
|
)
|
|
|
|
os.kill(proc.pid, 9)
|
|
|
|
wait_for_condition(
|
|
check_backlog, timeout=10, retry_interval_ms=1000, expect_backlog=False
|
|
)
|
|
|
|
|
|
def get_infeasible_queued(ray_ctx):
|
|
resources_batch = get_resource_usage(
|
|
gcs_address=ray_ctx.address_info["gcs_address"]
|
|
)
|
|
|
|
infeasible_queued = (
|
|
resources_batch.resource_load_by_shape.resource_demands[
|
|
0
|
|
].num_infeasible_requests_queued
|
|
if len(resources_batch.resource_load_by_shape.resource_demands) > 0
|
|
and hasattr(
|
|
resources_batch.resource_load_by_shape.resource_demands[0],
|
|
"num_infeasible_requests_queued",
|
|
)
|
|
else 0
|
|
)
|
|
|
|
return infeasible_queued
|
|
|
|
|
|
def check_infeasible(expect_infeasible, ray_ctx) -> bool:
|
|
infeasible_queued = get_infeasible_queued(ray_ctx)
|
|
if expect_infeasible:
|
|
return infeasible_queued > 0
|
|
else:
|
|
return infeasible_queued == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
["""ray start --head"""],
|
|
indirect=True,
|
|
)
|
|
def test_kill_driver_clears_infeasible(call_ray_start):
|
|
driver = """
|
|
import ray
|
|
|
|
@ray.remote
|
|
def f():
|
|
pass
|
|
|
|
ray.get(f.options(num_cpus=99999999).remote())
|
|
"""
|
|
proc = run_string_as_driver_nonblocking(driver)
|
|
ctx = ray.init(address=call_ray_start)
|
|
|
|
wait_for_condition(
|
|
check_infeasible,
|
|
timeout=10,
|
|
retry_interval_ms=1000,
|
|
expect_infeasible=True,
|
|
ray_ctx=ctx,
|
|
)
|
|
|
|
os.kill(proc.pid, 9)
|
|
|
|
wait_for_condition(
|
|
check_infeasible,
|
|
timeout=10,
|
|
retry_interval_ms=1000,
|
|
expect_infeasible=False,
|
|
ray_ctx=ctx,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
["ray start --head --ray-client-server-port=25555"],
|
|
indirect=True,
|
|
)
|
|
def test_exiting_driver_clears_infeasible(call_ray_start):
|
|
# Test that there is no leaking infeasible demands
|
|
# from an exited driver.
|
|
# See https://github.com/ray-project/ray/issues/43687
|
|
# for a bug where it happened.
|
|
|
|
driver = """
|
|
import ray
|
|
|
|
ray.init()
|
|
|
|
@ray.remote
|
|
def f():
|
|
pass
|
|
|
|
f.options(num_cpus=99999999).remote()
|
|
"""
|
|
proc = run_string_as_driver_nonblocking(driver)
|
|
proc.wait()
|
|
|
|
client_driver = """
|
|
import ray
|
|
|
|
ray.init("ray://127.0.0.1:25555")
|
|
|
|
@ray.remote
|
|
def f():
|
|
pass
|
|
|
|
f.options(num_cpus=99999999).remote()
|
|
"""
|
|
proc = run_string_as_driver_nonblocking(client_driver)
|
|
proc.wait()
|
|
|
|
ctx = ray.init(address=call_ray_start)
|
|
|
|
# Give gcs some time to update the load
|
|
time.sleep(1)
|
|
|
|
wait_for_condition(
|
|
check_infeasible,
|
|
timeout=10,
|
|
retry_interval_ms=1000,
|
|
expect_infeasible=False,
|
|
ray_ctx=ctx,
|
|
)
|
|
|
|
|
|
def test_kill_driver_keep_infeasible_detached_actor(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
address = cluster.address
|
|
|
|
cluster.add_node(num_cpus=1)
|
|
|
|
driver_script = """
|
|
import ray
|
|
|
|
@ray.remote
|
|
class A:
|
|
def fn(self):
|
|
pass
|
|
|
|
ray.init(address="{}", namespace="test_det")
|
|
|
|
ray.get(A.options(num_cpus=123, name="det", lifetime="detached").remote())
|
|
""".format(
|
|
cluster.address
|
|
)
|
|
|
|
proc = run_string_as_driver_nonblocking(driver_script)
|
|
|
|
ctx = ray.init(address=address, namespace="test_det")
|
|
|
|
wait_for_condition(
|
|
check_infeasible,
|
|
timeout=10,
|
|
retry_interval_ms=1000,
|
|
expect_infeasible=True,
|
|
ray_ctx=ctx,
|
|
)
|
|
|
|
os.kill(proc.pid, 9)
|
|
|
|
cluster.add_node(num_cpus=200)
|
|
|
|
det_actor = ray.get_actor("det")
|
|
|
|
ray.get(det_actor.fn.remote())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
["""ray start --head"""],
|
|
indirect=True,
|
|
)
|
|
def test_reference_global_import_does_not_leak_worker_upon_driver_exit(call_ray_start):
|
|
driver = """
|
|
import ray
|
|
import numpy as np
|
|
import tensorflow
|
|
|
|
@ray.remote(max_retries=0)
|
|
def leak_repro(obj):
|
|
tensorflow
|
|
return []
|
|
|
|
refs = []
|
|
for i in range(100_000):
|
|
refs.append(leak_repro.remote(i))
|
|
|
|
ray.get(refs)
|
|
"""
|
|
try:
|
|
run_string_as_driver(driver)
|
|
except subprocess.CalledProcessError:
|
|
pass
|
|
|
|
ray.init(address=call_ray_start)
|
|
|
|
def no_object_leaks():
|
|
objects = list_objects(_explain=True, timeout=3)
|
|
return len(objects) == 0
|
|
|
|
wait_for_condition(no_object_leaks, timeout=10, retry_interval_ms=1000)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="subprocess command only works for unix"
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
["""ray start --head --system-config={"enable_worker_prestart":true}"""],
|
|
indirect=True,
|
|
)
|
|
def test_worker_prestart_on_node_manager_start(call_ray_start, shutdown_only):
|
|
def num_idle_workers(count):
|
|
result = subprocess.check_output(
|
|
"ps aux | grep ray::IDLE | grep -v grep",
|
|
shell=True,
|
|
)
|
|
return len(result.splitlines()) == count
|
|
|
|
wait_for_condition(num_idle_workers, count=get_num_cpus())
|
|
|
|
with ray.init():
|
|
for _ in range(5):
|
|
workers = list_workers(
|
|
filters=[("worker_type", "=", "WORKER")], raise_on_missing_output=False
|
|
)
|
|
assert len(workers) == get_num_cpus(), workers
|
|
time.sleep(1)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"call_ray_start",
|
|
["""ray start --head"""],
|
|
indirect=True,
|
|
)
|
|
def test_jobs_prestart_worker_once(call_ray_start, shutdown_only):
|
|
with ray.init():
|
|
workers = list_workers(
|
|
filters=[("worker_type", "=", "WORKER")], raise_on_missing_output=False
|
|
)
|
|
assert len(workers) == get_num_cpus(), workers
|
|
with ray.init():
|
|
for _ in range(5):
|
|
workers = list_workers(
|
|
filters=[("worker_type", "=", "WORKER")], raise_on_missing_output=False
|
|
)
|
|
assert len(workers) == get_num_cpus(), workers
|
|
time.sleep(1)
|
|
|
|
|
|
def test_can_use_prestart_idle_workers(ray_start_cluster):
|
|
"""Test that actors and GPU tasks can use prestarted workers."""
|
|
cluster = ray_start_cluster
|
|
NUM_CPUS = 4
|
|
NUM_GPUS = 4
|
|
cluster.add_node(num_cpus=NUM_CPUS, num_gpus=NUM_GPUS)
|
|
ray.init(address=cluster.address)
|
|
|
|
wait_for_condition(
|
|
lambda: len(
|
|
list_workers(
|
|
filters=[("worker_type", "=", "WORKER")], raise_on_missing_output=False
|
|
)
|
|
)
|
|
== NUM_CPUS
|
|
)
|
|
|
|
# These workers don't have job_id or is_actor_worker.
|
|
workers = list_workers(
|
|
filters=[("worker_type", "=", "WORKER")],
|
|
detail=True,
|
|
raise_on_missing_output=False,
|
|
)
|
|
worker_pids = {worker.pid for worker in workers}
|
|
assert len(worker_pids) == NUM_CPUS
|
|
|
|
@ray.remote
|
|
class A:
|
|
def getpid(self):
|
|
return os.getpid()
|
|
|
|
@ray.remote
|
|
def f():
|
|
return os.getpid()
|
|
|
|
used_worker_pids = set()
|
|
cpu_actor = A.options(num_cpus=1).remote()
|
|
used_worker_pids.add(ray.get(cpu_actor.getpid.remote()))
|
|
|
|
gpu_actor = A.options(num_gpus=1).remote()
|
|
used_worker_pids.add(ray.get(gpu_actor.getpid.remote()))
|
|
|
|
used_worker_pids.add(ray.get(f.options(num_cpus=1).remote()))
|
|
used_worker_pids.add(ray.get(f.options(num_gpus=1).remote()))
|
|
|
|
assert used_worker_pids == worker_pids
|
|
|
|
|
|
MyPlugin = "HangOnSecondWorkerPlugin"
|
|
MY_PLUGIN_CLASS_PATH = "ray.tests.test_node_manager.HangOnSecondWorkerPlugin"
|
|
PLUGIN_TIMEOUT = 10
|
|
|
|
|
|
class HangOnSecondWorkerPlugin(RuntimeEnvPlugin):
|
|
"""
|
|
The first worker will start up normally, but all subsequent workers will hang at
|
|
start up indefinitely. How it works: Ray RuntimeEnvAgent caches the modified context
|
|
so we can't do it in modify_context. Instead, we use a bash command to read a file
|
|
and hang forever. We don't have a good file lock mechanism in bash (flock is not
|
|
installed by default in macos), so we also serialize the worker startup.
|
|
"""
|
|
|
|
name = MyPlugin
|
|
|
|
def __init__(self):
|
|
# Each URI has a temp dir, a counter file, and a hang.sh script.
|
|
self.uris = collections.defaultdict(dict)
|
|
|
|
def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]: # noqa: F821
|
|
return [runtime_env[self.name]]
|
|
|
|
async def create(
|
|
self,
|
|
uri: Optional[str],
|
|
runtime_env,
|
|
context: RuntimeEnvContext,
|
|
logger: logging.Logger,
|
|
) -> float:
|
|
d = self.uris[uri]
|
|
d["temp_dir"] = tempfile.mkdtemp()
|
|
logger.info(f"caching temp dir {d['temp_dir']} for uri {uri}")
|
|
d["counter_file"] = os.path.join(d["temp_dir"], "script_run_count")
|
|
with open(d["counter_file"], "w+") as f:
|
|
f.write("0")
|
|
d["hang_sh"] = os.path.join(d["temp_dir"], "hang.sh")
|
|
with open(d["hang_sh"], "w+") as f:
|
|
f.write(
|
|
f"""#!/bin/bash
|
|
|
|
counter_file="{d['counter_file']}"
|
|
|
|
count=$(cat "$counter_file")
|
|
|
|
if [ "$count" -eq "0" ]; then
|
|
echo "1" > "$counter_file"
|
|
echo "first time run"
|
|
exit 0
|
|
elif [ "$count" -eq "1" ]; then
|
|
echo "2" > "$counter_file"
|
|
echo "second time run, sleeping..."
|
|
sleep 1000
|
|
fi
|
|
"""
|
|
)
|
|
os.chmod(d["hang_sh"], 0o755)
|
|
return 0.1
|
|
|
|
def modify_context(
|
|
self,
|
|
uris: List[str],
|
|
runtime_env: "RuntimeEnv", # noqa: F821
|
|
ctx: RuntimeEnvContext,
|
|
logger: logging.Logger,
|
|
) -> None:
|
|
logger.info(f"Starting worker: {uris}, {runtime_env}")
|
|
if self.name not in runtime_env:
|
|
return
|
|
assert len(uris) == 1
|
|
uri = uris[0]
|
|
hang_sh = self.uris[uri]["hang_sh"]
|
|
ctx.command_prefix += ["bash", hang_sh, "&&"]
|
|
|
|
def delete_uri(self, uri: str, logger: logging.Logger) -> float:
|
|
temp_dir = self.uris[uri]["temp_dir"]
|
|
shutil.rmtree(temp_dir)
|
|
del self.uris[uri]
|
|
logger.info(f"temp_dir removed: {temp_dir}")
|
|
|
|
|
|
@pytest.fixture
|
|
def serialize_worker_startup(monkeypatch):
|
|
"""Only one worker starts up each time, since our bash script is not process-safe"""
|
|
monkeypatch.setenv("RAY_worker_maximum_startup_concurrency", "1")
|
|
yield
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"set_runtime_env_plugins",
|
|
[
|
|
'[{"class":"' + MY_PLUGIN_CLASS_PATH + '"}]',
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_can_reuse_released_workers(
|
|
serialize_worker_startup, set_runtime_env_plugins, ray_start_cluster
|
|
):
|
|
"""
|
|
Uses a runtime env plugin to make sure only 1 worker can start and all subsequent
|
|
workers will hang in runtime start up forever. We issue 10 tasks and test that
|
|
all the following tasks can still be scheduled on the first worker released from the
|
|
first task, i.e. tasks are not binded to the workers that they requested to start.
|
|
"""
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=2)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(runtime_env={"env_vars": {"HELLO": "WORLD"}, MyPlugin: "key"})
|
|
def f():
|
|
# Sleep for a while to make sure other tasks also request workers.
|
|
time.sleep(1)
|
|
print(f"pid={os.getpid()}, env HELLO={os.environ.get('HELLO')}")
|
|
return os.getpid()
|
|
|
|
objs = [f.remote() for i in range(10)]
|
|
|
|
pids = ray.get(objs)
|
|
for pid in pids:
|
|
assert pid == pids[0]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|