wehub-resource-sync
397 lines
12 KiB
Python
397 lines
12 KiB
Python
# coding: utf-8
|
|
|
|
import gc
|
|
import logging
|
|
import sys
|
|
import time
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray._private.test_utils import client_test_enabled
|
|
from ray._private.worker import _wait_generators_bulk
|
|
from ray.exceptions import ObjectRefStreamEndOfStreamError, RayTaskError
|
|
|
|
if client_test_enabled():
|
|
from ray.util.client import ray
|
|
else:
|
|
import ray
|
|
import ray.util.state
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def test_wait(ray_start_regular):
|
|
@ray.remote
|
|
def f(delay):
|
|
time.sleep(delay)
|
|
return
|
|
|
|
object_refs = [f.remote(0), f.remote(0), f.remote(0), f.remote(0)]
|
|
ready_ids, remaining_ids = ray.wait(object_refs)
|
|
assert len(ready_ids) == 1
|
|
assert len(remaining_ids) == 3
|
|
ready_ids, remaining_ids = ray.wait(object_refs, num_returns=4)
|
|
assert set(ready_ids) == set(object_refs)
|
|
assert remaining_ids == []
|
|
|
|
object_refs = [f.remote(0), f.remote(5)]
|
|
ready_ids, remaining_ids = ray.wait(object_refs, timeout=0.5, num_returns=2)
|
|
assert len(ready_ids) == 1
|
|
assert len(remaining_ids) == 1
|
|
|
|
# Verify that calling wait with duplicate object refs throws an
|
|
# exception.
|
|
x = ray.put(1)
|
|
with pytest.raises(Exception):
|
|
ray.wait([x, x])
|
|
|
|
# Make sure it is possible to call wait with an empty list.
|
|
ready_ids, remaining_ids = ray.wait([])
|
|
assert ready_ids == []
|
|
assert remaining_ids == []
|
|
|
|
# Test semantics of num_returns with no timeout.
|
|
obj_refs = [ray.put(i) for i in range(10)]
|
|
(found, rest) = ray.wait(obj_refs, num_returns=2)
|
|
assert len(found) == 2
|
|
assert len(rest) == 8
|
|
|
|
# Verify that incorrect usage raises a TypeError.
|
|
x = ray.put(1)
|
|
with pytest.raises(TypeError):
|
|
ray.wait(x)
|
|
with pytest.raises(TypeError):
|
|
ray.wait(1)
|
|
with pytest.raises(TypeError):
|
|
ray.wait([1])
|
|
|
|
|
|
def test_wait_timing(ray_start_2_cpus):
|
|
@ray.remote
|
|
def f():
|
|
time.sleep(1)
|
|
|
|
future = f.remote()
|
|
|
|
start = time.time()
|
|
ready, not_ready = ray.wait([future], timeout=0.2)
|
|
assert 0.2 < time.time() - start < 0.3
|
|
assert len(ready) == 0
|
|
assert len(not_ready) == 1
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test_wait_always_fetch_local(monkeypatch, ray_start_cluster):
|
|
monkeypatch.setenv("RAY_scheduler_report_pinned_bytes_only", "false")
|
|
cluster = ray_start_cluster
|
|
head_node = cluster.add_node(num_cpus=0, object_store_memory=300e6)
|
|
ray.init(address=cluster.address)
|
|
worker_node = cluster.add_node(num_cpus=1, object_store_memory=300e6)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def return_large_object():
|
|
# 100mb so will spill on worker, but not once on head
|
|
return np.zeros(100 * 1024 * 1024, dtype=np.uint8)
|
|
|
|
@ray.remote(num_cpus=0)
|
|
def small_local_task():
|
|
return 1
|
|
|
|
put_on_head = {ray._raylet.RAY_NODE_ID_KEY: head_node.node_id}
|
|
put_on_worker = {ray._raylet.RAY_NODE_ID_KEY: worker_node.node_id}
|
|
x = small_local_task.options(label_selector=put_on_head).remote()
|
|
y = return_large_object.options(label_selector=put_on_worker).remote()
|
|
z = return_large_object.options(label_selector=put_on_worker).remote()
|
|
|
|
# will return when tasks are done
|
|
ray.wait([x, y, z], num_returns=3, fetch_local=False)
|
|
assert (
|
|
ray._private.state.available_resources_per_node()[head_node.node_id][
|
|
"object_store_memory"
|
|
]
|
|
> 250e6
|
|
)
|
|
|
|
# x should be immediately available locally, start fetching y and z
|
|
ray.wait([x, y, z], num_returns=1, fetch_local=True)
|
|
assert (
|
|
ray._private.state.available_resources_per_node()[head_node.node_id][
|
|
"object_store_memory"
|
|
]
|
|
> 250e6
|
|
)
|
|
|
|
time.sleep(5)
|
|
# y, z should be pulled here
|
|
assert (
|
|
ray._private.state.available_resources_per_node()[head_node.node_id][
|
|
"object_store_memory"
|
|
]
|
|
< 150e6
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test__wait_generators_bulk_fetch_local(monkeypatch, ray_start_cluster):
|
|
monkeypatch.setenv("RAY_scheduler_report_pinned_bytes_only", "false")
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=0, object_store_memory=500e6)
|
|
ray.init(address=cluster.address)
|
|
worker_node = cluster.add_node(num_cpus=2, object_store_memory=500e6)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def gen_large_objects():
|
|
# 100mb so the object is stored in plasma.
|
|
yield np.zeros(100 * 1024 * 1024, dtype=np.uint8)
|
|
yield np.ones(100 * 1024 * 1024, dtype=np.uint8)
|
|
|
|
put_on_worker = {ray._raylet.RAY_NODE_ID_KEY: worker_node.node_id}
|
|
gen1 = gen_large_objects.options(label_selector=put_on_worker).remote()
|
|
gen2 = gen_large_objects.options(label_selector=put_on_worker).remote()
|
|
|
|
ready = _wait_generators_bulk(
|
|
[(gen1, [True, False]), (gen2, [False, True])],
|
|
num_return=2,
|
|
timeout=10,
|
|
)
|
|
assert len(ready) == 2
|
|
assert [gen for gen, _ in ready] == [gen1, gen2]
|
|
assert all(len(refs) == 2 for _, refs in ready)
|
|
|
|
assert np.all(ray.get(ready[0][1][0], timeout=0) == 0)
|
|
assert np.all(ray.get(ready[1][1][1], timeout=0) == 1)
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test__wait_generators_bulk_wait_for_at_most_num_return(ray_start_regular):
|
|
@ray.remote
|
|
def gen(base, delays):
|
|
for i, delay in enumerate(delays):
|
|
time.sleep(delay)
|
|
yield base + i
|
|
|
|
gen1 = gen.remote(10, [0, 0, 0])
|
|
gen2 = gen.remote(20, [0, 5])
|
|
|
|
ready = _wait_generators_bulk(
|
|
[(gen1, [True, False]), (gen2, [False, True])],
|
|
num_return=1,
|
|
timeout=2,
|
|
)
|
|
|
|
assert len(ready) == 1
|
|
ready_gen, refs = ready[0]
|
|
assert ready_gen is gen1
|
|
assert ray.get(refs) == [10, 11]
|
|
|
|
# The returned refs are consumed from the stream.
|
|
assert ray.get(next(gen1)) == 12
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test__wait_generators_bulk_timeout(ray_start_regular):
|
|
@ray.remote(num_cpus=0, max_concurrency=2)
|
|
class Signal:
|
|
def __init__(self):
|
|
self.ready = False
|
|
|
|
def wait(self):
|
|
while not self.ready:
|
|
time.sleep(0.01)
|
|
|
|
def send(self):
|
|
self.ready = True
|
|
|
|
@ray.remote
|
|
def slow_gen(signal):
|
|
ray.get(signal.wait.remote())
|
|
yield 1
|
|
|
|
signal = Signal.remote()
|
|
gen = slow_gen.remote(signal)
|
|
assert _wait_generators_bulk([(gen, [False])], timeout=0.01) == []
|
|
ray.get(signal.send.remote())
|
|
ready = _wait_generators_bulk([(gen, [False])], timeout=5)
|
|
assert len(ready) == 1
|
|
ready_gen, refs = ready[0]
|
|
assert ready_gen is gen
|
|
assert ray.get(refs) == [1]
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test__wait_generators_bulk_validation(ray_start_regular):
|
|
@ray.remote
|
|
def gen():
|
|
yield 1
|
|
|
|
gen = gen.remote()
|
|
|
|
with pytest.raises(TypeError):
|
|
_wait_generators_bulk({})
|
|
with pytest.raises(TypeError):
|
|
_wait_generators_bulk([(ray.put(1), [False])])
|
|
with pytest.raises(TypeError):
|
|
_wait_generators_bulk([(gen, False)])
|
|
with pytest.raises(ValueError):
|
|
_wait_generators_bulk([(gen, [])])
|
|
with pytest.raises(ValueError):
|
|
_wait_generators_bulk([(gen, [False])], num_return=2)
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test__consume_next_ref_n_rejects_unready(ray_start_regular):
|
|
"""Consuming before the last requested ref is ready must raise rather than
|
|
silently advancing past (and dropping) the not-yet-produced object."""
|
|
|
|
@ray.remote(num_cpus=0, max_concurrency=2)
|
|
class Signal:
|
|
def __init__(self):
|
|
self.ready = False
|
|
|
|
def wait(self):
|
|
while not self.ready:
|
|
time.sleep(0.01)
|
|
|
|
def send(self):
|
|
self.ready = True
|
|
|
|
@ray.remote
|
|
def slow_gen(signal):
|
|
ray.get(signal.wait.remote())
|
|
yield 1
|
|
|
|
signal = Signal.remote()
|
|
gen = slow_gen.remote(signal)
|
|
|
|
# Peek without waiting: the ref isn't produced yet, so consuming it is rejected.
|
|
gen._get_next_ref_n(1)
|
|
with pytest.raises(ValueError):
|
|
gen._consume_next_ref_n(1)
|
|
|
|
# After the value is produced, the same generator consumes normally.
|
|
ray.get(signal.send.remote())
|
|
ready = _wait_generators_bulk([(gen, [False])], timeout=10)
|
|
assert len(ready) == 1
|
|
assert ray.get(ready[0][1]) == [1]
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test__wait_generators_bulk_after_eof_raise_EndOfStreamError(ray_start_regular):
|
|
@ray.remote
|
|
def empty_gen():
|
|
if False:
|
|
yield 1
|
|
|
|
empty = empty_gen.remote()
|
|
ready = _wait_generators_bulk([(empty, [True, True, True])], timeout=1)
|
|
assert len(ready) == 1
|
|
ready_gen, refs = ready[0]
|
|
assert ready_gen is empty
|
|
assert len(set(refs)) == 3
|
|
for ref in refs:
|
|
with pytest.raises(ObjectRefStreamEndOfStreamError):
|
|
ray.get(ref)
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test__wait_generators_bulk_after_partial_eof(ray_start_regular):
|
|
@ray.remote
|
|
def one_item_gen():
|
|
yield 1
|
|
|
|
one_item = one_item_gen.remote()
|
|
ready = _wait_generators_bulk([(one_item, [False, False, False])], timeout=1)
|
|
assert len(ready) == 1
|
|
ready_gen, refs = ready[0]
|
|
assert ready_gen is one_item
|
|
assert len(set(refs)) == 3
|
|
assert ray.get(refs[0]) == 1
|
|
for ref in refs[1:]:
|
|
with pytest.raises(ObjectRefStreamEndOfStreamError):
|
|
ray.get(ref)
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test__wait_generators_bulk_after_partial_error(ray_start_regular):
|
|
@ray.remote
|
|
def one_item_then_error_gen():
|
|
yield 1
|
|
raise ValueError("expected test error")
|
|
|
|
one_item_then_error = one_item_then_error_gen.remote()
|
|
ready = _wait_generators_bulk(
|
|
[(one_item_then_error, [False, False, False])], timeout=1
|
|
)
|
|
assert len(ready) == 1
|
|
ready_gen, refs = ready[0]
|
|
assert ready_gen is one_item_then_error
|
|
assert len(set(refs)) == 3
|
|
assert ray.get(refs[0]) == 1
|
|
with pytest.raises(RayTaskError) as exc_info:
|
|
ray.get(refs[1])
|
|
assert isinstance(exc_info.value.as_instanceof_cause(), ValueError)
|
|
with pytest.raises(ObjectRefStreamEndOfStreamError):
|
|
ray.get(refs[2])
|
|
|
|
|
|
def _assert_no_owned_refs_leak():
|
|
"""Wait until the owner holds no live references and the store is empty."""
|
|
|
|
def check():
|
|
gc.collect()
|
|
core_worker = ray._private.worker.global_worker.core_worker
|
|
ref_counts = core_worker.get_all_reference_counts()
|
|
for rc in ref_counts.values():
|
|
if rc["local"] != 0 or rc["submitted"] != 0:
|
|
return False
|
|
return core_worker.get_memory_store_size() == 0
|
|
|
|
wait_for_condition(check, timeout=30)
|
|
|
|
|
|
@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
|
|
def test__wait_generators_bulk_no_ref_leak(ray_start_regular):
|
|
"""Draining a generator entirely via _wait_generators_bulk must not leak
|
|
owner-side references for the consumed objects.
|
|
|
|
The bulk peek (peek_object_ref_stream_n) hands back ObjectRefs that add their
|
|
own local reference, while the owner-side stream reference taken at peek/report
|
|
time is only released for *unconsumed* refs at stream teardown. This test
|
|
confirms whether consumed refs leave that owner-side reference dangling.
|
|
"""
|
|
|
|
@ray.remote
|
|
def gen():
|
|
for i in range(3):
|
|
yield i
|
|
|
|
g = gen.remote()
|
|
|
|
collected = []
|
|
saw_eof = False
|
|
while not saw_eof:
|
|
ready = _wait_generators_bulk([(g, [True, True, True])], timeout=10)
|
|
assert len(ready) == 1
|
|
# Avoid binding the generator object to a local (e.g. via tuple unpacking),
|
|
# which would keep its stream alive and prevent teardown.
|
|
refs = ready[0][1]
|
|
for ref in refs:
|
|
try:
|
|
collected.append(ray.get(ref))
|
|
except ObjectRefStreamEndOfStreamError:
|
|
saw_eof = True
|
|
break
|
|
|
|
assert collected == [0, 1, 2]
|
|
|
|
# Drop every handle to the generator and its (consumed) objects.
|
|
del g, ready, refs, ref
|
|
|
|
_assert_no_owned_refs_leak()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|