async def test_worker_keys(): """Ensure that redefining adapt with a lower maximum removes workers""" async with SpecCluster( scheduler={"cls": Scheduler, "options": {"dashboard_address": ":0"}}, workers={ "a-1": {"cls": Worker}, "a-2": {"cls": Worker}, "b-1": {"cls": Worker}, "b-2": {"cls": Worker}, }, asynchronous=True, ) as cluster: def key(ws): return ws.name.split("-")[0] cluster._adaptive_options = {"worker_key": key} adaptive = cluster.adapt(minimum=1) await adaptive.adapt() while len(cluster.scheduler.workers) == 4: await asyncio.sleep(0.01) names = {ws.name for ws in cluster.scheduler.workers.values()} assert names == {"a-1", "a-2"} or names == {"b-1", "b-2"}
def test_worker_keys(): """ Ensure that redefining adapt with a lower maximum removes workers """ cluster = yield SpecCluster( workers={ "a-1": {"cls": Worker}, "a-2": {"cls": Worker}, "b-1": {"cls": Worker}, "b-2": {"cls": Worker}, }, asynchronous=True, ) try: def key(ws): return ws.name.split("-")[0] cluster._adaptive_options = {"worker_key": key} adaptive = cluster.adapt(minimum=1) yield adaptive._adapt() while len(cluster.scheduler.workers) == 4: yield gen.sleep(0.01) names = {ws.name for ws in cluster.scheduler.workers.values()} assert names == {"a-1", "a-2"} or names == {"b-1", "b-2"} finally: yield cluster.close()