Add giant new-run to scaffold a config.toml + run dir ahead of training
Pulled forward from the not-yet-mergeable condor-gpu-train-rollout branch: `new-run` resolves CLI hyperparameter overrides into a full config.toml and run dir (reusing the existing default_out_dir_name collision-avoidance and a newly factored-out router-override helper shared with `train`), so a run can be prepared and reviewed before `giant train` actually kicks off. Also brings README up to date with the model/CLI as it actually stands (physical/embedding conditioning, WGAN/MoE-router modes, giant analyze, W&B, setup-stage caching), which had drifted back to describing the Phase-1 proof-of-concept. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
@@ -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" 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()
|
||||
Reference in New Issue
Block a user