[unsupervised AI] Preserve worker name restrictions across restarts

distributed/scheduler.py

@@ -1347,12 +1347,12 @@ class TaskState:
     #: certain hosts. A hostname may correspond to one or several connected workers.
     host_restrictions: set[str] | None
 
-    #: A set of complete worker addresses where this can be run (or ``None`` if empty).
+    #: A set of worker addresses or names where this can be run (or ``None`` if empty).
     #: Usually this is empty unless the task has been specifically restricted to only
     #: run on certain workers.
-    #: Note this is tracking worker addresses, not worker states, since the specific
-    #: workers may not be connected at this time.
-    worker_restrictions: set[str] | None
+    #: Note this is tracking worker addresses or names, not worker states, since the
+    #: specific workers may not be connected at this time.
+    worker_restrictions: set[Hashable] | None
 
     #: Resources required by this task, such as ``{'gpu': 1}`` or ``{'memory': 1e9}``
     #: These are user-defined names and are matched against the : contents of each
@@ -3196,6 +3196,15 @@ def get_comm_cost(self, ts: TaskState, ws: WorkerState) -> float:
         nbytes = sum(dts.get_nbytes() for dts in deps)
         return nbytes / self.bandwidth
 
+    def _resolve_worker_restrictions(
+        self, worker_restrictions: Collection[Hashable]
+    ) -> set[str]:
+        return {
+            address
+            for worker in worker_restrictions
+            if isinstance(address := self.aliases.get(worker, worker), str)
+        }
+
     def valid_workers(self, ts: TaskState) -> set[WorkerState] | None:
         """Return set of currently valid workers for key
 
@@ -3212,7 +3221,11 @@ def valid_workers(self, ts: TaskState) -> set[WorkerState] | None:
         s: set[str] | None = None
 
         if ts.worker_restrictions:
-            s = {addr for addr in ts.worker_restrictions if addr in self.workers}
+            s = {
+                addr
+                for addr in self._resolve_worker_restrictions(ts.worker_restrictions)
+                if addr in self.workers
+            }
 
         if ts.host_restrictions:
             # Resolve the alias here rather than early, for the worker
@@ -5285,6 +5298,9 @@ def _apply_annotations(
                     host_restrictions = set()
                     worker_restrictions = set()
                     for w in value:
+                        if w in self.aliases:
+                            worker_restrictions.add(w)
+                            continue
                         try:
                             w = self.coerce_address(w)
                         except ValueError:
@@ -9258,7 +9274,7 @@ def __init__(
         self,
         task: Key,
         host_restrictions: set[str],
-        worker_restrictions: set[str],
+        worker_restrictions: set[Hashable],
         resource_restrictions: dict[str, float],
         timeout: float,
     ):

distributed/shuffle/_scheduler_plugin.py

@@ -277,7 +277,9 @@ def _calculate_worker_for(self, spec: ShuffleSpec) -> dict[Any, str]:
         barrier = self.scheduler.tasks[barrier_key(shuffle_id)]
 
         if barrier.worker_restrictions:
-            workers = list(barrier.worker_restrictions)
+            workers = list(
+                self.scheduler._resolve_worker_restrictions(barrier.worker_restrictions)
+            )
         else:
             workers = list(self.scheduler.workers)
 

distributed/tests/test_client.py

@@ -4929,6 +4929,23 @@ async def test_restart_workers_by_name(c, s, a, b, by_name):
     assert results == {a.name if by_name else a_addr: "OK", b_addr: "OK"}
 
 
+@pytest.mark.slow
+@gen_cluster(client=True, Worker=Nanny, nthreads=[("", 1)])
+async def test_restart_worker_preserves_name_restrictions(c, s, a):
+    event = Event()
+    future = c.submit(block_on_event, event, workers=[a.name])
+    await wait_for_state(future.key, "processing", s)
+
+    old_address = a.worker_address
+    await c.restart_workers([a.name])
+    while a.worker_address == old_address:
+        await asyncio.sleep(0.01)
+
+    assert s.valid_workers(s.tasks[future.key]) == set(s.workers.values())
+    await event.set()
+    await future
+
+
 class MyException(Exception):
     pass