giant/tests/test_cli_train_overrides.py

"""Tests for `giant train`'s stage-prefixed CLI flags: --stage1-*/--stage2-*
must independently override each stage's config block, and must take precedence
over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still
apply the same value to both stages for backward compatibility."""

from __future__ import annotations

from pathlib import Path

from typer.testing import CliRunner

import giant.cli as cli

runner = CliRunner()


def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dict:
    captured: dict = {}

    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
        captured["cfg"] = cfg

    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)

    result = runner.invoke(
        cli.app,
        ["train", "dummy.parquet", "--out", str(tmp_path / "run")] + args,
    )
    assert result.exit_code == 0, result.output
    return captured["cfg"]


def test_stage_prefixed_generator_overrides_shared_mode(monkeypatch, tmp_path):
    cfg = _invoke_and_capture_cfg(
        monkeypatch,
        tmp_path,
        ["--mode", "wgan", "--stage1-generator", "flow"],
    )
    assert cfg["stage1_model"]["generator"] == "flow"
    assert cfg["stage2_model"]["generator"] == "wgan"


def test_stage2_only_knobs(monkeypatch, tmp_path):
    cfg = _invoke_and_capture_cfg(
        monkeypatch,
        tmp_path,
        [
            "--stage2-decoder",
            "one_shot",
            "--stage2-k-max",
            "8",
            "--stage2-hidden-dim",
            "32",
            "--stage2-context-dim",
            "16",
            "--stage2-stage1-context",
            "sampled",
        ],
    )
    assert cfg["stage2_model"]["decoder"] == "one_shot"
    assert cfg["stage2_model"]["k_max"] == 8
    assert cfg["stage2_model"]["hidden_dim"] == 32
    assert cfg["stage2_model"]["context_dim"] == 16
    assert cfg["stage2_model"]["stage1_context"] == "sampled"
    # untouched stage1 defaults
    assert cfg["stage1_model"]["hidden_dim"] == 256


def test_stage1_hidden_dim_flag_overrides_legacy_hidden_dim_flag(monkeypatch, tmp_path):
    cfg = _invoke_and_capture_cfg(
        monkeypatch,
        tmp_path,
        ["--hidden-dim", "64", "--stage1-hidden-dim", "128"],
    )
    assert cfg["stage1_model"]["hidden_dim"] == 128


def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
    cfg = _invoke_and_capture_cfg(
        monkeypatch,
        tmp_path,
        [
            "--mode",
            "wgan",
            "--n-critic",
            "5",
            "--stage1-n-critic",
            "3",
            "--stage2-gp-weight",
            "2.5",
        ],
    )
    assert cfg["stage1_model"]["wgan"]["n_critic"] == 3
    assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
    assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
    assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5


def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
    monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
    result = runner.invoke(
        cli.app,
        ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
    )
    assert result.exit_code == 1
    assert "--batch-size must be an integer or 'auto'" in result.output


def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_path):
    captured: dict = {}

    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
        captured["out_dir"] = out_dir

    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)

    resume_dir = tmp_path / "resumed_run"
    resume_dir.mkdir()
    (resume_dir / "last.pt").touch()
    explicit_out = tmp_path / "explicit_run"

    result = runner.invoke(
        cli.app,
        ["train", "dummy.parquet", "--out", str(explicit_out), "--resume", str(resume_dir / "last.pt")],
    )
    assert result.exit_code == 0, result.output
    assert captured["out_dir"] == explicit_out


def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
    captured: dict = {}

    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
        captured["out_dir"] = out_dir

    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)

    resume_dir = tmp_path / "resumed_run"
    resume_dir.mkdir()
    (resume_dir / "last.pt").touch()

    result = runner.invoke(cli.app, ["train", "dummy.parquet", "--resume", str(resume_dir / "last.pt")])
    assert result.exit_code == 0, result.output
    assert captured["out_dir"] == resume_dir


def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypatch, tmp_path):
    captured: dict = {}

    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
        captured["out_dir"] = out_dir

    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
    monkeypatch.chdir(tmp_path)

    result = runner.invoke(cli.app, ["train", "dummy.parquet"])
    assert result.exit_code == 0, result.output
    assert captured["out_dir"] == Path("checkpoints") / cli.gconfig.default_out_dir_name(cli.gconfig.DEFAULT_CONFIG)


def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
    captured: dict = {}

    def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
        captured["batch_size"] = cfg["train"]["batch_size"]

    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
    monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)

    result = runner.invoke(
        cli.app,
        ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "auto"],
    )
    assert result.exit_code == 0, result.output
    assert captured["batch_size"] == 123
    assert "batch_size: 123 (auto-estimated from free GPU memory)" in result.output