python312Packages.numpyro: reduce number of pytest parallel workers t…

…o reduce RAM usage
NixOS · Jan 8, 2025 · 25d3852 · 25d3852
1 parent 6d88f20
commit 25d3852
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/pkgs/development/python-modules/numpyro/default.nix b/pkgs/development/python-modules/numpyro/default.nix
@@ -63,7 +63,12 @@ buildPythonPackage rec {
 
   pythonImportsCheck = [ "numpyro" ];
 
+  # Required to override `--numprocesses` in `pytestFlagsArray`
+  dontUsePytestXdist = true;
   pytestFlagsArray = [
+    # Tests memory consumption grows significantly with the number of parallel processes (reaches ~200GB with 80 jobs)
+    "--numprocesses=$((NIX_BUILD_CORES > 8 ? 8 : NIX_BUILD_CORES))"
+
     # A few tests fail with:
     # UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1.
     # Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program.