From f226a459777036a49d19745c3807fe96e283467c Mon Sep 17 00:00:00 2001 From: Abhinav Gorrepati Date: Tue, 11 Aug 2026 16:31:57 -0700 Subject: [PATCH] Fix late worker alias restrictions Resolve a worker alias that connects after task submission to its exact address instead of expanding the restriction to every worker on the same host. Add a regression test for mixed present and late worker names. Signed-off-by: Abhinav Gorrepati --- distributed/scheduler.py | 19 +++++++++++-------- distributed/tests/test_scheduler.py | 24 ++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/distributed/scheduler.py b/distributed/scheduler.py index 8f4ba87da3..15437587dc 100644 --- a/distributed/scheduler.py +++ b/distributed/scheduler.py @@ -3215,17 +3215,20 @@ def valid_workers(self, ts: TaskState) -> set[WorkerState] | None: s = {addr for addr in ts.worker_restrictions if addr in self.workers} if ts.host_restrictions: - # Resolve the alias here rather than early, for the worker - # may not be connected when host_restrictions is populated - hr = [self.coerce_hostname(h) for h in ts.host_restrictions] - # XXX need HostState? - sl = [] - for h in hr: + ss: set[str] = set() + for h in ts.host_restrictions: + # Resolve the alias here rather than early, for the worker may not be + # connected when host_restrictions is populated. An alias identifies + # one worker, not every worker on the same host. + if (addr := self.aliases.get(h)) is not None: + ss.add(addr) + continue + + assert isinstance(h, str) dh = self.host_info.get(h) if dh is not None: - sl.append(dh["addresses"]) + ss.update(dh["addresses"]) - ss = set.union(*sl) if sl else set() if s is None: s = ss else: diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py index 8d92afbf18..892588960d 100644 --- a/distributed/tests/test_scheduler.py +++ b/distributed/tests/test_scheduler.py @@ -707,6 +707,30 @@ async def test_no_valid_workers(client, s, a, b, c): await wait_for(x, 0.05) +@gen_cluster(client=True, nthreads=[]) +async def test_valid_workers_with_late_worker_alias(client, s): + async with ( + Worker(s.address, name="a-0") as a0, + Worker(s.address, name="a-1"), + ): + x = client.submit(inc, 1, workers=["a-0", "b-0"]) + assert await x == 2 + + ts = s.tasks[x.key] + assert ts.worker_restrictions == {a0.address} + assert ts.host_restrictions == {"b-0"} + + async with Worker(s.address, name="b-0") as b0: + await async_poll_for( + lambda: s.workers[b0.address].status == Status.running, + timeout=1, + ) + assert s.valid_workers(ts) == { + s.workers[a0.address], + s.workers[b0.address], + } + + @gen_cluster(client=True, nthreads=[("127.0.0.1", 1)] * 3) async def test_no_valid_workers_loose_restrictions(client, s, a, b, c): x = client.submit(inc, 1, workers="127.0.0.5:9999", allow_other_workers=True)