99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from airfrans_frontier.remote.artifacts import verify_artifacts
|
|
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.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)
|
|
|
|
|
|
class ArtifactVerificationTests(unittest.TestCase):
|
|
def test_verify_artifacts_requires_contract_files_and_writes_manifest(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
(root / "final_metrics.json").write_text(json.dumps({"loss": 1.0}) + "\n")
|
|
(root / "metrics.jsonl").write_text(json.dumps({"step": 0}) + "\n")
|
|
(root / "checkpoint.pt").write_bytes(b"weights")
|
|
(root / "run_manifest.json").write_text(json.dumps({"exit_code": 0}) + "\n")
|
|
|
|
manifest = verify_artifacts(root)
|
|
|
|
self.assertEqual(manifest["file_count"], 4)
|
|
self.assertTrue((root / "artifact_manifest.json").is_file())
|
|
self.assertTrue((root / "checksums.txt").is_file())
|
|
|
|
|
|
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()
|