Files
ray-project--ray/python/ray/tests/test_advanced_3.py
T
2026-07-13 13:17:40 +08:00

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__]))