Parallelism in the Scientific Python Stack¶

The hidden parallelism problem¶

Many scientific Python libraries already use multiple threads internally. NumPy delegates linear algebra to compiled BLAS libraries. scikit-learn uses Joblib to parallelize model fitting. Polars runs queries on a Rust thread pool. These native threads bypass Python's GIL, so they do achieve true parallelism.

On your laptop, this is fine. The library detects your 8 cores and uses them.

On a shared cluster, it's a problem. Your Slurm job might request 4 CPUs, but NumPy sees the node's full 64 cores and spawns 64 threads. Those threads fight for 4 physical cores, slow everything down, and interfere with other users' jobs on the same node.

The fix is straightforward: tell each library how many threads to use. The rest of this recipe shows how for each major library.

NumPy, SciPy, and BLAS threads¶

NumPy and SciPy delegate heavy linear algebra (matrix multiply, SVD, eigenvalue decomposition, linear solves, FFT) to compiled BLAS/LAPACK implementations. On Unity, this is typically OpenBLAS or Intel MKL. These libraries spawn native threads that bypass the GIL entirely.

What triggers multi-threading¶

Not every NumPy operation is multi-threaded. The speedup comes from linear algebra routines, not element-wise math:

If your workload is mostly element-wise operations, adding more threads won't help.

Controlling thread count¶

Set environment variables in your sbatch script before Python starts:

export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
export MKL_NUM_THREADS=$SLURM_CPUS_PER_TASK
export OPENBLAS_NUM_THREADS=$SLURM_CPUS_PER_TASK

These variables control OpenMP (used by most BLAS implementations), Intel MKL, and OpenBLAS respectively. Setting all three covers whichever backend is linked.

Checking what's linked¶

To see which BLAS library NumPy is using:

import numpy as np
np.show_config()

For runtime thread inspection and control, use the threadpoolctl package:

from threadpoolctl import threadpool_info
import json
print(json.dumps(threadpool_info(), indent=2))

This shows each thread pool, its library, and how many threads it's using.

Example: matrix multiply scaling¶

Save this as blas_bench.py to see how BLAS thread count affects performance:

"""Benchmark matrix multiply with different thread counts."""

import os
import time
import numpy as np

# Read thread count from environment (set by sbatch script)
n_threads = os.environ.get("OMP_NUM_THREADS", "not set")
print(f"OMP_NUM_THREADS={n_threads}")

# Generate two large random matrices
rng = np.random.default_rng(42)
size = 4000
A = rng.standard_normal((size, size))
B = rng.standard_normal((size, size))

# Warm-up run (first call may include JIT/setup overhead)
_ = A @ B

# Timed run
t0 = time.perf_counter()
C = A @ B
elapsed = time.perf_counter() - t0

print(f"Matrix multiply ({size}x{size}): {elapsed:.2f}s")

Submit it with different thread counts to see the scaling curve:

#!/bin/bash
#SBATCH --job-name=blas-bench
#SBATCH --partition=cpu
#SBATCH --time=00:10:00
#SBATCH --cpus-per-task=8 (1)
#SBATCH --mem=8G
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err

# Load required modules
module load uv/latest

# Run your work
export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
export MKL_NUM_THREADS=$SLURM_CPUS_PER_TASK
export OPENBLAS_NUM_THREADS=$SLURM_CPUS_PER_TASK

uv run python blas_bench.py

1. Try 1, 2, 4, 8, 16 to see how performance scales with thread count.

You'll typically see good speedup up to 4-8 threads for a 4000×4000 matrix, with diminishing returns beyond that.

pandas¶

pandas is mostly single-threaded for DataFrame operations. Functions like df.apply(), df.groupby().agg(), and df.merge() run on one core regardless of how many CPUs you've requested.

There are a few exceptions:

• I/O with PyArrow: pd.read_parquet() and pd.read_csv(engine="pyarrow") can use multiple threads for decompression and parsing.
• Certain internal operations: Some recent pandas versions parallelize specific operations when using the PyArrow backend (pd.options.mode.dtype_backend = "pyarrow").

For genuinely parallel tabular data processing, consider Polars (covered below) or Dask.

scikit-learn and Joblib¶

Many scikit-learn estimators accept an n_jobs parameter that controls parallelism:

from sklearn.ensemble import RandomForestClassifier

# Use 8 parallel workers for fitting
clf = RandomForestClassifier(n_trees=100, n_jobs=8)
clf.fit(X_train, y_train)

Under the hood, scikit-learn uses Joblib, a lightweight library for parallel function execution. Joblib's default backend (loky) spawns worker processes, similar to multiprocessing.

The n_jobs=-1 trap¶

Setting n_jobs=-1 means "use all available CPUs." Joblib reads os.cpu_count() to decide how many that is. On a Slurm cluster, os.cpu_count() returns the node's total core count (e.g., 64), not the number of cores you requested (e.g., 8). Your 8-CPU job spawns 64 workers, oversubscribing the node.

The fix: set the LOKY_MAX_CPU_COUNT environment variable in your sbatch script:

export LOKY_MAX_CPU_COUNT=$SLURM_CPUS_PER_TASK

With this set, n_jobs=-1 correctly spawns workers matching your allocation. Alternatively, pass the count explicitly:

import os
n_cpus = int(os.environ.get("SLURM_CPUS_PER_TASK", os.cpu_count()))
clf = RandomForestClassifier(n_trees=100, n_jobs=n_cpus)

scikit-learn functions that support n_jobs¶

Not all scikit-learn operations parallelize. Here are the most common ones that do:

Check each estimator's documentation for whether it accepts n_jobs.

Example: random forest scaling¶

Save this as sklearn_bench.py:

"""Benchmark RandomForest fitting with different n_jobs values."""

import os
import time
import numpy as np
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier

n_cpus = int(os.environ.get("SLURM_CPUS_PER_TASK", os.cpu_count()))
print(f"SLURM_CPUS_PER_TASK={n_cpus}")

# Generate a synthetic dataset
X, y = make_classification(n_samples=50_000, n_features=50, random_state=42)

for n_jobs in [1, 2, 4, n_cpus]:
    clf = RandomForestClassifier(n_estimators=200, n_jobs=n_jobs, random_state=42)
    t0 = time.perf_counter()
    clf.fit(X, y)
    elapsed = time.perf_counter() - t0
    print(f"n_jobs={n_jobs:>2d}: {elapsed:.2f}s")

#!/bin/bash
#SBATCH --job-name=sklearn-bench
#SBATCH --partition=cpu
#SBATCH --time=00:10:00
#SBATCH --cpus-per-task=8
#SBATCH --mem=4G
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err

# Load required modules
module load uv/latest

# Run your work
export LOKY_MAX_CPU_COUNT=$SLURM_CPUS_PER_TASK

uv run python sklearn_bench.py

Polars¶

Polars is a DataFrame library written in Rust that automatically parallelizes operations using a thread pool. If your workload is data wrangling (filtering, grouping, joining, aggregating), Polars can be significantly faster than pandas without any explicit parallelism code.

Polars reads the POLARS_MAX_THREADS environment variable:

export POLARS_MAX_THREADS=$SLURM_CPUS_PER_TASK

Without this, Polars detects the node's full core count and oversubscribes your allocation.

Numba¶

Numba compiles Python functions to machine code using LLVM. By default, @numba.jit produces single-threaded code. Adding parallel=True and using prange (parallel range) enables automatic loop parallelism via OpenMP-style threads that bypass the GIL.

import numba
import numpy as np

@numba.jit(nopython=True, parallel=True)  # (1)!
def pairwise_distance(X):
    """Compute pairwise Euclidean distance matrix."""
    n = X.shape[0]
    D = np.empty((n, n))
    for i in numba.prange(n):  # (2)!
        for j in range(n):
            d = 0.0
            for k in range(X.shape[1]):
                d += (X[i, k] - X[j, k]) ** 2
            D[i, j] = np.sqrt(d)
    return D

1. parallel=True tells Numba to look for parallelizable loops. nopython=True (the default for @jit) forces compilation rather than falling back to the Python interpreter.
2. numba.prange is the parallel version of range. Numba distributes iterations across threads.

Control the thread count with:

export NUMBA_NUM_THREADS=$SLURM_CPUS_PER_TASK

Numba is a good fit for custom numerical kernels (inner loops, distance computations, simulations) that don't map cleanly to vectorized NumPy operations. If your hot loop can be expressed as NumPy operations, prefer that: BLAS-backed NumPy is typically faster and doesn't require a JIT compilation step.

Dask¶

Dask provides parallel and out-of-core versions of NumPy arrays (dask.array) and pandas DataFrames (dask.dataframe). It's useful when your data doesn't fit in memory or when you need lazy evaluation pipelines.

For single-node parallelism, use LocalCluster:

from dask.distributed import Client, LocalCluster
import os

n_workers = int(os.environ.get("SLURM_CPUS_PER_TASK", 4))
cluster = LocalCluster(
    n_workers=n_workers,
    threads_per_worker=1,  # (1)!
    memory_limit="auto",
)
client = Client(cluster)

1. For CPU-bound work, one thread per worker avoids GIL contention. For I/O-heavy pipelines, you might increase threads_per_worker.

The Slurm wrapper script¶

Here's a template sbatch script that sets all the relevant environment variables in one place. Copy it as a starting point for any scientific Python job:

#!/bin/bash
#SBATCH --job-name=sci-python
#SBATCH --partition=cpu
#SBATCH --time=01:00:00
#SBATCH --cpus-per-task=8 (1)
#SBATCH --mem-per-cpu=2G
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err

# Load required modules
module load uv/latest

# Run your work
# --- Thread control (set BEFORE Python starts) ---
export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
export MKL_NUM_THREADS=$SLURM_CPUS_PER_TASK
export OPENBLAS_NUM_THREADS=$SLURM_CPUS_PER_TASK
export NUMBA_NUM_THREADS=$SLURM_CPUS_PER_TASK
export POLARS_MAX_THREADS=$SLURM_CPUS_PER_TASK
export LOKY_MAX_CPU_COUNT=$SLURM_CPUS_PER_TASK

uv run python my_script.py

1. Adjust to match the parallelism your workload can use. Memory scales with this: 8 CPUs × 2 GB = 16 GB total.

All six variables point to $SLURM_CPUS_PER_TASK, which means you change the core count in one place (--cpus-per-task) and everything adapts.

Mixing multiprocessing with threaded libraries¶

If you use multiprocessing.Pool (or concurrent.futures.ProcessPoolExecutor) to parallelize work and each worker calls NumPy linear algebra, you can accidentally create N processes × M threads. Eight workers, each running NumPy with 8 BLAS threads, produces 64 threads fighting for 8 cores.

The fix: when using multiprocessing with NumPy/SciPy, set BLAS threads to 1:

export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1

This forces each worker process to use a single BLAS thread, and your process-level parallelism handles the scaling. This is usually the right tradeoff for embarrassingly parallel workloads where each task does moderate linear algebra.

Quick reference¶

The "Default (unset)" column is the one that causes trouble. Any library that defaults to "all node cores" will oversubscribe your Slurm allocation if you don't set the corresponding variable.

Common pitfalls¶

from threadpoolctl import threadpool_info
print(threadpool_info())