444 lines
13 KiB
Python
444 lines
13 KiB
Python
# coding: utf-8
|
|
import importlib
|
|
import logging
|
|
import os
|
|
import pickle
|
|
import socket
|
|
import sys
|
|
import time
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import ray
|
|
import ray._private.ray_constants
|
|
import ray._private.utils
|
|
from ray._private.test_utils import check_call_ray, wait_for_num_actors
|
|
from ray.util.state import list_actors
|
|
|
|
import psutil
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def test_global_state_api(shutdown_only):
|
|
|
|
ray.init(
|
|
num_cpus=5, num_gpus=3, resources={"CustomResource": 1}, include_dashboard=True
|
|
)
|
|
|
|
assert ray.cluster_resources()["CPU"] == 5
|
|
assert ray.cluster_resources()["GPU"] == 3
|
|
assert ray.cluster_resources()["CustomResource"] == 1
|
|
|
|
job_id = ray._private.utils.compute_job_id_from_driver(
|
|
ray.WorkerID(ray._private.worker.global_worker.worker_id)
|
|
)
|
|
|
|
client_table = ray.nodes()
|
|
node_ip_address = ray._private.worker.global_worker.node_ip_address
|
|
|
|
assert len(client_table) == 1
|
|
assert client_table[0]["NodeManagerAddress"] == node_ip_address
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
_ = Actor.options(name="test_actor").remote() # noqa: F841
|
|
# Wait for actor to be created
|
|
wait_for_num_actors(1)
|
|
|
|
actor_table = list_actors() # should be using this API now for fetching actors
|
|
assert len(actor_table) == 1
|
|
|
|
actor_info = actor_table[0]
|
|
assert actor_info.job_id == job_id.hex()
|
|
assert actor_info.name == "test_actor"
|
|
|
|
job_table = ray._private.state.jobs()
|
|
|
|
assert len(job_table) == 1
|
|
assert job_table[0]["JobID"] == job_id.hex()
|
|
assert job_table[0]["DriverIPAddress"] == node_ip_address
|
|
|
|
|
|
def test_logging_to_driver(capsys, shutdown_only):
|
|
ray.init(num_cpus=1, log_to_driver=True)
|
|
|
|
@ray.remote
|
|
def f():
|
|
# It's important to make sure that these print statements occur even
|
|
# without calling sys.stdout.flush() and sys.stderr.flush().
|
|
for i in range(10):
|
|
print(i, end=" ")
|
|
print(100 + i, end=" ", file=sys.stderr)
|
|
|
|
ray.get(f.remote())
|
|
time.sleep(1)
|
|
|
|
out, err = capsys.readouterr()
|
|
for i in range(10):
|
|
assert str(i) in out
|
|
|
|
for i in range(100, 110):
|
|
assert str(i) in err
|
|
|
|
|
|
def test_not_logging_to_driver_via_env_var(monkeypatch, capsys, shutdown_only):
|
|
monkeypatch.setenv("RAY_LOG_TO_DRIVER", "0")
|
|
importlib.reload(ray._private.ray_constants)
|
|
ray.init(num_cpus=1)
|
|
|
|
@ray.remote
|
|
def f():
|
|
for i in range(100):
|
|
print(i)
|
|
print(100 + i, file=sys.stderr)
|
|
sys.stdout.flush()
|
|
sys.stderr.flush()
|
|
|
|
capsys.readouterr()
|
|
ray.get(f.remote())
|
|
time.sleep(1)
|
|
|
|
out, err = capsys.readouterr()
|
|
assert len(out) == 0
|
|
assert len(err) == 0
|
|
|
|
|
|
def test_not_logging_to_driver(capsys, shutdown_only):
|
|
ray.init(num_cpus=1, log_to_driver=False)
|
|
|
|
@ray.remote
|
|
def f():
|
|
for i in range(100):
|
|
print(i)
|
|
print(100 + i, file=sys.stderr)
|
|
sys.stdout.flush()
|
|
sys.stderr.flush()
|
|
|
|
capsys.readouterr()
|
|
ray.get(f.remote())
|
|
time.sleep(1)
|
|
|
|
out, err = capsys.readouterr()
|
|
assert len(out) == 0
|
|
assert len(err) == 0
|
|
|
|
|
|
def test_workers(shutdown_only):
|
|
num_workers = 3
|
|
ray.init(num_cpus=num_workers)
|
|
|
|
@ray.remote
|
|
def f():
|
|
return id(ray._private.worker.global_worker), os.getpid()
|
|
|
|
# Wait until all of the workers have started.
|
|
worker_ids = set()
|
|
while len(worker_ids) != num_workers:
|
|
worker_ids = set(ray.get([f.remote() for _ in range(10)]))
|
|
|
|
|
|
def test_object_ref_properties():
|
|
id_bytes = b"0011223344556677889900001111"
|
|
object_ref = ray.ObjectRef(id_bytes)
|
|
assert object_ref.binary() == id_bytes
|
|
object_ref = ray.ObjectRef.nil()
|
|
assert object_ref.is_nil()
|
|
with pytest.raises(ValueError, match=r".*needs to have length.*"):
|
|
ray.ObjectRef(id_bytes + b"1234")
|
|
with pytest.raises(ValueError, match=r".*needs to have length.*"):
|
|
ray.ObjectRef(b"0123456789")
|
|
object_ref = ray.ObjectRef.from_random()
|
|
assert not object_ref.is_nil()
|
|
assert object_ref.binary() != id_bytes
|
|
id_dumps = pickle.dumps(object_ref)
|
|
id_from_dumps = pickle.loads(id_dumps)
|
|
assert id_from_dumps == object_ref
|
|
|
|
|
|
def test_wait_reconstruction(shutdown_only):
|
|
ray.init(num_cpus=1, object_store_memory=int(10**8))
|
|
|
|
@ray.remote
|
|
def f():
|
|
return np.zeros(6 * 10**7, dtype=np.uint8)
|
|
|
|
x_id = f.remote()
|
|
ray.wait([x_id])
|
|
ray.wait([f.remote()])
|
|
assert not ray._private.worker.global_worker.core_worker.object_exists(x_id)
|
|
ready_ids, _ = ray.wait([x_id])
|
|
assert len(ready_ids) == 1
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="Windows doesn't support changing process title."
|
|
)
|
|
def test_ray_setproctitle(ray_start_2_cpus):
|
|
@ray.remote
|
|
class UniqueName:
|
|
def __init__(self):
|
|
assert psutil.Process().cmdline()[0] == "ray::UniqueName.__init__"
|
|
|
|
def f(self):
|
|
assert psutil.Process().cmdline()[0] == "ray::UniqueName.f"
|
|
|
|
@ray.remote
|
|
def unique_1():
|
|
assert psutil.Process().cmdline()[0] == "ray::unique_1"
|
|
|
|
actor = UniqueName.remote()
|
|
ray.get(actor.f.remote())
|
|
ray.get(unique_1.remote())
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="Windows doesn't support changing process title."
|
|
)
|
|
def test_ray_task_name_setproctitle(ray_start_2_cpus):
|
|
method_task_name = "foo"
|
|
|
|
@ray.remote
|
|
class UniqueName:
|
|
def __init__(self):
|
|
assert psutil.Process().cmdline()[0] == "ray::UniqueName.__init__"
|
|
|
|
def f(self):
|
|
assert psutil.Process().cmdline()[0] == f"ray::{method_task_name}"
|
|
|
|
task_name = "bar"
|
|
|
|
@ray.remote
|
|
def unique_1():
|
|
assert psutil.Process().cmdline()[0] == f"ray::{task_name}"
|
|
|
|
actor = UniqueName.remote()
|
|
ray.get(actor.f.options(name=method_task_name).remote())
|
|
ray.get(unique_1.options(name=task_name).remote())
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="Windows doesn't support changing process title."
|
|
)
|
|
def test_ray_task_generator_setproctitle(ray_start_2_cpus):
|
|
@ray.remote
|
|
def generator_task():
|
|
for i in range(4):
|
|
assert psutil.Process().cmdline()[0] == "ray::generator_task"
|
|
yield i
|
|
|
|
ray.get(generator_task.options(num_returns=2).remote()[0])
|
|
ray.get(generator_task.options(num_returns="dynamic").remote())
|
|
generator = generator_task.remote()
|
|
for _ in range(4):
|
|
ray.get(next(generator))
|
|
|
|
@ray.remote
|
|
class UniqueName:
|
|
def f(self):
|
|
for i in range(4):
|
|
assert psutil.Process().cmdline()[0] == "ray::UniqueName.f"
|
|
yield i
|
|
|
|
actor = UniqueName.remote()
|
|
ray.get(actor.f.options(num_returns=2).remote()[0])
|
|
ray.get(actor.f.options(num_returns="dynamic").remote())
|
|
generator = actor.f.remote()
|
|
for _ in range(4):
|
|
ray.get(next(generator))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux",
|
|
reason=(
|
|
"Test specifically targets the Linux /proc/self/stat fallback in "
|
|
"src/ray/thirdparty/setproctitle/spt_setup.c."
|
|
),
|
|
)
|
|
def test_setproctitle_falls_back_to_proc_stat_when_environ_broken(tmp_path):
|
|
"""Regression test for the vendored-setproctitle deflake.
|
|
|
|
spt_setup() in src/ray/thirdparty/setproctitle/spt_setup.c originally
|
|
relied on find_argv_from_env(), which walks backward from environ[0] to
|
|
locate the original argv memory region. That walk silently fails once
|
|
libc setenv()/putenv() has moved environ, leaving every subsequent
|
|
setproctitle() call as a no-op. We added find_argv_from_proc_stat() as
|
|
a fallback that reads the kernel-recorded arg_start/arg_end from
|
|
/proc/self/stat (set at execve() time, immune to env mutation).
|
|
|
|
This test deliberately clears os.environ in a subprocess BEFORE the
|
|
first setproctitle() call. Without the fallback, env[0] is NULL,
|
|
find_argv_from_env() fails, and the title write is silently dropped.
|
|
With the fallback in place, the kernel-recorded bounds are used and
|
|
/proc/self/cmdline reflects the new title.
|
|
"""
|
|
import subprocess
|
|
import textwrap
|
|
|
|
marker = "ray::PROC_STAT_FALLBACK_OK"
|
|
script = textwrap.dedent(
|
|
f"""
|
|
import os, sys
|
|
# Import _raylet BEFORE wiping environ — the .so load itself may
|
|
# need a sane environment, and spt_setup() is lazy (only fires on
|
|
# the first setproctitle() call), so we can safely break environ
|
|
# afterwards.
|
|
import ray._raylet
|
|
|
|
# Break find_argv_from_env(): environ[0] becomes NULL once cleared,
|
|
# so the backward walk bails out immediately.
|
|
os.environ.clear()
|
|
|
|
ray._raylet.setproctitle({marker!r})
|
|
|
|
with open("/proc/self/cmdline", "rb") as f:
|
|
cmdline = f.read()
|
|
# /proc/self/cmdline is NUL-separated; the title is written into
|
|
# argv[0]'s buffer and padded with NULs.
|
|
first_arg = cmdline.split(b"\\x00", 1)[0].decode("utf-8", "replace")
|
|
print("FIRST_ARG=" + first_arg)
|
|
"""
|
|
)
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", script],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=60,
|
|
)
|
|
assert (
|
|
result.returncode == 0
|
|
), f"subprocess failed: stdout={result.stdout!r} stderr={result.stderr!r}"
|
|
# Locate our marker line in stdout (other Ray imports may emit chatter).
|
|
first_arg_lines = [
|
|
line for line in result.stdout.splitlines() if line.startswith("FIRST_ARG=")
|
|
]
|
|
assert first_arg_lines, (
|
|
f"FIRST_ARG line missing from subprocess output. "
|
|
f"stdout={result.stdout!r} stderr={result.stderr!r}"
|
|
)
|
|
first_arg = first_arg_lines[-1][len("FIRST_ARG=") :]
|
|
assert first_arg == marker, (
|
|
f"setproctitle did not take effect after os.environ.clear(): "
|
|
f"expected {marker!r}, got {first_arg!r}. "
|
|
f"This typically means the /proc/self/stat fallback in "
|
|
f"src/ray/thirdparty/setproctitle/spt_setup.c regressed."
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
os.getenv("TRAVIS") is None, reason="This test should only be run on Travis."
|
|
)
|
|
def test_ray_stack(ray_start_2_cpus):
|
|
def unique_name_1():
|
|
time.sleep(1000)
|
|
|
|
@ray.remote
|
|
def unique_name_2():
|
|
time.sleep(1000)
|
|
|
|
@ray.remote
|
|
def unique_name_3():
|
|
unique_name_1()
|
|
|
|
unique_name_2.remote()
|
|
unique_name_3.remote()
|
|
|
|
success = False
|
|
start_time = time.time()
|
|
while time.time() - start_time < 30:
|
|
# Attempt to parse the "ray stack" call.
|
|
output = ray._common.utils.decode(
|
|
check_call_ray(["stack"], capture_stdout=True)
|
|
)
|
|
if (
|
|
"unique_name_1" in output
|
|
and "unique_name_2" in output
|
|
and "unique_name_3" in output
|
|
):
|
|
success = True
|
|
break
|
|
|
|
if not success:
|
|
raise Exception("Failed to find necessary information with 'ray stack'")
|
|
|
|
|
|
def test_raylet_is_robust_to_random_messages(ray_start_regular):
|
|
node_manager_address = None
|
|
node_manager_port = None
|
|
for client in ray.nodes():
|
|
if "NodeManagerAddress" in client:
|
|
node_manager_address = client["NodeManagerAddress"]
|
|
node_manager_port = client["NodeManagerPort"]
|
|
assert node_manager_address
|
|
assert node_manager_port
|
|
# Try to bring down the node manager:
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
s.connect((node_manager_address, node_manager_port))
|
|
s.send(1000 * b"asdf")
|
|
|
|
@ray.remote
|
|
def f():
|
|
return 1
|
|
|
|
assert ray.get(f.remote()) == 1
|
|
|
|
|
|
def test_non_ascii_comment(ray_start_regular):
|
|
@ray.remote
|
|
def f():
|
|
# 日本語 Japanese comment
|
|
return 1
|
|
|
|
assert ray.get(f.remote()) == 1
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_start_object_store_memory", [150 * 1024 * 1024], indirect=True
|
|
)
|
|
def test_put_pins_object(ray_start_object_store_memory):
|
|
obj = np.ones(200 * 1024, dtype=np.uint8)
|
|
x_id = ray.put(obj)
|
|
x_binary = x_id.binary()
|
|
assert (ray.get(ray.ObjectRef(x_binary)) == obj).all()
|
|
|
|
# x cannot be evicted since x_id pins it
|
|
for _ in range(10):
|
|
ray.put(np.zeros(10 * 1024 * 1024))
|
|
assert (ray.get(x_id) == obj).all()
|
|
assert (ray.get(ray.ObjectRef(x_binary)) == obj).all()
|
|
|
|
# now it can be evicted since x_id pins it but x_binary does not
|
|
del x_id
|
|
for _ in range(10):
|
|
ray.put(np.zeros(10 * 1024 * 1024))
|
|
assert not ray._private.worker.global_worker.core_worker.object_exists(
|
|
ray.ObjectRef(x_binary)
|
|
)
|
|
|
|
|
|
def test_decorated_function(ray_start_regular):
|
|
def function_invocation_decorator(f):
|
|
def new_f(args, kwargs):
|
|
# Reverse the arguments.
|
|
return f(args[::-1], {"d": 5}), kwargs
|
|
|
|
return new_f
|
|
|
|
def f(a, b, c, d=None):
|
|
return a, b, c, d
|
|
|
|
f.__ray_invocation_decorator__ = function_invocation_decorator
|
|
f = ray.remote(f)
|
|
|
|
result_id, kwargs = f.remote(1, 2, 3, d=4)
|
|
assert kwargs == {"d": 4}
|
|
assert ray.get(result_id) == (3, 2, 1, 5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|