airfRANS-model-exploration/tests/test_remote_run.py

from __future__ import annotations

import json
import tempfile
import shutil
import unittest
from pathlib import Path

import torch

from airfrans_frontier.remote.artifacts import verify_artifacts
from airfrans_frontier.remote.cli import _classify_artifacts, _stage_resume_checkpoint
from airfrans_frontier.remote.config import load_remote_run_config
from airfrans_frontier.remote.skypilot import render_skypilot_yaml
from airfrans_frontier.remote.vast import VastOffer, choose_offer


class RemoteRunConfigTests(unittest.TestCase):
    def test_loads_remote_smoke_config(self) -> None:
        config = load_remote_run_config("configs/remote_smoke.toml")

        self.assertEqual(config.provider.kind, "vastai")
        self.assertEqual(config.provider.gpu.name, "RTX 4090")
        self.assertEqual(config.job.artifact_dir.as_posix(), "artifacts/current_run")
        self.assertIn("checkpoint_latest.pt", config.artifacts.required)


class VastSelectionTests(unittest.TestCase):
    def test_selection_filters_bad_hosts_and_drops_suspiciously_cheap_tail(self) -> None:
        config = load_remote_run_config("configs/remote_smoke.toml")
        offers = [
            offer(1, price=0.10, host=1),
            offer(2, price=0.20, host=2),
            offer(3, price=0.30, host=3),
            offer(4, price=0.40, host=4),
            offer(5, price=0.50, host=59017),
            offer(6, price=0.25, host=6, geo="CN"),
            offer(7, price=0.26, host=7, verification="deverified"),
        ]

        result = choose_offer(offers, config, query={"test": True})

        # Four reachable RTX 4090 offers remain; drop_cheap_frac=0.30 drops floor(1.2)=1 cheapest.
        self.assertEqual(result.selected_offer_id, 2)
        self.assertEqual(result.candidate_count, 7)
        self.assertEqual(result.survivor_count, 3)

    def test_rendered_yaml_injects_selected_offer(self) -> None:
        config = load_remote_run_config("configs/remote_smoke.toml")
        result = choose_offer([offer(123, price=0.30, host=22), offer(124, price=0.40, host=23)], config, query={})

        yaml = render_skypilot_yaml(config, result, run_id="airfrans-test")

        self.assertIn("selected_offer_id: 123", yaml)
        self.assertNotIn("sky launch", yaml)
        self.assertIn("remote-run smoke-train", yaml)
        self.assertIn("configs/aggressive_smoke.toml", yaml)

    def test_rendered_yaml_can_pass_resume_checkpoint(self) -> None:
        config = load_remote_run_config("configs/remote_smoke.toml")
        result = choose_offer([offer(123, price=0.30, host=22), offer(124, price=0.40, host=23)], config, query={})

        yaml = render_skypilot_yaml(
            config,
            result,
            run_id="airfrans-test",
            resume_checkpoint=".airfrans_resume/airfrans-test/checkpoint_latest.pt",
        )

        self.assertIn("AIRFRANS_RESUME_CHECKPOINT: .airfrans_resume/airfrans-test/checkpoint_latest.pt", yaml)


class ArtifactVerificationTests(unittest.TestCase):
    def test_verify_artifacts_requires_contract_files_and_writes_manifest(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            _write_contract_artifacts(root, success=True)

            manifest = verify_artifacts(root)

            self.assertGreaterEqual(manifest["file_count"], 8)
            self.assertTrue((root / "artifact_manifest.json").is_file())
            self.assertTrue((root / "checksums.txt").is_file())

    def test_verify_artifacts_accepts_failure_report_terminal_state(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            _write_contract_artifacts(root, success=False)
            (root / "failure_report.json").write_text(json.dumps({"error_type": "NonFiniteLoss"}) + "\n")

            manifest = verify_artifacts(root)

            self.assertGreaterEqual(manifest["file_count"], 7)

    def test_classifies_and_stages_restart_checkpoint(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self.assertEqual(_classify_artifacts(root), "incomplete")
            _write_contract_artifacts(root, success=False)
            self.assertEqual(_classify_artifacts(root), "restartable")

            staged = _stage_resume_checkpoint(root, "test-run")

            self.assertIsNotNone(staged)
            assert staged is not None
            self.assertTrue(staged.is_file())
            self.assertEqual(staged.as_posix(), ".airfrans_resume/test-run/checkpoint_latest.pt")
            shutil.rmtree(".airfrans_resume")


def _write_contract_artifacts(root: Path, *, success: bool) -> None:
    (root / "config.toml").write_text("[run]\nname = 'test'\n")
    (root / "metrics.jsonl").write_text(json.dumps({"step": 0}) + "\n")
    (root / "latest_metrics.json").write_text(json.dumps({"step": 0}) + "\n")
    (root / "heartbeat.json").write_text(json.dumps({"phase": "training"}) + "\n")
    checkpoint = {
        "schema_version": 1,
        "step": 0,
        "model_state_dict": {},
        "optimizer_state_dict": {},
        "normalization": {},
    }
    torch.save(checkpoint, root / "checkpoint_latest.pt")
    torch.save(checkpoint, root / "checkpoint_best.pt")
    if success:
        torch.save(checkpoint, root / "checkpoint_final.pt")
        (root / "final_metrics.json").write_text(json.dumps({"loss": 1.0}) + "\n")


def offer(
    offer_id: int,
    *,
    price: float,
    host: int,
    geo: str = "US",
    verification: str = "verified",
) -> VastOffer:
    return VastOffer(
        id=offer_id,
        gpu_name="RTX 4090",
        dph_total=price,
        gpu_ram=24_000,
        geolocation=geo,
        inet_down_cost_per_tb=0.0,
        inet_up_cost_per_tb=0.0,
        host_id=host,
        verification=verification,
        reliability2=0.99,
        cuda_max_good=12.8,
        direct_port_count=1,
        inet_down=500.0,
        inet_up=100.0,
        verified=True,
    )


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