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()
