chore: import upstream snapshot with attribution

python/ray/tests/test_client_references.py

@@ -0,0 +1,334 @@
+import sys
+from concurrent.futures import Future
+
+import pytest
+
+import ray as real_ray
+from ray._common.test_utils import wait_for_condition
+from ray._private.test_utils import object_memory_usage
+from ray._raylet import ActorID, ObjectRef
+from ray.util.client import _ClientContext
+from ray.util.client.common import ClientActorRef, ClientObjectRef
+from ray.util.client.ray_client_helpers import (
+    ray_start_client_server,
+    ray_start_client_server_pair,
+    ray_start_cluster_client_server_pair,
+)
+
+
+def test_client_object_ref_basics(ray_start_regular):
+    with ray_start_client_server_pair() as pair:
+        ray, server = pair
+        ref = ray.put("Hello World")
+        # Make sure ClientObjectRef is a subclass of ObjectRef
+        assert isinstance(ref, ClientObjectRef)
+        assert isinstance(ref, ObjectRef)
+
+        # Invalid ref format.
+        with pytest.raises(Exception):
+            ClientObjectRef(b"\0")
+
+        obj_id = b"\0" * 28
+        fut = Future()
+        fut.set_result(obj_id)
+        server_ref = ObjectRef(obj_id)
+        for client_ref in [ClientObjectRef(obj_id), ClientObjectRef(fut)]:
+            client_members = set(client_ref.__dir__())
+            server_members = set(server_ref.__dir__())
+            client_members = {m for m in client_ref.__dir__() if not m.startswith("_")}
+            server_members = {m for m in server_ref.__dir__() if not m.startswith("_")}
+            assert client_members.difference(server_members) == {"id"}
+            assert server_members.difference(client_members) == set()
+
+            # Test __eq__()
+            assert client_ref == ClientObjectRef(obj_id)
+            assert client_ref != ref
+            assert client_ref != server_ref
+
+            # Test other methods
+            assert client_ref.__repr__() == f"ClientObjectRef({obj_id.hex()})"
+            assert client_ref.binary() == obj_id
+            assert client_ref.hex() == obj_id.hex()
+            assert not client_ref.is_nil()
+            assert client_ref.task_id() == server_ref.task_id()
+            assert client_ref.job_id() == server_ref.job_id()
+
+
+def test_client_actor_ref_basics(ray_start_regular):
+    with ray_start_client_server_pair() as pair:
+        ray, server = pair
+
+        @ray.remote
+        class Counter:
+            def __init__(self):
+                self.acc = 0
+
+            def inc(self):
+                self.acc += 1
+
+            def get(self):
+                return self.acc
+
+        counter = Counter.remote()
+        ref = counter.actor_ref
+
+        # Make sure ClientActorRef is a subclass of ActorID
+        assert isinstance(ref, ClientActorRef)
+        assert isinstance(ref, ActorID)
+
+        # Invalid ref format.
+        with pytest.raises(Exception):
+            ClientActorRef(b"\0")
+
+        actor_id = b"\0" * 16
+        fut = Future()
+        fut.set_result(actor_id)
+        server_ref = ActorID(actor_id)
+        for client_ref in [ClientActorRef(actor_id), ClientActorRef(fut)]:
+            client_members = {m for m in client_ref.__dir__() if not m.startswith("_")}
+            server_members = {m for m in server_ref.__dir__() if not m.startswith("_")}
+            assert client_members.difference(server_members) == {"id"}
+            assert server_members.difference(client_members) == set()
+
+            # Test __eq__()
+            assert client_ref == ClientActorRef(actor_id)
+            assert client_ref != ref
+            assert client_ref != server_ref
+
+            # Test other methods
+            assert client_ref.__repr__() == f"ClientActorRef({actor_id.hex()})"
+            assert client_ref.binary() == actor_id
+            assert client_ref.hex() == actor_id.hex()
+            assert not client_ref.is_nil()
+
+
+def server_object_ref_count(server, n):
+    assert server is not None
+
+    def test_cond():
+        if len(server.task_servicer.object_refs) == 0:
+            # No open clients
+            return n == 0
+        client_id = list(server.task_servicer.object_refs.keys())[0]
+        return len(server.task_servicer.object_refs[client_id]) == n
+
+    return test_cond
+
+
+def server_actor_ref_count(server, n):
+    assert server is not None
+
+    def test_cond():
+        if len(server.task_servicer.actor_refs) == 0:
+            # No running actors
+            return n == 0
+        return len(server.task_servicer.actor_refs) == n
+
+    return test_cond
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_nodes": 1,
+            "do_init": False,
+        }
+    ],
+    indirect=True,
+)
+def test_delete_refs_on_disconnect(ray_start_cluster):
+    cluster = ray_start_cluster
+    with ray_start_cluster_client_server_pair(cluster.address) as pair:
+        ray, server = pair
+
+        @ray.remote
+        def f(x):
+            return x + 2
+
+        thing1 = f.remote(6)  # noqa
+        thing2 = ray.put("Hello World")  # noqa
+
+        # One put, one function -- the function result thing1 is
+        # in a different category, according to the raylet.
+        # But we're maintaining the reference
+        assert server_object_ref_count(server, 3)()
+        # And can get the data
+        assert ray.get(thing1) == 8
+
+        # Close the client.
+        ray.close()
+
+        wait_for_condition(server_object_ref_count(server, 0), timeout=5)
+
+        # Connect to the real ray again, since we disconnected
+        # upon num_clients = 0.
+        real_ray.init(address=cluster.address, namespace="default_test_namespace")
+
+        def test_cond():
+            return object_memory_usage() == 0
+
+        wait_for_condition(test_cond, timeout=5)
+
+
+def test_delete_ref_on_object_deletion(ray_start_regular):
+    with ray_start_client_server_pair() as pair:
+        ray, server = pair
+        vals = {
+            "ref": ray.put("Hello World"),
+            "ref2": ray.put("This value stays"),
+        }
+
+        del vals["ref"]
+
+        wait_for_condition(server_object_ref_count(server, 1), timeout=5)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_nodes": 1,
+            "do_init": False,
+            "include_dashboard": True,
+        }
+    ],
+    indirect=True,
+)
+def test_delete_actor_on_disconnect(ray_start_cluster):
+    cluster = ray_start_cluster
+    with ray_start_cluster_client_server_pair(cluster.address) as pair:
+        ray, server = pair
+
+        @ray.remote
+        class Accumulator:
+            def __init__(self):
+                self.acc = 0
+
+            def inc(self):
+                self.acc += 1
+
+            def get(self):
+                return self.acc
+
+        actor = Accumulator.remote()
+        actor.inc.remote()
+
+        assert server_actor_ref_count(server, 1)()
+
+        assert ray.get(actor.get.remote()) == 1
+
+        ray.close()
+
+        wait_for_condition(server_actor_ref_count(server, 0), timeout=5)
+
+        def test_cond():
+            alive_actors = [
+                v for v in real_ray.util.state.list_actors() if v.state != "DEAD"
+            ]
+            return len(alive_actors) == 0
+
+        # Connect to the real ray again, since we disconnected
+        # upon num_clients = 0.
+        real_ray.init(address=cluster.address, namespace="default_test_namespace")
+
+        wait_for_condition(test_cond, timeout=10)
+
+
+def test_delete_actor(ray_start_regular):
+    with ray_start_client_server_pair() as pair:
+        ray, server = pair
+
+        @ray.remote
+        class Accumulator:
+            def __init__(self):
+                self.acc = 0
+
+            def inc(self):
+                self.acc += 1
+
+        actor = Accumulator.remote()
+        actor.inc.remote()
+        actor2 = Accumulator.remote()
+        actor2.inc.remote()
+
+        assert server_actor_ref_count(server, 2)()
+
+        del actor
+
+        wait_for_condition(server_actor_ref_count(server, 1), timeout=5)
+
+
+def test_simple_multiple_references(ray_start_regular):
+    with ray_start_client_server() as ray:
+
+        @ray.remote
+        class A:
+            def __init__(self):
+                self.x = ray.put("hi")
+
+            def get(self):
+                return [self.x]
+
+        a = A.remote()
+        ref1 = ray.get(a.get.remote())[0]
+        ref2 = ray.get(a.get.remote())[0]
+        del a
+        assert ray.get(ref1) == "hi"
+        del ref1
+        assert ray.get(ref2) == "hi"
+        del ref2
+
+
+def test_named_actor_refcount(ray_start_regular):
+    with ray_start_client_server_pair() as (ray, server):
+
+        @ray.remote
+        class ActorTest:
+            def __init__(self):
+                self._counter = 0
+
+            def bump(self):
+                self._counter += 1
+
+            def check(self):
+                return self._counter
+
+        ActorTest.options(name="actor", lifetime="detached").remote()
+
+        def connect_api():
+            api = _ClientContext()
+            api.connect("localhost:50051", namespace="default_test_namespace")
+            api.get_actor("actor")
+            return api
+
+        def check_owners(size):
+            return size == sum(
+                len(x) for x in server.task_servicer.actor_owners.values()
+            )
+
+        apis = [connect_api() for i in range(3)]
+        assert check_owners(3)
+        assert len(server.task_servicer.actor_refs) == 1
+        assert len(server.task_servicer.named_actors) == 1
+
+        apis.pop(0).disconnect()
+        assert check_owners(2)
+        assert len(server.task_servicer.actor_refs) == 1
+        assert len(server.task_servicer.named_actors) == 1
+
+        apis.pop(0).disconnect()
+        assert check_owners(1)
+        assert len(server.task_servicer.actor_refs) == 1
+        assert len(server.task_servicer.named_actors) == 1
+
+        apis.pop(0).disconnect()
+        # no more owners should be seen
+        assert check_owners(0)
+        # actor refs shouldn't be removed
+        assert len(server.task_servicer.actor_refs) == 1
+        assert len(server.task_servicer.named_actors) == 1
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))