change dask default client

Author: stevebachmeier

src/pseudopeople/interface.py

@@ -23,6 +23,7 @@
     configure_logging_to_terminal,
     get_engine_from_string,
     get_state_abbreviation,
+    set_up_dask_client,
 )
 
 
@@ -111,7 +112,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 +127,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.

src/pseudopeople/utilities.py

@@ -303,6 +303,27 @@ def get_engine_from_string(engine: str) -> Engine:
     DataFrame = pd.DataFrame  # type: ignore [misc]
 
 
+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()
+
+
 ##########################
 # Data utility functions #
 ##########################

tests/unit/test_utils.py

@@ -1,18 +1,25 @@
+import os
+
 import numpy as np
 import pandas as pd
+import psutil
 import pytest
+from dask.distributed import LocalCluster, get_client
+from dask.system import CPU_COUNT
 
 from pseudopeople.dataset import Dataset
 from pseudopeople.noise_functions import _corrupt_tokens
 from pseudopeople.schema_entities import DATASET_SCHEMAS, DtypeNames
 from pseudopeople.utilities import (
     get_hash,
     get_index_to_noise,
+    set_up_dask_client,
     to_string_as_integer,
     two_d_array_choice,
     vectorized_choice,
 )
 from tests.conftest import FuzzyChecker
+from tests.utilities import is_on_slurm
 
 
 @pytest.fixture()
@@ -268,3 +275,50 @@ def test_two_d_array_choice(
                 target_proportion=1 / num_choices,
                 name_additional=f"team {team} for sport {sport}",
             )
+
+
+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 sum(worker["memory_limit"] / 1024**3 for worker in workers.values()) == available_memory
+
+
+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/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