chore: import upstream snapshot with attribution

python/ray/tests/test_client_multi.py

@@ -0,0 +1,211 @@
+import sys
+
+import pytest
+
+import ray
+
+
+def test_multi_raise_exception_on_non_client(shutdown_only):
+    with pytest.raises(
+        RuntimeError,
+        match=("Do not pass the `allow_multiple` to `ray.init` to fix the issue."),
+    ):
+        ray.init(allow_multiple=True)
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="PSUtil does not work the same on windows."
+)
+@pytest.mark.parametrize(
+    "call_ray_start",
+    ["ray start --head --ray-client-server-port 25001 --port 0"],
+    indirect=True,
+)
+def test_multi_cli_basic(call_ray_start):
+    ray.init("ray://localhost:25001")
+    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)
+    cli2 = ray.init("ray://localhost:25001", allow_multiple=True)
+    with cli1:
+        a = ray.put(10)
+
+    with cli2:
+        b = ray.put(20)
+
+    # TODO better error message.
+    # Right now, it's EOFError actually
+    with pytest.raises(Exception):
+        ray.get(a)
+
+    with pytest.raises(Exception), cli2:
+        ray.get(a)
+
+    with pytest.raises(Exception), cli1:
+        ray.get(b)
+
+    c = ray.put(30)
+
+    with cli1:
+        assert 10 == ray.get(a)
+
+    with cli2:
+        assert 20 == ray.get(b)
+
+    with pytest.raises(Exception), cli1:
+        ray.get(c)
+
+    with pytest.raises(Exception), cli2:
+        ray.get(c)
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="PSUtil does not work the same on windows."
+)
+@pytest.mark.parametrize(
+    "call_ray_start",
+    ["ray start --head --ray-client-server-port 25001 --port 0"],
+    indirect=True,
+)
+def test_multi_cli_init(call_ray_start):
+    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)  # noqa
+    with pytest.raises(
+        ValueError,
+        match="The client has already connected to the cluster "
+        "with allow_multiple=True. Please set allow_multiple=True"
+        " to proceed",
+    ):
+        ray.init("ray://localhost:25001")
+    cli2 = ray.init("ray://localhost:25001", allow_multiple=True)  # noqa
+
+    cli1.disconnect()
+    cli2.disconnect()
+
+    ray.init("ray://localhost:25001")
+    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)  # noqa
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="PSUtil does not work the same on windows."
+)
+@pytest.mark.parametrize(
+    "call_ray_start",
+    ["ray start --head --ray-client-server-port 25001 --port 0"],
+    indirect=True,
+)
+def test_multi_cli_func(call_ray_start):
+    @ray.remote
+    def hello():
+        return "world"
+
+    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)
+    cli2 = ray.init("ray://localhost:25001", allow_multiple=True)
+
+    # TODO better error message.
+    # Right now, it's EOFError actually
+    with pytest.raises(Exception):
+        ray.get(hello.remote())
+
+    with cli1:
+        o1 = hello.remote()
+        assert "world" == ray.get(o1)
+
+    with cli2:
+        o2 = hello.remote()
+        assert "world" == ray.get(o2)
+
+    with pytest.raises(Exception), cli1:
+        ray.get(o2)
+
+    with pytest.raises(Exception), cli2:
+        ray.get(o1)
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="PSUtil does not work the same on windows."
+)
+@pytest.mark.parametrize(
+    "call_ray_start",
+    ["ray start --head --ray-client-server-port 25001 --port 0"],
+    indirect=True,
+)
+def test_multi_cli_actor(call_ray_start):
+    @ray.remote
+    class Actor:
+        def __init__(self, v):
+            self.v = v
+
+        def double(self):
+            return self.v * 2
+
+    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)
+    cli2 = ray.init("ray://localhost:25001", allow_multiple=True)
+
+    # TODO better error message.
+    # Right now, it's EOFError actually
+    with pytest.raises(Exception):
+        a = Actor.remote(10)
+        ray.get(a.double.remote())
+
+    with cli1:
+        a1 = Actor.remote(10)
+        o1 = a1.double.remote()
+        assert 20 == ray.get(o1)
+
+    with cli2:
+        a2 = Actor.remote(20)
+        o2 = a2.double.remote()
+        assert 40 == ray.get(o2)
+
+    with pytest.raises(Exception), cli1:
+        ray.get(a2.double.remote())
+
+    with pytest.raises(Exception), cli1:
+        ray.get(o2)
+
+    with pytest.raises(Exception), cli2:
+        ray.get(a1.double.remote())
+
+    with pytest.raises(Exception), cli2:
+        ray.get(o1)
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="PSUtil does not work the same on windows."
+)
+@pytest.mark.parametrize(
+    "call_ray_start",
+    ["ray start --head --ray-client-server-port 25001 --port 0"],
+    indirect=True,
+)
+def test_multi_cli_threading(call_ray_start):
+    import threading
+
+    b = threading.Barrier(2)
+    ret = [None, None]
+
+    def get(idx):
+        cli = ray.init("ray://localhost:25001", allow_multiple=True)
+        with cli:
+            a = ray.put(idx)
+            b.wait()
+            v = ray.get(a)
+            assert idx == v
+            b.wait()
+            ret[idx] = v
+
+    t1 = threading.Thread(target=get, args=(0,))
+    t2 = threading.Thread(target=get, args=(1,))
+    t1.start()
+    t2.start()
+    t1.join()
+    t2.join()
+    assert ret == [0, 1]
+
+
+if __name__ == "__main__":
+    # The following should be removed after
+    # https://github.com/ray-project/ray/issues/20355
+    # is fixed.
+    from ray._private import auto_init_hook
+
+    auto_init_hook.enable_auto_connect = False
+    sys.exit(pytest.main(["-sv", __file__]))