279 lines
13 KiB
Python
279 lines
13 KiB
Python
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()
|