async def test_GatewayCluster_cleanup_atexit():
async with temp_gateway() as g:
def test():
return GatewayCluster(address=g.address,
proxy_address=g.proxy_address)
loop = get_running_loop()
cluster = await loop.run_in_executor(None, test)
assert len(GatewayCluster._instances) == 1
def test_cleanup():
# No warnings raised by cleanup function
with pytest.warns(None) as rec:
cleanup_lingering_clusters()
for r in rec:
assert not issubclass(r.category, UserWarning)
# Cluster is now closed
assert cluster.status == "closed"
# No harm in double running
with pytest.warns(None) as rec:
cleanup_lingering_clusters()
for r in rec:
assert not issubclass(r.category, UserWarning)
await loop.run_in_executor(None, test_cleanup)
async with g.gateway_client() as gateway:
# No cluster running
clusters = await gateway.list_clusters()
assert not clusters
async def test_cluster_widget():
pytest.importorskip("ipywidgets")
def test():
with GatewayCluster(address=g.address,
proxy_address=g.proxy_address) as cluster:
# Smoke test widget
cluster._widget()
template = "<tr><th>Workers</th> <td>%d</td></tr>"
assert (template % 0) in cluster._widget_status()
cluster.scale(1)
timeout = time.time() + 30
while time.time() < timeout:
if cluster.scheduler_info.get("workers"):
break
time.sleep(0.25)
else:
assert False, "didn't scale up in time"
assert (template % 1) in cluster._widget_status()
async with temp_gateway() as g:
loop = get_running_loop()
await loop.run_in_executor(None, test)
async def test_GatewayCluster_shutdown_on_close(tmpdir):
async with temp_gateway(cluster_manager_class=InProcessClusterManager,
temp_dir=str(tmpdir)) as gateway_proc:
def test():
cluster = GatewayCluster(
address=gateway_proc.public_urls.connect_url,
proxy_address=gateway_proc.gateway_urls.connect_url,
)
assert cluster.shutdown_on_close
assert cluster in GatewayCluster._instances
loop = get_running_loop()
await loop.run_in_executor(None, test)
assert len(GatewayCluster._instances) == 0
async with Gateway(
address=gateway_proc.public_urls.connect_url,
proxy_address=gateway_proc.gateway_urls.connect_url,
asynchronous=True,
) as gateway:
# No cluster running
clusters = await gateway.list_clusters()
assert not clusters
async def test_cluster_widget(tmpdir):
pytest.importorskip("ipywidgets")
def test():
with GatewayCluster(
address=gateway_proc.public_urls.connect_url,
proxy_address=gateway_proc.gateway_urls.connect_url,
) as cluster:
# Smoke test widget
cluster._widget()
template = "<tr><th>Workers</th> <td>%d</td></tr>"
assert (template % 0) in cluster._widget_status()
cluster.scale(1)
timeout = time.time() + 30
while time.time() < timeout:
if cluster.scheduler_info.get("workers"):
break
time.sleep(0.25)
else:
assert False, "didn't scale up in time"
assert (template % 1) in cluster._widget_status()
async with temp_gateway(cluster_manager_class=InProcessClusterManager,
temp_dir=str(tmpdir)) as gateway_proc:
loop = get_running_loop()
await loop.run_in_executor(None, test)
async def start_cluster(self):
# Initiate state
loop = get_running_loop()
self.failed = loop.create_future()
self.stop_cluster_called = loop.create_future()
# Start the cluster
async for state in super().start_cluster():
yield state
# Launch a task to kill the cluster soon
self.task_pool.create_task(self.delay_fail_cluster())
async def delete_pod(self, manager, pod_name):
loop = get_running_loop()
if pod_name is not None:
try:
await loop.run_in_executor(
None,
manager.kube_client.delete_namespaced_pod,
pod_name,
manager.namespace,
)
except ApiException as e:
if e.status == 404:
return
raise
async def test_sync_constructors():
def test():
with g.gateway_client(asynchronous=False) as gateway:
with gateway.new_cluster() as cluster:
cluster.scale(1)
client = cluster.get_client()
res = client.submit(lambda x: x + 1, 1).result()
assert res == 2
with gateway.connect(cluster.name) as cluster2:
client2 = cluster2.get_client()
res = client2.submit(lambda x: x + 1, 1).result()
assert res == 2
async with temp_gateway() as g:
loop = get_running_loop()
await loop.run_in_executor(None, test)
async def test_GatewayCluster_shutdown_on_close():
async with temp_gateway() as g:
def test():
cluster = GatewayCluster(address=g.address, proxy_address=g.proxy_address)
assert cluster.shutdown_on_close
assert cluster in GatewayCluster._instances
loop = get_running_loop()
await loop.run_in_executor(None, test)
assert len(GatewayCluster._instances) == 0
async with g.gateway_client() as gateway:
# No cluster running
clusters = await gateway.list_clusters()
assert not clusters
async def cleanup_cluster(self, manager, cluster_state, worker_states):
for state in worker_states + [cluster_state]:
await self.delete_pod(manager, state.get("pod_name"))
secret_name = cluster_state.get("secret_name")
if secret_name is not None:
try:
loop = get_running_loop()
await loop.run_in_executor(
None,
manager.kube_client.delete_namespaced_secret,
secret_name,
manager.namespace,
)
except ApiException as e:
if e.status == 404:
return
raise
async def test_GatewayCluster_cleanup_atexit(tmpdir):
async with temp_gateway(
cluster_manager_class=InProcessClusterManager,
temp_dir=str(tmpdir.join("dask-gateway")),
) as gateway_proc:
def test():
return GatewayCluster(
address=gateway_proc.public_urls.connect_url,
proxy_address=gateway_proc.gateway_urls.connect_url,
)
loop = get_running_loop()
cluster = await loop.run_in_executor(None, test)
assert len(GatewayCluster._instances) == 1
def test_cleanup():
# No warnings raised by cleanup function
with pytest.warns(None) as rec:
cleanup_lingering_clusters()
for r in rec:
assert not issubclass(r.category, UserWarning)
# Cluster is now closed
assert cluster.status == "closed"
# No harm in double running
with pytest.warns(None) as rec:
cleanup_lingering_clusters()
for r in rec:
assert not issubclass(r.category, UserWarning)
await loop.run_in_executor(None, test_cleanup)
async with Gateway(
address=gateway_proc.public_urls.connect_url,
proxy_address=gateway_proc.gateway_urls.connect_url,
asynchronous=True,
) as gateway:
# No cluster running
clusters = await gateway.list_clusters()
assert not clusters
async def test_GatewayCluster_client_error_doesnt_prevent_cleanup():
"""Check that an error on closing clients doesn't prevent cluster shutdown"""
async with temp_gateway() as g:
class BadGatewayCluster(GatewayCluster):
async def _stop_async(self):
await super()._stop_async()
raise ValueError("OH NO")
def test():
cluster = BadGatewayCluster(address=g.address,
proxy_address=g.proxy_address)
assert cluster in GatewayCluster._instances
loop = get_running_loop()
await loop.run_in_executor(None, test)
assert len(GatewayCluster._instances) == 0
async with g.gateway_client() as gateway:
# No cluster running
clusters = await gateway.list_clusters()
assert not clusters
async def do_setup(self):
await super().do_setup()
loop = get_running_loop()
self.stop_cluster_called = loop.create_future()
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
loop = get_running_loop()
self.stop_worker_called = loop.create_future()
self.worker_connected = loop.create_future()
