[unsupervised AI] Fix late worker alias restrictions

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:

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)