train/rollout: submit as remote-GPU HTCondor jobs on TOpAS/NEMO2

Adds `giant train-submit`/`giant rollout-submit`, mirroring `giant analyze submit`'s CPU-job pattern but for single remote-GPU jobs: +RemoteJob/ RequestGPUs, TARGET.ProvidesEtpCeph instead of the local-only ProvidesETPResources, and a self-contained condor/ run dir (wrapper, submit description, and a CondorJobMeta sidecar recording what was submitted and the assigned cluster id) so a run stays traceable after the fact. Training jobs re-check for last.pt on every wrapper invocation so a preempted job resumes instead of restarting. Also adds `giant new-run` to scaffold a run's config.toml + run dir (with collision-free naming via the new shared `default_out_dir`) ahead of submission, and factors router-flag parsing into `_router_cli_overrides` so `train` and `new-run` resolve it identically. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-24 13:12:01 +02:00
parent 1115301eb9
commit 430917d8f2
7 changed files with 1046 additions and 26 deletions
@@ -0,0 +1,118 @@
+"""Tests for `giant new-run` (config.toml + run-dir scaffolding)."""
+
+from __future__ import annotations
+
+import tomllib
+from pathlib import Path
+
+from typer.testing import CliRunner
+
+from giant.cli import app
+
+runner = CliRunner()
+
+
+def test_writes_config_with_overrides_applied(tmp_path: Path):
+    out_dir = tmp_path / "run1"
+    result = runner.invoke(
+        app,
+        [
+            "new-run",
+            "--out",
+            str(out_dir),
+            "--mode",
+            "ddpm",
+            "--hidden-dim",
+            "128",
+            "--n-blocks",
+            "4",
+            "--lr",
+            "0.0005",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+
+    config_path = out_dir / "config.toml"
+    assert config_path.exists()
+    with open(config_path, "rb") as f:
+        cfg = tomllib.load(f)
+
+    assert cfg["train"]["mode"] == "ddpm"
+    assert cfg["train"]["lr"] == 0.0005
+    assert cfg["model"]["hidden_dim"] == 128
+    assert cfg["model"]["n_blocks"] == 4
+    # untouched defaults still present
+    assert cfg["train"]["epochs"] == 100
+    assert "router" in cfg["model"]
+
+    assert str(out_dir) in result.output
+    assert "<data.parquet>" in result.output
+    assert "giant train-submit" in result.output
+
+
+def test_comment_and_provenance_recorded_in_meta(tmp_path: Path):
+    out_dir = tmp_path / "run2"
+    result = runner.invoke(
+        app,
+        ["new-run", "--out", str(out_dir), "--comment", "quick test"],
+    )
+    assert result.exit_code == 0, result.output
+
+    with open(out_dir / "config.toml", "rb") as f:
+        cfg = tomllib.load(f)
+
+    assert cfg["meta"]["comment"] == "quick test"
+    assert cfg["meta"]["created_by"] == "giant new-run"
+    assert "created_at" in cfg["meta"]
+    assert "git_hash" in cfg["meta"]
+
+
+def test_data_flag_fills_printed_next_step_commands(tmp_path: Path):
+    out_dir = tmp_path / "run3"
+    result = runner.invoke(
+        app,
+        ["new-run", "--out", str(out_dir), "--data", "/ceph/lbogner/train.parquet"],
+    )
+    assert result.exit_code == 0, result.output
+    assert "/ceph/lbogner/train.parquet" in result.output
+    assert "<data.parquet>" not in result.output
+
+
+def test_dry_run_writes_nothing(tmp_path: Path):
+    out_dir = tmp_path / "run4"
+    result = runner.invoke(
+        app,
+        ["new-run", "--out", str(out_dir), "--hidden-dim", "512", "--dry-run"],
+    )
+    assert result.exit_code == 0, result.output
+    assert "dry-run" in result.output
+    assert "hidden_dim = 512" in result.output
+    assert not out_dir.exists()
+
+
+def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
+    out_dir = tmp_path / "run5"
+    out_dir.mkdir()
+    (out_dir / "last.pt").touch()
+
+    result = runner.invoke(app, ["new-run", "--out", str(out_dir), "--mode", "ddpm"])
+    assert result.exit_code != 0
+    assert "already has last.pt" in result.output
+    assert not (out_dir / "config.toml").exists()
+
+    result = runner.invoke(
+        app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"]
+    )
+    assert result.exit_code == 0, result.output
+    assert (out_dir / "config.toml").exists()
+
+
+def test_default_out_dir_used_when_out_omitted(tmp_path: Path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+    result = runner.invoke(app, ["new-run", "--hidden-dim", "64"])
+    assert result.exit_code == 0, result.output
+
+    checkpoints_dir = tmp_path / "checkpoints"
+    run_dirs = list(checkpoints_dir.iterdir()) if checkpoints_dir.exists() else []
+    assert len(run_dirs) == 1
+    assert (run_dirs[0] / "config.toml").exists()
@@ -58,6 +58,17 @@ def test_each_call_produces_a_distinct_uuid(tmp_path):
    assert uuid1 != uuid2


+def test_pinned_pred_uuid_is_used_as_is(tmp_path):
+    # `rollout-submit` pins the uuid at submit time so its condor run-dir,
+    # the output filename, and the eventual YAML sidecar all agree.
+    data = tmp_path / "data.parquet"
+    pinned = "abcd1234-abcd-4abc-9abc-abcdabcdabcd"
+    out, _, pred_uuid = _resolve_prediction_output(data, None, pred_uuid=pinned)
+
+    assert pred_uuid == pinned
+    assert out.name == f"{pinned}.parquet"
+
+
 def test_dataset_path_is_resolved(tmp_path):
    data = tmp_path / "data.parquet"
    _, dataset_path, _ = _resolve_prediction_output(data, None)
@@ -0,0 +1,211 @@
+"""Tests for GPU-job HTCondor submission (`giant train-submit` / `giant rollout-submit`)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from typer.testing import CliRunner
+
+from giant import config as gconfig
+from giant.cli import app
+from giant.condor import CondorJobMeta, GpuSubmitConfig, parse_cluster_id, write_gpu_submit
+
+runner = CliRunner()
+
+_MINIMAL_TOML = """\
+[train]
+mode = "flow"
+epochs = 1
+
+[model]
+hidden_dim = 8
+n_blocks = 2
+"""
+
+
+# ---------------------------------------------------------------------------
+# default_out_dir
+# ---------------------------------------------------------------------------
+
+
+def test_default_out_dir_encodes_hyperparams():
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}, {})
+    out_dir = gconfig.default_out_dir(cfg)
+    name = out_dir.name
+    assert f"_h{cfg['model']['hidden_dim']}_" in name
+    assert f"_b{cfg['model']['n_blocks']}_" in name
+    assert f"_c{cfg['model']['conditioning']}_" in name
+
+
+def test_default_out_dir_avoids_collisions():
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}, {})
+    a = gconfig.default_out_dir(cfg)
+    b = gconfig.default_out_dir(cfg)
+    assert a != b
+    # same hyperparam-derived prefix, differing only in the uuid tag
+    assert a.name.rsplit("_", 1)[0] == b.name.rsplit("_", 1)[0]
+
+
+# ---------------------------------------------------------------------------
+# giant/condor.py: GpuSubmitConfig / write_gpu_submit
+# ---------------------------------------------------------------------------
+
+
+def test_write_gpu_submit_description(tmp_path: Path):
+    cfg = GpuSubmitConfig(
+        run_dir=tmp_path / "condor",
+        accounting_group="cms",
+        repo_dir=tmp_path,
+        command="uv run giant train data.parquet --config c.toml --out out --device cuda",
+    )
+    sub = write_gpu_submit(cfg, "#!/bin/bash\necho hi\n")
+    txt = sub.read_text()
+
+    assert "universe = docker" in txt
+    assert "docker_image = mschnepf/slc7-condocker" in txt
+    assert "+RemoteJob = True" in txt
+    assert "RequestGPUs = 1" in txt
+    assert "accounting_group = cms" in txt
+    assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in txt
+    assert "ProvidesETPResources" not in txt
+    assert "queue 1" in txt
+
+    wrapper = cfg.run_dir / "run.sh"
+    assert wrapper.exists()
+    assert wrapper.stat().st_mode & 0o111
+    assert wrapper.read_text() == "#!/bin/bash\necho hi\n"
+
+
+def test_gpu_requirements_include_type_and_memory_pins(tmp_path: Path):
+    cfg = GpuSubmitConfig(
+        run_dir=tmp_path / "condor",
+        accounting_group="cms",
+        repo_dir=tmp_path,
+        command="uv run giant train ...",
+        request_gpus=2,
+        gpu_type="Tesla V100-PCIE-32GB",
+        gpu_memory_mb=16000,
+    )
+    txt = write_gpu_submit(cfg, "#!/bin/bash\n").read_text()
+
+    assert "RequestGPUs = 2" in txt
+    assert 'TARGET.GPUs_DeviceName =?= "Tesla V100-PCIE-32GB"' in txt
+    assert "TARGET.GPUs_GlobalMemoryMb >= 16000" in txt
+
+
+# ---------------------------------------------------------------------------
+# CondorJobMeta / parse_cluster_id
+# ---------------------------------------------------------------------------
+
+
+def test_condor_job_meta_round_trip(tmp_path: Path):
+    meta = CondorJobMeta(
+        accounting_group="cms",
+        request_gpus=1,
+        gpu_type=None,
+        gpu_memory_mb=None,
+        request_memory_mb=16384,
+        request_walltime_s=172800,
+        docker_image="mschnepf/slc7-condocker",
+        repo_dir="/ceph/lbogner/giant",
+        command="uv run giant train ...",
+        submitted_at="2026-07-24T00:00:00+00:00",
+    )
+    path = tmp_path / "submission.json"
+    meta.save(path)
+    loaded = CondorJobMeta.load(path)
+    assert loaded == meta
+    assert loaded.cluster_id is None
+
+    loaded.cluster_id = 123456
+    loaded.save(path)
+    assert CondorJobMeta.load(path).cluster_id == 123456
+
+
+def test_parse_cluster_id():
+    assert parse_cluster_id("1 job(s) submitted to cluster 123456.\n") == 123456
+    assert parse_cluster_id("ERROR: something went wrong\n") is None
+
+
+# ---------------------------------------------------------------------------
+# CLI: giant train-submit / giant rollout-submit (--dry-run only, no cluster contact)
+# ---------------------------------------------------------------------------
+
+
+def test_train_submit_dry_run_writes_condor_dir(tmp_path: Path):
+    config = tmp_path / "config.toml"
+    config.write_text(_MINIMAL_TOML)
+    data = tmp_path / "train.parquet"
+    out_dir = tmp_path / "ckpt"
+
+    result = runner.invoke(
+        app,
+        [
+            "train-submit",
+            str(data),
+            "--config",
+            str(config),
+            "--accounting-group",
+            "cms",
+            "--out",
+            str(out_dir),
+            "--dry-run",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+
+    run_dir = out_dir / "condor"
+    sub_txt = (run_dir / "job.sub").read_text()
+    assert "+RemoteJob = True" in sub_txt
+    assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in sub_txt
+
+    wrapper = (run_dir / "run.sh").read_text()
+    assert "--device cuda" in wrapper
+    assert "last.pt" in wrapper
+    assert "RESUME" in wrapper
+
+    meta = CondorJobMeta.load(run_dir / "submission.json")
+    assert meta.accounting_group == "cms"
+    assert meta.cluster_id is None
+
+
+def test_rollout_submit_dry_run_writes_condor_dir(tmp_path: Path):
+    data = tmp_path / "seed.parquet"
+    checkpoint = tmp_path / "best.pt"
+    geometry = tmp_path / "oracle.pkl"
+    out = tmp_path / "rollout.parquet"
+
+    result = runner.invoke(
+        app,
+        [
+            "rollout-submit",
+            str(data),
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            str(geometry),
+            "--accounting-group",
+            "cms",
+            "--out",
+            str(out),
+            "--dry-run",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+
+    condor_dirs = list(tmp_path.glob("condor_*"))
+    assert len(condor_dirs) == 1
+    run_dir = condor_dirs[0]
+
+    sub_txt = (run_dir / "job.sub").read_text()
+    assert "+RemoteJob = True" in sub_txt
+    assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in sub_txt
+
+    wrapper = (run_dir / "run.sh").read_text()
+    assert "--device cuda" in wrapper
+    assert "--prediction-id" in wrapper
+    assert "last.pt" not in wrapper
+    assert "RESUME" not in wrapper
+
+    meta = CondorJobMeta.load(run_dir / "submission.json")
+    assert meta.accounting_group == "cms"