airfRANS-model-exploration/tests/test_sweep.py

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from airfrans_frontier.sweep import (
    BudgetLedger,
    PoolConfig,
    can_expand_pool,
    collect_job_status,
    generate_jobs,
    read_jobs,
    utilization_summary,
)
from airfrans_frontier.training.config import load_training_config


class SweepFoundationTests(unittest.TestCase):
    def test_generate_jobs_writes_deterministic_manifests_and_coordinate_encoding_configs(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            jobs = generate_jobs(
                output_dir=root,
                data_root="artifacts/data_cache/airfrans_processed/processed/full",
                bands=("100m",),
                families=("film_fourier_inr",),
                encodings=("nerf_multires", "raw"),
                group="test_sweep",
            )

            persisted = read_jobs(root / "jobs.jsonl")
            self.assertEqual([job["job_id"] for job in persisted], [job["job_id"] for job in jobs])
            self.assertTrue((root / "pool.toml").is_file())
            self.assertTrue((root / "budget_ledger.json").is_file())

            config = load_training_config(root / "configs" / "100m_film_fourier_inr_nerf_multires.toml")
            self.assertEqual(config.coordinate_encoding.type, "nerf_multires")
            self.assertEqual(config.coordinate_encoding.features, ("x", "y", "sdf"))
            self.assertEqual(config.model.encoding_levels, 16)
            self.assertEqual(config.model.fourier_scales, ())
            self.assertEqual(config.model.hidden_width, 2048)
            self.assertEqual(config.model.depth, 16)

    def test_generate_jobs_only_crosses_encoding_axis_for_compatible_families(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            generate_jobs(
                output_dir=root,
                data_root="data/full",
                bands=("100m",),
                families=("film_fourier_inr", "siren_conditioned_inr"),
                encodings=("raw", "random_fourier"),
            )

            job_ids = [job["job_id"] for job in read_jobs(root / "jobs.jsonl")]

            self.assertIn("100m_film_fourier_inr_random_fourier", job_ids)
            self.assertIn("100m_siren_conditioned_inr_raw", job_ids)
            self.assertNotIn("100m_siren_conditioned_inr_random_fourier", job_ids)

    def test_collect_reconstructs_success_from_required_job_artifacts(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            run_root = root / "runs"
            jobs = generate_jobs(
                output_dir=root / "sweep",
                data_root="data/full",
                artifact_dir=str(run_root),
                bands=("100m",),
                families=("mlp",),
                encodings=("raw",),
            )
            run_dir = run_root / "20260727T000000Z_100m_mlp_raw"
            run_dir.mkdir(parents=True)
            for name in (
                "config.toml",
                "job_manifest.json",
                "metrics.jsonl",
                "latest_metrics.json",
                "final_metrics.json",
                "run_manifest.json",
                "environment_manifest.json",
                "utilization.jsonl",
            ):
                (run_dir / name).write_text("{}\n")

            status = collect_job_status(jobs_path=root / "sweep" / "jobs.jsonl")

            self.assertEqual(status["counts"], {"succeeded": 1})
            self.assertEqual(status["jobs"][0]["job_id"], jobs[0]["job_id"])
            self.assertEqual(status["jobs"][0]["missing_artifacts"], [])

    def test_expansion_gate_requires_utilization_backlog_stability_and_budget(self) -> None:
        pool = PoolConfig(nodes=(), budget_usd=25.0)
        good_utilization = {"gpu_util_median": 91.0, "gpu_idle_fraction": 0.04}

        allowed, reasons = can_expand_pool(
            pool=pool,
            utilization=good_utilization,
            backlog=3,
            recent_failures=0,
            ledger=BudgetLedger(budget_usd=25.0, spent_usd=5.0),
            next_pool_cost_usd=10.0,
        )
        self.assertTrue(allowed)
        self.assertEqual(reasons, ())

        blocked, reasons = can_expand_pool(
            pool=pool,
            utilization={"gpu_util_median": 70.0, "gpu_idle_fraction": 0.20},
            backlog=0,
            recent_failures=1,
            ledger=BudgetLedger(budget_usd=25.0, spent_usd=24.0),
            next_pool_cost_usd=2.0,
        )
        self.assertFalse(blocked)
        self.assertIn("median GPU utilization below gate", reasons)
        self.assertIn("GPU idle fraction above gate", reasons)
        self.assertIn("no job backlog", reasons)
        self.assertIn("recent failures are not isolated", reasons)
        self.assertIn("remaining budget does not support larger pool", reasons)

    def test_utilization_summary_reports_gate_metrics(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "utilization.jsonl"
            path.write_text(
                "".join(
                    json.dumps(sample) + "\n"
                    for sample in (
                        {"gpu_util_percent": 90, "memory_used_mb": 1000},
                        {"gpu_util_percent": 95, "memory_used_mb": 1500},
                        {"gpu_util_percent": 0, "memory_used_mb": 1200},
                    )
                )
            )

            summary = utilization_summary(path)

            self.assertEqual(summary["gpu_util_median"], 90.0)
            self.assertAlmostEqual(summary["gpu_idle_fraction"], 1 / 3)
            self.assertEqual(summary["memory_used_peak"], 1500.0)


if __name__ == "__main__":
    unittest.main()