chore: import upstream snapshot with attribution

python/ray/tests/test_worker_capping.py

@@ -0,0 +1,296 @@
+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__]))