"""Tests for giant.training.plots (gitea #75) — render smoke tests skipped where plotstyle/LaTeX is unavailable, plus pure-function column-classification coverage that needs neither.""" from __future__ import annotations import csv from pathlib import Path import pytest pytest.importorskip("plotstyle") from giant.training import plots as plots_mod # noqa: E402 from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics # noqa: E402 # --- fixtures ---------------------------------------------------------- _RICH_HEADER = [ "epoch", "stage1/train/loss", "stage1/train/loss_gen", "stage1/train/nsec_acc", "stage1/train/grad_norm", "stage1/val/loss", "stage1/val/loss_gen", "stage1/val/nsec_acc", "stage1/lr", "stage1/router/entropy", "stage1/router/util_min", "stage1/router/util_max", "stage1/router/util_std", "stage2/train/d_loss", "stage2/train/g_loss", "stage2/train/wasserstein", "stage2/train/gp_loss", "stage2/train/loss_nsec", "stage2/train/nsec_acc", "stage2/train/grad_norm_d", "stage2/train/grad_norm_g", "stage2/lr", "stage2/critic_lr", "val/loss", "val/marginal_kl", "grad_norm", "gpu_mem_mb", "samples_per_sec", "is_best", "epoch_time_s", ] _RICH_ROWS = [ [ 1, 1.0, 0.8, 0.5, 1.2, 0.9, 0.7, 0.6, 3e-4, 1.5, 0.05, 0.3, 0.1, -0.2, 0.3, 0.5, 0.1, 0.4, 0.4, 0.9, 1.1, 3e-4, 1e-4, 0.85, 0.4, 2.1, 512.0, 100.0, 1, 5.0, ], [ 2, 0.8, 0.6, 0.6, 1.0, 0.7, 0.5, 0.7, 2e-4, 1.6, 0.06, 0.28, 0.09, -0.1, 0.25, 0.4, 0.09, 0.3, 0.5, 0.8, 1.0, 2e-4, 8e-5, 0.7, 0.35, 1.9, 520.0, 105.0, 0, 5.1, ], ] _MINIMAL_HEADER = [ "epoch", "stage1/train/loss", "stage1/train/loss_gen", "stage1/val/loss", "stage1/val/loss_gen", "stage1/lr", "val/loss", "grad_norm", "gpu_mem_mb", "samples_per_sec", "is_best", "epoch_time_s", ] _MINIMAL_ROWS = [ [1, 1.0, 0.8, 0.9, 0.7, 3e-4, 0.85, 0.4, 0.0, 100.0, 0, 5.0], [2, 0.8, 0.6, 0.7, 0.5, 2e-4, 0.7, 0.35, 0.0, 105.0, 1, 5.1], ] def _write_csv(path: Path, header: list[str], rows: list[list]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", newline="") as f: writer = csv.writer(f) writer.writerow(header) writer.writerows(rows) # --- MetricsTable -------------------------------------------------------- def test_metrics_table_load_round_trips(tmp_path: Path): csv_path = tmp_path / "metrics.csv" _write_csv(csv_path, _MINIMAL_HEADER, _MINIMAL_ROWS) table = MetricsTable.load(csv_path) assert table.epochs == [1, 2] assert table.columns["stage1/train/loss"] == [1.0, 0.8] assert "epoch" not in table.columns assert table.best_epochs() == [2] def test_metrics_table_best_epochs_empty_without_is_best_column(): table = MetricsTable(epochs=[1, 2], columns={"stage1/train/loss": [1.0, 0.5]}) assert table.best_epochs() == [] # --- column classification (pure functions, no matplotlib) -------------- def _rich_columns() -> dict[str, list]: return {name: [0.0] for name in _RICH_HEADER if name != "epoch"} def test_stages_detects_only_stages_present(): assert plots_mod._stages(_rich_columns()) == ["stage1", "stage2"] assert plots_mod._stages({"stage2/train/loss": [0.0]}) == ["stage2"] assert plots_mod._stages({"val/loss": [0.0]}) == [] def test_split_matches_stage_and_split_prefix_only(): cols = _rich_columns() train = plots_mod._split(cols, "stage1", "train") assert train == { "loss": "stage1/train/loss", "loss_gen": "stage1/train/loss_gen", "nsec_acc": "stage1/train/nsec_acc", "grad_norm": "stage1/train/grad_norm", } assert plots_mod._split(cols, "stage2", "val") == {} def test_point_in_time_excludes_train_val_router(): cols = _rich_columns() pit = plots_mod._point_in_time(cols, "stage1") assert pit == {"lr": "stage1/lr"} pit2 = plots_mod._point_in_time(cols, "stage2") assert pit2 == {"lr": "stage2/lr", "critic_lr": "stage2/critic_lr"} def test_router_columns(): cols = _rich_columns() assert plots_mod._router(cols, "stage1") == { "entropy": "stage1/router/entropy", "util_min": "stage1/router/util_min", "util_max": "stage1/router/util_max", "util_std": "stage1/router/util_std", } assert plots_mod._router(cols, "stage2") == {} def test_run_level_excludes_stage_prefixed_columns_including_val_loss_lookalike(): cols = _rich_columns() run_level = plots_mod._run_level(cols) assert set(run_level) == { "val/loss", "val/marginal_kl", "grad_norm", "gpu_mem_mb", "samples_per_sec", "is_best", "epoch_time_s", } # stage-prefixed "val/loss" lookalike (stage1/val/loss) must not leak in assert "stage1/val/loss" not in run_level def test_loss_keys_excludes_acc_and_wgan_and_grad_norm(): train = {"loss": "x", "loss_gen": "x", "nsec_acc": "x", "grad_norm": "x", "d_loss": "x"} val = {"loss": "x", "loss_gen": "x"} assert plots_mod._loss_keys(train, val) == ["loss", "loss_gen"] # --- derive_metrics_dir --------------------------------------------------- def test_derive_metrics_dir_explicit_out_dir_wins(): assert derive_metrics_dir("runs/my-run", out_dir="/somewhere") == Path("/somewhere") def test_derive_metrics_dir_default_base(): assert derive_metrics_dir("runs/my-run", default_base="/data/analysis_runs") == Path( "/data/analysis_runs/metrics_my-run" ) def test_derive_metrics_dir_falls_back_to_cwd_analysis_runs(monkeypatch, tmp_path): monkeypatch.chdir(tmp_path) assert derive_metrics_dir("runs/my-run") == tmp_path / "analysis_runs" / "metrics_my-run" # --- render_metrics end to end ------------------------------------------- def _try_render(run_dir: Path, out_dir: Path) -> list[Path]: try: return render_metrics(run_dir, out_dir) except RuntimeError as e: # LaTeX missing at render time pytest.skip(f"LaTeX rendering unavailable: {e}") def test_render_metrics_rich_run_produces_expected_plots_outside_run_dir(tmp_path: Path): run_dir = tmp_path / "run" out_dir = tmp_path / "out" _write_csv(run_dir / "metrics.csv", _RICH_HEADER, _RICH_ROWS) paths = _try_render(run_dir, out_dir) names = {p.stem for p in paths} assert names == { "overview", "stage1_loss", "stage2_loss", "lr", "stage1_accuracy", "stage2_accuracy", "grad_norm", "stage1_router", "stage2_wgan_balance", "throughput", } assert all(p.exists() for p in paths) assert all(p.is_relative_to(out_dir) for p in paths) # nothing written into the training run directory itself assert not any(run_dir.rglob("*.pdf")) def test_render_metrics_minimal_run_omits_router_wgan_accuracy(tmp_path: Path): run_dir = tmp_path / "run" out_dir = tmp_path / "out" _write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS) paths = _try_render(run_dir, out_dir) names = {p.stem for p in paths} assert names == {"overview", "stage1_loss", "lr", "grad_norm", "throughput"} assert "stage1_accuracy" not in names assert "stage1_router" not in names assert "stage1_wgan_balance" not in names def test_render_metrics_default_out_dir_uses_analysis_runs_convention(tmp_path: Path): run_dir = tmp_path / "runs" / "my-run" _write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS) default_base = tmp_path / "analysis_runs" try: paths = render_metrics(run_dir, default_base=default_base) except RuntimeError as e: pytest.skip(f"LaTeX rendering unavailable: {e}") assert paths assert all(p.is_relative_to(default_base / "metrics_my-run") for p in paths) # --- CLI ------------------------------------------------------------------- def test_cli_analyze_metrics_smoke(tmp_path: Path): from typer.testing import CliRunner from giant.cli import app run_dir = tmp_path / "run" out_dir = tmp_path / "out" _write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS) runner = CliRunner() result = runner.invoke(app, ["analyze", "metrics", str(run_dir), "--out", str(out_dir)]) if result.exit_code != 0 and "latex" in (str(result.output) + str(result.exception)).lower(): pytest.skip("LaTeX rendering unavailable") assert result.exit_code == 0, result.output or result.exception assert any(out_dir.glob("*.pdf"))