Tibo De Peuter
23f0b45614
- env/hpc/modules.txt: curated list of HPC EasyBuild modules loaded by vsc-venv - env/hpc/requirements.txt: pip packages not covered by HPC modules (auto-generated) - scripts/export_hpc_requirements.py: developer utility to regenerate requirements.txt from pyproject.toml; to be called by CI when pyproject.toml changes
76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
||
"""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.
|
||
|
||
Usage:
|
||
uv run scripts/export_hpc_requirements.py
|
||
"""
|
||
|
||
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."""
|
||
return re.sub(r"[-_.]+", "-", name).lower()
|
||
|
||
|
||
def pkg_name(dep: str) -> str:
|
||
"""Extract the bare package name from a PEP 508 dependency string."""
|
||
return re.split(r"[\[=><~!;]", dep)[0].strip()
|
||
|
||
|
||
def main() -> None:
|
||
try:
|
||
import tomllib # Python ≥ 3.11
|
||
except ImportError:
|
||
import tomli as tomllib # type: ignore[no-redef]
|
||
|
||
pyproject_path = ROOT / "pyproject.toml"
|
||
with pyproject_path.open("rb") as f:
|
||
data = tomllib.load(f)
|
||
|
||
deps: list[str] = data.get("project", {}).get("dependencies", [])
|
||
|
||
missing: list[str] = []
|
||
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)
|
||
continue
|
||
missing.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)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|