| |
| """Extract NequIP demo launch metadata and its Hydra training config.""" |
|
|
| from __future__ import annotations |
|
|
| import re |
| import shlex |
| import sys |
| from pathlib import Path |
|
|
| import yaml |
|
|
|
|
| META_KEYS = {"name", "launch", "slurm", "env"} |
|
|
|
|
| def _assignment(name: str, value) -> None: |
| print(f"{name}={shlex.quote(str(value))}") |
|
|
|
|
| def _positive_int(value, field: str) -> int: |
| try: |
| parsed = int(value) |
| except (TypeError, ValueError) as error: |
| raise ValueError(f"{field} must be a positive integer") from error |
| if parsed < 1: |
| raise ValueError(f"{field} must be a positive integer") |
| return parsed |
|
|
|
|
| def _config(path: str) -> dict: |
| text = Path(path).read_text(encoding="utf-8") |
| text = re.sub( |
| r"\$\{demo_dir:([^}]+)\}", |
| lambda match: str(Path(__file__).parent.resolve() / match.group(1)), |
| text, |
| ) |
| return yaml.safe_load(text) or {} |
|
|
|
|
| def _print_launch(cfg: dict) -> None: |
| launch = cfg.get("launch", {}) or {} |
| trainer = cfg.get("trainer", {}) or {} |
| mode = launch.get("mode", "local") |
| if mode not in {"auto", "local", "submit"}: |
| raise ValueError("launch.mode must be 'auto', 'local', or 'submit'") |
|
|
| nodes = _positive_int(launch.get("num_nodes", 1), "launch.num_nodes") |
| devices = _positive_int(launch.get("num_gpus", 1), "launch.num_gpus") |
| trainer_nodes = _positive_int(trainer.get("num_nodes", 1), "trainer.num_nodes") |
| trainer_devices = _positive_int(trainer.get("devices", 1), "trainer.devices") |
| if nodes != trainer_nodes: |
| raise ValueError("launch.num_nodes must equal trainer.num_nodes") |
| if devices != trainer_devices: |
| raise ValueError("launch.num_gpus must equal trainer.devices") |
| _assignment("RUN_MODE", mode) |
| _assignment("NODES", nodes) |
| _assignment("GPUS_PER_NODE", devices) |
| _assignment("WORLD_SIZE", nodes * devices) |
|
|
|
|
| def _print_slurm(cfg: dict) -> None: |
| slurm = cfg.get("slurm", {}) or {} |
| _assignment("PARTITION", slurm.get("partition", "hx1hdnormal01")) |
| _assignment("TIME", slurm.get("time", "01:00:00")) |
| _assignment( |
| "CPUS_PER_TASK", |
| _positive_int(slurm.get("cpus_per_task", 8), "slurm.cpus_per_task"), |
| ) |
| _assignment("NODELIST", slurm.get("nodelist", "")) |
|
|
|
|
| def _print_env(cfg: dict) -> None: |
| env = cfg.get("env", {}) or {} |
| for name, value in env.items(): |
| if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name): |
| raise ValueError(f"invalid environment variable name: {name}") |
| _assignment(f"export {name}", value) |
|
|
|
|
| def _print_training_config(cfg: dict) -> None: |
| training = {key: value for key, value in cfg.items() if key not in META_KEYS} |
| yaml.safe_dump( |
| training, |
| sys.stdout, |
| sort_keys=False, |
| default_flow_style=False, |
| allow_unicode=True, |
| ) |
|
|
|
|
| def main() -> None: |
| if len(sys.argv) != 3: |
| raise SystemExit( |
| "usage: _parse_config.py <config.yaml> " |
| "<name|launch|slurm|env|training-config>" |
| ) |
| cfg = _config(sys.argv[1]) |
| action = sys.argv[2] |
| actions = { |
| "name": lambda: print(cfg.get("name", "nequip_run")), |
| "launch": lambda: _print_launch(cfg), |
| "slurm": lambda: _print_slurm(cfg), |
| "env": lambda: _print_env(cfg), |
| "training-config": lambda: _print_training_config(cfg), |
| "finetune-config": lambda: _print_training_config(cfg), |
| } |
| try: |
| actions[action]() |
| except KeyError as error: |
| raise SystemExit(f"unknown action: {action}") from error |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
