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