wehub-resource-sync
297 lines
7.7 KiB
Python
297 lines
7.7 KiB
Python
import asyncio
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import Semaphore
|
|
|
|
|
|
def test_nested_tasks(shutdown_only):
|
|
ray.init(num_cpus=1)
|
|
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self):
|
|
self.count = 0
|
|
|
|
def inc(self):
|
|
self.count += 1
|
|
# Since we relex the cap after a timeout we can have slightly more
|
|
# than 1 task. We should never have 20 though since that takes 2^20
|
|
# * 10ms time.
|
|
assert self.count < 20
|
|
|
|
def dec(self):
|
|
self.count -= 1
|
|
|
|
counter = Counter.remote()
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def g():
|
|
return None
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def f():
|
|
ray.get(counter.inc.remote())
|
|
res = ray.get(g.remote())
|
|
ray.get(counter.dec.remote())
|
|
return res
|
|
|
|
ready, _ = ray.wait(
|
|
[f.remote() for _ in range(1000)], timeout=60.0, num_returns=1000
|
|
)
|
|
assert len(ready) == 1000, len(ready)
|
|
# Ensure the assertion in `inc` didn't fail.
|
|
ray.get(ready)
|
|
|
|
|
|
def test_recursion(shutdown_only):
|
|
ray.init(num_cpus=1)
|
|
|
|
@ray.remote
|
|
def summer(n):
|
|
if n == 0:
|
|
return 0
|
|
return n + ray.get(summer.remote(n - 1))
|
|
|
|
assert ray.get(summer.remote(10)) == sum(range(11))
|
|
|
|
|
|
def test_out_of_order_scheduling(shutdown_only):
|
|
"""Ensure that when a task runs before its dependency, and they're of the same
|
|
scheduling class, the dependency is eventually able to run."""
|
|
ray.init(num_cpus=1)
|
|
|
|
@ray.remote
|
|
def foo(arg, path):
|
|
(ref,) = arg
|
|
should_die = not os.path.exists(path)
|
|
with open(path, "w") as f:
|
|
f.write("")
|
|
if should_die:
|
|
print("dying!!!")
|
|
os._exit(-1)
|
|
|
|
if ref:
|
|
print("hogging the only available slot for a while")
|
|
ray.get(ref)
|
|
return "done!"
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
path = f"{tmpdir}/temp.txt"
|
|
first = foo.remote((None,), path)
|
|
second = foo.remote((first,), path)
|
|
print(ray.get(second))
|
|
|
|
|
|
def test_limit_concurrency(shutdown_only):
|
|
ray.init(num_cpus=1)
|
|
|
|
block_task = Semaphore.remote(0)
|
|
block_driver = Semaphore.remote(0)
|
|
|
|
ray.get([block_task.locked.remote(), block_driver.locked.remote()])
|
|
|
|
@ray.remote(num_cpus=1)
|
|
def foo():
|
|
ray.get(block_driver.release.remote())
|
|
ray.get(block_task.acquire.remote())
|
|
|
|
refs = [foo.remote() for _ in range(20)]
|
|
|
|
block_driver_refs = [block_driver.acquire.remote() for _ in range(20)]
|
|
|
|
# Some of the tasks will run since we relax the cap, but not all because it
|
|
# should take exponentially long for the cap to be increased.
|
|
ready, not_ready = ray.wait(block_driver_refs, timeout=10, num_returns=20)
|
|
assert len(not_ready) >= 1
|
|
|
|
# Now the first instance of foo finishes, so the second starts to run.
|
|
ray.get([block_task.release.remote() for _ in range(19)])
|
|
|
|
ready, not_ready = ray.wait(block_driver_refs, timeout=10, num_returns=20)
|
|
assert len(not_ready) == 0
|
|
|
|
ready, not_ready = ray.wait(refs, num_returns=20, timeout=15)
|
|
assert len(ready) == 19
|
|
assert len(not_ready) == 1
|
|
|
|
|
|
def test_zero_cpu_scheduling(shutdown_only):
|
|
ray.init(num_cpus=1)
|
|
|
|
block_task = Semaphore.remote(0)
|
|
block_driver = Semaphore.remote(0)
|
|
|
|
@ray.remote(num_cpus=0)
|
|
def foo():
|
|
ray.get(block_driver.release.remote())
|
|
ray.get(block_task.acquire.remote())
|
|
|
|
foo.remote()
|
|
foo.remote()
|
|
|
|
ray.get(block_driver.acquire.remote())
|
|
|
|
block_driver_ref = block_driver.acquire.remote()
|
|
|
|
# Both tasks should be running, so the driver should be unblocked.
|
|
timeout_value = 5 if sys.platform == "win32" else 1
|
|
_, not_ready = ray.wait([block_driver_ref], timeout=timeout_value)
|
|
assert len(not_ready) == 0
|
|
|
|
|
|
def test_exponential_wait(shutdown_only):
|
|
ray.init(num_cpus=2)
|
|
|
|
num_tasks = 6
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class Barrier:
|
|
def __init__(self, limit):
|
|
self.i = 0
|
|
self.limit = limit
|
|
|
|
async def join(self):
|
|
self.i += 1
|
|
while self.i < self.limit:
|
|
await asyncio.sleep(1)
|
|
|
|
b = Barrier.remote(num_tasks)
|
|
|
|
@ray.remote
|
|
def f(i, start):
|
|
delta = time.time() - start
|
|
print("Launch", i, delta)
|
|
ray.get(b.join.remote())
|
|
return delta
|
|
|
|
start = time.time()
|
|
results = ray.get([f.remote(i, start) for i in range(num_tasks)])
|
|
|
|
last_wait = results[-1] - results[-2]
|
|
second_last = results[-2] - results[-3]
|
|
|
|
# Assert that last_wwait / second_last ~= 2, with a healthy buffer since ci
|
|
# is noisy.
|
|
assert second_last < last_wait < 4 * second_last
|
|
assert 7 < last_wait
|
|
|
|
|
|
def test_spillback(ray_start_cluster):
|
|
"""Ensure that we can spillback without waiting for the worker cap to be lifed"""
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(
|
|
num_cpus=1,
|
|
resources={"head": 1},
|
|
_system_config={
|
|
"worker_cap_initial_backoff_delay_ms": 36000_000,
|
|
"worker_cap_max_backoff_delay_ms": 36000_000,
|
|
},
|
|
)
|
|
cluster.wait_for_nodes()
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class Counter:
|
|
def __init__(self):
|
|
self.i = 0
|
|
|
|
def inc(self):
|
|
self.i = self.i + 1
|
|
|
|
def get(self):
|
|
return self.i
|
|
|
|
counter = Counter.remote()
|
|
|
|
@ray.remote
|
|
def get_node_id():
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
@ray.remote
|
|
def func(i, counter):
|
|
if i == 0:
|
|
counter.inc.remote()
|
|
while True:
|
|
time.sleep(1)
|
|
else:
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
refs = [func.remote(i, counter) for i in range(2)]
|
|
|
|
# Make sure the first task is running
|
|
# and the second task is in the dispatch queue hitting the worker cap
|
|
while ray.get(counter.get.remote()) != 1:
|
|
time.sleep(0.1)
|
|
time.sleep(1)
|
|
|
|
# A new node is added,
|
|
# the second task should be spilled back to it
|
|
# instead of waiting for the cap to be lifted on the head node after 10h.
|
|
cluster.add_node(
|
|
num_cpus=1,
|
|
resources={"worker": 1},
|
|
)
|
|
worker_node_id = ray.get(
|
|
get_node_id.options(num_cpus=0, resources={"worker": 1}).remote()
|
|
)
|
|
assert ray.get(refs[1]) == worker_node_id
|
|
|
|
ray.cancel(refs[0], force=True)
|
|
|
|
|
|
def test_idle_workers(shutdown_only):
|
|
ray.init(
|
|
num_cpus=2,
|
|
_system_config={
|
|
"idle_worker_killing_time_threshold_ms": 10,
|
|
},
|
|
)
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class Actor:
|
|
def get(self):
|
|
pass
|
|
|
|
@ray.remote
|
|
def getpid():
|
|
time.sleep(0.1)
|
|
return os.getpid()
|
|
|
|
# We start exactly as many workers as there are CPUs.
|
|
for _ in range(3):
|
|
pids = set(ray.get([getpid.remote() for _ in range(4)]))
|
|
assert len(pids) <= 2, pids
|
|
# Wait for at least the idle worker timeout.
|
|
time.sleep(0.1)
|
|
|
|
# Now test with two actors that uses 1 process each but 0 CPUs.
|
|
a1 = Actor.remote()
|
|
a2 = Actor.remote()
|
|
ray.get([a1.get.remote(), a2.get.remote()])
|
|
|
|
for _ in range(3):
|
|
pids = set(ray.get([getpid.remote() for _ in range(4)]))
|
|
assert len(pids) <= 2, pids
|
|
# Wait for at least the idle worker timeout.
|
|
time.sleep(0.1)
|
|
|
|
# Kill the actors and test again.
|
|
del a1
|
|
del a2
|
|
for _ in range(3):
|
|
pids = set(ray.get([getpid.remote() for _ in range(4)]))
|
|
assert len(pids) <= 2, pids
|
|
# Wait for at least the idle worker timeout.
|
|
time.sleep(0.1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
os.environ["RAY_worker_cap_enabled"] = "true"
|
|
sys.exit(pytest.main(["-sv", __file__]))
|