fix(hpc): modules loading

docs/CONTRIBUTING.md

@@ -32,4 +32,4 @@ Code readability is paramount, as code is read far more frequently than it is wr
 * **Algorithms & Frameworks:** Proximal Policy Optimization (PPO) is the recommended baseline algorithm. CleanRL should be used as a starting point and adapted for continuous action spaces. All Artificial Neural Network (ANN) controller architectures must be implemented using Flax.
 * **Simulation:** The simulation environment utilizes a MuJoCo brittle star. XML MuJoCo structures must remain realistic and respect morphological constraints.
 * **Experiment Tracking:** Weights & Biases (wandb) must be utilized for tracking and logging all experiments.
-* **Code Styling:** All code must conform to the chosen style guide (i.e. Google standard). This is enforced using build tools and pre-commit hooks such as flake8, black, or isort.
+* **Code Styling:** All code must conform to the chosen style guide (Google standard). This is enforced via `uv` using **ruff** and pre-commit hooks.

env/hpc/modules.txt

@@ -1,8 +1,6 @@
-PyTorch/2.1.2-foss-2023a-CUDA-12.1.1
-jax/0.4.25-gfbf-2023a-CUDA-12.1.1
-Flax/0.8.4-gfbf-2023a-CUDA-12.1.1
-Optax/0.2.2-gfbf-2023a-CUDA-12.1.1
-wandb/0.16.1-GCC-12.3.0
-matplotlib/3.7.2-gfbf-2023a
-PyYAML/6.0-GCCcore-12.3.0
-FFmpeg/6.0-GCCcore-12.3.0
+gfbf/2024a
+GCCcore/13.3.0
+Python/3.12.3-GCCcore-13.3.0
+PyTorch/2.6.0-foss-2024a-CUDA-12.6.0
+FFmpeg/7.0.2-GCCcore-13.3.0
+PyYAML/6.0.2-GCCcore-13.3.0

env/hpc/requirements.txt

@@ -1,9 +1,14 @@
 biorobot==0.4.2
 cleanrl>=0.4.8
 evosax==0.2.0
+flax>=0.12.2
 gymnasium>=1.2.3
 ipykernel==7.2.0
+jax==0.9.0.1
+matplotlib==3.10.8
 mediapy==1.2.6
+optax>=0.2.6
 pyopengl>=3.1.10
 pyopengl-accelerate>=3.1.10
 tyro>=1.0.10
+wandb==0.24.2

scripts/export_hpc_requirements.py

@@ -2,10 +2,9 @@
 """Export HPC pip requirements from pyproject.toml.
 
 This is a LOCAL DEVELOPER UTILITY — run it on your own machine before pushing
-code whenever pyproject.toml dependencies change. It reads the full dependency
-list from pyproject.toml and subtracts packages already provided by HPC modules
-(listed in env/hpc/modules.txt), then writes the remainder to
-env/hpc/requirements.txt for use by vsc-venv on the cluster.
+code whenever pyproject.toml dependencies change. It reads the modules from
+env/hpc/modules.txt and the full dependency list from pyproject.toml, then
+writes the remainder to env/hpc/requirements.txt.
 
 Usage:
     uv run scripts/export_hpc_requirements.py
@@ -13,27 +12,12 @@ Usage:
 
 from __future__ import annotations
 
-import importlib.util
 import re
 import sys
 from pathlib import Path
 
 ROOT = Path(__file__).parent.parent
 
-# Packages provided by HPC modules (PyPI name → Python import name).
-# We skip any pyproject.toml dep whose import can be found after loading
-# the HPC modules. As this script runs locally (without those modules), we
-# maintain an explicit exclusion list keyed by normalised PyPI package name.
-EXCLUDED_BY_MODULE = {
-    "jax",
-    "flax",
-    "optax",
-    "wandb",
-    "matplotlib",
-    "pyyaml",  # PyYAML module
-    "ffmpeg",  # FFmpeg is a system tool, not a Python package
-}
-
 
 def normalise(name: str) -> str:
     """Normalise a PyPI package name for comparison."""
@@ -46,10 +30,19 @@ def pkg_name(dep: str) -> str:
 
 
 def main() -> None:
-    try:
-        import tomllib  # Python ≥ 3.11
-    except ImportError:
-        import tomli as tomllib  # type: ignore[no-redef]
+    import tomllib
+
+    modules_path = ROOT / "env" / "hpc" / "modules.txt"
+    if not modules_path.exists():
+        print(f"Error: {modules_path} not found.", file=sys.stderr)
+        sys.exit(1)
+
+    # Read normalized module names from modules.txt
+    module_names = [
+        normalise(line.split()[0].split("/")[0])
+        for line in modules_path.read_text().splitlines()
+        if line.strip() and not line.startswith("#")
+    ]
 
     pyproject_path = ROOT / "pyproject.toml"
     with pyproject_path.open("rb") as f:
@@ -57,19 +50,23 @@ def main() -> None:
 
     deps: list[str] = data.get("project", {}).get("dependencies", [])
 
-    missing: list[str] = []
+    final_deps: list[str] = []
+    print(f"Checking dependencies against {modules_path}...", file=sys.stderr)
     for dep in deps:
         name = normalise(pkg_name(dep))
-        if name in EXCLUDED_BY_MODULE:
-            print(f"  [skip – provided by HPC module] {dep}", file=sys.stderr)
+        # Smart check: if the package name is a substring of any loaded module name
+        # (e.g. 'torch' in 'pytorch', 'scipy' in 'scipy-bundle')
+        if any(name in mod for mod in module_names):
+            print(f"  [skip – module provider found] {dep}", file=sys.stderr)
             continue
-        missing.append(dep)
+
+        final_deps.append(dep)
         print(f"  [pip] {dep}", file=sys.stderr)
 
     output_path = ROOT / "env" / "hpc" / "requirements.txt"
     output_path.parent.mkdir(parents=True, exist_ok=True)
-    output_path.write_text("\n".join(missing) + "\n")
-    print(f"\nWrote {len(missing)} requirement(s) to {output_path}", file=sys.stderr)
+    output_path.write_text("\n".join(final_deps) + "\n")
+    print(f"\nWrote {len(final_deps)} requirement(s) to {output_path}", file=sys.stderr)
 
 
 if __name__ == "__main__":