Skip to content

Commit

Permalink
Avoid a crash when there is a timeout when shutting down the Dask clu…
Browse files Browse the repository at this point in the history
…ster (#2580)
  • Loading branch information
bouweandela authored Dec 3, 2024
1 parent 0f9e888 commit 4c36a0c
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 3 deletions.
9 changes: 8 additions & 1 deletion esmvalcore/config/_dask.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,4 +80,11 @@ def get_distributed_client():
if client is not None:
client.close()
if cluster is not None:
cluster.close()
try:
cluster.close()
except TimeoutError:
logger.warning(
"Timeout while trying to shut down the cluster at %s, "
"you may want to check it was stopped.",
cluster.scheduler_address,
)
7 changes: 5 additions & 2 deletions tests/unit/config/test_dask.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,8 @@ def test_get_distributed_client_external(mocker, tmp_path, warn_unused_args):
mock_client.close.assert_called()


def test_get_distributed_client_slurm(mocker, tmp_path):
@pytest.mark.parametrize("shutdown_timeout", [False, True])
def test_get_distributed_client_slurm(mocker, tmp_path, shutdown_timeout):
cfg = {
"cluster": {
"type": "dask_jobqueue.SLURMCluster",
Expand Down Expand Up @@ -66,10 +67,12 @@ def test_get_distributed_client_slurm(mocker, tmp_path):
create_autospec=True,
return_value=mock_module,
)
mock_cluster = mock_cluster_cls.return_value
if shutdown_timeout:
mock_cluster.close.side_effect = TimeoutError
with _dask.get_distributed_client() as client:
assert client is mock_client
mock_client.close.assert_called()
mock_cluster = mock_cluster_cls.return_value
_dask.Client.assert_called_with(address=mock_cluster.scheduler_address)
args = {k: v for k, v in cfg["cluster"].items() if k != "type"}
mock_cluster_cls.assert_called_with(**args)
Expand Down

0 comments on commit 4c36a0c

Please sign in to comment.