From bfd5c2606aa96083b3ab941ae3a071690dfe5451 Mon Sep 17 00:00:00 2001
From: Martin Weinelt <hexa@darmstadt.ccc.de>
Date: Sat, 30 Sep 2023 16:18:45 +0200
Subject: [PATCH] Revert "python3Packages.jax: remove pytest-xdist"

This reverts commit bddafba6a54a4ea07428a6ff0713898f81b25cca.

Running without pytest-xdist leads to reliable crashes on both
aarch64-linux (Ampere Q80-30) and x86_64-linux (Zen 3 Epyc) on hydra,
which prevents the package from being cached.

We revert this change to prevent the churn on downstream maintainers
and consumers, that would otherwise have to build this expensive
package locally, which does not work reliably.
---
 pkgs/development/python-modules/jax/default.nix | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix
index dbdf655e4007..25ff005569a5 100644
--- a/pkgs/development/python-modules/jax/default.nix
+++ b/pkgs/development/python-modules/jax/default.nix
@@ -12,6 +12,7 @@
 , numpy
 , opt-einsum
 , pytestCheckHook
+, pytest-xdist
 , pythonOlder
 , scipy
 , stdenv
@@ -61,13 +62,18 @@ buildPythonPackage rec {
     jaxlib'
     matplotlib
     pytestCheckHook
+    pytest-xdist
   ];
 
+  # high parallelism will result in the tests getting stuck
+  dontUsePytestXdist = true;
+
   # NOTE: Don't run the tests in the expiremental directory as they require flax
   # which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2.
   # Not a big deal, this is how the JAX docs suggest running the test suite
   # anyhow.
   pytestFlagsArray = [
+    "--numprocesses=4"
     "-W ignore::DeprecationWarning"
     "tests/"
   ];