From 868489cfebddc8893248560fcc16b2fc2d15cc3d Mon Sep 17 00:00:00 2001 From: Gernot Bauer Date: Sat, 14 Mar 2026 19:44:32 +0100 Subject: [PATCH] Add and expose Rayon thread pool handling to the Python module. --- CHANGELOG.md | 2 + docs/api/index.md | 3 +- docs/api/thread_pool.md | 38 ++++ examples/managing_threads.ipynb | 314 ++++++++++++++++++++++++++++++++ py-feos/src/lib.rs | 98 ++++++++++ 5 files changed, 454 insertions(+), 1 deletion(-) create mode 100644 docs/api/thread_pool.md create mode 100644 examples/managing_threads.ipynb diff --git a/CHANGELOG.md b/CHANGELOG.md index f7998e383..998556d0f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added +- Add Rayon global thread pool control via `FEOS_MAX_THREADS` and `set_num_threads()`/ `get_num_threads()` to Python. [#346](https://github.com/feos-org/feos/pull/346) ## [0.9.4] - 2026-03-09 ### Changed diff --git a/docs/api/index.md b/docs/api/index.md index 85df22f77..f3fae50e3 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -10,4 +10,5 @@ All functions and classes in FeOS are exported at the package root. Here, they a eos dft ad -``` \ No newline at end of file + thread_pool +``` diff --git a/docs/api/thread_pool.md b/docs/api/thread_pool.md new file mode 100644 index 000000000..4a3ef95d2 --- /dev/null +++ b/docs/api/thread_pool.md @@ -0,0 +1,38 @@ +# Global thread pool + +Several functions in `feos` use [Rayon](https://github.com/rayon-rs/rayon) for parallelism. +By default, Rayon uses all logical CPUs available on your machine, which is usually what you want when working on your local machine. +In other environments, for example HPC clusters, you may want to limit the number of threads to match your job allocation. + +There are three ways to configure this, in order of priority: + +- `FEOS_MAX_THREADS` environment variable: for HPC or "script" environments, defined before launching Python +- `feos.set_num_threads()`: for interactive use, at the top of a script or notebook +- Do nothing: local machines where using all cores is fine + +You can get the number of threads configured via `feos.get_num_threads()`. + +## Important +- The thread pool can only be configured **once** per Python process. +- Whichever method runs first wins. Any later attempt to change it will have no effect and a warning will be emitted. +- Calling `get_num_threads` without setting `FEOS_MAX_THREADS` or `set_num_threads` will initialze the thread pool with the default (all logical CPUs). +- To test the different behaviour in a notebook, you have to restart the kernel and start from the respective cell you want to test. + +## Example Usage + +```python +import feos + +feos.set_num_threads(4) +print(f"Active threads: {feos.get_num_threads()}") +``` + +```{eval-rst} +.. currentmodule:: feos + +.. autosummary:: + :toctree: generated/ + + set_num_threads + get_num_threads +``` diff --git a/examples/managing_threads.ipynb b/examples/managing_threads.ipynb new file mode 100644 index 000000000..0ae8a20fe --- /dev/null +++ b/examples/managing_threads.ipynb @@ -0,0 +1,314 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "b9685fcb", + "metadata": {}, + "source": [ + "# Controlling the Thread Pool in `feos`\n", + "\n", + "Several functions in `feos` use [Rayon](https://github.com/rayon-rs/rayon) for parallelism.\n", + "By default, Rayon uses all logical CPUs available on your machine, which is usually what you want when working on your local machine. \n", + "In other environments, for example HPC clusters, you may want to limit the number of threads to match your job allocation.\n", + "\n", + "There are three ways to configure this, in order of priority:\n", + "\n", + "- `FEOS_MAX_THREADS` environment variable: for HPC or \"script\" environments, defined before launching Python\n", + "- `feos.set_num_threads()`: for interactive use, at the top of a script or notebook\n", + "- Do nothing: local machines where using all cores is fine\n", + "\n", + "You can get the number of threads configured via `feos.get_num_threads()`.\n", + "\n", + "## Important\n", + "- The thread pool can only be configured **once** per Python process.\n", + "- Whichever method runs first wins. Any later attempt to change it will have no effect and a warning will be emitted.\n", + "- Calling `get_num_threads` without setting `FEOS_MAX_THREADS` or `set_num_threads` will initialze the thread pool with the default (all logical CPUs).\n", + "- To test the different behaviour in this notebook, restart the kernel and start from the respective cell you want to test." + ] + }, + { + "cell_type": "markdown", + "id": "34e0be01", + "metadata": {}, + "source": [ + "## Method 1: Environment variable (recommended for HPC)\n", + "\n", + "Set `FEOS_MAX_THREADS` **before** starting Python or launching your Jupyter kernel.\n", + "The thread pool is initialized automatically when `feos` is imported.\n", + "\n", + "In a Slurm job script:\n", + "\n", + "```bash\n", + "#!/bin/bash\n", + "#SBATCH --cpus-per-task=8\n", + "\n", + "export FEOS_MAX_THREADS=$SLURM_CPUS_PER_TASK\n", + "python my_script.py\n", + "```\n", + "\n", + "Or in a terminal before starting Jupyter:\n", + "\n", + "```bash\n", + "export FEOS_MAX_THREADS=4\n", + "jupyter lab\n", + "```\n", + "\n", + "You can verify it was picked up after importing:" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "3e161b43", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Active threads: 10\n" + ] + } + ], + "source": [ + "import feos\n", + "\n", + "# If FEOS_MAX_THREADS was set before starting Python, the pool\n", + "# was already configured at import time.\n", + "print(f\"Active threads: {feos.get_num_threads()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "b60da8bf", + "metadata": {}, + "source": [ + "## Method 2: `set_num_threads()` (interactive use)\n", + "\n", + "If you did not set the environment variable, you can configure the thread pool\n", + "programmatically. This must be done **before any parallel computation is triggered**.\n", + "\n", + "In a notebook or script, call it immediately after importing `feos`: (if the code below emits a warning, restart the kernel and run the code below as first cell)" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "d5fce92f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Active threads: 4\n" + ] + } + ], + "source": [ + "import feos\n", + "\n", + "feos.set_num_threads(4)\n", + "\n", + "print(f\"Active threads: {feos.get_num_threads()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "b8a862cd", + "metadata": {}, + "source": [ + "You can also read the thread count from `SLURM_CPUS_PER_TASK` manually if you prefer\n", + "to keep configuration in Python rather than in your shell environment:" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "a8e08786", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Active threads: 4\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/var/folders/3s/t93ws1md04qdbbq5d1jdz8640000gn/T/ipykernel_9962/4108200440.py:5: UserWarning: set_num_threads(10) without effect: The thread pool was already initialized with 4 thread(s) Call set_num_threads() before any parallel work or set FEOS_MAX_THREADS before starting Python.\n", + " feos.set_num_threads(n_threads)\n" + ] + } + ], + "source": [ + "import os\n", + "import feos\n", + "\n", + "n_threads = int(os.environ.get(\"SLURM_CPUS_PER_TASK\", os.cpu_count()))\n", + "feos.set_num_threads(n_threads)\n", + "\n", + "print(f\"Active threads: {feos.get_num_threads()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "9c891bdc", + "metadata": {}, + "source": [ + "## Method 3: Do nothing (Rayon default)\n", + "\n", + "If you neither set `FEOS_MAX_THREADS` nor call `set_num_threads()`, Rayon will initialize the thread pool lazily the first time a parallel function is called, using all available logical CPUs. This is usually the right choice on a local workstation.\n", + "Note that calling `get_num_threads` will set the threads to the default." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "61c0c96d", + "metadata": {}, + "outputs": [], + "source": [ + "import feos\n", + "\n", + "# No configuration — Rayon will use all logical CPUs.\n", + "# get_num_threads() triggers lazy initialization if not already done.\n", + "print(f\"Active threads: {feos.get_num_threads()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "142bc70d", + "metadata": {}, + "source": [ + "## What happens if you call `set_num_threads()` too late?\n", + "\n", + "If the pool is already initialized — because `FEOS_MAX_THREADS` was set, or because\n", + "a parallel function (or `get_num_threads()`) has already run — `set_num_threads()`\n", + "has no effect and emits a warning:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cba243c9", + "metadata": {}, + "outputs": [], + "source": [ + "import feos\n", + "\n", + "print(feos.get_num_threads()) # triggers lazy initialization\n", + "\n", + "feos.set_num_threads(2) # UserWarning: had no effect" + ] + }, + { + "cell_type": "markdown", + "id": "5ba53f01", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "```\n", + " import feos\n", + " │\n", + " FEOS_MAX_THREADS set?\n", + " ┌──────┴───────┐\n", + " Yes No\n", + " │ │\n", + " Pool initialized set_num_threads() called?\n", + " with env var value ┌──────┴───────┐\n", + " Yes No\n", + " │ │\n", + " Pool initialized Pool initialized lazily\n", + " with given value on first parallel call\n", + " (all logical CPUs)\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "fe867ce1-10f6-43b4-8382-b9f5bc9b387d", + "metadata": {}, + "source": [ + "## Example\n", + "\n", + "The following example calculates vapor pressures and derivatives w.r.t. the model's parameters in parallel.\n", + "Set the number of threads to check the impact on calculation time (restart kernel as before to change the threads)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "3ee76e3a-6727-45c4-8ca1-58f6c3b27231", + "metadata": {}, + "outputs": [], + "source": [ + "import feos\n", + "import numpy as np\n", + "import timeit\n", + "\n", + "# modify the number of threads and rerun the cell below to see impact\n", + "feos.set_num_threads(0) \n", + "\n", + "n = 1_000_000\n", + "fit_params = [\"m\", \"sigma\", \"epsilon_k\"]\n", + "\n", + "# order: m, sigma, epsilon_k, mu\n", + "parameters = np.array([[1.5, 3.4, 230.0, 2.3]] * n)\n", + "temperature = np.expand_dims(np.linspace(250.0, 400.0, n), 1)\n", + "eos = feos.EquationOfStateAD.PcSaftNonAssoc" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "4315f002-ce50-4e01-a2aa-49f28eebd84d", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Mean of 5 runs (1000000 VLEs each): 0.559s using 10 thread(s)\n" + ] + } + ], + "source": [ + "n_runs = 5\n", + "n_threads = feos.get_num_threads()\n", + "elapsed = np.mean(timeit.repeat(\n", + " lambda: feos.vapor_pressure_derivatives(eos, fit_params, parameters, temperature),\n", + " number=1,\n", + " repeat=n_runs\n", + "))\n", + "\n", + "print(f\"Mean of {n_runs} runs ({n} VLEs each): {elapsed:.3f}s using {n_threads} thread(s)\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/py-feos/src/lib.rs b/py-feos/src/lib.rs index b4f94441a..0886c3eb0 100644 --- a/py-feos/src/lib.rs +++ b/py-feos/src/lib.rs @@ -51,10 +51,108 @@ impl From for Verbosity { } } +#[cfg(feature = "rayon")] +mod rayon_features { + use pyo3::exceptions::{PyRuntimeError, PyUserWarning}; + use pyo3::prelude::*; + use std::ffi::CString; + + /// Reads the `FEOS_MAX_THREADS` environment variable and, if present, + /// initializes the global Rayon thread pool with that many threads. + /// Called automatically at module import time. + pub fn rayon_threads_from_env() { + if let Some(n) = std::env::var("FEOS_MAX_THREADS") + .ok() + .and_then(|s| s.parse::().ok()) + { + let _ = rayon::ThreadPoolBuilder::new() + .num_threads(n) + .build_global(); + } + } + + #[pyfunction] + /// Set the number of threads used for any parallel calculations. + /// + /// Must be called before any parallel computation is performed and + /// before the `FEOS_MAX_THREADS` environment variable takes effect. + /// If the thread pool has already been initialized — either + /// because `FEOS_MAX_THREADS` was set at import time or because a + /// parallel function has already run — this call has no effect and + /// a warning is emitted. + /// + /// Args: + /// n (int): Number of threads. Pass `0` to use the default + /// (number of logical CPUs). + /// + /// Example: + /// >>> import feos + /// >>> feos.set_num_threads(4) + pub fn set_num_threads(py: Python<'_>, n: usize) -> PyResult<()> { + match rayon::ThreadPoolBuilder::new() + .num_threads(n) + .build_global() + { + Ok(_) => Ok(()), + Err(_) => { + // build useful warning + let current = rayon::current_num_threads(); + let reason = if std::env::var("FEOS_MAX_THREADS").is_ok() { + format!( + "FEOS_MAX_THREADS is set. \ + The thread pool was already initialized with {} thread(s) \ + (probably configured at import time). \ + To change this, set FEOS_MAX_THREADS before starting Python.", + current + ) + } else { + format!( + "The thread pool was already initialized with {} thread(s) \ + Call set_num_threads() before any parallel work or set \ + FEOS_MAX_THREADS before starting Python.", + current + ) + }; + + let msg = + CString::new(format!("set_num_threads({}) without effect: {}", n, reason)) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + PyErr::warn(py, &py.get_type::(), &msg, 1) + } + } + } + + #[pyfunction] + /// Return the number of threads in the thread pool. + /// + /// If the thread pool has not yet been initialized, calling this + /// function will trigger initialization with the default + /// (number of logical CPUs), making any subsequent call to + /// `set_num_threads()` ineffective. + /// + /// Returns: + /// int: Number of threads currently configured. + /// + /// Example: + /// >>> import feos + /// >>> feos.get_num_threads() + /// 8 + pub fn get_num_threads() -> usize { + rayon::current_num_threads() + } +} + #[pymodule] fn feos(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add("__version__", env!("CARGO_PKG_VERSION"))?; + #[cfg(feature = "rayon")] + { + rayon_features::rayon_threads_from_env(); + m.add_function(wrap_pyfunction!(rayon_features::set_num_threads, m)?)?; + m.add_function(wrap_pyfunction!(rayon_features::get_num_threads, m)?)?; + } + // Utility m.add_class::()?;