chore: import upstream snapshot with attribution

python/ray/tests/test_ray_init.py

@@ -0,0 +1,377 @@
+import json
+import os
+import signal
+import subprocess
+import sys
+import tempfile
+import unittest.mock
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import grpc
+import pytest
+
+import ray
+from ray._common.network_utils import build_address, parse_address
+from ray._private import ray_constants
+from ray._private.test_utils import persistent_gcs_test_enabled
+from ray.client_builder import ClientContext
+from ray.cluster_utils import Cluster
+from ray.runtime_env.runtime_env import RuntimeEnv
+from ray.util.client.common import ClientObjectRef
+from ray.util.client.ray_client_helpers import ray_start_client_server
+from ray.util.client.worker import Worker
+
+
+@pytest.mark.skipif(
+    os.environ.get("CI") and sys.platform == "win32",
+    reason="Flaky when run on windows CI",
+)
+@pytest.mark.parametrize("input", [None, "auto"])
+def test_ray_address(input, call_ray_start):
+    address = call_ray_start
+    with unittest.mock.patch.dict(os.environ, {"RAY_ADDRESS": address}):
+        res = ray.init(input)
+        # Ensure this is not a client.connect()
+        assert not isinstance(res, ClientContext)
+        assert res.address_info["gcs_address"] == address
+        ray.shutdown()
+
+    addr = f"localhost:{parse_address(address)[-1]}"
+    with unittest.mock.patch.dict(os.environ, {"RAY_ADDRESS": addr}):
+        res = ray.init(input)
+        # Ensure this is not a client.connect()
+        assert not isinstance(res, ClientContext)
+        assert res.address_info["gcs_address"] == address
+        ray.shutdown()
+
+
+@pytest.mark.parametrize("address", [None, "auto"])
+def test_ray_init_no_local_instance(shutdown_only, address):
+    # Starts a new Ray instance.
+    if address is None:
+        ray.init(address=address)
+    else:
+        # Throws an error if we explicitly want to connect to an existing
+        # instance and none exists.
+        with pytest.raises(ConnectionError):
+            ray.init(address=address)
+
+
+@pytest.mark.skipif(
+    os.environ.get("CI") and sys.platform == "win32",
+    reason="Flaky when run on windows CI",
+)
+@pytest.mark.parametrize("address", [None, "auto"])
+def test_ray_init_existing_instance(call_ray_start, address):
+    ray_address = call_ray_start
+    # If no address is specified, we will default to an existing cluster.
+    res = ray.init(address=address)
+    assert res.address_info["gcs_address"] == ray_address
+    ray.shutdown()
+
+    # Start a second local Ray instance.
+    try:
+        subprocess.check_output("ray start --head", shell=True)
+        # If there are multiple local instances, connect to the latest.
+        res = ray.init(address=address)
+        assert res.address_info["gcs_address"] != ray_address
+        ray.shutdown()
+
+        # If there are multiple local instances and we specify an address
+        # explicitly, it works.
+        with unittest.mock.patch.dict(os.environ, {"RAY_ADDRESS": ray_address}):
+            res = ray.init(address=address)
+            assert res.address_info["gcs_address"] == ray_address
+    finally:
+        ray.shutdown()
+        subprocess.check_output("ray stop --force", shell=True)
+
+
+@pytest.mark.skipif(
+    os.environ.get("CI") and sys.platform == "win32",
+    reason="Flaky when run on windows CI",
+)
+def test_ray_init_existing_instance_via_blocked_ray_start():
+    """Run a blocked ray start command and check that ray.init() connects to it."""
+    blocked_start_cmd = subprocess.Popen(
+        ["ray", "start", "--head", "--block", "--num-cpus", "1999"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+
+    def _wait_for_startup_msg():
+        for line in blocked_start_cmd.stdout:
+            l = line.decode("utf-8")
+            print(l)
+            if "Ray runtime started." in l:
+                return
+
+    try:
+        # Wait for the blocked start command's output to indicate that the local Ray
+        # instance has been started successfully. This is done in a background thread
+        # because there is no direct way to read the process' stdout with a timeout.
+        tp = ThreadPoolExecutor(max_workers=1)
+        fut = tp.submit(_wait_for_startup_msg)
+        fut.result(timeout=30)
+
+        # Verify that `ray.init()` connects to the existing cluster
+        # (verified by checking the resources specified to the `ray start` command).
+        ray.init()
+        assert ray.cluster_resources().get("CPU", 0) == 1999
+    finally:
+        ray.shutdown()
+        blocked_start_cmd.terminate()
+        blocked_start_cmd.wait()
+        tp.shutdown()
+        subprocess.check_output("ray stop --force", shell=True)
+
+
+@pytest.mark.skipif(
+    os.environ.get("CI") and sys.platform == "win32",
+    reason="Flaky when run on windows CI",
+)
+@pytest.mark.parametrize("address", [None, "auto"])
+def test_ray_init_existing_instance_crashed(address):
+    ray._private.utils.write_ray_address("localhost:6379")
+    try:
+        # If no address is specified, we will default to an existing cluster.
+        ray_constants.NUM_REDIS_GET_RETRIES = 1
+        with pytest.raises(ConnectionError):
+            ray.init(address=address)
+    finally:
+        ray._common.utils.reset_ray_address()
+
+
+class Credentials(grpc.ChannelCredentials):
+    def __init__(self, name):
+        self.name = name
+
+
+class Stop(Exception):
+    def __init__(self, credentials):
+        self.credentials = credentials
+
+
+def test_ray_init_credentials_with_client(monkeypatch):
+    def mock_init(
+        self,
+        conn_str="",
+        secure=False,
+        metadata=None,
+        connection_retries=3,
+        _credentials=None,
+    ):
+        raise (Stop(_credentials))
+
+    monkeypatch.setattr(Worker, "__init__", mock_init)
+    with pytest.raises(Stop) as stop:
+        with ray_start_client_server(_credentials=Credentials("test")):
+            pass
+
+    assert stop.value.credentials.name == "test"
+
+
+def test_ray_init_credential(monkeypatch):
+    def mock_secure_channel(conn_str, credentials, options=None, compression=None):
+        raise (Stop(credentials))
+
+    monkeypatch.setattr(grpc, "secure_channel", mock_secure_channel)
+
+    with pytest.raises(Stop) as stop:
+        ray.init("ray://127.0.0.1", _credentials=Credentials("test"))
+
+    ray.util.disconnect()
+    assert stop.value.credentials.name == "test"
+
+
+def test_auto_init_non_client(call_ray_start):
+    address = call_ray_start
+    with unittest.mock.patch.dict(os.environ, {"RAY_ADDRESS": address}):
+        res = ray.put(300)
+        # Ensure this is not a client.connect()
+        assert not isinstance(res, ClientObjectRef)
+        ray.shutdown()
+
+    addr = f"localhost:{parse_address(address)[-1]}"
+    with unittest.mock.patch.dict(os.environ, {"RAY_ADDRESS": addr}):
+        res = ray.put(300)
+        # Ensure this is not a client.connect()
+        assert not isinstance(res, ClientObjectRef)
+
+
+@pytest.mark.parametrize(
+    "call_ray_start",
+    ["ray start --head --ray-client-server-port 25036 --port 0"],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "function", [lambda: ray.put(300), lambda: ray.remote(ray.nodes).remote()]
+)
+def test_auto_init_client(call_ray_start, function):
+    address = parse_address(call_ray_start)[0]
+
+    with unittest.mock.patch.dict(
+        os.environ, {"RAY_ADDRESS": f"ray://{build_address(address, 25036)}"}
+    ):
+        res = function()
+        # Ensure this is a client connection.
+        assert isinstance(res, ClientObjectRef)
+        ray.shutdown()
+
+    with unittest.mock.patch.dict(os.environ, {"RAY_ADDRESS": "ray://localhost:25036"}):
+        res = function()
+        # Ensure this is a client connection.
+        assert isinstance(res, ClientObjectRef)
+
+
+@pytest.mark.skipif(
+    os.environ.get("CI") and sys.platform != "linux",
+    reason="This test is only run on linux CI machines.",
+)
+def test_ray_init_using_hostname(ray_start_cluster):
+    import socket
+
+    hostname = socket.gethostname()
+    cluster = Cluster(
+        initialize_head=True,
+        head_node_args={
+            "node_ip_address": hostname,
+        },
+    )
+
+    # Use `ray.init` to test the connection.
+    ray.init(address=cluster.address, _node_ip_address=hostname)
+
+    node_table = cluster.global_state.node_table()
+    assert len(node_table) == 1
+    assert node_table[0].get("NodeManagerHostname", "") == hostname
+
+
+def test_new_ray_instance_new_session_dir(shutdown_only):
+    ray.init()
+    session_dir = ray._private.worker._global_node.get_session_dir_path()
+    ray.shutdown()
+    ray.init()
+    if persistent_gcs_test_enabled():
+        assert ray._private.worker._global_node.get_session_dir_path() == session_dir
+    else:
+        assert ray._private.worker._global_node.get_session_dir_path() != session_dir
+
+
+def test_new_cluster_new_session_dir(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node()
+    ray.init(address=cluster.address)
+    session_dir = ray._private.worker._global_node.get_session_dir_path()
+    ray.shutdown()
+    cluster.shutdown()
+    cluster.add_node()
+    ray.init(address=cluster.address)
+    if persistent_gcs_test_enabled():
+        assert ray._private.worker._global_node.get_session_dir_path() == session_dir
+    else:
+        assert ray._private.worker._global_node.get_session_dir_path() != session_dir
+    ray.shutdown()
+    cluster.shutdown()
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGTERM only on posix")
+def test_ray_init_sigterm_handler():
+    TEST_FILENAME = "sigterm.txt"
+
+    def sigterm_handler_cmd(ray_init=False):
+        return f"""
+import os
+import sys
+import signal
+def sigterm_handler(signum, frame):
+    f = open("{TEST_FILENAME}", "w")
+    sys.exit(0)
+signal.signal(signal.SIGTERM, sigterm_handler)
+
+import ray
+{"ray.init()" if ray_init else ""}
+os.kill(os.getpid(), signal.SIGTERM)
+"""
+
+    # test if sigterm handler is not overwritten by import ray
+    test_child = subprocess.run(["python", "-c", sigterm_handler_cmd()])
+    assert test_child.returncode == 0 and os.path.exists(TEST_FILENAME)
+    os.remove(TEST_FILENAME)
+
+    # test if sigterm handler is overwritten by ray.init
+    test_child = subprocess.run(["python", "-c", sigterm_handler_cmd(ray_init=True)])
+    assert test_child.returncode == signal.SIGTERM and not os.path.exists(TEST_FILENAME)
+
+
+@pytest.fixture
+def runtime_env_working_dir():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        path = Path(tmp_dir)
+        working_dir = path / "working_dir"
+        working_dir.mkdir(parents=True)
+        yield working_dir
+
+
+@pytest.fixture
+def py_module_whl():
+    f = tempfile.NamedTemporaryFile(suffix=".whl", delete=False)
+    f.close()
+    yield f.name
+    os.unlink(f.name)
+
+
+@pytest.fixture
+def ray_shutdown():
+    yield
+    ray.shutdown()
+
+
+def test_ray_init_with_runtime_env_as_dict(
+    runtime_env_working_dir, py_module_whl, ray_shutdown
+):
+    working_dir_path = runtime_env_working_dir
+    working_dir_str = str(working_dir_path)
+    ray.init(
+        runtime_env={"working_dir": working_dir_str, "py_modules": [py_module_whl]}
+    )
+    worker = ray._private.worker.global_worker.core_worker
+    parsed_runtime_env = json.loads(worker.get_current_runtime_env())
+    assert "gcs://" in parsed_runtime_env["working_dir"]
+    assert len(parsed_runtime_env["py_modules"]) == 1
+    assert "gcs://" in parsed_runtime_env["py_modules"][0]
+
+
+def test_ray_init_with_runtime_env_as_object(
+    runtime_env_working_dir, py_module_whl, ray_shutdown
+):
+    working_dir_path = runtime_env_working_dir
+    working_dir_str = str(working_dir_path)
+    ray.init(
+        runtime_env=RuntimeEnv(working_dir=working_dir_str, py_modules=[py_module_whl])
+    )
+    worker = ray._private.worker.global_worker.core_worker
+    parsed_runtime_env = json.loads(worker.get_current_runtime_env())
+    assert "gcs://" in parsed_runtime_env["working_dir"]
+    assert len(parsed_runtime_env["py_modules"]) == 1
+    assert "gcs://" in parsed_runtime_env["py_modules"][0]
+
+
+def test_shutdown_wait_for_processes(shutdown_only):
+    ray.init()
+    node = ray._private.worker._global_node
+    procs = [proc for _, proc in node.live_processes()]
+    assert procs, "ray.init() should have started subprocesses"
+
+    with unittest.mock.patch.object(
+        node, "kill_all_processes", wraps=node.kill_all_processes
+    ) as mock_kill:
+        ray.shutdown(wait_for_processes=True)
+
+    assert mock_kill.call_args.kwargs.get("wait") is True
+    assert all(proc.poll() is not None for proc in procs)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))