wehub-resource-sync
335 lines
9.8 KiB
Python
335 lines
9.8 KiB
Python
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__]))
|