from __future__ import annotations import json import shutil import sys import tempfile import types import unittest from pathlib import Path from unittest.mock import Mock, patch from airfrans_frontier.runtime import remove_pythonpath_entries remove_pythonpath_entries() from airfrans_frontier.remote.cleanup import reconcile_cleanup from airfrans_frontier.remote.collection import ARTIFACT_COLLECTION_REPORT, collect_artifact_paths, required_collection_failures from airfrans_frontier.remote.config import load_remote_run_config from airfrans_frontier.remote.launch_group import LaunchGroupScheduler, LaunchRunSpec from airfrans_frontier.remote.selection import require_fresh_selection, selection_freshness_report from airfrans_frontier.remote.skypilot import render_skypilot_yaml from airfrans_frontier.remote.vast import VastOffer, choose_offer from airfrans_frontier.training.hf_upload import HfArtifactUploader from airfrans_frontier.training.streaming_data import StreamingEventRecorder class LaunchGroupSchedulingTests(unittest.TestCase): def test_healthy_runs_release_fragile_launch_capacity_without_serializing_training(self) -> None: with tempfile.TemporaryDirectory() as tmp: state_path = Path(tmp) / "launch_state.json" scheduler = LaunchGroupScheduler( [ LaunchRunSpec("run-a", "configs/a.toml"), LaunchRunSpec("run-b", "configs/b.toml"), LaunchRunSpec("run-c", "configs/c.toml"), ], max_active=3, max_fragile=1, state_path=state_path, group_id="group-local", ) self.assertTrue(scheduler.try_start("run-a", selected_offer_id=101, selected_host_id=11)) self.assertFalse(scheduler.try_start("run-b", selected_offer_id=102, selected_host_id=12)) self.assertEqual(scheduler.capacity_snapshot()["fragile"], 1) scheduler.mark_training_healthy("run-a") self.assertTrue(scheduler.try_start("run-b", selected_offer_id=102, selected_host_id=12)) self.assertEqual(scheduler.capacity_snapshot()["active"], 2) self.assertEqual(scheduler.capacity_snapshot()["fragile"], 1) payload = json.loads(state_path.read_text()) self.assertEqual(payload["launch_group_id"], "group-local") self.assertEqual(payload["healthy_runs"], ["run-a"]) self.assertEqual(payload["running_runs"], ["run-b"]) self.assertEqual(payload["runs"]["run-b"]["selected_host_id"], 12) self.assertIn("capacity_blocked", [event["event"] for event in payload["events"]]) class HostAntiCollisionTests(unittest.TestCase): def test_active_launches_avoid_duplicate_hosts_unless_allowed(self) -> None: scheduler = LaunchGroupScheduler( [LaunchRunSpec("run-a", "a.toml"), LaunchRunSpec("run-b", "b.toml")], max_active=2, max_fragile=2, ) self.assertTrue(scheduler.try_start("run-a", selected_offer_id=1, selected_host_id=9)) self.assertFalse(scheduler.try_start("run-b", selected_offer_id=2, selected_host_id=9)) self.assertEqual(scheduler.to_payload()["runs"]["run-b"]["blocked_reason"], "host_collision") allowed = LaunchGroupScheduler( [LaunchRunSpec("run-a", "a.toml"), LaunchRunSpec("run-b", "b.toml")], max_active=2, max_fragile=2, allow_duplicate_hosts=True, ) self.assertTrue(allowed.try_start("run-a", selected_offer_id=1, selected_host_id=9)) self.assertTrue(allowed.try_start("run-b", selected_offer_id=2, selected_host_id=9)) def test_offer_selection_skips_reserved_active_hosts(self) -> None: config = load_remote_run_config("configs/remote_smoke.toml") result = choose_offer( [offer(10, price=0.20, host=1), offer(11, price=0.22, host=2)], config, query={"test": True}, reserved_host_ids=(1,), ) self.assertEqual(result.selected_offer.host_id, 2) self.assertEqual(result.policy["reserved_host_ids"], [1]) class SelectionFreshnessTests(unittest.TestCase): def test_selection_artifacts_record_and_enforce_freshness(self) -> None: fresh = {"selected_offer_id": 1, "created_at": 1000.0} report = selection_freshness_report(fresh, max_age_seconds=60, now=1020.0) self.assertTrue(report["is_fresh"]) self.assertEqual(report["age_seconds"], 20.0) stale = {"selected_offer_id": 1, "created_at": 1000.0} with self.assertRaisesRegex(ValueError, "stale"): require_fresh_selection(stale, max_age_seconds=60, now=1100.0, path="selection.json") config = load_remote_run_config("configs/remote_smoke.toml") manifest = choose_offer([offer(20, price=0.20, host=3)], config, query={}).to_manifest() self.assertIn("created_at", manifest) self.assertIn("created_at_iso", manifest) self.assertIn("age_seconds", manifest) class CleanupReconciliationTests(unittest.TestCase): def test_reconciliation_uses_vast_ground_truth_for_orphans_and_records_actions(self) -> None: destroyed: list[int] = [] report = reconcile_cleanup( sky_state={"clusters": [{"name": "known-run", "instance_id": 77}]}, vast_instances=[ {"id": 77, "actual_status": "running", "gpu_name": "RTX 4090", "num_gpus": 1, "dph_total": 0.40}, {"id": 88, "host_id": 123, "actual_status": "running", "gpu_name": "RTX 4090", "num_gpus": 1, "dph_total": 0.45, "label": "orphan-run"}, ], known_run_ids=("known-run", "orphan-run"), destroy_orphans=True, destroy_instance=lambda instance_id: destroyed.append(instance_id), now=1234.0, ) orphan = next(item for item in report["instances"] if item["vast_instance_id"] == 88) self.assertEqual(report["unexpected_live_count"], 1) self.assertEqual(destroyed, [88]) self.assertEqual(orphan["cleanup_action_attempted"], "destroy_orphan") self.assertEqual(orphan["cleanup_result"], "destroy_requested") self.assertEqual(orphan["hourly_cost"], 0.45) class HfSafetyTests(unittest.TestCase): def test_rate_limit_suppression_preserves_training_success_as_hf_incomplete(self) -> None: class FakeRateLimitError(RuntimeError): def __init__(self) -> None: super().__init__("429 Too Many Requests") self.response = types.SimpleNamespace(headers={"Retry-After": "600"}) class FakeCommitOperationAdd: def __init__(self, *, path_in_repo: str, path_or_fileobj: str) -> None: self.path_in_repo = path_in_repo self.path_or_fileobj = path_or_fileobj fake_module = types.SimpleNamespace(CommitOperationAdd=FakeCommitOperationAdd) with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}): run_dir = Path(tmp) (run_dir / "metrics.jsonl").write_text("{}\n") uploader = HfArtifactUploader( enabled=True, run_dir=run_dir, repo_id="owner/repo", repo_type="model", path_in_repo="runs/run-1", max_rate_limit_sleep_seconds=0, ) uploader._api = types.SimpleNamespace(create_commit=Mock(side_effect=FakeRateLimitError())) with self.assertRaises(FakeRateLimitError): uploader.upload_files(("metrics.jsonl",), commit_message="upload metrics") suppressed = uploader.upload_files(("metrics.jsonl",), commit_message="retry metrics") final = uploader.finalize(training_success=True) self.assertTrue(suppressed["rate_limited"]) self.assertEqual(final["hf_publication_status"], "training_succeeded_hf_incomplete") manifest = json.loads((run_dir / "hf_upload_manifest.json").read_text()) self.assertTrue(manifest["training_success"]) self.assertFalse(manifest["publication_complete"]) self.assertGreater(manifest["rate_limit_until"], 0) def test_final_reporting_distinguishes_training_failure_from_hf_success(self) -> None: with tempfile.TemporaryDirectory() as tmp: disabled = HfArtifactUploader(enabled=False, run_dir=Path(tmp)) self.assertEqual(disabled.finalize(training_success=False)["hf_publication_status"], "disabled") with tempfile.TemporaryDirectory() as tmp: uploader = HfArtifactUploader(enabled=True, run_dir=Path(tmp), repo_id="owner/repo", repo_type="model", path_in_repo="run") self.assertEqual(uploader.finalize(training_success=False)["hf_publication_status"], "training_failed") with tempfile.TemporaryDirectory() as tmp: uploader = HfArtifactUploader(enabled=True, run_dir=Path(tmp), repo_id="owner/repo", repo_type="model", path_in_repo="run") self.assertEqual(uploader.finalize(training_success=True)["hf_publication_status"], "hf_publication_succeeded") class ArtifactCollectionReportTests(unittest.TestCase): def test_collection_report_classifies_produced_missing_partial_and_failed_copy(self) -> None: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) def copy_one(relative_path: str) -> int | None: if relative_path == "produced.json": (root / relative_path).write_text("{}\n") return 0 if relative_path == "missing.json": return 0 if relative_path == "partial.pt": partial = root / ".rsync-partial" / relative_path partial.parent.mkdir(parents=True) partial.write_bytes(b"partial") return 0 raise RuntimeError("rsync failed") report = collect_artifact_paths( local_dir=root, remote_dir="remote:~/artifacts", paths=("produced.json", "missing.json", "partial.pt", "failed.json"), required=("produced.json", "partial.pt", "failed.json"), collection_kind="terminal", copy_one=copy_one, ) by_path = {attempt["expected_path"]: attempt for attempt in report["attempts"]} self.assertEqual(by_path["produced.json"]["final_status"], "success") self.assertEqual(by_path["missing.json"]["likely_reason"], "remote_missing_or_not_produced") self.assertEqual(by_path["partial.pt"]["final_status"], "partial") self.assertEqual(by_path["failed.json"]["likely_reason"], "collection_command_failed") self.assertEqual( {item["expected_path"] for item in required_collection_failures(report)}, {"partial.pt", "failed.json"}, ) saved = json.loads((root / ARTIFACT_COLLECTION_REPORT).read_text()) self.assertFalse(saved["summary"]["ok"]) class DiskPhilosophyTests(unittest.TestCase): def test_disk_paths_record_telemetry_and_backpressure_state_instead_of_capacity_mismatch_hard_fail(self) -> None: config = load_remote_run_config("configs/remote_smoke.toml") yaml = render_skypilot_yaml(config, choose_offer([offer(30, price=0.20, host=4)], config, query={}), run_id="disk-check") self.assertIn("disk_telemetry.json", yaml) self.assertIn("backpressure_adaptive", yaml) self.assertIn("airfrans_disk_capacity_status=below_requested", yaml) self.assertNotIn("exit 74", yaml) with tempfile.TemporaryDirectory() as tmp: run_dir = Path(tmp) recorder = StreamingEventRecorder(run_dir) usage = shutil._ntuple_diskusage(total=1000, used=900, free=100) with patch("airfrans_frontier.training.streaming_data.shutil.disk_usage", return_value=usage): recorder.observe_cache(run_dir, cache_bytes=950) recorder.emit("cache_high_water", phase="data", cache_bytes=950, high_water_bytes=900) recorder.emit("producer_paused", phase="data", reason="cache_high_water") recorder.emit("cache_low_water", phase="data", cache_bytes=500, low_water_bytes=600) recorder.emit("producer_resumed", phase="data", reason="cache_low_water", idle_seconds=1.25) summary = recorder.to_dict() self.assertEqual(summary["minimum_free_disk_bytes"], 100) self.assertEqual(summary["cache_high_water_events"], 1) self.assertEqual(summary["cache_low_water_events"], 1) self.assertEqual(summary["producer_pause_events"], 1) self.assertEqual(summary["producer_resume_events"], 1) self.assertGreater(summary["producer_idle_backpressure_seconds"], 0) def offer(offer_id: int, *, price: float, host: int) -> VastOffer: return VastOffer( id=offer_id, gpu_name="RTX 4090", dph_total=price, gpu_ram=24_000, disk_space=256.0, geolocation="US", inet_down_cost_per_tb=0.0, inet_up_cost_per_tb=0.0, host_id=host, verification="verified", 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()