Albrja/mic 6801/pytest xdist (#273)

Albrja/mic 6801/pytest xdist

Add pytest xdist feature to parallelize test suites
- *Category*: Test
- *JIRA issue*: https://jira.ihme.washington.edu/browse/MIC-6801
- *Research reference*: <!--Link to research documentation for code -->

Changes and notes
-add pytest-xdist to install requirements
-add conftest hook to automatically discover number of CPUs allocated on srun on slurm
-allow for parallelization on local builds but default to running tests linearly

Verification and Testing
-built new environment locally and on the cluster. Saw expected default behavior and tested -n flag


*** REMINDER ***
 model results directory is up to date
- [X ] all tests pass (`pytest --runslow` with both *vivarium_gates_mncnh_artifact* and *vivarium_gates_mncnh_simulation*)

Author: albrja

pyproject.toml

@@ -3,6 +3,7 @@ requires = ["packaging", "setuptools"]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
+addopts = "-nauto"
 
 [tool.black]
 line_length = 94

setup.py

@@ -63,6 +63,7 @@
         "papermill",
         "jupyterlab",
         "vivarium_testing_utils",
+        "pytest-xdist",
     ]
     validation_requirements = ["vivarium_testing_utils[validation]"]
     lint_requirements = [

tests/conftest.py

@@ -87,6 +87,28 @@ def pytest_collection_modifyitems(config, items):
                 item.add_marker(skip_slow)
 
 
+def pytest_xdist_auto_num_workers(config):
+    """Automatically determine the number of workers for pytest-xdist.
+
+    - On SLURM: Use CPUs allocated to the job (via SLURM environment variables)
+    - Not on SLURM: Return 1 (no parallelization by default)
+    - Users can override by explicitly passing -n flag to pytest
+    """
+    cpus = 1
+    if IS_ON_SLURM:
+        # On SLURM clusters, use the number of CPUs allocated to the job
+        # Check SLURM environment variables in order of preference
+        slurm_cpus = os.environ.get("SLURM_CPUS_PER_TASK") or os.environ.get(
+            "SLURM_CPUS_ON_NODE"
+        )
+        if slurm_cpus:
+            cpus = int(slurm_cpus)
+        # Fallback to total CPUs if SLURM vars not found (shouldn't happen)
+        cpus = os.cpu_count()
+
+    return cpus
+
+
 @pytest.fixture(scope="session")
 def model_spec_path() -> Path:
     repo_path = paths.BASE_DIR