From 7aac2109aa0bb1bd2b0debf86d37584e681ac36e Mon Sep 17 00:00:00 2001 From: vedjaw Date: Thu, 6 Aug 2026 02:02:56 +0530 Subject: [PATCH] Schedule resource-restricted tasks against free resources valid_workers() compared a task's resource requirements against each worker's declared total, ignoring what that worker had already been given. The scheduler therefore committed more of a resource than exists: with a worker declaring A: 10 and fifteen tasks asking for A: 3 each, ws.used_resources["A"] reaches 45 while the worker itself only ever runs three of them. Subtracting used_resources is not sufficient on its own. Tasks held back that way land in no-worker, and the only event that revisited no-worker tasks was a worker joining, so they stayed there forever. This adds stimulus_resources_maybe_released(), a sibling of stimulus_queue_slots_maybe_opened() called from the same sites, which retries those tasks one at a time as running tasks release what they hold. no-workers-timeout also needed teaching: a task waiting for a busy resource is not a task with unsatisfiable restrictions, and would otherwise be failed with NoValidWorkerError. valid_workers() grows an only_available argument so the timeout can ask whether the restrictions are satisfiable at all. Resolves #9108 Assisted-by: Claude Fable 5 --- distributed/scheduler.py | 86 ++++++++++++++++++++- distributed/tests/test_resources.py | 112 +++++++++++++++++++++++++++- distributed/tests/test_steal.py | 11 ++- 3 files changed, 203 insertions(+), 6 deletions(-) diff --git a/distributed/scheduler.py b/distributed/scheduler.py index 8f4ba87da3..bed0dfce87 100644 --- a/distributed/scheduler.py +++ b/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 diff --git a/distributed/tests/test_resources.py b/distributed/tests/test_resources.py index 8a6be6090b..0976a293c3 100644 --- a/distributed/tests/test_resources.py +++ b/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 diff --git a/distributed/tests/test_steal.py b/distributed/tests/test_steal.py index 93aba1d5ad..dc1800c4b3 100644 --- a/distributed/tests/test_steal.py +++ b/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)