[unsupervised AI] Schedule resource-restricted tasks against free resources

distributed/scheduler.py

@@ -3196,7 +3196,9 @@ def get_comm_cost(self, ts: TaskState, ws: WorkerState) -> float:
         nbytes = sum(dts.get_nbytes() for dts in deps)
         return nbytes / self.bandwidth
 
-    def valid_workers(self, ts: TaskState) -> set[WorkerState] | None:
+    def valid_workers(
+        self, ts: TaskState, *, only_available: bool = True
+    ) -> set[WorkerState] | None:
         """Return set of currently valid workers for key
 
         If all workers are valid then this returns ``None``, in which case
@@ -3208,6 +3210,15 @@ def valid_workers(self, ts: TaskState) -> set[WorkerState] | None:
         *  worker_restrictions
         *  host_restrictions
         *  resource_restrictions
+
+        Parameters
+        ----------
+        only_available
+            If True (default), a worker only counts for a resource restriction when
+            it has enough of the resource *unused*, so the returned workers can run
+            the task right now. If False, the worker's declared total is used
+            instead, which answers whether the restrictions are satisfiable at all
+            regardless of what is running.
         """
         s: set[str] | None = None
 
@@ -3240,7 +3251,16 @@ def valid_workers(self, ts: TaskState) -> set[WorkerState] | None:
 
                 sw = set()
                 for addr, supplied in dr.items():
-                    if supplied >= required:
+                    available = supplied
+                    if only_available:
+                        # Comparing against the declared total lets the scheduler
+                        # commit more of a resource than the worker has, so
+                        # ws.used_resources overshoots ws.resources while the worker
+                        # holds the surplus tasks back anyway.
+                        ws = self.workers.get(addr)
+                        if ws is not None:
+                            available -= ws.used_resources.get(resource, 0)
+                    if available >= required:
                         sw.add(addr)
 
                 dw[resource] = sw
@@ -4777,6 +4797,7 @@ async def add_worker(
                 self.bulk_schedule_unrunnable_after_adding_worker(ws), stimulus_id
             )
             self.stimulus_queue_slots_maybe_opened(stimulus_id=stimulus_id)
+            self.stimulus_resources_maybe_released(stimulus_id=stimulus_id)
 
         logger.info("Register worker addr: %s name: %s", ws.address, ws.name)
 
@@ -5400,6 +5421,42 @@ def stimulus_queue_slots_maybe_opened(self, *, stimulus_id: str) -> None:
                 assert qts.state == "processing"
                 assert not self.queued or self.queued.peek() != qts
 
+    def stimulus_resources_maybe_released(self, *, stimulus_id: str) -> None:
+        """Respond to an event which may have released resources on workers
+
+        Transitions ``no-worker`` tasks whose resource restrictions can now be
+        satisfied to ``processing``.
+
+        A resource is only freed when a task holding it leaves ``processing``, so
+        tasks that `Scheduler.valid_workers` held back for want of a free resource
+        have to be reconsidered at that point. Without this they would stay
+        unrunnable until a new worker joined, which is the only other event that
+        revisits ``no-worker`` tasks.
+
+        Notes
+        -----
+        Tasks are transitioned one at a time so that each one takes its resources
+        before the next is considered; recommending them in bulk would let tasks
+        that no longer fit bounce straight back to ``no-worker``.
+
+        Other transitions related to this stimulus should be fully processed
+        beforehand, for the same reason as in
+        `Scheduler.stimulus_queue_slots_maybe_opened`.
+        """
+        if not self.unrunnable:
+            return
+
+        # Snapshot first: transitioning mutates self.unrunnable
+        candidates = [ts for ts in self.unrunnable if ts.resource_restrictions]
+        if not candidates:
+            return
+
+        candidates.sort(key=operator.attrgetter("priority"))
+        for ts in candidates:
+            if not self.valid_workers(ts):
+                continue
+            self.transitions({ts.key: "processing"}, stimulus_id)
+
     def stimulus_task_finished(
         self,
         worker: str,
@@ -5835,6 +5892,7 @@ def client_releases_keys(
         self.transitions(recommendations, stimulus_id)
 
         self.stimulus_queue_slots_maybe_opened(stimulus_id=stimulus_id)
+        self.stimulus_resources_maybe_released(stimulus_id=stimulus_id)
 
     def client_heartbeat(self, client: str) -> None:
         """Handle heartbeats from Client"""
@@ -6021,6 +6079,7 @@ def handle_task_finished(
         self.send_all(client_msgs, worker_msgs)
 
         self.stimulus_queue_slots_maybe_opened(stimulus_id=stimulus_id)
+        self.stimulus_resources_maybe_released(stimulus_id=stimulus_id)
 
     def handle_task_erred(self, key: Key, stimulus_id: str, **msg: Any) -> None:
         r: tuple = self.stimulus_task_erred(key=key, stimulus_id=stimulus_id, **msg)
@@ -6029,6 +6088,7 @@ def handle_task_erred(self, key: Key, stimulus_id: str, **msg: Any) -> None:
         self.send_all(client_msgs, worker_msgs)
 
         self.stimulus_queue_slots_maybe_opened(stimulus_id=stimulus_id)
+        self.stimulus_resources_maybe_released(stimulus_id=stimulus_id)
 
     def release_worker_data(self, key: Key, worker: str, stimulus_id: str) -> None:
         ts = self.tasks.get(key)
@@ -6093,6 +6153,7 @@ def handle_long_running(
         self.check_idle_saturated(ws)
 
         self.stimulus_queue_slots_maybe_opened(stimulus_id=stimulus_id)
+        self.stimulus_resources_maybe_released(stimulus_id=stimulus_id)
 
     def handle_worker_status_change(
         self, status: str | Status, worker: str | WorkerState, stimulus_id: str
@@ -6123,6 +6184,7 @@ def handle_worker_status_change(
                 self.bulk_schedule_unrunnable_after_adding_worker(ws), stimulus_id
             )
             self.stimulus_queue_slots_maybe_opened(stimulus_id=stimulus_id)
+            self.stimulus_resources_maybe_released(stimulus_id=stimulus_id)
         else:
             self.running.discard(ws)
             self.idle.pop(ws.address, None)
@@ -8735,6 +8797,21 @@ def _check_no_workers(self) -> None:
                 {"action": "no-workers-timeout-exceeded", "keys": affected},
             )
 
+    def _waiting_for_busy_resources(self, ts: TaskState) -> bool:
+        """Whether *ts* is unrunnable only because the resources it asks for are
+        currently held by other tasks
+
+        Such a task is not waiting on unsatisfiable restrictions: some worker
+        declares enough of the resource, and the task becomes runnable as soon as
+        that resource is released.
+        """
+        if not ts.resource_restrictions:
+            return False
+        if self.valid_workers(ts):
+            # A worker can take it right now, so it is not blocked on resources.
+            return False
+        return bool(self.valid_workers(ts, only_available=False))
+
     def _check_unrunnable_task_timeouts(
         self, timestamp: float, recommendations: Recs, stimulus_id: str
     ) -> set[Key]:
@@ -8746,6 +8823,11 @@ def _check_unrunnable_task_timeouts(
                 # unrunnable is insertion-ordered, which means that unrunnable_since will
                 # be monotonically increasing in this loop.
                 break
+            if self._waiting_for_busy_resources(ts):
+                # The restrictions can be satisfied; the task is only waiting for
+                # other tasks to release the resources it needs. Failing it here
+                # would kill work that is about to become runnable.
+                continue
             if (
                 self._no_workers_since is None
                 or self._no_workers_since >= unrunnable_since

distributed/tests/test_resources.py

@@ -7,9 +7,19 @@
 import dask
 from dask import delayed
 
-from distributed import Lock, Worker
+from distributed import Event, Lock, Worker
 from distributed.client import wait
-from distributed.utils_test import NO_AMM, gen_cluster, inc, lock_inc, slowadd, slowinc
+from distributed.utils_test import (
+    NO_AMM,
+    async_poll_for,
+    block_on_event,
+    gen_cluster,
+    inc,
+    lock_inc,
+    slowadd,
+    slowinc,
+    wait_for_state,
+)
 from distributed.worker_state_machine import (
     ComputeTaskEvent,
     Execute,
@@ -557,3 +567,101 @@ def test_resumed_with_different_resources(ws_with_running_task, done_ev_cls):
 
     ws.handle_stimulus(done_ev_cls.dummy(key="x", stimulus_id="s3"))
     assert ws.available_resources == {"R": 1}
+
+
+@gen_cluster(
+    client=True,
+    nthreads=[("127.0.0.1", 20, {"resources": {"A": 10}})],
+)
+async def test_resources_not_over_allocated(c, s, a):
+    """The scheduler must schedule against the resources a worker still has free,
+    not against the total it declared.
+
+    See https://github.com/dask/distributed/issues/9108
+    """
+    ev = Event()
+    futs = [
+        c.submit(block_on_event, ev, resources={"A": 3}, pure=False, key=f"x-{i}")
+        for i in range(15)
+    ]
+
+    ws = s.workers[a.address]
+    # 3 tasks * 3 A = 9 <= 10; a fourth would exceed what the worker declared.
+    await async_poll_for(lambda: len(ws.processing) == 3, timeout=5)
+    # Give the scheduler a chance to over-commit if it is going to.
+    await asyncio.sleep(0.2)
+
+    assert len(ws.processing) == 3
+    assert ws.used_resources["A"] == 9
+    assert ws.used_resources["A"] <= ws.resources["A"]
+    # The rest are waiting for the resource, not silently reported as running.
+    assert len(s.unrunnable) == 12
+
+    await ev.set()
+    await c.gather(futs)
+
+    assert ws.used_resources["A"] == 0
+    assert not s.unrunnable
+
+
+@gen_cluster(
+    client=True,
+    nthreads=[("127.0.0.1", 20, {"resources": {"A": 10}})],
+)
+async def test_resource_tasks_rescheduled_when_resources_released(c, s, a):
+    """Tasks held back for want of a free resource must be picked up again once
+    running tasks release it, rather than staying in ``no-worker`` forever.
+    """
+    ev = Event()
+    blockers = [
+        c.submit(block_on_event, ev, resources={"A": 5}, pure=False, key=f"b-{i}")
+        for i in range(2)
+    ]
+    await async_poll_for(lambda: len(s.workers[a.address].processing) == 2, timeout=5)
+
+    waiter = c.submit(inc, 1, resources={"A": 5}, key="waiter")
+    await wait_for_state("waiter", "no-worker", s)
+
+    # Releasing the blockers must make the waiter runnable without any new worker
+    # joining the cluster.
+    await ev.set()
+    assert await waiter == 2
+    await c.gather(blockers)
+    assert s.workers[a.address].used_resources["A"] == 0
+
+
+@gen_cluster(
+    client=True,
+    nthreads=[("127.0.0.1", 20, {"resources": {"A": 10}})],
+    config={"distributed.scheduler.no-workers-timeout": "100ms"},
+)
+async def test_no_workers_timeout_does_not_fail_tasks_awaiting_resources(c, s, a):
+    """``no-workers-timeout`` must not fail a task whose restrictions are
+    satisfiable and which is only waiting for a busy resource to be released.
+    """
+    ev = Event()
+    blocker = c.submit(block_on_event, ev, resources={"A": 10}, key="blocker")
+    await async_poll_for(lambda: len(s.workers[a.address].processing) == 1, timeout=5)
+
+    waiter = c.submit(inc, 1, resources={"A": 10}, key="waiter")
+    await wait_for_state("waiter", "no-worker", s)
+
+    # Well past no-workers-timeout: the task must still be alive.
+    await asyncio.sleep(0.5)
+    assert s.tasks["waiter"].state == "no-worker"
+
+    await ev.set()
+    assert await waiter == 2
+    await blocker
+
+
+@gen_cluster(
+    client=True,
+    nthreads=[("127.0.0.1", 20, {"resources": {"A": 10}})],
+    config={"distributed.scheduler.no-workers-timeout": "100ms"},
+)
+async def test_no_workers_timeout_still_fails_unsatisfiable_resources(c, s, a):
+    """A resource restriction that no worker can ever satisfy must still time out."""
+    fut = c.submit(inc, 1, resources={"A": 100}, key="impossible")
+    with pytest.raises(Exception, match="impossible"):
+        await fut

distributed/tests/test_steal.py

@@ -527,13 +527,20 @@ async def test_dont_steal_resource_restrictions(c, s, a, b):
     assert len(b.state.tasks) == 0
 
 
-@gen_cluster(client=True, nthreads=[("", 1, {"resources": {"A": 2}})])
+@gen_cluster(client=True, nthreads=[("", 1, {"resources": {"A": 100}})])
 async def test_steal_resource_restrictions(c, s, a):
+    # Both workers declare enough of A to hold all 100 tasks at once. Previously
+    # this test used A: 2 and A: 4 and still expected all 100 tasks to be assigned
+    # to `a`, which only held because the scheduler committed more of a resource
+    # than the worker had. The point of the test is that stealing rebalances
+    # resource-restricted tasks, so the capacities are now large enough for the
+    # assignment to be legitimate. test_dont_steal_resource_restrictions still
+    # covers a thief that cannot satisfy the restriction at all.
     futures = c.map(slowinc, range(100), delay=0.05, resources={"A": 1})
     while len(a.state.tasks) < 100:
         await asyncio.sleep(0.01)
 
-    async with Worker(s.address, nthreads=1, resources={"A": 4}) as b:
+    async with Worker(s.address, nthreads=1, resources={"A": 100}) as b:
         while s.workers[b.address].status != Status.running:
             await asyncio.sleep(0.01)