chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,471 @@
|
||||
"""Reference counting tests that require their own custom fixture.
|
||||
|
||||
The other reference counting tests use a shared Ray instance across the test module
|
||||
to reduce overheads & overall test runtime.
|
||||
"""
|
||||
# coding: utf-8
|
||||
import logging
|
||||
import platform
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
import ray.cluster_utils
|
||||
from ray._common.test_utils import (
|
||||
SignalActor,
|
||||
fetch_prometheus_metrics,
|
||||
wait_for_condition,
|
||||
)
|
||||
from ray._private.internal_api import memory_summary
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _fill_object_store_and_get(obj, succeed=True, object_MiB=20, num_objects=5):
|
||||
for _ in range(num_objects):
|
||||
ray.put(np.zeros(object_MiB * 1024 * 1024, dtype=np.uint8))
|
||||
|
||||
if type(obj) is bytes:
|
||||
obj = ray.ObjectRef(obj)
|
||||
|
||||
if succeed:
|
||||
wait_for_condition(
|
||||
lambda: ray._private.worker.global_worker.core_worker.object_exists(obj)
|
||||
)
|
||||
else:
|
||||
wait_for_condition(
|
||||
lambda: not ray._private.worker.global_worker.core_worker.object_exists(obj)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(platform.system() in ["Windows"], reason="Failing on Windows.")
|
||||
def test_object_unpin(ray_start_cluster):
|
||||
nodes = []
|
||||
cluster = ray_start_cluster
|
||||
head_node = cluster.add_node(
|
||||
num_cpus=0,
|
||||
object_store_memory=100 * 1024 * 1024,
|
||||
_system_config={
|
||||
"subscriber_timeout_ms": 100,
|
||||
"health_check_initial_delay_ms": 0,
|
||||
"health_check_period_ms": 1000,
|
||||
"health_check_failure_threshold": 5,
|
||||
},
|
||||
)
|
||||
ray.init(address=cluster.address)
|
||||
|
||||
# Add worker nodes.
|
||||
for i in range(2):
|
||||
nodes.append(
|
||||
cluster.add_node(
|
||||
num_cpus=1,
|
||||
resources={f"node_{i}": 1},
|
||||
object_store_memory=100 * 1024 * 1024,
|
||||
)
|
||||
)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
one_mb_array = np.ones(1 * 1024 * 1024, dtype=np.uint8)
|
||||
ten_mb_array = np.ones(10 * 1024 * 1024, dtype=np.uint8)
|
||||
|
||||
@ray.remote
|
||||
class ObjectsHolder:
|
||||
def __init__(self):
|
||||
self.ten_mb_objs = []
|
||||
self.one_mb_objs = []
|
||||
|
||||
def put_10_mb(self):
|
||||
self.ten_mb_objs.append(ray.put(ten_mb_array))
|
||||
|
||||
def put_1_mb(self):
|
||||
self.one_mb_objs.append(ray.put(one_mb_array))
|
||||
|
||||
def pop_10_mb(self):
|
||||
if len(self.ten_mb_objs) == 0:
|
||||
return False
|
||||
self.ten_mb_objs.pop()
|
||||
return True
|
||||
|
||||
def pop_1_mb(self):
|
||||
if len(self.one_mb_objs) == 0:
|
||||
return False
|
||||
self.one_mb_objs.pop()
|
||||
return True
|
||||
|
||||
# Head node contains 11MB of data.
|
||||
one_mb_arrays = []
|
||||
ten_mb_arrays = []
|
||||
|
||||
one_mb_arrays.append(ray.put(one_mb_array))
|
||||
ten_mb_arrays.append(ray.put(ten_mb_array))
|
||||
|
||||
def check_memory(mb):
|
||||
return f"Plasma memory usage {mb} MiB" in memory_summary(
|
||||
address=head_node.address, stats_only=True
|
||||
)
|
||||
|
||||
def wait_until_node_dead(node):
|
||||
for n in ray.nodes():
|
||||
if n["ObjectStoreSocketName"] == node.address_info["object_store_address"]:
|
||||
return not n["Alive"]
|
||||
return False
|
||||
|
||||
wait_for_condition(lambda: check_memory(11))
|
||||
|
||||
# Pop one mb array and see if it works.
|
||||
one_mb_arrays.pop()
|
||||
wait_for_condition(lambda: check_memory(10))
|
||||
|
||||
# Pop 10 MB.
|
||||
ten_mb_arrays.pop()
|
||||
wait_for_condition(lambda: check_memory(0))
|
||||
|
||||
# Put 11 MB for each actor.
|
||||
# actor 1: 1MB + 10MB
|
||||
# actor 2: 1MB + 10MB
|
||||
actor_on_node_1 = ObjectsHolder.options(resources={"node_0": 1}).remote()
|
||||
actor_on_node_2 = ObjectsHolder.options(resources={"node_1": 1}).remote()
|
||||
ray.get(actor_on_node_1.put_1_mb.remote())
|
||||
ray.get(actor_on_node_1.put_10_mb.remote())
|
||||
ray.get(actor_on_node_2.put_1_mb.remote())
|
||||
ray.get(actor_on_node_2.put_10_mb.remote())
|
||||
wait_for_condition(lambda: check_memory(22))
|
||||
|
||||
# actor 1: 10MB
|
||||
# actor 2: 1MB
|
||||
ray.get(actor_on_node_1.pop_1_mb.remote())
|
||||
ray.get(actor_on_node_2.pop_10_mb.remote())
|
||||
wait_for_condition(lambda: check_memory(11))
|
||||
|
||||
# The second node is dead, and actor 2 is dead.
|
||||
cluster.remove_node(nodes[1], allow_graceful=False)
|
||||
wait_for_condition(lambda: wait_until_node_dead(nodes[1]))
|
||||
wait_for_condition(lambda: check_memory(10))
|
||||
|
||||
# The first actor is dead, so object should be GC'ed.
|
||||
ray.kill(actor_on_node_1)
|
||||
wait_for_condition(lambda: check_memory(0))
|
||||
|
||||
|
||||
@pytest.mark.skipif(platform.system() in ["Windows"], reason="Failing on Windows.")
|
||||
def test_object_unpin_stress(ray_start_cluster):
|
||||
nodes = []
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(
|
||||
num_cpus=1, resources={"head": 1}, object_store_memory=1000 * 1024 * 1024
|
||||
)
|
||||
ray.init(address=cluster.address)
|
||||
|
||||
# Add worker nodes.
|
||||
for i in range(2):
|
||||
nodes.append(
|
||||
cluster.add_node(
|
||||
num_cpus=1,
|
||||
resources={f"node_{i}": 1},
|
||||
object_store_memory=1000 * 1024 * 1024,
|
||||
)
|
||||
)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
one_mb_array = np.ones(1 * 1024 * 1024, dtype=np.uint8)
|
||||
ten_mb_array = np.ones(10 * 1024 * 1024, dtype=np.uint8)
|
||||
|
||||
@ray.remote
|
||||
class ObjectsHolder:
|
||||
def __init__(self):
|
||||
self.ten_mb_objs = []
|
||||
self.one_mb_objs = []
|
||||
|
||||
def put_10_mb(self):
|
||||
self.ten_mb_objs.append(ray.put(ten_mb_array))
|
||||
|
||||
def put_1_mb(self):
|
||||
self.one_mb_objs.append(ray.put(one_mb_array))
|
||||
|
||||
def pop_10_mb(self):
|
||||
if len(self.ten_mb_objs) == 0:
|
||||
return False
|
||||
self.ten_mb_objs.pop()
|
||||
return True
|
||||
|
||||
def pop_1_mb(self):
|
||||
if len(self.one_mb_objs) == 0:
|
||||
return False
|
||||
self.one_mb_objs.pop()
|
||||
return True
|
||||
|
||||
def get_obj_size(self):
|
||||
return len(self.ten_mb_objs) * 10 + len(self.one_mb_objs)
|
||||
|
||||
actor_on_node_1 = ObjectsHolder.options(resources={"node_0": 1}).remote()
|
||||
actor_on_node_2 = ObjectsHolder.options(resources={"node_1": 1}).remote()
|
||||
actor_on_head_node = ObjectsHolder.options(resources={"head": 1}).remote()
|
||||
|
||||
ray.get(actor_on_node_1.get_obj_size.remote())
|
||||
ray.get(actor_on_node_2.get_obj_size.remote())
|
||||
ray.get(actor_on_head_node.get_obj_size.remote())
|
||||
|
||||
def random_ops(actors):
|
||||
r = random.random()
|
||||
for actor in actors:
|
||||
if r <= 0.25:
|
||||
actor.put_10_mb.remote()
|
||||
elif r <= 0.5:
|
||||
actor.put_1_mb.remote()
|
||||
elif r <= 0.75:
|
||||
actor.pop_10_mb.remote()
|
||||
else:
|
||||
actor.pop_1_mb.remote()
|
||||
|
||||
total_iter = 15
|
||||
for _ in range(total_iter):
|
||||
random_ops([actor_on_node_1, actor_on_node_2, actor_on_head_node])
|
||||
|
||||
# Simulate node dead.
|
||||
cluster.remove_node(nodes[1])
|
||||
for _ in range(total_iter):
|
||||
random_ops([actor_on_node_1, actor_on_head_node])
|
||||
|
||||
total_size = sum(
|
||||
[
|
||||
ray.get(actor_on_node_1.get_obj_size.remote()),
|
||||
ray.get(actor_on_head_node.get_obj_size.remote()),
|
||||
]
|
||||
)
|
||||
|
||||
wait_for_condition(
|
||||
lambda: (
|
||||
(f"Plasma memory usage {total_size} MiB") in memory_summary(stats_only=True)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("inline_args", [True, False])
|
||||
def test_inlined_nested_refs(ray_start_cluster, inline_args):
|
||||
cluster = ray_start_cluster
|
||||
config = {}
|
||||
if not inline_args:
|
||||
config["max_direct_call_object_size"] = 0
|
||||
cluster.add_node(
|
||||
num_cpus=2, object_store_memory=100 * 1024 * 1024, _system_config=config
|
||||
)
|
||||
ray.init(address=cluster.address)
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def nested(self):
|
||||
return ray.put("x")
|
||||
|
||||
@ray.remote
|
||||
def nested_nested(a):
|
||||
return a.nested.remote()
|
||||
|
||||
@ray.remote
|
||||
def foo(ref):
|
||||
time.sleep(1)
|
||||
return ray.get(ref)
|
||||
|
||||
a = Actor.remote()
|
||||
nested_nested_ref = nested_nested.remote(a)
|
||||
# We get nested_ref's value directly from its owner.
|
||||
nested_ref = ray.get(nested_nested_ref)
|
||||
|
||||
del nested_nested_ref
|
||||
x = foo.remote(nested_ref)
|
||||
del nested_ref
|
||||
ray.get(x)
|
||||
|
||||
|
||||
# https://github.com/ray-project/ray/issues/17553
|
||||
@pytest.mark.parametrize("inline_args", [True, False])
|
||||
def test_return_nested_ids(shutdown_only, inline_args):
|
||||
config = dict()
|
||||
if inline_args:
|
||||
config["max_direct_call_object_size"] = 100 * 1024 * 1024
|
||||
else:
|
||||
config["max_direct_call_object_size"] = 0
|
||||
ray.init(object_store_memory=100 * 1024 * 1024, _system_config=config)
|
||||
|
||||
class Nested:
|
||||
def __init__(self, blocks):
|
||||
self._blocks = blocks
|
||||
|
||||
@ray.remote
|
||||
def echo(fn):
|
||||
return fn()
|
||||
|
||||
@ray.remote
|
||||
def create_nested():
|
||||
refs = [ray.put(np.random.random(1024 * 1024)) for _ in range(10)]
|
||||
return Nested(refs)
|
||||
|
||||
@ray.remote
|
||||
def test():
|
||||
ref = create_nested.remote()
|
||||
result1 = ray.get(ref)
|
||||
del ref
|
||||
result = echo.remote(lambda: result1) # noqa
|
||||
del result1
|
||||
|
||||
time.sleep(5)
|
||||
block = ray.get(result)._blocks[0]
|
||||
print(ray.get(block))
|
||||
|
||||
ray.get(test.remote())
|
||||
|
||||
|
||||
def _check_refcounts(expected):
|
||||
actual = ray._private.worker.global_worker.core_worker.get_all_reference_counts()
|
||||
assert len(expected) == len(actual)
|
||||
for object_ref, (local, submitted) in expected.items():
|
||||
hex_id = object_ref.hex().encode("ascii")
|
||||
assert hex_id in actual
|
||||
assert local == actual[hex_id]["local"]
|
||||
assert submitted == actual[hex_id]["submitted"]
|
||||
|
||||
|
||||
def test_out_of_band_serialized_object_ref(ray_start_regular):
|
||||
assert (
|
||||
len(ray._private.worker.global_worker.core_worker.get_all_reference_counts())
|
||||
== 0
|
||||
)
|
||||
obj_ref = ray.put("hello")
|
||||
_check_refcounts({obj_ref: (1, 0)})
|
||||
obj_ref_str = ray.cloudpickle.dumps(obj_ref)
|
||||
_check_refcounts({obj_ref: (2, 0)})
|
||||
del obj_ref
|
||||
assert (
|
||||
len(ray._private.worker.global_worker.core_worker.get_all_reference_counts())
|
||||
== 1
|
||||
)
|
||||
assert ray.get(ray.cloudpickle.loads(obj_ref_str)) == "hello"
|
||||
|
||||
|
||||
def test_captured_object_ref(ray_start_regular):
|
||||
captured_id = ray.put(np.zeros(1024, dtype=np.uint8))
|
||||
|
||||
@ray.remote
|
||||
def f(signal):
|
||||
ray.get(signal.wait.remote())
|
||||
ray.get(captured_id) # noqa: F821
|
||||
|
||||
signal = SignalActor.remote()
|
||||
obj_ref = f.remote(signal)
|
||||
|
||||
# Delete local references.
|
||||
del f
|
||||
del captured_id
|
||||
|
||||
# Test that the captured object ref is pinned despite having no local
|
||||
# references.
|
||||
ray.get(signal.send.remote())
|
||||
_fill_object_store_and_get(obj_ref)
|
||||
|
||||
captured_id = ray.put(np.zeros(1024, dtype=np.uint8))
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def get(self, signal):
|
||||
ray.get(signal.wait.remote())
|
||||
ray.get(captured_id) # noqa: F821
|
||||
|
||||
signal = SignalActor.remote()
|
||||
actor = Actor.remote()
|
||||
obj_ref = actor.get.remote(signal)
|
||||
|
||||
# Delete local references.
|
||||
del Actor
|
||||
del captured_id
|
||||
|
||||
# Test that the captured object ref is pinned despite having no local
|
||||
# references.
|
||||
ray.get(signal.send.remote())
|
||||
_fill_object_store_and_get(obj_ref)
|
||||
|
||||
|
||||
def test_borrowed_id_failure_while_pulling(ray_start_cluster):
|
||||
"""The driver creates an object and passes the ref to actor A via an actor
|
||||
task. That task passes the ref on to B, then A kills itself before
|
||||
finishing the task, so the task never reports the borrower B to the
|
||||
driver. The driver can therefore erase the ref while B is still pulling
|
||||
the object, and B's get must then fail promptly instead of hanging.
|
||||
"""
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(
|
||||
num_cpus=1,
|
||||
resources={"head_node": 1},
|
||||
object_store_memory=100 * 1024 * 1024,
|
||||
_system_config={
|
||||
"testing_asio_delay_us": (
|
||||
"ObjectManagerService.grpc_server.Pull=5000000000:5000000000"
|
||||
),
|
||||
"metrics_report_interval_ms": 200,
|
||||
},
|
||||
)
|
||||
ray.init(address=cluster.address)
|
||||
cluster.add_node(
|
||||
num_cpus=1,
|
||||
resources={"worker_node": 1},
|
||||
object_store_memory=100 * 1024 * 1024,
|
||||
)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
@ray.remote(resources={"head_node": 1})
|
||||
class A:
|
||||
def pass_ref(self, ref, b):
|
||||
ray.get(b.receive_ref.remote(ref))
|
||||
sys.exit(-1)
|
||||
|
||||
@ray.remote(resources={"worker_node": 1})
|
||||
class B:
|
||||
def __init__(self):
|
||||
self.ref = None
|
||||
|
||||
def receive_ref(self, ref):
|
||||
self.ref = ref[0]
|
||||
|
||||
def resolve_ref(self):
|
||||
with pytest.raises(ray.exceptions.ObjectLostError):
|
||||
ray.get(self.ref)
|
||||
return True
|
||||
|
||||
def ping(self):
|
||||
return
|
||||
|
||||
a = A.remote()
|
||||
b = B.remote()
|
||||
ray.get(b.ping.remote())
|
||||
|
||||
obj = ray.put(np.zeros(1024 * 1024, dtype=np.uint8))
|
||||
with pytest.raises(ray.exceptions.RayActorError):
|
||||
ray.get(a.pass_ref.remote([obj], b))
|
||||
|
||||
resolved = b.resolve_ref.remote()
|
||||
|
||||
def pull_in_flight():
|
||||
(worker_node,) = [n for n in ray.nodes() if "worker_node" in n["Resources"]]
|
||||
address = (
|
||||
f"{worker_node['NodeManagerAddress']}:{worker_node['MetricsExportPort']}"
|
||||
)
|
||||
samples = fetch_prometheus_metrics([address]).get(
|
||||
"ray_pull_manager_usage_bytes", []
|
||||
)
|
||||
return any(
|
||||
s.labels.get("Type") == "BeingPulled" and s.value > 0 for s in samples
|
||||
)
|
||||
|
||||
# make sure to only del the last ref to let ref count go to 0 after actor B's raylet starts pulling that object.
|
||||
wait_for_condition(pull_in_flight, timeout=30)
|
||||
del obj
|
||||
assert ray.get(resolved, timeout=30)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
Reference in New Issue
Block a user