chore: import upstream snapshot with attribution

python/ray/tests/test_client_init.py

@@ -0,0 +1,291 @@
+"""Client tests that run their own init (as with init_and_serve) live here"""
+
+import json
+import random
+import subprocess
+import sys
+import time
+from unittest.mock import patch
+
+import pytest
+
+import ray
+import ray.core.generated.ray_client_pb2 as ray_client_pb2
+import ray.util.client.server.server as ray_client_server
+from ray._private.client_mode_hook import _explicitly_disable_client_mode
+from ray._private.ray_logging.logging_config import LoggingConfig
+from ray.cloudpickle.compat import pickle
+from ray.cluster_utils import cluster_not_supported
+from ray.job_config import JobConfig
+from ray.util.client import _ClientContext, ray as ray_client
+from ray.util.client.worker import prepare_init_request_args
+
+
+@ray.remote
+def hello_world():
+    c1 = complex_task.remote(random.randint(1, 10))
+    c2 = complex_task.remote(random.randint(1, 10))
+    return sum(ray.get([c1, c2]))
+
+
+@ray.remote
+def complex_task(value):
+    time.sleep(1)
+    return value * 10
+
+
+@ray.remote
+class C:
+    def __init__(self, x):
+        self.val = x
+
+    def double(self):
+        self.val += self.val
+
+    def get(self):
+        return self.val
+
+
+@pytest.mark.xfail(cluster_not_supported, reason="cluster not supported")
+@pytest.fixture
+def init_and_serve_lazy(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=1, num_gpus=0)
+    cluster.wait_for_nodes(1)
+    address = cluster.address
+
+    def connect(job_config=None, **ray_init_kwargs):
+        ray.init(address=address, job_config=job_config, **ray_init_kwargs)
+
+    server_handle = ray_client_server.serve("localhost", 50051, connect)
+    yield server_handle
+    ray_client_server.shutdown_with_server(server_handle.grpc_server)
+
+
+def test_validate_port():
+    """Check that ports outside of 1024-65535 are rejected."""
+    for port in [1000, 1023, 65536, 700000]:
+        with pytest.raises(subprocess.CalledProcessError) as excinfo:
+            subprocess.check_output(
+                [
+                    "ray",
+                    "start",
+                    "--head",
+                    "--num-cpus",
+                    "8",
+                    "--ray-client-server-port",
+                    f"{port}",
+                ]
+            )
+            assert "ValueError" in str(excinfo.traceback)
+            assert "65535" in str(excinfo.traceback)
+
+
+def test_basic_preregister(init_and_serve):
+    """Tests conversion of Ray actors and remote functions to client actors
+    and client remote functions.
+
+    Checks that the conversion works when disconnecting and reconnecting client
+    sessions.
+    """
+    from ray.util.client import ray
+
+    for _ in range(2):
+        ray.connect("localhost:50051")
+        val = ray.get(hello_world.remote())
+        print(val)
+        assert val >= 20
+        assert val <= 200
+        c = C.remote(3)
+        x = c.double.remote()
+        y = c.double.remote()
+        ray.wait([x, y])
+        val = ray.get(c.get.remote())
+        assert val == 12
+        ray.disconnect()
+
+
+def test_idempotent_disconnect(init_and_serve):
+    from ray.util.client import ray
+
+    ray.disconnect()
+    ray.disconnect()
+    ray.connect("localhost:50051")
+    ray.disconnect()
+    ray.disconnect()
+
+
+def test_public_shutdown_wait_for_processes_in_client_mode(init_and_serve):
+    ray_client.connect("localhost:50051")
+    try:
+        assert ray_client.is_connected()
+        # In client mode, wait_for_processes is a no-op
+        ray.shutdown(wait_for_processes=True)
+        assert not ray_client.is_connected()
+    finally:
+        _explicitly_disable_client_mode()
+
+
+def test_num_clients(init_and_serve_lazy):
+    # Tests num clients reporting; useful if you want to build an app that
+    # load balances clients between Ray client servers.
+
+    def get_job_id(api):
+        return api.get_runtime_context().worker.current_job_id
+
+    api1 = _ClientContext()
+    info1 = api1.connect("localhost:50051")
+    job_id_1 = get_job_id(api1)
+    assert info1["num_clients"] == 1, info1
+    api2 = _ClientContext()
+    info2 = api2.connect("localhost:50051")
+    job_id_2 = get_job_id(api2)
+    assert info2["num_clients"] == 2, info2
+
+    assert job_id_1 == job_id_2
+
+    # Disconnect the first two clients.
+    api1.disconnect()
+    api2.disconnect()
+    time.sleep(1)
+
+    api3 = _ClientContext()
+    info3 = api3.connect("localhost:50051")
+    job_id_3 = get_job_id(api3)
+    assert info3["num_clients"] == 1, info3
+    assert job_id_1 != job_id_3
+
+    # Check info contains ray and python version.
+    assert isinstance(info3["ray_version"], str), info3
+    assert isinstance(info3["ray_commit"], str), info3
+    assert isinstance(info3["python_version"], str), info3
+    api3.disconnect()
+
+
+def test_python_version(init_and_serve):
+    server_handle = init_and_serve
+    ray = _ClientContext()
+    info1 = ray.connect("localhost:50051")
+    assert info1["python_version"] == ".".join(
+        [str(x) for x in list(sys.version_info)[:3]]
+    )
+    ray.disconnect()
+    time.sleep(1)
+
+    def mock_connection_response():
+        return ray_client_pb2.ConnectionInfoResponse(
+            num_clients=1,
+            python_version="2.7.12",
+            ray_version="",
+            ray_commit="",
+        )
+
+    # inject mock connection function
+    server_handle.data_servicer._build_connection_response = mock_connection_response
+
+    ray = _ClientContext()
+    with pytest.raises(RuntimeError):
+        _ = ray.connect("localhost:50051")
+
+    ray = _ClientContext()
+    info3 = ray.connect("localhost:50051", ignore_version=True)
+    assert info3["num_clients"] == 1, info3
+    ray.disconnect()
+
+
+@patch("ray.util.client.server.dataservicer.CLIENT_SERVER_MAX_THREADS", 6)
+@patch("ray.util.client.server.logservicer.CLIENT_SERVER_MAX_THREADS", 6)
+def test_max_clients(init_and_serve):
+    # Tests max clients. Should support up to CLIENT_SERVER_MAX_THREADS / 2.
+    def get_job_id(api):
+        return api.get_runtime_context().worker.current_job_id
+
+    for i in range(3):
+        api1 = _ClientContext()
+        info1 = api1.connect("localhost:50051")
+
+        assert info1["num_clients"] == i + 1, info1
+
+    with pytest.raises(ConnectionError):
+        api = _ClientContext()
+        _ = api.connect("localhost:50051")
+
+
+@pytest.mark.parametrize(
+    "call_ray_start",
+    ["ray start --head --ray-client-server-port=50056"],
+    indirect=True,
+)
+def test_logging_config_in_client_mode(call_ray_start):
+    """logging_config reaches task workers (JobConfig) when using Ray Client."""
+
+    @ray.remote
+    def job_logging_encoding():
+        w = ray.get_runtime_context().worker
+        jlc = w.job_logging_config
+        return None if jlc is None else jlc.encoding
+
+    ray.init(
+        "ray://localhost:50056",
+        logging_config=LoggingConfig(encoding="TEXT", log_level="INFO"),
+        log_to_driver=False,
+        configure_logging=False,
+    )
+
+    assert ray.get(job_logging_encoding.remote()) == "TEXT"
+
+
+def test_prepare_init_request_args():
+    """``prepare_init_request_args`` copies kwargs, JSON-serializes logging config,
+    propagates it onto ``JobConfig``, and pickles the job for ``InitRequest``."""
+    # no job config, no ray init kwargs
+    serialized, kwargs = prepare_init_request_args(None, None)
+    assert serialized is None and kwargs == {}
+
+    # invalid logging config type
+    with pytest.raises(
+        TypeError, match="logging_config must be a dict or LoggingConfig"
+    ):
+        prepare_init_request_args(JobConfig(), {"logging_config": 42})
+
+    # ray init kwargs contains serialized logging config
+    lc = LoggingConfig(encoding="JSON", log_level="DEBUG")
+    orig = {"logging_config": lc, "num_cpus": 1}
+    job = JobConfig()
+    _, out = prepare_init_request_args(job, orig)
+    assert isinstance(orig["logging_config"], LoggingConfig)
+    assert isinstance(out["logging_config"], dict)
+    assert out["logging_config"]["encoding"] == "JSON"
+    assert orig["num_cpus"] == 1
+    json.dumps(out)
+    assert job.py_logging_config is not None
+    assert job.py_logging_config.encoding == "JSON"
+
+    # write logging config to JobConfig
+    job_dict = JobConfig()
+    prepare_init_request_args(
+        job_dict,
+        {"logging_config": {"encoding": "TEXT", "log_level": "ERROR"}},
+    )
+    assert job_dict.py_logging_config.encoding == "TEXT"
+    assert job_dict.py_logging_config.log_level == "ERROR"
+
+    # don't overwrite existing logging config on JobConfig
+    existing = LoggingConfig(encoding="TEXT")
+    job_keep = JobConfig()
+    job_keep.set_py_logging_config(existing)
+    prepare_init_request_args(
+        job_keep,
+        {"logging_config": LoggingConfig(encoding="JSON")},
+    )
+    assert job_keep.py_logging_config is existing
+
+    # pickle the job for InitRequest
+    job_pickle = JobConfig()
+    blob, _ = prepare_init_request_args(job_pickle, {})
+    assert isinstance(blob, bytes) and len(blob) > 0
+    assert isinstance(pickle.loads(blob), JobConfig)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))