Ensure retire_workers works if AMM extension hasn't been loaded (#…

…7863)
dask · May 30, 2023 · 61efdfe · 61efdfe
1 parent 8dad1fc
commit 61efdfe
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 4 deletions.
diff --git a/distributed/active_memory_manager.py b/distributed/active_memory_manager.py
@@ -108,7 +108,6 @@ def __init__(
         self.measure = measure
 
         if register:
-            scheduler.extensions["amm"] = self
             scheduler.handlers["amm_handler"] = self.amm_handler
 
         if interval is None:

diff --git a/distributed/diagnostics/memory_sampler.py b/distributed/diagnostics/memory_sampler.py
@@ -185,7 +185,6 @@ class MemorySamplerExtension:
 
     def __init__(self, scheduler: Scheduler):
         self.scheduler = scheduler
-        self.scheduler.extensions["memory_sampler"] = self
         self.scheduler.handlers["memory_sampler_start"] = self.start
         self.scheduler.handlers["memory_sampler_stop"] = self.stop
         self.samples = {}

diff --git a/distributed/scheduler.py b/distributed/scheduler.py
@@ -6876,8 +6876,8 @@ async def retire_workers(
                 return {}
 
             stop_amm = False
-            amm: ActiveMemoryManagerExtension = self.extensions["amm"]
-            if not amm.running:
+            amm: ActiveMemoryManagerExtension | None = self.extensions.get("amm")
+            if not amm or not amm.running:
                 amm = ActiveMemoryManagerExtension(
                     self, policies=set(), register=False, start=True, interval=2.0
                 )

diff --git a/distributed/tests/test_active_memory_manager.py b/distributed/tests/test_active_memory_manager.py
@@ -759,6 +759,15 @@ async def test_RetireWorker_amm_on_off(c, s, a, b, start_amm):
     assert "x" in b.data
 
 
+@gen_cluster(client=True, scheduler_kwargs={"extensions": {}})
+async def test_RetireWorker_no_extension(c, s, a, b):
+    """retire_workers must work when the AMM extension is not loaded"""
+    futures = await c.scatter({"x": 1}, workers=[a.address])
+    await c.retire_workers([a.address])
+    assert a.address not in s.workers
+    assert "x" in b.data
+
+
 @gen_cluster(
     client=True,
     config={