lars
87e37ebe14
CI / Format (ruff format) (push) Successful in 29s
CI / Lint (ruff check) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 38s
CI / Lint (ruff check) (pull_request) Successful in 31s
CI / Format (ruff format) (pull_request) Successful in 44s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 47s
CI / Tests (push) Successful in 4m32s
CI / Tests (pull_request) Successful in 4m25s
stage{1,2}_model.active = false already trains one stage alone, but the
checkpoint it writes holds only that stage, so giant rollout refuses it --
the "retrain stage 2 alone against a fixed, known-good stage 1" experiment
the 2026-08-03 species failure calls for wasn't runnable end to end.
Adds stage{1,2}_model.init_from (a checkpoint .pt to load this stage's
weights from before training) and .freeze (never update them), symmetric
across both stages. Both stages stay active = true, so both get built and
both land in the output checkpoint -- the frozen stage is merely
initialized from disk instead of from scratch.
Decisions made during planning:
- Soft freeze: forward/backward still run every batch (loss/grad_norm stay
meaningful, no autograd special-casing), only optimizer.step() (and, for
the frozen stage, lr_sched.step()/EMA update) is skipped -- weights are
byte-identical for the whole run. This is StageTrainer._step_optimizer,
shared by the non-adversarial path and both halves (generator + critic)
of the WGAN path, so a frozen WGAN stage's critic freezes too.
- validate_config requires init_from whenever freeze = true, unless the run
is a --resume (a resumed frozen stage's weights come from the resume
checkpoint instead) -- freezing a randomly-initialized model is almost
certainly a mistake.
- CLI flags on both `giant train` and `giant new-run`
(--stage{1,2}-init-from/--stage{1,2}-freeze), matching every other
per-stage model knob's existing treatment.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
143 lines
4.3 KiB
Python
143 lines
4.3 KiB
Python
"""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["stage1_model"]["generator"] == "ddpm"
|
|
assert cfg["stage2_model"]["generator"] == "ddpm"
|
|
assert cfg["train"]["lr"] == 0.0005
|
|
assert cfg["stage1_model"]["hidden_dim"] == 128
|
|
assert cfg["stage1_model"]["n_res_blocks"] == 4
|
|
# untouched defaults still present
|
|
assert cfg["train"]["epochs"] == 100
|
|
assert "router" in cfg["stage1_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_stage1_init_from_and_freeze_flags_scaffold_a_partial_retrain_config(tmp_path: Path):
|
|
"""gitea #42."""
|
|
out_dir = tmp_path / "run5"
|
|
result = runner.invoke(
|
|
app,
|
|
[
|
|
"new-run",
|
|
"--out",
|
|
str(out_dir),
|
|
"--stage1-init-from",
|
|
"ckpt/stage1_good/best.pt",
|
|
"--stage1-freeze",
|
|
],
|
|
)
|
|
assert result.exit_code == 0, result.output
|
|
|
|
with open(out_dir / "config.toml", "rb") as f:
|
|
cfg = tomllib.load(f)
|
|
|
|
assert cfg["stage1_model"]["init_from"] == "ckpt/stage1_good/best.pt"
|
|
assert cfg["stage1_model"]["freeze"] is True
|
|
assert cfg["stage2_model"]["init_from"] == ""
|
|
assert cfg["stage2_model"]["freeze"] is False
|
|
|
|
|
|
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()
|