chore: import upstream snapshot with attribution

python/ray/tests/test_advanced_3.py

@@ -0,0 +1,443 @@
+# 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__]))