from __future__ import annotations from contextlib import redirect_stdout from io import StringIO import json import tempfile import shutil import unittest from unittest.mock import patch from pathlib import Path from airfrans_frontier.runtime import remove_pythonpath_entries remove_pythonpath_entries() import torch from airfrans_frontier.remote.artifacts import verify_artifacts from airfrans_frontier.remote.cli import _classify_artifacts, _stage_resume_checkpoint, _terminal_artifact_names, main as remote_main 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) def test_loads_remote_hf_smoke_config_for_cheap_upload(self) -> None: config = load_remote_run_config("configs/remote_hf_smoke.toml") self.assertEqual(config.provider.kind, "vastai") self.assertEqual(config.provider.gpu.name, "RTX 3060") self.assertLessEqual(config.provider.max_price_per_hour or 999.0, 0.08) self.assertEqual(config.artifacts.mode, "object_store_upload") self.assertIn("remote-run hf-smoke", config.job.command) self.assertIn("hf_upload_manifest.json", config.artifacts.required) def test_loads_remote_wandb_smoke_config_for_observability(self) -> None: config = load_remote_run_config("configs/remote_wandb_smoke.toml") self.assertEqual(config.provider.kind, "vastai") self.assertEqual(config.provider.gpu.name, "RTX 3060 Ti") self.assertLessEqual(config.provider.max_price_per_hour or 999.0, 0.09) self.assertEqual(config.artifacts.mode, "rsync") self.assertIn("remote-run wandb-smoke", config.job.command) self.assertIn("wandb_smoke_manifest.json", 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) self.assertIn("df -h .", yaml) self.assertIn("airfrans_disk_requested_gb=128", yaml) self.assertIn("AIRFRANS_STARTUP_TIMELINE: artifacts/current_run/startup_timeline.jsonl", yaml) self.assertIn("airfrans_timeline 'setup' 'started'", yaml) self.assertIn("airfrans_timeline 'data_validation' 'started'", yaml) self.assertIn("airfrans_timeline 'training_command' 'started'", yaml) def test_terminal_collection_includes_startup_timeline(self) -> None: self.assertIn("startup_timeline.jsonl", _terminal_artifact_names(())) 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 VastInstanceCliTests(unittest.TestCase): def test_vast_instances_reports_api_ground_truth(self) -> None: stdout = StringIO() with patch.dict("os.environ", {"VAST_API_KEY": "token"}), patch( "airfrans_frontier.remote.cli.list_instances", return_value=[{"id": 123, "actual_status": "running", "gpu_name": "RTX 4090"}], ), redirect_stdout(stdout): code = remote_main(["vast-instances"]) self.assertEqual(code, 0) payload = json.loads(stdout.getvalue()) self.assertEqual(payload["instance_count"], 1) self.assertEqual(payload["instances"][0]["id"], 123) self.assertEqual(payload["instances"][0]["actual_status"], "running") 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()) self.assertTrue((root / "verification_report.json").is_file()) report = json.loads((root / "verification_report.json").read_text()) self.assertTrue(report["ok"]) 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_remote_success_before_checkpoint_collection(self) -> None: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) (root / "final_metrics.json").write_text(json.dumps({"loss": 1.0}) + "\n") self.assertEqual(_classify_artifacts(root), "success") 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, disk_space=256.0, 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()