change default dask client

Author: stevebachmeier

src/pseudopeople/interface.py

@@ -111,7 +111,7 @@ def _generate_dataset(
         noised_dataset = pd.concat(noised_datasets_list, ignore_index=True)
 
         noised_dataset = coerce_dtypes(noised_dataset, dataset_schema)
-    else:
+    else:  # dask
         try:
             from distributed.client import default_client
 
@@ -126,6 +126,8 @@ def _generate_dataset(
         import dask
         import dask.dataframe as dd
 
+        set_up_dask_client()
+
         # Our work depends on the particulars of how dtypes work, and is only
         # built to work with NumPy dtypes, so we turn off the Dask default behavior
         # of using PyArrow dtypes.
@@ -931,3 +933,24 @@ def get_dataset_filepaths(source: Path, dataset_name: str) -> list[Path]:
     dataset_paths = [x for x in directory.glob(f"{dataset_name}*")]
     sorted_dataset_paths = sorted(dataset_paths)
     return sorted_dataset_paths
+
+
+def set_up_dask_client() -> None:
+    """Sets up a Dask client if one is not already running."""
+    from dask.distributed import get_client
+
+    # Determine whether or not a Dask client is already running. If not,
+    # create a new one.
+    try:
+        client = get_client()
+    except ValueError:
+        # No Dask client is running so we create one.
+        from dask.distributed import LocalCluster
+        from dask.system import CPU_COUNT
+
+        # extract the memory limit from the environment variable
+        cluster = LocalCluster(
+            n_workers=CPU_COUNT,
+            threads_per_worker=1,
+        )
+        client = cluster.get_client()

tests/unit/test_interface.py

@@ -1,17 +1,24 @@
+import os
 from pathlib import Path
 
+import numpy as np
+import psutil
 import pytest
 from _pytest.tmpdir import TempPathFactory
+from dask.distributed import LocalCluster, get_client
+from dask.system import CPU_COUNT
 from packaging.version import parse
 from pytest_mock import MockerFixture
 
 from pseudopeople.constants.metadata import DatasetNames
 from pseudopeople.exceptions import DataSourceError
 from pseudopeople.interface import (
     _get_data_changelog_version,
+    set_up_dask_client,
     validate_source_compatibility,
 )
 from pseudopeople.schema_entities import DATASET_SCHEMAS
+from tests.utilities import is_on_slurm
 
 CENSUS = DATASET_SCHEMAS.get_dataset_schema(DatasetNames.CENSUS)
 
@@ -94,3 +101,50 @@ def test_validate_source_compatibility_wrong_directory(tmp_path: Path) -> None:
     bad_path.mkdir()
     with pytest.raises(FileNotFoundError, match="Could not find 'decennial_census' in"):
         validate_source_compatibility(bad_path, CENSUS)
+
+
+def test_set_up_dask_client_default() -> None:
+
+    # There should be no dask client yet
+    with pytest.raises(ValueError):
+        client = get_client()
+
+    set_up_dask_client()
+    client = get_client()
+    workers = client.scheduler_info()["workers"]
+    assert len(workers) == CPU_COUNT
+    assert all(worker["nthreads"] == 1 for worker in workers.values())
+    if is_on_slurm():
+        try:
+            available_memory = float(os.environ["SLURM_MEM_PER_NODE"]) / 1024
+        except KeyError:
+            raise RuntimeError(
+                "You are on Slurm but SLURM_MEM_PER_NODE is not set. "
+                "It is likely that you are SSHed onto a node (perhaps using VSCode). "
+                "In this case, dask will assign the total memory of the node to each "
+                "worker instead of the allocated memory from the srun call. "
+                "Pseudopeople should only be used on Slurm directly on the node "
+                "assigned via an srun (both for pytests as well as actual work)."
+            )
+    else:
+        available_memory = psutil.virtual_memory().total / (1024 ** 3)
+    assert np.isclose(sum(worker["memory_limit"] / 1024**3 for worker in workers.values()), available_memory, rtol=0.01)
+
+
+def test_set_up_dask_client_custom() -> None:
+    memory_limit = 1  # gb
+    n_workers = 3
+    cluster = LocalCluster(
+        name="custom",
+        n_workers=n_workers,
+        threads_per_worker=2,
+        memory_limit=memory_limit * 1024**3,
+    )
+    client = cluster.get_client()
+    set_up_dask_client()
+    client = get_client()
+    assert client.cluster.name == "custom"
+    workers = client.scheduler_info()["workers"]
+    assert len(workers) == 3
+    assert all(worker["nthreads"] == 2 for worker in workers.values())
+    assert sum(worker["memory_limit"] / 1024**3 for worker in workers.values()) == memory_limit * n_workers

tests/unit/test_utils.py

@@ -1,3 +1,4 @@
+
 import numpy as np
 import pandas as pd
 import pytest

tests/utilities.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 import math
+import shutil
 from collections.abc import Callable
-from functools import partial
 from typing import Any
 
 import numpy as np
@@ -190,3 +190,15 @@ def get_single_noise_type_config(
                     ] = new_probability
 
     return config_dict
+
+
+def is_on_slurm() -> bool:
+    """Returns True if the current environment is a SLURM cluster.
+
+    Notes
+    -----
+    This function simply checks for the presence of the `sbatch` command to _infer_
+    if SLURM is installed. It does _not_ check if SLURM is currently active or
+    managing jobs.
+    """
+    return shutil.which("sbatch") is not None