Add support for closing n workers

Matt711 · Matt711 · commit 15ad9923b57f · 2022-05-16T13:34:17.000-04:00
diff --git a/distributed/http/scheduler/api.py b/distributed/http/scheduler/api.py
@@ -17,8 +17,13 @@ async def post(self):
         scheduler = self.server
         try:
             params = json.loads(self.request.body)
-            workers = params["workers"]
-            workers_info = await scheduler.retire_workers(workers)
+            n_workers = params.get("n", 0)
+            if n_workers:
+                workers = scheduler.workers_to_close(n=n_workers)
+                workers_info = await scheduler.retire_workers(workers=workers)
+            else:
+                workers = params.get("workers", {})
+                workers_info = await scheduler.retire_workers(workers=workers)
             self.write(json.dumps(workers_info))
         except Exception as e:
             self.set_status(500, str(e))
