feat: first remote run
This commit is contained in:
parent
99b122999a
commit
05476f1724
30 changed files with 7501 additions and 73 deletions
6
.skyignore
Normal file
6
.skyignore
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
/artifacts
|
||||
/data/raw
|
||||
/.venv
|
||||
/notebooks
|
||||
__pycache__
|
||||
*.pyc
|
||||
32
configs/remote_mlp_tiny.toml
Normal file
32
configs/remote_mlp_tiny.toml
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
[run]
|
||||
name = "remote_mlp_tiny"
|
||||
seed = 0
|
||||
artifact_dir = "artifacts/current_run/training_runs"
|
||||
|
||||
[data]
|
||||
root = "data/processed/minimal"
|
||||
train_cases = 4
|
||||
val_cases = 1
|
||||
test_cases = 1
|
||||
points_per_case = 128
|
||||
batch_size = 128
|
||||
|
||||
[model]
|
||||
type = "mlp"
|
||||
hidden_width = 128
|
||||
depth = 4
|
||||
activation = "gelu"
|
||||
|
||||
[optim]
|
||||
lr = 0.001
|
||||
weight_decay = 0.0
|
||||
steps = 100
|
||||
log_interval = 25
|
||||
|
||||
[device]
|
||||
type = "cuda"
|
||||
allow_cpu_fallback = false
|
||||
benchmark_kernels = true
|
||||
|
||||
[loss]
|
||||
type = "normalized_mse"
|
||||
73
configs/remote_smoke.toml
Normal file
73
configs/remote_smoke.toml
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
[run]
|
||||
name = "airfrans-smoke"
|
||||
timeout_minutes = 45
|
||||
local_artifact_dir = "artifacts/remote_runs"
|
||||
max_attempts = 2
|
||||
|
||||
[provider]
|
||||
kind = "vastai"
|
||||
disk_gb = 64
|
||||
max_price_per_hour = 0.60
|
||||
image = "vastai/base:0.0.2"
|
||||
|
||||
[provider.gpu]
|
||||
name = "RTX 4090"
|
||||
count = 1
|
||||
min_vram_gb = 16
|
||||
|
||||
[selection]
|
||||
min_reliability = 0.95
|
||||
min_down_mbps = 100
|
||||
min_up_mbps = 25
|
||||
require_verified = true
|
||||
blocked_geos = ["CN"]
|
||||
blacklist_hosts = [59017]
|
||||
drop_cheap_frac = 0.30
|
||||
image_size_gb = 5.0
|
||||
base_url = "https://cloud.vast.ai"
|
||||
|
||||
[workspace]
|
||||
workdir = "."
|
||||
exclude = [
|
||||
"/artifacts",
|
||||
"/data/raw",
|
||||
"/.venv",
|
||||
"/notebooks",
|
||||
"__pycache__",
|
||||
"*.pyc",
|
||||
]
|
||||
|
||||
[bootstrap]
|
||||
command = """
|
||||
uv sync --no-dev
|
||||
uv run --no-dev python -c "import torch; assert torch.cuda.is_available(); print(torch.cuda.get_device_name(0))"
|
||||
"""
|
||||
|
||||
[data]
|
||||
validation_command = """
|
||||
uv run --no-dev python -c "from pathlib import Path; files=sorted(Path('data/processed/minimal').glob('*.npz')); assert len(files) >= 6; print(f'processed_minimal_cases={len(files)}')"
|
||||
"""
|
||||
|
||||
[job]
|
||||
command = """
|
||||
uv run --no-dev remote-run smoke-train configs/remote_mlp_tiny.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID"
|
||||
"""
|
||||
artifact_dir = "artifacts/current_run"
|
||||
heartbeat_file = "artifacts/current_run/heartbeat.json"
|
||||
metrics_file = "artifacts/current_run/metrics.jsonl"
|
||||
|
||||
[artifacts]
|
||||
mode = "rsync"
|
||||
required = [
|
||||
"final_metrics.json",
|
||||
"metrics.jsonl",
|
||||
"checkpoint.pt",
|
||||
"run_manifest.json",
|
||||
"environment_manifest.json",
|
||||
"artifact_manifest.json",
|
||||
"checksums.txt",
|
||||
]
|
||||
|
||||
[cleanup]
|
||||
on_success = "sky_down"
|
||||
on_failure = "collect_then_keep"
|
||||
File diff suppressed because one or more lines are too long
1152
notebooks/explore_test_run.ipynb
Normal file
1152
notebooks/explore_test_run.ipynb
Normal file
File diff suppressed because one or more lines are too long
|
|
@ -10,6 +10,7 @@ dependencies = [
|
|||
|
||||
[project.scripts]
|
||||
airfrans-frontier = "airfrans_frontier.cli:main"
|
||||
remote-run = "airfrans_frontier.remote.cli:main"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["src/airfrans_frontier"]
|
||||
|
|
@ -25,4 +26,5 @@ dev = [
|
|||
"numpy>=2.4.0",
|
||||
"nbclient>=0.10.2",
|
||||
"nbformat>=5.10.4",
|
||||
"skypilot[vast]>=0.12.3.post1",
|
||||
]
|
||||
|
|
|
|||
1
src/airfrans_frontier/remote/__init__.py
Normal file
1
src/airfrans_frontier/remote/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""Remote GPU smoke-run orchestration helpers."""
|
||||
71
src/airfrans_frontier/remote/artifacts.py
Normal file
71
src/airfrans_frontier/remote/artifacts.py
Normal file
|
|
@ -0,0 +1,71 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
DEFAULT_REQUIRED = ("final_metrics.json", "metrics.jsonl", "checkpoint.pt", "run_manifest.json")
|
||||
|
||||
|
||||
def verify_artifacts(artifact_dir: str | Path, required: Iterable[str] = DEFAULT_REQUIRED) -> dict[str, Any]:
|
||||
root = Path(artifact_dir)
|
||||
if not root.exists():
|
||||
raise FileNotFoundError(f"Artifact directory not found: {root}")
|
||||
if not root.is_dir():
|
||||
raise ValueError(f"Artifact path is not a directory: {root}")
|
||||
|
||||
missing = [name for name in required if not (root / name).is_file()]
|
||||
if missing:
|
||||
raise ValueError(f"Artifact directory missing required files: {', '.join(missing)}")
|
||||
|
||||
_validate_json(root / "final_metrics.json")
|
||||
_validate_json(root / "run_manifest.json")
|
||||
_validate_jsonl(root / "metrics.jsonl")
|
||||
|
||||
files = sorted(path for path in root.rglob("*") if path.is_file())
|
||||
manifest = {
|
||||
"artifact_dir": str(root),
|
||||
"file_count": len(files),
|
||||
"files": [
|
||||
{
|
||||
"path": str(path.relative_to(root)),
|
||||
"bytes": path.stat().st_size,
|
||||
"sha256": sha256_file(path),
|
||||
}
|
||||
for path in files
|
||||
],
|
||||
}
|
||||
(root / "artifact_manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
||||
(root / "checksums.txt").write_text(
|
||||
"".join(f"{item['sha256']} {item['path']}\n" for item in manifest["files"])
|
||||
)
|
||||
return manifest
|
||||
|
||||
|
||||
def sha256_file(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as file:
|
||||
for chunk in iter(lambda: file.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _validate_json(path: Path) -> None:
|
||||
try:
|
||||
json.loads(path.read_text())
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"Invalid JSON artifact {path}: {exc}") from exc
|
||||
|
||||
|
||||
def _validate_jsonl(path: Path) -> None:
|
||||
with path.open() as file:
|
||||
for line_number, line in enumerate(file, start=1):
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
continue
|
||||
try:
|
||||
json.loads(stripped)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"Invalid JSONL artifact {path}:{line_number}: {exc}") from exc
|
||||
241
src/airfrans_frontier/remote/cli.py
Normal file
241
src/airfrans_frontier/remote/cli.py
Normal file
|
|
@ -0,0 +1,241 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from airfrans_frontier.remote.artifacts import verify_artifacts
|
||||
from airfrans_frontier.remote.config import RemoteRunConfig, load_remote_run_config
|
||||
from airfrans_frontier.remote.skypilot import render_skypilot_yaml, write_skyignore
|
||||
from airfrans_frontier.remote.skypilot_patch import apply_patch, patch_status, require_patch
|
||||
from airfrans_frontier.remote.smoke import run_smoke_training
|
||||
from airfrans_frontier.remote.vast import SelectionResult, select_offer
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(prog="remote-run")
|
||||
subparsers = parser.add_subparsers(required=True)
|
||||
|
||||
doctor = subparsers.add_parser("doctor", help="check local remote-run and SkyPilot readiness")
|
||||
doctor.add_argument("--apply-skypilot-patch", action="store_true")
|
||||
doctor.set_defaults(command="doctor")
|
||||
|
||||
select = subparsers.add_parser("select", help="select a Vast.ai offer from a remote config")
|
||||
select.add_argument("config")
|
||||
select.add_argument("--out")
|
||||
select.set_defaults(command="select")
|
||||
|
||||
render = subparsers.add_parser("render", help="render patched SkyPilot YAML")
|
||||
render.add_argument("config")
|
||||
render.add_argument("--selection", required=True)
|
||||
render.add_argument("--run-id", required=True)
|
||||
render.add_argument("--out")
|
||||
render.set_defaults(command="render")
|
||||
|
||||
verify = subparsers.add_parser("verify-artifacts", help="verify a collected artifact directory")
|
||||
verify.add_argument("artifact_dir")
|
||||
verify.set_defaults(command="verify-artifacts")
|
||||
|
||||
smoke = subparsers.add_parser("smoke-train", help="run the configured smoke training job and flatten artifacts")
|
||||
smoke.add_argument("training_config")
|
||||
smoke.add_argument("--artifact-dir", required=True)
|
||||
smoke.add_argument("--run-id", required=True)
|
||||
smoke.set_defaults(command="smoke-train")
|
||||
|
||||
run = subparsers.add_parser("run", help="select a Vast offer and execute a SkyPilot run")
|
||||
run.add_argument("config")
|
||||
run.add_argument("--dry-run", action="store_true")
|
||||
run.add_argument("--skip-down", action="store_true", help="leave the SkyPilot cluster running after collection")
|
||||
run.set_defaults(command="run")
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
try:
|
||||
if args.command == "doctor":
|
||||
return _doctor(apply=args.apply_skypilot_patch)
|
||||
if args.command == "select":
|
||||
config = load_remote_run_config(args.config)
|
||||
result = select_offer(config)
|
||||
_emit_json(result.to_manifest(), args.out)
|
||||
return 0
|
||||
if args.command == "render":
|
||||
config = load_remote_run_config(args.config)
|
||||
selection = _selection_from_manifest(Path(args.selection))
|
||||
text = render_skypilot_yaml(config, selection, run_id=args.run_id)
|
||||
if args.out:
|
||||
Path(args.out).write_text(text)
|
||||
else:
|
||||
print(text)
|
||||
return 0
|
||||
if args.command == "verify-artifacts":
|
||||
manifest = verify_artifacts(args.artifact_dir)
|
||||
print(json.dumps({"status": "ok", "file_count": manifest["file_count"]}, sort_keys=True))
|
||||
return 0
|
||||
if args.command == "smoke-train":
|
||||
output = run_smoke_training(args.training_config, artifact_dir=args.artifact_dir, run_id=args.run_id)
|
||||
print(f"artifact_dir: {output}")
|
||||
return 0
|
||||
if args.command == "run":
|
||||
return _run(args.config, dry_run=args.dry_run, skip_down=args.skip_down)
|
||||
except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError, subprocess.CalledProcessError) as exc:
|
||||
print(f"error: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
parser.error(f"unknown command: {args.command}")
|
||||
return 2
|
||||
|
||||
|
||||
def _doctor(*, apply: bool) -> int:
|
||||
problems: list[str] = []
|
||||
if not os.environ.get("VAST_API_KEY"):
|
||||
problems.append("VAST_API_KEY is not set")
|
||||
sky = shutil.which("sky")
|
||||
if not sky:
|
||||
problems.append("sky executable not found on PATH")
|
||||
if apply:
|
||||
status = apply_patch()
|
||||
else:
|
||||
status = patch_status()
|
||||
if not status.installed:
|
||||
problems.append(status.message)
|
||||
elif not status.patched:
|
||||
problems.append(status.message)
|
||||
payload = {
|
||||
"vast_api_key_present": bool(os.environ.get("VAST_API_KEY")),
|
||||
"sky_executable": sky,
|
||||
"skypilot_patch": {
|
||||
"path": str(status.path) if status.path else None,
|
||||
"installed": status.installed,
|
||||
"patched": status.patched,
|
||||
"message": status.message,
|
||||
},
|
||||
"problems": problems,
|
||||
}
|
||||
print(json.dumps(payload, indent=2, sort_keys=True))
|
||||
return 1 if problems else 0
|
||||
|
||||
|
||||
def _run(config_path: str | Path, *, dry_run: bool, skip_down: bool) -> int:
|
||||
config = load_remote_run_config(config_path)
|
||||
if not dry_run:
|
||||
require_patch()
|
||||
run_id = f"{config.run.name}-{time.strftime('%Y%m%dT%H%M%SZ', time.gmtime())}"
|
||||
local_run_dir = config.run.local_artifact_dir / run_id
|
||||
local_run_dir.mkdir(parents=True, exist_ok=False)
|
||||
state_path = local_run_dir / "orchestrator_state.json"
|
||||
|
||||
def state(phase: str, **extra: Any) -> None:
|
||||
payload = {
|
||||
"run_id": run_id,
|
||||
"phase": phase,
|
||||
"updated_at": time.time(),
|
||||
"config_path": str(config.path),
|
||||
**extra,
|
||||
}
|
||||
state_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
||||
|
||||
state("SELECTING_OFFER")
|
||||
selection = select_offer(config)
|
||||
selection_path = local_run_dir / "selection_manifest.json"
|
||||
selection_path.write_text(json.dumps(selection.to_manifest(), indent=2, sort_keys=True) + "\n")
|
||||
|
||||
state("RENDERING_SKYPILOT_CONFIG", selected_offer_id=selection.selected_offer_id)
|
||||
sky_yaml = render_skypilot_yaml(config, selection, run_id=run_id)
|
||||
sky_yaml_path = local_run_dir / "sky.yaml"
|
||||
sky_yaml_path.write_text(sky_yaml)
|
||||
(local_run_dir / "config.toml").write_text(config.raw_text)
|
||||
write_skyignore(config)
|
||||
|
||||
if dry_run:
|
||||
state("DRY_RUN", selected_offer_id=selection.selected_offer_id, sky_yaml=str(sky_yaml_path))
|
||||
print(f"run_id: {run_id}")
|
||||
print(f"selection: {selection_path}")
|
||||
print(f"sky_yaml: {sky_yaml_path}")
|
||||
return 0
|
||||
|
||||
env = _subprocess_env()
|
||||
try:
|
||||
state("PROVISIONING", selected_offer_id=selection.selected_offer_id)
|
||||
_run_checked(["sky", "launch", "-c", run_id, str(sky_yaml_path), "-y"], env=env, timeout=config.run.timeout_minutes * 60)
|
||||
|
||||
state("COLLECTING", selected_offer_id=selection.selected_offer_id)
|
||||
_collect_with_rsync(cluster=run_id, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, env=env)
|
||||
|
||||
state("VERIFYING_ARTIFACTS", selected_offer_id=selection.selected_offer_id)
|
||||
verify_artifacts(local_run_dir, required=config.artifacts.required)
|
||||
except Exception as exc:
|
||||
state("FAILED", selected_offer_id=selection.selected_offer_id, error=str(exc))
|
||||
if config.cleanup.on_failure == "sky_down" and not skip_down:
|
||||
_run_best_effort(["sky", "down", run_id, "-y"], env=env)
|
||||
raise
|
||||
else:
|
||||
if config.cleanup.on_success == "sky_down" and not skip_down:
|
||||
state("CLEANING_UP", selected_offer_id=selection.selected_offer_id)
|
||||
_run_checked(["sky", "down", run_id, "-y"], env=env, timeout=300)
|
||||
state("SUCCEEDED", selected_offer_id=selection.selected_offer_id)
|
||||
print(f"run_id: {run_id}")
|
||||
print(f"artifacts: {local_run_dir}")
|
||||
return 0
|
||||
|
||||
|
||||
def _collect_with_rsync(*, cluster: str, remote_dir: Path, local_dir: Path, env: dict[str, str]) -> None:
|
||||
local_dir.mkdir(parents=True, exist_ok=True)
|
||||
source = f"{cluster}:~/sky_workdir/{remote_dir}/"
|
||||
_run_checked(["rsync", "-Pavz", source, f"{local_dir}/"], env=env, timeout=600)
|
||||
|
||||
|
||||
def _run_checked(argv: list[str], *, env: dict[str, str], timeout: int) -> None:
|
||||
subprocess.run(argv, check=True, env=env, timeout=timeout)
|
||||
|
||||
|
||||
def _run_best_effort(argv: list[str], *, env: dict[str, str]) -> None:
|
||||
try:
|
||||
subprocess.run(argv, check=False, env=env, timeout=300)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _subprocess_env() -> dict[str, str]:
|
||||
env = dict(os.environ)
|
||||
env.pop("PYTHONPATH", None)
|
||||
return env
|
||||
|
||||
|
||||
def _emit_json(data: dict[str, Any], out: str | None) -> None:
|
||||
text = json.dumps(data, indent=2, sort_keys=True) + "\n"
|
||||
if out:
|
||||
Path(out).write_text(text)
|
||||
else:
|
||||
print(text, end="")
|
||||
|
||||
|
||||
def _selection_from_manifest(path: Path) -> SelectionResult:
|
||||
from airfrans_frontier.remote.vast import VastOffer, effective_price
|
||||
|
||||
data = json.loads(path.read_text())
|
||||
raw_offer = data.get("selected_offer")
|
||||
if not isinstance(raw_offer, dict):
|
||||
raise ValueError(f"Selection manifest missing selected_offer object: {path}")
|
||||
offer = VastOffer.from_mapping({**raw_offer, "id": data.get("selected_offer_id", raw_offer.get("id"))})
|
||||
# Preserve manifest values by building a minimal SelectionResult. Effective price is already stored.
|
||||
return SelectionResult(
|
||||
selected_offer=offer,
|
||||
candidate_count=int(data.get("candidate_count", 0)),
|
||||
survivor_count=int(data.get("survivor_count", 0)),
|
||||
effective_price=float(raw_offer.get("effective_price", data.get("effective_price", 0.0))) if raw_offer else 0.0,
|
||||
query=data.get("query", {}) if isinstance(data.get("query"), dict) else {},
|
||||
policy=data.get("policy", {}) if isinstance(data.get("policy"), dict) else {},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
331
src/airfrans_frontier/remote/config.py
Normal file
331
src/airfrans_frontier/remote/config.py
Normal file
|
|
@ -0,0 +1,331 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import tomllib
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RunConfig:
|
||||
name: str
|
||||
timeout_minutes: int
|
||||
local_artifact_dir: Path
|
||||
max_attempts: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GpuConfig:
|
||||
name: str | None
|
||||
count: int
|
||||
min_vram_gb: int | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProviderConfig:
|
||||
kind: str
|
||||
gpu: GpuConfig
|
||||
disk_gb: int
|
||||
max_price_per_hour: float | None
|
||||
image: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SelectionConfig:
|
||||
min_reliability: float
|
||||
min_down_mbps: float
|
||||
min_up_mbps: float | None
|
||||
require_verified: bool
|
||||
blocked_geos: tuple[str, ...]
|
||||
blacklist_hosts: tuple[int, ...]
|
||||
drop_cheap_frac: float
|
||||
image_size_gb: float | None
|
||||
base_url: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkspaceConfig:
|
||||
workdir: Path
|
||||
exclude: tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BootstrapConfig:
|
||||
command: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DataConfig:
|
||||
validation_command: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class JobConfig:
|
||||
command: str
|
||||
artifact_dir: Path
|
||||
heartbeat_file: Path
|
||||
metrics_file: Path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArtifactConfig:
|
||||
mode: str
|
||||
required: tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CleanupConfig:
|
||||
on_success: str
|
||||
on_failure: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RemoteRunConfig:
|
||||
path: Path
|
||||
raw_text: str
|
||||
run: RunConfig
|
||||
provider: ProviderConfig
|
||||
selection: SelectionConfig
|
||||
workspace: WorkspaceConfig
|
||||
bootstrap: BootstrapConfig
|
||||
data: DataConfig
|
||||
job: JobConfig
|
||||
artifacts: ArtifactConfig
|
||||
cleanup: CleanupConfig
|
||||
|
||||
|
||||
_REQUIRED_SECTIONS = ("run", "provider", "selection", "workspace", "bootstrap", "job", "artifacts")
|
||||
|
||||
|
||||
def load_remote_run_config(path: str | Path) -> RemoteRunConfig:
|
||||
config_path = Path(path).expanduser()
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"Remote run config not found: {config_path}")
|
||||
if not config_path.is_file():
|
||||
raise ValueError(f"Remote run config is not a file: {config_path}")
|
||||
|
||||
text = config_path.read_text()
|
||||
try:
|
||||
raw = tomllib.loads(text)
|
||||
except tomllib.TOMLDecodeError as exc:
|
||||
raise ValueError(f"Invalid TOML config {config_path}: {exc}") from exc
|
||||
|
||||
for section in _REQUIRED_SECTIONS:
|
||||
if section not in raw or not isinstance(raw[section], dict):
|
||||
raise ValueError(f"Remote run config missing [{section}] section")
|
||||
|
||||
run_raw = raw["run"]
|
||||
provider_raw = raw["provider"]
|
||||
gpu_raw = _table(provider_raw, "gpu")
|
||||
selection_raw = raw["selection"]
|
||||
workspace_raw = raw["workspace"]
|
||||
bootstrap_raw = raw["bootstrap"]
|
||||
data_raw = raw.get("data", {})
|
||||
if data_raw is None:
|
||||
data_raw = {}
|
||||
if not isinstance(data_raw, dict):
|
||||
raise ValueError("Remote run [data] section must be a table")
|
||||
job_raw = raw["job"]
|
||||
artifacts_raw = raw["artifacts"]
|
||||
cleanup_raw = raw.get("cleanup", {})
|
||||
if cleanup_raw is None:
|
||||
cleanup_raw = {}
|
||||
if not isinstance(cleanup_raw, dict):
|
||||
raise ValueError("Remote run [cleanup] section must be a table")
|
||||
|
||||
gpu = GpuConfig(
|
||||
name=_optional_string(gpu_raw, "name"),
|
||||
count=_integer(gpu_raw, "count", minimum=1),
|
||||
min_vram_gb=_optional_integer(gpu_raw, "min_vram_gb", minimum=1),
|
||||
)
|
||||
provider = ProviderConfig(
|
||||
kind=_choice(_string(provider_raw, "kind"), {"vastai"}, "provider.kind"),
|
||||
gpu=gpu,
|
||||
disk_gb=_integer(provider_raw, "disk_gb", minimum=16),
|
||||
max_price_per_hour=_optional_number(provider_raw, "max_price_per_hour", minimum=0.0, exclusive_minimum=True),
|
||||
image=_optional_string(provider_raw, "image"),
|
||||
)
|
||||
|
||||
drop_cheap_frac = _number(selection_raw, "drop_cheap_frac", default=0.30, minimum=0.0)
|
||||
if drop_cheap_frac >= 1.0:
|
||||
raise ValueError("selection.drop_cheap_frac must be less than 1.0")
|
||||
selection = SelectionConfig(
|
||||
min_reliability=_number(selection_raw, "min_reliability", default=0.95, minimum=0.0),
|
||||
min_down_mbps=_number(selection_raw, "min_down_mbps", default=100.0, minimum=0.0),
|
||||
min_up_mbps=_optional_number(selection_raw, "min_up_mbps", minimum=0.0, exclusive_minimum=True),
|
||||
require_verified=_boolean(selection_raw, "require_verified", default=False),
|
||||
blocked_geos=_string_tuple(selection_raw, "blocked_geos", default=("CN",)),
|
||||
blacklist_hosts=tuple(int(value) for value in _integer_list(selection_raw, "blacklist_hosts", default=(59017,))),
|
||||
drop_cheap_frac=drop_cheap_frac,
|
||||
image_size_gb=_optional_number(selection_raw, "image_size_gb", minimum=0.0, exclusive_minimum=True),
|
||||
base_url=_string(selection_raw, "base_url", default="https://cloud.vast.ai"),
|
||||
)
|
||||
|
||||
workspace = WorkspaceConfig(
|
||||
workdir=_path(workspace_raw, "workdir", default="."),
|
||||
exclude=_string_tuple(workspace_raw, "exclude", default=()),
|
||||
)
|
||||
bootstrap = BootstrapConfig(command=_string(bootstrap_raw, "command"))
|
||||
data = DataConfig(validation_command=_optional_string(data_raw, "validation_command"))
|
||||
job = JobConfig(
|
||||
command=_string(job_raw, "command"),
|
||||
artifact_dir=Path(_string(job_raw, "artifact_dir")),
|
||||
heartbeat_file=Path(_string(job_raw, "heartbeat_file", default="artifacts/current_run/heartbeat.json")),
|
||||
metrics_file=Path(_string(job_raw, "metrics_file", default="artifacts/current_run/metrics.jsonl")),
|
||||
)
|
||||
artifacts = ArtifactConfig(
|
||||
mode=_choice(_string(artifacts_raw, "mode", default="rsync"), {"rsync", "object_store_upload"}, "artifacts.mode"),
|
||||
required=_string_tuple(artifacts_raw, "required", default=("final_metrics.json", "metrics.jsonl", "checkpoint.pt", "run_manifest.json")),
|
||||
)
|
||||
cleanup = CleanupConfig(
|
||||
on_success=_choice(_string(cleanup_raw, "on_success", default="sky_down"), {"sky_down", "keep"}, "cleanup.on_success"),
|
||||
on_failure=_choice(_string(cleanup_raw, "on_failure", default="collect_then_keep"), {"collect_then_keep", "sky_down"}, "cleanup.on_failure"),
|
||||
)
|
||||
|
||||
return RemoteRunConfig(
|
||||
path=config_path,
|
||||
raw_text=text,
|
||||
run=RunConfig(
|
||||
name=_string(run_raw, "name"),
|
||||
timeout_minutes=_integer(run_raw, "timeout_minutes", minimum=1, default=60),
|
||||
local_artifact_dir=_path(run_raw, "local_artifact_dir", default="artifacts/remote_runs"),
|
||||
max_attempts=_integer(run_raw, "max_attempts", minimum=1, default=2),
|
||||
),
|
||||
provider=provider,
|
||||
selection=selection,
|
||||
workspace=workspace,
|
||||
bootstrap=bootstrap,
|
||||
data=data,
|
||||
job=job,
|
||||
artifacts=artifacts,
|
||||
cleanup=cleanup,
|
||||
)
|
||||
|
||||
|
||||
def _table(section: dict[str, Any], key: str) -> dict[str, Any]:
|
||||
value = section.get(key)
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"Remote run config missing table: {key}")
|
||||
return value
|
||||
|
||||
|
||||
def _path(section: dict[str, Any], key: str, *, default: str | None = None) -> Path:
|
||||
return Path(_string(section, key, default=default)).expanduser()
|
||||
|
||||
|
||||
def _string(section: dict[str, Any], key: str, *, default: str | None = None) -> str:
|
||||
if key not in section:
|
||||
if default is not None:
|
||||
return default
|
||||
raise ValueError(f"Remote run config missing key: {key}")
|
||||
value = section[key]
|
||||
if not isinstance(value, str) or not value:
|
||||
raise ValueError(f"Remote run config key must be a non-empty string: {key}")
|
||||
return value
|
||||
|
||||
|
||||
def _optional_string(section: dict[str, Any], key: str) -> str | None:
|
||||
if key not in section:
|
||||
return None
|
||||
return _string(section, key)
|
||||
|
||||
|
||||
def _integer(section: dict[str, Any], key: str, *, minimum: int | None = None, default: int | None = None) -> int:
|
||||
if key not in section:
|
||||
if default is not None:
|
||||
return default
|
||||
raise ValueError(f"Remote run config missing key: {key}")
|
||||
value = section[key]
|
||||
if not isinstance(value, int):
|
||||
raise ValueError(f"Remote run config key must be an integer: {key}")
|
||||
if minimum is not None and value < minimum:
|
||||
raise ValueError(f"Remote run config key must be >= {minimum}: {key}")
|
||||
return value
|
||||
|
||||
|
||||
def _optional_integer(section: dict[str, Any], key: str, *, minimum: int | None = None) -> int | None:
|
||||
if key not in section:
|
||||
return None
|
||||
return _integer(section, key, minimum=minimum)
|
||||
|
||||
|
||||
def _integer_list(section: dict[str, Any], key: str, *, default: tuple[int, ...]) -> tuple[int, ...]:
|
||||
if key not in section:
|
||||
return default
|
||||
value = section[key]
|
||||
if not isinstance(value, list):
|
||||
raise ValueError(f"Remote run config key must be a list of integers: {key}")
|
||||
result: list[int] = []
|
||||
for item in value:
|
||||
if not isinstance(item, int):
|
||||
raise ValueError(f"Remote run config key must be a list of integers: {key}")
|
||||
result.append(item)
|
||||
return tuple(result)
|
||||
|
||||
|
||||
def _number(
|
||||
section: dict[str, Any],
|
||||
key: str,
|
||||
*,
|
||||
default: float | None = None,
|
||||
minimum: float | None = None,
|
||||
exclusive_minimum: bool = False,
|
||||
) -> float:
|
||||
if key not in section:
|
||||
if default is not None:
|
||||
return default
|
||||
raise ValueError(f"Remote run config missing key: {key}")
|
||||
value = section[key]
|
||||
if not isinstance(value, (int, float)):
|
||||
raise ValueError(f"Remote run config key must be a number: {key}")
|
||||
result = float(value)
|
||||
if minimum is not None:
|
||||
invalid = result <= minimum if exclusive_minimum else result < minimum
|
||||
if invalid:
|
||||
comparator = ">" if exclusive_minimum else ">="
|
||||
raise ValueError(f"Remote run config key must be {comparator} {minimum}: {key}")
|
||||
return result
|
||||
|
||||
|
||||
def _optional_number(
|
||||
section: dict[str, Any],
|
||||
key: str,
|
||||
*,
|
||||
minimum: float | None = None,
|
||||
exclusive_minimum: bool = False,
|
||||
) -> float | None:
|
||||
if key not in section:
|
||||
return None
|
||||
return _number(section, key, minimum=minimum, exclusive_minimum=exclusive_minimum)
|
||||
|
||||
|
||||
def _boolean(section: dict[str, Any], key: str, *, default: bool | None = None) -> bool:
|
||||
if key not in section:
|
||||
if default is not None:
|
||||
return default
|
||||
raise ValueError(f"Remote run config missing key: {key}")
|
||||
value = section[key]
|
||||
if not isinstance(value, bool):
|
||||
raise ValueError(f"Remote run config key must be a boolean: {key}")
|
||||
return value
|
||||
|
||||
|
||||
def _string_tuple(section: dict[str, Any], key: str, *, default: tuple[str, ...]) -> tuple[str, ...]:
|
||||
if key not in section:
|
||||
return default
|
||||
value = section[key]
|
||||
if not isinstance(value, list):
|
||||
raise ValueError(f"Remote run config key must be a list of strings: {key}")
|
||||
result: list[str] = []
|
||||
for item in value:
|
||||
if not isinstance(item, str):
|
||||
raise ValueError(f"Remote run config key must be a list of strings: {key}")
|
||||
result.append(item)
|
||||
return tuple(result)
|
||||
|
||||
|
||||
def _choice(value: str, allowed: set[str], key: str) -> str:
|
||||
if value not in allowed:
|
||||
allowed_values = ", ".join(sorted(allowed))
|
||||
raise ValueError(f"Unsupported {key}: {value}; expected one of {allowed_values}")
|
||||
return value
|
||||
105
src/airfrans_frontier/remote/skypilot.py
Normal file
105
src/airfrans_frontier/remote/skypilot.py
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from airfrans_frontier.remote.config import RemoteRunConfig
|
||||
from airfrans_frontier.remote.vast import SelectionResult
|
||||
|
||||
|
||||
def render_skypilot_yaml(config: RemoteRunConfig, selection: SelectionResult, *, run_id: str) -> str:
|
||||
setup = _compose_setup(config)
|
||||
run = _compose_run(config, run_id=run_id)
|
||||
accelerator = _accelerator(config)
|
||||
lines = [
|
||||
f"name: {run_id}",
|
||||
"",
|
||||
"resources:",
|
||||
" infra: vast",
|
||||
f" accelerators: {accelerator}",
|
||||
f" disk_size: {config.provider.disk_gb}",
|
||||
]
|
||||
if config.provider.max_price_per_hour is not None:
|
||||
lines.append(f" max_hourly_cost: {config.provider.max_price_per_hour}")
|
||||
if config.provider.image:
|
||||
image = config.provider.image
|
||||
if not image.startswith("docker:"):
|
||||
image = f"docker:{image}"
|
||||
lines.append(f" image_id: {image}")
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
f"workdir: {_yaml_scalar(str(config.workspace.workdir))}",
|
||||
"",
|
||||
"envs:",
|
||||
f" AIRFRANS_REMOTE_RUN_ID: {run_id}",
|
||||
"",
|
||||
"setup: |",
|
||||
*_indent_block(setup),
|
||||
"",
|
||||
"run: |",
|
||||
*_indent_block(run),
|
||||
"",
|
||||
"config:",
|
||||
" vast:",
|
||||
" create_instance_kwargs:",
|
||||
f" selected_offer_id: {selection.selected_offer_id}",
|
||||
"",
|
||||
]
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def write_skyignore(config: RemoteRunConfig, path: str | Path = ".skyignore") -> Path:
|
||||
skyignore_path = Path(path)
|
||||
entries = list(config.workspace.exclude)
|
||||
content = "\n".join(entries).rstrip() + "\n"
|
||||
skyignore_path.write_text(content)
|
||||
return skyignore_path
|
||||
|
||||
|
||||
def _compose_setup(config: RemoteRunConfig) -> str:
|
||||
return "\n".join(
|
||||
[
|
||||
"set -euo pipefail",
|
||||
"export PATH=\"$HOME/.local/bin:$PATH\"",
|
||||
"if ! command -v uv >/dev/null 2>&1; then curl -LsSf https://astral.sh/uv/install.sh | sh; fi",
|
||||
"export PATH=\"$HOME/.local/bin:$PATH\"",
|
||||
config.bootstrap.command.strip(),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _compose_run(config: RemoteRunConfig, *, run_id: str) -> str:
|
||||
lines = [
|
||||
"set -euo pipefail",
|
||||
"export PATH=\"$HOME/.local/bin:$PATH\"",
|
||||
f"mkdir -p {_sh_quote(str(config.job.artifact_dir))}",
|
||||
f"nvidia-smi | tee {_sh_quote(str(config.job.artifact_dir / 'nvidia_smi.txt'))}",
|
||||
]
|
||||
if config.data.validation_command:
|
||||
lines.append(config.data.validation_command.strip())
|
||||
lines.append(config.job.command.strip())
|
||||
lines.append(f"uv run --no-dev remote-run verify-artifacts {_sh_quote(str(config.job.artifact_dir))}")
|
||||
lines.append(f"echo 'remote run {run_id} complete'")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _accelerator(config: RemoteRunConfig) -> str:
|
||||
name = config.provider.gpu.name or "T4"
|
||||
# SkyPilot accelerator names omit spaces for common Vast names.
|
||||
sky_name = name.replace(" ", "")
|
||||
return f"{sky_name}:{config.provider.gpu.count}"
|
||||
|
||||
|
||||
def _indent_block(text: str) -> list[str]:
|
||||
return [f" {line}" if line else "" for line in text.splitlines()]
|
||||
|
||||
|
||||
def _yaml_scalar(value: str) -> str:
|
||||
if value and all(ch.isalnum() or ch in "./_-" for ch in value):
|
||||
return value
|
||||
return repr(value)
|
||||
|
||||
|
||||
def _sh_quote(value: str) -> str:
|
||||
return "'" + value.replace("'", "'\\''") + "'"
|
||||
83
src/airfrans_frontier/remote/skypilot_patch.py
Normal file
83
src/airfrans_frontier/remote/skypilot_patch.py
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
PATCH_MARKER = "AIRFRANS_SELECTED_OFFER_ID_PATCH"
|
||||
SDK_API_KEY_MARKER = "AIRFRANS_VAST_SDK_API_KEY_PATCH"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PatchStatus:
|
||||
path: Path | None
|
||||
installed: bool
|
||||
patched: bool
|
||||
message: str
|
||||
|
||||
|
||||
def locate_vast_utils() -> Path:
|
||||
try:
|
||||
module = importlib.import_module("sky.provision.vast.utils")
|
||||
except ModuleNotFoundError as exc:
|
||||
raise RuntimeError("SkyPilot is not importable in this Python environment") from exc
|
||||
path = getattr(module, "__file__", None)
|
||||
if not path:
|
||||
raise RuntimeError("Could not locate sky.provision.vast.utils source file")
|
||||
return Path(path)
|
||||
|
||||
|
||||
def patch_status() -> PatchStatus:
|
||||
try:
|
||||
path = locate_vast_utils()
|
||||
except RuntimeError as exc:
|
||||
return PatchStatus(path=None, installed=False, patched=False, message=str(exc))
|
||||
text = path.read_text()
|
||||
patched = PATCH_MARKER in text and SDK_API_KEY_MARKER in text
|
||||
if patched:
|
||||
return PatchStatus(path=path, installed=True, patched=True, message="SkyPilot Vast selected-offer patch is installed")
|
||||
return PatchStatus(path=path, installed=True, patched=False, message="SkyPilot Vast selected-offer patch is missing")
|
||||
|
||||
|
||||
def apply_patch() -> PatchStatus:
|
||||
path = locate_vast_utils()
|
||||
text = path.read_text()
|
||||
already_selected = PATCH_MARKER in text
|
||||
already_api_key = SDK_API_KEY_MARKER in text
|
||||
if already_selected and already_api_key:
|
||||
return PatchStatus(path=path, installed=True, patched=True, message="SkyPilot Vast selected-offer patch already installed")
|
||||
|
||||
if not already_selected:
|
||||
old = ''' instance_list = vast.vast().search_offers(query=query_str)\n\n if isinstance(instance_list, int) or len(instance_list) == 0:\n raise RuntimeError('Failed to create instances, could not find an '\n 'offer that satisfies the requirements '\n f'"{query_str}".')\n\n instance_touse = instance_list[0]\n\n # Start with user-provided kwargs as the base\n launch_params: Dict[str, Any] = dict(create_instance_kwargs or {})\n # Remove None values to avoid overriding defaults\n launch_params = {k: v for k, v in launch_params.items() if v is not None}\n'''
|
||||
new = f''' # {PATCH_MARKER}: allow callers to bypass SkyPilot's Vast offer\n # search after selecting a vetted Vast offer themselves. This keeps\n # SkyPilot's create/wait/bootstrap/run lifecycle while preserving\n # external marketplace-quality filtering.\n launch_params: Dict[str, Any] = dict(create_instance_kwargs or {{}})\n launch_params = {{k: v for k, v in launch_params.items() if v is not None}}\n selected_offer_id = launch_params.pop('selected_offer_id', None)\n\n if selected_offer_id is not None:\n logger.info(f'Using externally selected Vast offer {{selected_offer_id}}.')\n try:\n instance_touse = {{'id': int(selected_offer_id)}}\n except (TypeError, ValueError) as e:\n raise RuntimeError(\n f'Invalid selected_offer_id for Vast: {{selected_offer_id!r}}'\n ) from e\n else:\n instance_list = vast.vast().search_offers(query=query_str)\n\n if isinstance(instance_list, int) or len(instance_list) == 0:\n raise RuntimeError('Failed to create instances, could not find an '\n 'offer that satisfies the requirements '\n f'"{{query_str}}".')\n\n instance_touse = instance_list[0]\n\n'''
|
||||
if old not in text:
|
||||
raise RuntimeError(
|
||||
f"Could not apply SkyPilot Vast selected-offer patch; expected source block not found in {path}"
|
||||
)
|
||||
text = text.replace(old, new, 1)
|
||||
|
||||
if not already_api_key:
|
||||
old_key = " f'echo \"{vast.vast().client.api_key}\" > ~/.vast_api_key',"
|
||||
new_key = (
|
||||
f" # {SDK_API_KEY_MARKER}: vastai-sdk exposes api_key directly in current releases.\n"
|
||||
" f'echo \"{vast.vast().api_key}\" > ~/.vast_api_key',"
|
||||
)
|
||||
if old_key not in text:
|
||||
raise RuntimeError(
|
||||
f"Could not apply SkyPilot Vast API-key patch; expected source line not found in {path}"
|
||||
)
|
||||
text = text.replace(old_key, new_key, 1)
|
||||
|
||||
path.write_text(text)
|
||||
return PatchStatus(path=path, installed=True, patched=True, message="SkyPilot Vast selected-offer patch installed")
|
||||
|
||||
|
||||
def require_patch() -> Path:
|
||||
status = patch_status()
|
||||
if not status.installed:
|
||||
raise RuntimeError(status.message)
|
||||
if not status.patched:
|
||||
raise RuntimeError(f"{status.message}; run `remote-run doctor --apply-skypilot-patch`")
|
||||
assert status.path is not None
|
||||
return status.path
|
||||
127
src/airfrans_frontier/remote/smoke.py
Normal file
127
src/airfrans_frontier/remote/smoke.py
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import platform
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from airfrans_frontier.remote.artifacts import verify_artifacts
|
||||
from airfrans_frontier.runtime import remove_pythonpath_entries
|
||||
|
||||
|
||||
def run_smoke_training(config_path: str | Path, *, artifact_dir: str | Path, run_id: str) -> Path:
|
||||
remove_pythonpath_entries()
|
||||
from airfrans_frontier.training.loop import train_from_config_path
|
||||
|
||||
output_dir = Path(artifact_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
heartbeat_path = output_dir / "heartbeat.json"
|
||||
_write_json(
|
||||
heartbeat_path,
|
||||
{
|
||||
"run_id": run_id,
|
||||
"phase": "starting",
|
||||
"timestamp": time.time(),
|
||||
},
|
||||
)
|
||||
_write_json(output_dir / "environment_manifest.json", environment_manifest())
|
||||
|
||||
started = time.time()
|
||||
_write_json(
|
||||
heartbeat_path,
|
||||
{
|
||||
"run_id": run_id,
|
||||
"phase": "training",
|
||||
"started_at": started,
|
||||
"timestamp": time.time(),
|
||||
},
|
||||
)
|
||||
result = train_from_config_path(config_path)
|
||||
finished = time.time()
|
||||
|
||||
training_dir = result.run_dir
|
||||
required_from_training = [
|
||||
"final_metrics.json",
|
||||
"metrics.jsonl",
|
||||
"checkpoint.pt",
|
||||
"config.toml",
|
||||
"normalization.json",
|
||||
"split_manifest.json",
|
||||
]
|
||||
for name in required_from_training:
|
||||
source = training_dir / name
|
||||
if source.is_file():
|
||||
shutil.copy2(source, output_dir / name)
|
||||
|
||||
run_manifest: dict[str, Any] = {
|
||||
"run_id": run_id,
|
||||
"command": f"remote-run smoke-train {config_path}",
|
||||
"started_at": started,
|
||||
"finished_at": finished,
|
||||
"elapsed_seconds": finished - started,
|
||||
"exit_code": 0,
|
||||
"training_run_dir": str(training_dir),
|
||||
"artifact_dir": str(output_dir),
|
||||
"final_metrics_path": str(output_dir / "final_metrics.json"),
|
||||
"checkpoint_path": str(output_dir / "checkpoint.pt"),
|
||||
}
|
||||
_write_json(output_dir / "run_manifest.json", run_manifest)
|
||||
_write_json(
|
||||
heartbeat_path,
|
||||
{
|
||||
"run_id": run_id,
|
||||
"phase": "completed",
|
||||
"started_at": started,
|
||||
"finished_at": finished,
|
||||
"timestamp": time.time(),
|
||||
},
|
||||
)
|
||||
verify_artifacts(output_dir)
|
||||
return output_dir
|
||||
|
||||
|
||||
def environment_manifest() -> dict[str, Any]:
|
||||
manifest: dict[str, Any] = {
|
||||
"python": sys.version,
|
||||
"platform": platform.platform(),
|
||||
"executable": sys.executable,
|
||||
}
|
||||
try:
|
||||
import torch
|
||||
|
||||
manifest.update(
|
||||
{
|
||||
"torch_version": torch.__version__,
|
||||
"cuda_available": torch.cuda.is_available(),
|
||||
"cuda_version": torch.version.cuda,
|
||||
"gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
|
||||
}
|
||||
)
|
||||
except Exception as exc: # pragma: no cover - environment reporting should not hide root run errors.
|
||||
manifest["torch_error"] = repr(exc)
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["nvidia-smi", "--query-gpu=name,memory.total,driver_version", "--format=csv,noheader"],
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
timeout=15,
|
||||
check=False,
|
||||
)
|
||||
manifest["nvidia_smi"] = {
|
||||
"returncode": result.returncode,
|
||||
"stdout": result.stdout.strip(),
|
||||
"stderr": result.stderr.strip(),
|
||||
}
|
||||
except OSError as exc:
|
||||
manifest["nvidia_smi"] = {"error": str(exc)}
|
||||
return manifest
|
||||
|
||||
|
||||
def _write_json(path: Path, data: Any) -> None:
|
||||
path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n")
|
||||
258
src/airfrans_frontier/remote/vast.py
Normal file
258
src/airfrans_frontier/remote/vast.py
Normal file
|
|
@ -0,0 +1,258 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any, Mapping
|
||||
|
||||
from airfrans_frontier.remote.config import RemoteRunConfig, SelectionConfig
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class VastOffer:
|
||||
id: int
|
||||
gpu_name: str
|
||||
dph_total: float
|
||||
gpu_ram: float | None
|
||||
geolocation: str | None
|
||||
inet_down_cost_per_tb: float
|
||||
inet_up_cost_per_tb: float
|
||||
host_id: int | None
|
||||
verification: str | None
|
||||
reliability2: float | None
|
||||
cuda_max_good: float | None
|
||||
direct_port_count: int | None
|
||||
inet_down: float | None
|
||||
inet_up: float | None
|
||||
verified: bool | None
|
||||
|
||||
@classmethod
|
||||
def from_mapping(cls, data: Mapping[str, Any]) -> VastOffer:
|
||||
return cls(
|
||||
id=_int(data, "id"),
|
||||
gpu_name=_string(data, "gpu_name"),
|
||||
dph_total=_float(data, "dph_total"),
|
||||
gpu_ram=_optional_float(data, "gpu_ram"),
|
||||
geolocation=_optional_string(data, "geolocation"),
|
||||
inet_down_cost_per_tb=_optional_float(data, "internet_down_cost_per_tb") or 0.0,
|
||||
inet_up_cost_per_tb=_optional_float(data, "internet_up_cost_per_tb") or 0.0,
|
||||
host_id=_optional_int(data, "host_id"),
|
||||
verification=_optional_string(data, "verification"),
|
||||
reliability2=_optional_float(data, "reliability2"),
|
||||
cuda_max_good=_optional_float(data, "cuda_max_good"),
|
||||
direct_port_count=_optional_int(data, "direct_port_count"),
|
||||
inet_down=_optional_float(data, "inet_down"),
|
||||
inet_up=_optional_float(data, "inet_up"),
|
||||
verified=_optional_bool(data, "verified"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SelectionResult:
|
||||
selected_offer: VastOffer
|
||||
candidate_count: int
|
||||
survivor_count: int
|
||||
effective_price: float
|
||||
query: dict[str, Any]
|
||||
policy: dict[str, Any]
|
||||
|
||||
@property
|
||||
def selected_offer_id(self) -> int:
|
||||
return self.selected_offer.id
|
||||
|
||||
def to_manifest(self) -> dict[str, Any]:
|
||||
offer = asdict(self.selected_offer)
|
||||
offer["effective_price"] = self.effective_price
|
||||
return {
|
||||
"selected_offer_id": self.selected_offer_id,
|
||||
"selected_offer": offer,
|
||||
"candidate_count": self.candidate_count,
|
||||
"survivor_count": self.survivor_count,
|
||||
"query": self.query,
|
||||
"policy": self.policy,
|
||||
}
|
||||
|
||||
|
||||
def select_offer(config: RemoteRunConfig, *, api_key: str | None = None) -> SelectionResult:
|
||||
if config.provider.kind != "vastai":
|
||||
raise ValueError(f"Unsupported provider: {config.provider.kind}")
|
||||
resolved_key = api_key or os.environ.get("VAST_API_KEY")
|
||||
if not resolved_key:
|
||||
raise RuntimeError("VAST_API_KEY is required for Vast.ai offer selection")
|
||||
|
||||
query = build_query(config)
|
||||
offers = search_offers(
|
||||
base_url=config.selection.base_url,
|
||||
api_key=resolved_key,
|
||||
query=query,
|
||||
)
|
||||
return choose_offer(offers, config, query=query)
|
||||
|
||||
|
||||
def build_query(config: RemoteRunConfig) -> dict[str, Any]:
|
||||
selection = config.selection
|
||||
provider = config.provider
|
||||
query: dict[str, Any] = {
|
||||
"rentable": {"eq": True},
|
||||
"rented": {"eq": False},
|
||||
"reliability2": {"gte": selection.min_reliability},
|
||||
"cuda_max_good": {"gte": 12.6},
|
||||
"direct_port_count": {"gte": 1},
|
||||
"num_gpus": {"eq": provider.gpu.count},
|
||||
"inet_down": {"gte": selection.min_down_mbps},
|
||||
"limit": 5000,
|
||||
}
|
||||
if selection.min_up_mbps is not None:
|
||||
query["inet_up"] = {"gte": selection.min_up_mbps}
|
||||
if selection.require_verified:
|
||||
query["verified"] = {"eq": True}
|
||||
if provider.gpu.min_vram_gb is not None:
|
||||
query["gpu_ram"] = {"gte": provider.gpu.min_vram_gb * 1024}
|
||||
if provider.gpu.name:
|
||||
query["gpu_name"] = {"eq": provider.gpu.name}
|
||||
return query
|
||||
|
||||
|
||||
def search_offers(*, base_url: str, api_key: str, query: Mapping[str, Any]) -> list[VastOffer]:
|
||||
encoded = urllib.parse.quote(json.dumps(query, separators=(",", ":")))
|
||||
url = f"{base_url.rstrip('/')}/api/v0/bundles/?q={encoded}"
|
||||
request = urllib.request.Request(url, headers={"Authorization": f"Bearer {api_key}"})
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=45) as response:
|
||||
payload = json.loads(response.read().decode("utf-8"))
|
||||
except urllib.error.HTTPError as exc:
|
||||
body = exc.read().decode("utf-8", errors="replace")
|
||||
raise RuntimeError(f"Vast offer search HTTP {exc.code}: {body}") from exc
|
||||
except OSError as exc:
|
||||
raise RuntimeError(f"Vast offer search failed: {exc}") from exc
|
||||
|
||||
raw_offers = payload.get("offers")
|
||||
if not isinstance(raw_offers, list):
|
||||
raise RuntimeError("Vast offer search response missing offers list")
|
||||
return [VastOffer.from_mapping(item) for item in raw_offers if isinstance(item, Mapping)]
|
||||
|
||||
|
||||
def choose_offer(offers: list[VastOffer], config: RemoteRunConfig, *, query: Mapping[str, Any]) -> SelectionResult:
|
||||
survivors = reachable_offers(offers, config.selection)
|
||||
ranked = rank_survivors(survivors, config.selection)
|
||||
if config.provider.max_price_per_hour is not None:
|
||||
ranked = [offer for offer in ranked if effective_price(offer, config.selection) <= config.provider.max_price_per_hour]
|
||||
if not ranked:
|
||||
raise RuntimeError("No Vast offers survived quality filters and price cap")
|
||||
selected = ranked[0]
|
||||
return SelectionResult(
|
||||
selected_offer=selected,
|
||||
candidate_count=len(offers),
|
||||
survivor_count=len(ranked),
|
||||
effective_price=effective_price(selected, config.selection),
|
||||
query=dict(query),
|
||||
policy=selection_policy_manifest(config),
|
||||
)
|
||||
|
||||
|
||||
def reachable_offers(offers: list[VastOffer], selection: SelectionConfig) -> list[VastOffer]:
|
||||
blacklist = set(selection.blacklist_hosts)
|
||||
blocked = tuple(item.upper() for item in selection.blocked_geos)
|
||||
result: list[VastOffer] = []
|
||||
for offer in offers:
|
||||
geo = (offer.geolocation or "").upper()
|
||||
if blocked and any(token and token in geo for token in blocked):
|
||||
continue
|
||||
if offer.host_id is not None and offer.host_id in blacklist:
|
||||
continue
|
||||
if offer.verification == "deverified":
|
||||
continue
|
||||
result.append(offer)
|
||||
return result
|
||||
|
||||
|
||||
def rank_survivors(offers: list[VastOffer], selection: SelectionConfig) -> list[VastOffer]:
|
||||
by_model: dict[str, list[VastOffer]] = {}
|
||||
for offer in offers:
|
||||
by_model.setdefault(offer.gpu_name, []).append(offer)
|
||||
|
||||
survivors: list[VastOffer] = []
|
||||
for group in by_model.values():
|
||||
group.sort(key=lambda offer: effective_price(offer, selection))
|
||||
drop = math.floor(selection.drop_cheap_frac * len(group))
|
||||
survivors.extend(group[drop:])
|
||||
survivors.sort(key=lambda offer: effective_price(offer, selection))
|
||||
return survivors
|
||||
|
||||
|
||||
def effective_price(offer: VastOffer, selection: SelectionConfig) -> float:
|
||||
image_pull = 0.0
|
||||
if selection.image_size_gb is not None:
|
||||
image_pull = selection.image_size_gb * offer.inet_down_cost_per_tb / 1000.0
|
||||
return offer.dph_total + image_pull
|
||||
|
||||
|
||||
def selection_policy_manifest(config: RemoteRunConfig) -> dict[str, Any]:
|
||||
return {
|
||||
"gpu_name": config.provider.gpu.name,
|
||||
"gpu_count": config.provider.gpu.count,
|
||||
"min_vram_gb": config.provider.gpu.min_vram_gb,
|
||||
"max_price_per_hour": config.provider.max_price_per_hour,
|
||||
"min_reliability": config.selection.min_reliability,
|
||||
"min_down_mbps": config.selection.min_down_mbps,
|
||||
"min_up_mbps": config.selection.min_up_mbps,
|
||||
"require_verified": config.selection.require_verified,
|
||||
"blocked_geos": list(config.selection.blocked_geos),
|
||||
"blacklist_hosts": list(config.selection.blacklist_hosts),
|
||||
"drop_cheap_frac": config.selection.drop_cheap_frac,
|
||||
"image_size_gb": config.selection.image_size_gb,
|
||||
}
|
||||
|
||||
|
||||
def _string(data: Mapping[str, Any], key: str) -> str:
|
||||
value = data.get(key)
|
||||
if not isinstance(value, str):
|
||||
raise ValueError(f"Vast offer missing string field: {key}")
|
||||
return value
|
||||
|
||||
|
||||
def _optional_string(data: Mapping[str, Any], key: str) -> str | None:
|
||||
value = data.get(key)
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
def _int(data: Mapping[str, Any], key: str) -> int:
|
||||
value = data.get(key)
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"Vast offer missing integer field: {key}")
|
||||
return value
|
||||
|
||||
|
||||
def _optional_int(data: Mapping[str, Any], key: str) -> int | None:
|
||||
value = data.get(key)
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
if isinstance(value, float) and value.is_integer():
|
||||
return int(value)
|
||||
return None
|
||||
|
||||
|
||||
def _float(data: Mapping[str, Any], key: str) -> float:
|
||||
value = data.get(key)
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
raise ValueError(f"Vast offer missing numeric field: {key}")
|
||||
return float(value)
|
||||
|
||||
|
||||
def _optional_float(data: Mapping[str, Any], key: str) -> float | None:
|
||||
value = data.get(key)
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
return None
|
||||
|
||||
|
||||
def _optional_bool(data: Mapping[str, Any], key: str) -> bool | None:
|
||||
value = data.get(key)
|
||||
return value if isinstance(value, bool) else None
|
||||
99
tests/test_remote_run.py
Normal file
99
tests/test_remote_run.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
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()
|
||||
15
tools/vastai/Cargo.toml
Normal file
15
tools/vastai/Cargo.toml
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
[package]
|
||||
name = "swactor-vastai"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
reqwest = { version = "0.12", features = ["json"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
tokio = { version = "1", features = ["time"] }
|
||||
urlencoding = "2"
|
||||
|
||||
[dev-dependencies]
|
||||
wiremock = "0.6"
|
||||
76
tools/vastai/src/client.rs
Normal file
76
tools/vastai/src/client.rs
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
use crate::types::{
|
||||
LabeledInstance, LifecyclePolicy, Offer, ProvisionRequest, ProvisionedFleet, RunningInstance,
|
||||
};
|
||||
|
||||
/// Small convenience wrapper around a reqwest client + vast.ai endpoint.
|
||||
#[derive(Clone)]
|
||||
pub struct VastClient {
|
||||
http: reqwest::Client,
|
||||
base_url: String,
|
||||
api_key: String,
|
||||
}
|
||||
|
||||
impl VastClient {
|
||||
pub fn new(api_key: impl Into<String>) -> Self {
|
||||
Self::with_base_url("https://cloud.vast.ai", api_key)
|
||||
}
|
||||
|
||||
pub fn with_base_url(base_url: impl Into<String>, api_key: impl Into<String>) -> Self {
|
||||
Self {
|
||||
http: reqwest::Client::new(),
|
||||
base_url: base_url.into(),
|
||||
api_key: api_key.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn http(&self) -> &reqwest::Client {
|
||||
&self.http
|
||||
}
|
||||
|
||||
pub fn base_url(&self) -> &str {
|
||||
&self.base_url
|
||||
}
|
||||
|
||||
pub fn api_key(&self) -> &str {
|
||||
&self.api_key
|
||||
}
|
||||
|
||||
pub async fn search_offers(
|
||||
&self,
|
||||
policy: &crate::types::SelectionPolicy,
|
||||
target_count: u32,
|
||||
) -> Result<Vec<Offer>, String> {
|
||||
crate::search::select_offer_pool_with_policy(
|
||||
&self.http,
|
||||
&self.base_url,
|
||||
&self.api_key,
|
||||
policy,
|
||||
target_count,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn provision(&self, req: ProvisionRequest) -> Result<ProvisionedFleet, String> {
|
||||
crate::lease::provision_fleet(&self.http, &self.base_url, &self.api_key, req).await
|
||||
}
|
||||
|
||||
pub async fn wait_for_running(
|
||||
&self,
|
||||
contract_id: u64,
|
||||
policy: &LifecyclePolicy,
|
||||
) -> Result<RunningInstance, String> {
|
||||
crate::monitor::wait_for_running_with_policy(
|
||||
&self.http,
|
||||
&self.base_url,
|
||||
&self.api_key,
|
||||
contract_id,
|
||||
policy,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn list_by_label(&self, label: &str) -> Result<Vec<LabeledInstance>, String> {
|
||||
crate::teardown::list_instances_by_label(&self.http, &self.base_url, &self.api_key, label)
|
||||
.await
|
||||
}
|
||||
}
|
||||
87
tools/vastai/src/config.rs
Normal file
87
tools/vastai/src/config.rs
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use crate::types::{LifecyclePolicy, SelectionPolicy};
|
||||
|
||||
pub const ENV_IMAGE_SIZE_GB: &str = "PP_IMAGE_SIZE_GB";
|
||||
pub const ENV_GPU_MIN_RAM_MB: &str = "PP_GPU_MIN_RAM_MB";
|
||||
pub const ENV_MIN_INET_DOWN_MBPS: &str = "PP_MIN_INET_DOWN_MBPS";
|
||||
pub const ENV_MIN_INET_UP_MBPS: &str = "PP_MIN_INET_UP_MBPS";
|
||||
pub const ENV_MIN_RELIABILITY: &str = "PP_MIN_RELIABILITY";
|
||||
pub const ENV_REQUIRE_VERIFIED: &str = "PP_REQUIRE_VERIFIED";
|
||||
pub const ENV_DROP_CHEAP_FRAC: &str = "PP_DROP_CHEAP_FRAC";
|
||||
pub const ENV_LEASE_PACE_MS: &str = "PP_LEASE_PACE_MS";
|
||||
pub const ENV_BLACKLIST_HOSTS: &str = "PP_BLACKLIST_HOSTS";
|
||||
pub const ENV_ASSUME_YES: &str = "PP_ASSUME_YES";
|
||||
|
||||
pub fn truthy_env(name: &str) -> bool {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.map(|s| matches!(s.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes"))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn env_positive_u64(name: &str) -> Option<u64> {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|s| s.trim().parse::<u64>().ok())
|
||||
.filter(|&n| n > 0)
|
||||
}
|
||||
|
||||
fn env_nonnegative_f64(name: &str, default: f64) -> f64 {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|s| s.trim().parse::<f64>().ok())
|
||||
.filter(|&v| v >= 0.0)
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn env_optional_positive_f64(name: &str) -> Option<f64> {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|s| s.trim().parse::<f64>().ok())
|
||||
.filter(|&v| v > 0.0)
|
||||
}
|
||||
|
||||
impl SelectionPolicy {
|
||||
/// Build selection policy from the historical `PP_*` environment knobs.
|
||||
pub fn from_env() -> Self {
|
||||
let mut policy = Self::default();
|
||||
policy.min_gpu_ram_mb = env_positive_u64(ENV_GPU_MIN_RAM_MB);
|
||||
policy.min_down_mbps = env_nonnegative_f64(ENV_MIN_INET_DOWN_MBPS, 100.0);
|
||||
policy.min_reliability = std::env::var(ENV_MIN_RELIABILITY)
|
||||
.ok()
|
||||
.and_then(|s| s.trim().parse::<f64>().ok())
|
||||
.filter(|&v| (0.0..=1.0).contains(&v))
|
||||
.unwrap_or(0.95);
|
||||
policy.require_verified = truthy_env(ENV_REQUIRE_VERIFIED);
|
||||
policy.min_up_mbps = env_optional_positive_f64(ENV_MIN_INET_UP_MBPS);
|
||||
policy.drop_cheap_frac = std::env::var(ENV_DROP_CHEAP_FRAC)
|
||||
.ok()
|
||||
.and_then(|s| s.trim().parse::<f64>().ok())
|
||||
.filter(|v| v.is_finite())
|
||||
.map(|v| v.clamp(0.0, 0.99))
|
||||
.unwrap_or(0.30);
|
||||
policy.image_size_gb = env_optional_positive_f64(ENV_IMAGE_SIZE_GB);
|
||||
if let Ok(raw) = std::env::var(ENV_BLACKLIST_HOSTS) {
|
||||
policy
|
||||
.blacklist_hosts
|
||||
.extend(raw.split(',').filter_map(|s| s.trim().parse::<u64>().ok()));
|
||||
}
|
||||
policy
|
||||
}
|
||||
}
|
||||
|
||||
impl LifecyclePolicy {
|
||||
/// Build lifecycle policy from environment, using caller-provided poll cadence.
|
||||
pub fn from_env(poll_interval: Duration) -> Self {
|
||||
let mut policy = Self::default();
|
||||
policy.lease_pace = Duration::from_millis(
|
||||
std::env::var(ENV_LEASE_PACE_MS)
|
||||
.ok()
|
||||
.and_then(|s| s.trim().parse::<u64>().ok())
|
||||
.unwrap_or(600),
|
||||
);
|
||||
policy.poll_interval = poll_interval;
|
||||
policy
|
||||
}
|
||||
}
|
||||
17
tools/vastai/src/filters.rs
Normal file
17
tools/vastai/src/filters.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use std::collections::HashSet;
|
||||
|
||||
use crate::types::{Offer, SelectionPolicy};
|
||||
|
||||
pub(crate) fn reachable_offers(offers: Vec<Offer>, policy: &SelectionPolicy) -> Vec<Offer> {
|
||||
let blacklist: HashSet<u64> = policy.blacklist_hosts.iter().copied().collect();
|
||||
offers
|
||||
.into_iter()
|
||||
.filter(|o| {
|
||||
o.geolocation
|
||||
.as_deref()
|
||||
.map_or(false, |g| !g.to_uppercase().contains("CN"))
|
||||
})
|
||||
.filter(|o| o.host_id.map_or(true, |h| !blacklist.contains(&h)))
|
||||
.filter(|o| o.verification.as_deref() != Some("deverified"))
|
||||
.collect()
|
||||
}
|
||||
288
tools/vastai/src/lease.rs
Normal file
288
tools/vastai/src/lease.rs
Normal file
|
|
@ -0,0 +1,288 @@
|
|||
use std::collections::{BTreeMap, HashSet};
|
||||
use std::io::{IsTerminal, Write};
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::config::{ENV_ASSUME_YES, truthy_env};
|
||||
use crate::monitor::wait_for_running_with_policy;
|
||||
use crate::pricing::{CostModel, plan_picks};
|
||||
use crate::provision::create_instance;
|
||||
use crate::search::select_offer_pool_with_policy;
|
||||
use crate::teardown::{destroy_instance_with_retry, rollback};
|
||||
use crate::types::{
|
||||
CreateInstanceRequest, InstanceInfo, Offer, ProvisionRequest, ProvisionedFleet,
|
||||
ProvisionedInstance,
|
||||
};
|
||||
|
||||
/// Print the planned lease + hourly cost and, on TTY, require y/N confirmation.
|
||||
pub fn confirm_lease(pool: &[Offer], num_instances: u32, cost: &CostModel) -> Result<(), String> {
|
||||
let picks = plan_picks(pool, num_instances);
|
||||
let total_dph: f64 = picks.iter().map(|o| o.dph_total).sum();
|
||||
let total_eff: f64 = picks.iter().map(|o| cost.effective_price(o)).sum();
|
||||
|
||||
eprintln!("vastai: lease plan — {num_instances} instance(s), cheapest on distinct hosts:");
|
||||
for (i, o) in picks.iter().enumerate() {
|
||||
eprintln!(
|
||||
" node {i} {:<14} {:>8} ${:.3}/hr [{}] host {}",
|
||||
o.gpu_name,
|
||||
o.gpu_ram
|
||||
.map(|r| format!("{:.0}MB", r))
|
||||
.unwrap_or_else(|| "?MB".into()),
|
||||
o.dph_total,
|
||||
o.geolocation.as_deref().unwrap_or("?"),
|
||||
o.host_id
|
||||
.map(|h| h.to_string())
|
||||
.unwrap_or_else(|| "?".into()),
|
||||
);
|
||||
}
|
||||
if picks.len() < num_instances as usize {
|
||||
eprintln!(
|
||||
" WARNING: only {} distinct-host offer(s) available for {num_instances} instance(s)",
|
||||
picks.len(),
|
||||
);
|
||||
}
|
||||
let eff_note = if (total_eff - total_dph).abs() > 1e-6 {
|
||||
format!(" (image-pull priced in: ${total_eff:.3}/hr eff)")
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
eprintln!(
|
||||
" TOTAL ${total_dph:.3}/hr (~${:.2}/day){eff_note}",
|
||||
total_dph * 24.0,
|
||||
);
|
||||
|
||||
if truthy_env(ENV_ASSUME_YES) {
|
||||
eprintln!("vastai: {ENV_ASSUME_YES} set — proceeding without confirmation");
|
||||
return Ok(());
|
||||
}
|
||||
if !std::io::stdin().is_terminal() {
|
||||
eprintln!(
|
||||
"vastai: stdin is not a TTY — proceeding without interactive confirmation \
|
||||
(set {ENV_ASSUME_YES}=1 to silence this)"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
eprint!("Proceed with renting these {num_instances} instance(s)? [y/N]: ");
|
||||
let _ = std::io::stderr().flush();
|
||||
let mut line = String::new();
|
||||
std::io::stdin()
|
||||
.read_line(&mut line)
|
||||
.map_err(|e| format!("failed to read lease confirmation: {e}"))?;
|
||||
let ans = line.trim().to_ascii_lowercase();
|
||||
if ans == "y" || ans == "yes" {
|
||||
Ok(())
|
||||
} else {
|
||||
Err("operator declined the lease (cost not confirmed); no instances were created".into())
|
||||
}
|
||||
}
|
||||
|
||||
fn next_eligible_offer<'a>(
|
||||
pool: &'a [Offer],
|
||||
tried_offer_ids: &[u64],
|
||||
used_host_ids: &HashSet<u64>,
|
||||
) -> Option<&'a Offer> {
|
||||
pool.iter().find(|o| {
|
||||
!tried_offer_ids.contains(&o.id) && o.host_id.map_or(true, |h| !used_host_ids.contains(&h))
|
||||
})
|
||||
}
|
||||
|
||||
fn env_for_index(req: &ProvisionRequest, index: u32) -> BTreeMap<String, String> {
|
||||
let mut env = req.env.clone();
|
||||
if let Some(extra) = req.per_instance_env.get(index as usize) {
|
||||
env.extend(extra.clone());
|
||||
}
|
||||
env
|
||||
}
|
||||
|
||||
async fn provision_one(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
req: &ProvisionRequest,
|
||||
pool: &[Offer],
|
||||
index: u32,
|
||||
tried_offer_ids: &mut Vec<u64>,
|
||||
used_host_ids: &mut HashSet<u64>,
|
||||
) -> Result<ProvisionedInstance, String> {
|
||||
let mut attempt = 1_u64;
|
||||
loop {
|
||||
let offer = match next_eligible_offer(pool, tried_offer_ids, used_host_ids) {
|
||||
Some(o) => o.clone(),
|
||||
None => {
|
||||
return Err(format!(
|
||||
"pool exhausted for index {index} (no untried offer on an unused host)"
|
||||
));
|
||||
}
|
||||
};
|
||||
tried_offer_ids.push(offer.id);
|
||||
let cost = CostModel::from_policy(&req.selection);
|
||||
eprintln!(
|
||||
"lease_chain: index {index} → offer {} — {} {} @ ${:.3}/hr [{}] host {} eff ${:.3}/hr",
|
||||
offer.id,
|
||||
offer.gpu_name,
|
||||
offer
|
||||
.gpu_ram
|
||||
.map(|r| format!("{:.0}MB", r))
|
||||
.unwrap_or_else(|| "?MB".into()),
|
||||
offer.dph_total,
|
||||
offer.geolocation.as_deref().unwrap_or("?"),
|
||||
offer
|
||||
.host_id
|
||||
.map(|h| h.to_string())
|
||||
.unwrap_or_else(|| "?".into()),
|
||||
cost.effective_price(&offer),
|
||||
);
|
||||
|
||||
let create = CreateInstanceRequest {
|
||||
offer_id: offer.id,
|
||||
image: req.image.clone(),
|
||||
disk_gb: req.disk_gb,
|
||||
label: req.label.clone(),
|
||||
env: env_for_index(req, index),
|
||||
onstart: req.onstart.clone(),
|
||||
};
|
||||
match create_instance(client, base_url, api_key, &create).await {
|
||||
Ok(info) => {
|
||||
if let Some(h) = offer.host_id {
|
||||
used_host_ids.insert(h);
|
||||
}
|
||||
return Ok(ProvisionedInstance {
|
||||
index,
|
||||
contract_id: info.contract_id,
|
||||
offer_id: offer.id,
|
||||
host_id: offer.host_id,
|
||||
gpu_name: offer.gpu_name,
|
||||
gpu_ram: offer.gpu_ram,
|
||||
dph_total: offer.dph_total,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"lease_chain: index {index} create on offer {} failed (attempt {attempt}): {e}",
|
||||
offer.id,
|
||||
);
|
||||
let is_429 = e.contains("429") || e.contains("Too Many Requests");
|
||||
let backoff = if is_429 {
|
||||
std::cmp::min(
|
||||
Duration::from_millis(2_000_u64.saturating_mul(attempt)),
|
||||
Duration::from_secs(30),
|
||||
)
|
||||
} else {
|
||||
Duration::from_millis(400)
|
||||
};
|
||||
tokio::time::sleep(backoff).await;
|
||||
attempt = attempt.saturating_add(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn as_instance_infos(instances: &[ProvisionedInstance]) -> Vec<InstanceInfo> {
|
||||
instances
|
||||
.iter()
|
||||
.map(|i| InstanceInfo {
|
||||
contract_id: i.contract_id,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Rent, monitor, replace, and roll back an N-instance fleet.
|
||||
pub async fn provision_fleet(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
req: ProvisionRequest,
|
||||
) -> Result<ProvisionedFleet, String> {
|
||||
let pool = select_offer_pool_with_policy(client, base_url, api_key, &req.selection, req.count)
|
||||
.await
|
||||
.map_err(|e| format!("lease_chain: {e}"))?;
|
||||
|
||||
if req.confirm_lease {
|
||||
confirm_lease(&pool, req.count, &CostModel::from_policy(&req.selection))?;
|
||||
}
|
||||
|
||||
let mut tried_offer_ids = Vec::new();
|
||||
let mut created: Vec<ProvisionedInstance> = Vec::with_capacity(req.count as usize);
|
||||
let mut used_host_ids = HashSet::new();
|
||||
|
||||
for index in 0..req.count {
|
||||
match provision_one(
|
||||
client,
|
||||
base_url,
|
||||
api_key,
|
||||
&req,
|
||||
&pool,
|
||||
index,
|
||||
&mut tried_offer_ids,
|
||||
&mut used_host_ids,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(info) => created.push(info),
|
||||
Err(e) => {
|
||||
rollback(client, base_url, api_key, &as_instance_infos(&created)).await;
|
||||
return Err(format!("lease_chain: {e}"));
|
||||
}
|
||||
}
|
||||
if index + 1 < req.count {
|
||||
tokio::time::sleep(req.lifecycle.lease_pace).await;
|
||||
}
|
||||
}
|
||||
|
||||
for index in 0..req.count {
|
||||
let idx = index as usize;
|
||||
loop {
|
||||
let cid = created[idx].contract_id;
|
||||
match wait_for_running_with_policy(client, base_url, api_key, cid, &req.lifecycle).await
|
||||
{
|
||||
Ok(_) => break,
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"lease_chain: index {index} contract {cid} did not reach running: {e}"
|
||||
);
|
||||
if let Err(de) =
|
||||
destroy_instance_with_retry(client, base_url, api_key, cid).await
|
||||
{
|
||||
eprintln!(
|
||||
"lease_chain: WARNING could not destroy dead contract {cid}: {de}"
|
||||
);
|
||||
}
|
||||
eprintln!("lease_chain: replacing index {index}");
|
||||
match provision_one(
|
||||
client,
|
||||
base_url,
|
||||
api_key,
|
||||
&req,
|
||||
&pool,
|
||||
index,
|
||||
&mut tried_offer_ids,
|
||||
&mut used_host_ids,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(info) => created[idx] = info,
|
||||
Err(pe) => {
|
||||
let survivors: Vec<InstanceInfo> = created
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(i, _)| *i != idx)
|
||||
.map(|(_, c)| InstanceInfo {
|
||||
contract_id: c.contract_id,
|
||||
})
|
||||
.collect();
|
||||
rollback(client, base_url, api_key, &survivors).await;
|
||||
return Err(format!(
|
||||
"lease_chain: index {index} replacement could not be provisioned: {pe}"
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(ProvisionedFleet {
|
||||
label: req.label,
|
||||
instances: created,
|
||||
})
|
||||
}
|
||||
34
tools/vastai/src/lib.rs
Normal file
34
tools/vastai/src/lib.rs
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
//! Peripheral vast.ai provisioning utility.
|
||||
//!
|
||||
//! This crate is intentionally outside `crates/`: swactor itself stays a
|
||||
//! general-purpose actor runtime, while this utility rents, monitors, and tears
|
||||
//! down vast.ai machines for apps that choose to use it.
|
||||
|
||||
pub mod client;
|
||||
pub mod config;
|
||||
pub mod filters;
|
||||
pub mod lease;
|
||||
pub mod logs;
|
||||
pub mod monitor;
|
||||
pub mod pricing;
|
||||
pub mod provision;
|
||||
pub mod search;
|
||||
pub mod state;
|
||||
pub mod teardown;
|
||||
pub mod types;
|
||||
|
||||
pub use client::VastClient;
|
||||
pub use lease::{confirm_lease, provision_fleet};
|
||||
pub use logs::{fetch_logs, request_logs};
|
||||
pub use monitor::{wait_for_running, wait_for_running_with_policy};
|
||||
pub use pricing::CostModel;
|
||||
pub use provision::create_instance;
|
||||
pub use search::{select_offer_pool, select_offer_pool_with_policy};
|
||||
pub use teardown::{
|
||||
destroy_all_instances, destroy_instance, destroy_instance_with_retry, list_instances_by_label,
|
||||
};
|
||||
pub use types::{
|
||||
ContractRef, CreateInstanceRequest, FleetState, InstanceInfo, LabeledInstance, LifecyclePolicy,
|
||||
Offer, ProvisionRequest, ProvisionedFleet, ProvisionedInstance, RunningInstance,
|
||||
SelectionPolicy,
|
||||
};
|
||||
36
tools/vastai/src/logs.rs
Normal file
36
tools/vastai/src/logs.rs
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
use std::time::Duration;
|
||||
|
||||
pub async fn request_logs(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
contract_id: u64,
|
||||
) -> Result<String, String> {
|
||||
let url = format!("{base_url}/api/v0/instances/request_logs/{contract_id}/");
|
||||
let resp = client
|
||||
.put(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("request_logs failed: {e}"))?;
|
||||
let body: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("request_logs parse failed: {e}"))?;
|
||||
body["result_url"]
|
||||
.as_str()
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| "no result_url in log response".to_string())
|
||||
}
|
||||
|
||||
pub async fn fetch_logs(client: &reqwest::Client, log_url: &str) -> Result<String, String> {
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
let resp = client
|
||||
.get(log_url)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("fetch_logs failed: {e}"))?;
|
||||
resp.text()
|
||||
.await
|
||||
.map_err(|e| format!("fetch_logs read failed: {e}"))
|
||||
}
|
||||
114
tools/vastai/src/monitor.rs
Normal file
114
tools/vastai/src/monitor.rs
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
use crate::types::{InstanceResponse, LifecyclePolicy, RunningInstance};
|
||||
|
||||
/// Historical env-backed polling wrapper.
|
||||
pub async fn wait_for_running(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
contract_id: u64,
|
||||
poll_interval: std::time::Duration,
|
||||
) -> Result<RunningInstance, String> {
|
||||
let policy = LifecyclePolicy::from_env(poll_interval);
|
||||
wait_for_running_with_policy(client, base_url, api_key, contract_id, &policy).await
|
||||
}
|
||||
|
||||
/// Poll vast.ai until an instance reaches `running`, or fail on terminal provider state.
|
||||
pub async fn wait_for_running_with_policy(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
contract_id: u64,
|
||||
policy: &LifecyclePolicy,
|
||||
) -> Result<RunningInstance, String> {
|
||||
let url = format!("{base_url}/api/v0/instances/{contract_id}/");
|
||||
let mut state_since = std::time::Instant::now();
|
||||
let mut last_state: Option<String> = None;
|
||||
|
||||
let mut poll = 0_u64;
|
||||
loop {
|
||||
poll += 1;
|
||||
let resp = match client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
eprintln!(" contract {contract_id} poll {poll}: request error: {e} (retrying)");
|
||||
tokio::time::sleep(policy.poll_interval).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
eprintln!(
|
||||
" contract {contract_id} poll {poll}: HTTP {status} (retrying): {}",
|
||||
body.chars().take(80).collect::<String>(),
|
||||
);
|
||||
tokio::time::sleep(policy.poll_interval).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
let wrapper: InstanceResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("wait_for_running parse failed: {e}"))?;
|
||||
let status = wrapper.instances;
|
||||
|
||||
let actual = status.actual_status.as_deref().unwrap_or("unknown");
|
||||
let intended = status.intended_status.as_deref().unwrap_or("unknown");
|
||||
|
||||
let msg = status.status_msg.clone();
|
||||
let disk = status.disk_usage;
|
||||
|
||||
if last_state.as_deref() != Some(actual) {
|
||||
state_since = std::time::Instant::now();
|
||||
}
|
||||
last_state = Some(actual.to_string());
|
||||
let in_state = state_since.elapsed().as_secs();
|
||||
let msg_disp = match msg.as_deref() {
|
||||
Some(m) if !m.is_empty() => format!(" msg=\"{m}\""),
|
||||
_ => String::new(),
|
||||
};
|
||||
let disk_disp = match disk {
|
||||
Some(d) if d >= 0.0 => format!(" disk={d:.2}GB"),
|
||||
_ => String::new(),
|
||||
};
|
||||
eprintln!(
|
||||
" contract {contract_id} poll {poll}: status={actual} in-state={in_state}s{msg_disp}{disk_disp}",
|
||||
);
|
||||
|
||||
if let Some(m) = &msg {
|
||||
if m.contains("Error") || m.contains("failed") {
|
||||
return Err(format!("instance {contract_id} error: {m}"));
|
||||
}
|
||||
}
|
||||
if intended == "stopped" && actual != "running" {
|
||||
return Err(format!(
|
||||
"instance {contract_id} stopped: {}",
|
||||
msg.unwrap_or_default()
|
||||
));
|
||||
}
|
||||
|
||||
match actual {
|
||||
"running" => {
|
||||
let ip = status
|
||||
.public_ipaddr
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
let port = status.ssh_port.unwrap_or(0);
|
||||
return Ok(RunningInstance { ip, port });
|
||||
}
|
||||
"exited" | "error" => {
|
||||
return Err(format!(
|
||||
"instance {contract_id} reached terminal status: {actual}"
|
||||
));
|
||||
}
|
||||
_ => {
|
||||
tokio::time::sleep(policy.poll_interval).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
67
tools/vastai/src/pricing.rs
Normal file
67
tools/vastai/src/pricing.rs
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
use crate::types::{Offer, SelectionPolicy};
|
||||
|
||||
/// Cost model for ranking offers on true lease cost rather than listed $/hr.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CostModel {
|
||||
/// Deploy-image size in GB; `None` → image pull is not priced in.
|
||||
pub image_gb: Option<f64>,
|
||||
}
|
||||
|
||||
impl CostModel {
|
||||
pub fn from_policy(policy: &SelectionPolicy) -> Self {
|
||||
Self {
|
||||
image_gb: policy.image_size_gb,
|
||||
}
|
||||
}
|
||||
|
||||
/// One-time cost of pulling the deploy image to this offer's host.
|
||||
pub fn pull_cost(&self, o: &Offer) -> f64 {
|
||||
self.image_gb
|
||||
.map_or(0.0, |gb| gb * o.inet_down_cost_per_tb / 1000.0)
|
||||
}
|
||||
|
||||
/// Effective price used for ranking.
|
||||
pub fn effective_price(&self, o: &Offer) -> f64 {
|
||||
o.dph_total + self.pull_cost(o)
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop the suspiciously-cheap tail within each GPU model, then merge by price.
|
||||
pub(crate) fn rank_survivors(offers: Vec<Offer>, cost: &CostModel, drop_frac: f64) -> Vec<Offer> {
|
||||
let mut by_model: std::collections::HashMap<String, Vec<Offer>> =
|
||||
std::collections::HashMap::new();
|
||||
for o in offers {
|
||||
by_model.entry(o.gpu_name.clone()).or_default().push(o);
|
||||
}
|
||||
let price = |o: &Offer| cost.effective_price(o);
|
||||
let by_price = |a: &Offer, b: &Offer| {
|
||||
price(a)
|
||||
.partial_cmp(&price(b))
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
};
|
||||
let mut survivors = Vec::new();
|
||||
for (_model, mut group) in by_model {
|
||||
group.sort_by(&by_price);
|
||||
let drop = (drop_frac * group.len() as f64).floor() as usize;
|
||||
survivors.extend(group.into_iter().skip(drop));
|
||||
}
|
||||
survivors.sort_by(&by_price);
|
||||
survivors
|
||||
}
|
||||
|
||||
pub(crate) fn plan_picks(pool: &[Offer], num_instances: u32) -> Vec<&Offer> {
|
||||
let mut picks = Vec::with_capacity(num_instances as usize);
|
||||
let mut used = std::collections::HashSet::new();
|
||||
for o in pool {
|
||||
if picks.len() == num_instances as usize {
|
||||
break;
|
||||
}
|
||||
if let Some(h) = o.host_id {
|
||||
if !used.insert(h) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
picks.push(o);
|
||||
}
|
||||
picks
|
||||
}
|
||||
52
tools/vastai/src/provision.rs
Normal file
52
tools/vastai/src/provision.rs
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
use serde_json::Value;
|
||||
|
||||
use crate::types::{CreateInstanceRequest, CreateResponse, InstanceInfo};
|
||||
|
||||
/// Create one vast.ai instance from a fully generic payload.
|
||||
pub async fn create_instance(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
req: &CreateInstanceRequest,
|
||||
) -> Result<InstanceInfo, String> {
|
||||
let url = format!("{base_url}/api/v0/asks/{}/", req.offer_id);
|
||||
let env = req
|
||||
.env
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), Value::String(v.clone())))
|
||||
.collect::<serde_json::Map<String, Value>>();
|
||||
|
||||
let mut body = serde_json::json!({
|
||||
"image": req.image,
|
||||
"env": env,
|
||||
"disk": req.disk_gb,
|
||||
});
|
||||
if let Some(onstart) = req.onstart.as_deref() {
|
||||
body["onstart"] = Value::String(onstart.to_string());
|
||||
}
|
||||
if let Some(label) = req.label.as_deref() {
|
||||
body["label"] = Value::String(label.to_string());
|
||||
}
|
||||
|
||||
let resp = client
|
||||
.put(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("create_instance request failed: {e}"))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("create_instance HTTP {status}: {body}"));
|
||||
}
|
||||
|
||||
let parsed: CreateResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("create_instance parse failed: {e}"))?;
|
||||
Ok(InstanceInfo {
|
||||
contract_id: parsed.new_contract,
|
||||
})
|
||||
}
|
||||
94
tools/vastai/src/search.rs
Normal file
94
tools/vastai/src/search.rs
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
use crate::filters::reachable_offers;
|
||||
use crate::pricing::{CostModel, rank_survivors};
|
||||
use crate::types::{Offer, SearchResponse, SelectionPolicy};
|
||||
|
||||
/// Historical env-backed offer search wrapper.
|
||||
pub async fn select_offer_pool(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
gpu_name: &str,
|
||||
target_count: u32,
|
||||
) -> Result<Vec<Offer>, String> {
|
||||
let mut policy = SelectionPolicy::from_env();
|
||||
if !gpu_name.is_empty() {
|
||||
policy.gpu_name = Some(gpu_name.to_string());
|
||||
}
|
||||
select_offer_pool_with_policy(client, base_url, api_key, &policy, target_count).await
|
||||
}
|
||||
|
||||
/// Search vast.ai offers, apply quality filters, and rank survivors.
|
||||
pub async fn select_offer_pool_with_policy(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
policy: &SelectionPolicy,
|
||||
target_count: u32,
|
||||
) -> Result<Vec<Offer>, String> {
|
||||
// Hard gates expressed server-side. Network speed cannot be probed before
|
||||
// renting, so this trusts vast.ai's measured inet figures.
|
||||
let mut query = serde_json::json!({
|
||||
"rentable": {"eq": true},
|
||||
"rented": {"eq": false},
|
||||
"reliability2": {"gte": policy.min_reliability},
|
||||
"cuda_max_good": {"gte": 12.6},
|
||||
"direct_port_count": {"gte": 1},
|
||||
"num_gpus": {"eq": 1},
|
||||
"inet_down": {"gte": policy.min_down_mbps},
|
||||
// vast.ai treats `limit` as a scan budget, not a simple result cap.
|
||||
"limit": 5000,
|
||||
});
|
||||
if let Some(up) = policy.min_up_mbps {
|
||||
query["inet_up"] = serde_json::json!({"gte": up});
|
||||
}
|
||||
if policy.require_verified {
|
||||
query["verified"] = serde_json::json!({"eq": true});
|
||||
}
|
||||
if let Some(min_ram) = policy.min_gpu_ram_mb {
|
||||
query["gpu_ram"] = serde_json::json!({"gte": min_ram});
|
||||
}
|
||||
if let Some(gpu_name) = policy.gpu_name.as_deref().filter(|s| !s.is_empty()) {
|
||||
query["gpu_name"] = serde_json::json!({"eq": gpu_name});
|
||||
}
|
||||
|
||||
let url = format!(
|
||||
"{base_url}/api/v0/bundles/?q={}",
|
||||
urlencoding::encode(&query.to_string())
|
||||
);
|
||||
let resp = client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("select_offer_pool request failed: {e}"))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("select_offer_pool HTTP {status}: {body}"));
|
||||
}
|
||||
|
||||
let body: SearchResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("select_offer_pool parse failed: {e}"))?;
|
||||
|
||||
let reachable = reachable_offers(body.offers, policy);
|
||||
let cost = CostModel::from_policy(policy);
|
||||
let pool = rank_survivors(reachable, &cost, policy.drop_cheap_frac);
|
||||
|
||||
if pool.is_empty() {
|
||||
return Err(
|
||||
"no offers available (after quality/geo/host-blacklist filters and cheap-tail drop)"
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
eprintln!(
|
||||
"select_offer_pool: {} survivor(s) for {target_count} instance(s) after \
|
||||
per-model {:.0}% cheap-drop (cheapest ${:.3}/hr eff)",
|
||||
pool.len(),
|
||||
policy.drop_cheap_frac * 100.0,
|
||||
cost.effective_price(&pool[0]),
|
||||
);
|
||||
Ok(pool)
|
||||
}
|
||||
19
tools/vastai/src/state.rs
Normal file
19
tools/vastai/src/state.rs
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
use std::path::Path;
|
||||
|
||||
use crate::types::FleetState;
|
||||
|
||||
impl FleetState {
|
||||
pub fn load(path: &Path) -> Result<Self, String> {
|
||||
let raw = std::fs::read_to_string(path)
|
||||
.map_err(|e| format!("cannot read fleet state {}: {e}", path.display()))?;
|
||||
serde_json::from_str(&raw)
|
||||
.map_err(|e| format!("cannot parse fleet state {}: {e}", path.display()))
|
||||
}
|
||||
|
||||
pub fn save(&self, path: &Path) -> Result<(), String> {
|
||||
let raw = serde_json::to_string_pretty(self)
|
||||
.map_err(|e| format!("cannot serialize fleet state: {e}"))?;
|
||||
std::fs::write(path, raw)
|
||||
.map_err(|e| format!("cannot write fleet state {}: {e}", path.display()))
|
||||
}
|
||||
}
|
||||
115
tools/vastai/src/teardown.rs
Normal file
115
tools/vastai/src/teardown.rs
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use crate::types::{InstanceInfo, InstanceListResponse, LabeledInstance};
|
||||
|
||||
/// Destroy one vast.ai instance by contract id.
|
||||
pub async fn destroy_instance(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
contract_id: u64,
|
||||
) -> Result<(), String> {
|
||||
let url = format!("{base_url}/api/v0/instances/{contract_id}/");
|
||||
let resp = client
|
||||
.delete(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("destroy_instance request failed: {e}"))?;
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(format!(
|
||||
"destroy_instance {contract_id} HTTP {status}: {body}"
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Destroy every contract and return per-id results in the same order.
|
||||
pub async fn destroy_all_instances(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
contract_ids: &[u64],
|
||||
) -> Vec<Result<(), String>> {
|
||||
let mut results = Vec::with_capacity(contract_ids.len());
|
||||
for &id in contract_ids {
|
||||
results.push(destroy_instance(client, base_url, api_key, id).await);
|
||||
}
|
||||
results
|
||||
}
|
||||
|
||||
/// Destroy one contract, retrying transient failures so rollback does not strand billing instances.
|
||||
pub async fn destroy_instance_with_retry(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
contract_id: u64,
|
||||
) -> Result<(), String> {
|
||||
let mut attempt = 1_u64;
|
||||
loop {
|
||||
match destroy_instance(client, base_url, api_key, contract_id).await {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(_) => {
|
||||
let backoff = std::cmp::min(
|
||||
Duration::from_millis(500_u64.saturating_mul(attempt)),
|
||||
Duration::from_secs(30),
|
||||
);
|
||||
tokio::time::sleep(backoff).await;
|
||||
attempt = attempt.saturating_add(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn rollback(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
created: &[InstanceInfo],
|
||||
) {
|
||||
for info in created {
|
||||
if let Err(e) =
|
||||
destroy_instance_with_retry(client, base_url, api_key, info.contract_id).await
|
||||
{
|
||||
eprintln!(
|
||||
"lease_chain: WARNING rollback could not destroy {}: {e}",
|
||||
info.contract_id
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// List every instance on the account tagged with `label`, sorted by contract id.
|
||||
pub async fn list_instances_by_label(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
label: &str,
|
||||
) -> Result<Vec<LabeledInstance>, String> {
|
||||
let url = format!("{base_url}/api/v0/instances/");
|
||||
let resp = client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("list_instances request failed: {e}"))?;
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(format!("list_instances HTTP {status}: {body}"));
|
||||
}
|
||||
let body: InstanceListResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("list_instances parse failed: {e}"))?;
|
||||
let mut out: Vec<LabeledInstance> = body
|
||||
.instances
|
||||
.into_iter()
|
||||
.filter(|e| e.label.as_deref() == Some(label))
|
||||
.map(Into::into)
|
||||
.collect();
|
||||
out.sort_by_key(|i| i.contract_id);
|
||||
Ok(out)
|
||||
}
|
||||
226
tools/vastai/src/types.rs
Normal file
226
tools/vastai/src/types.rs
Normal file
|
|
@ -0,0 +1,226 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Identifier for a rented vast.ai instance.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct InstanceInfo {
|
||||
pub contract_id: u64,
|
||||
}
|
||||
|
||||
/// A vast.ai offer (GPU rental option) returned by the bundle search endpoint.
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct Offer {
|
||||
pub id: u64,
|
||||
pub gpu_name: String,
|
||||
pub dph_total: f64,
|
||||
#[serde(default)]
|
||||
pub gpu_ram: Option<f64>,
|
||||
#[serde(default)]
|
||||
pub geolocation: Option<String>,
|
||||
/// Inbound bandwidth price ($/TB). vast.ai bills Docker image pulls here.
|
||||
#[serde(default, rename = "internet_down_cost_per_tb")]
|
||||
pub inet_down_cost_per_tb: f64,
|
||||
/// Outbound bandwidth price ($/TB), surfaced so callers can account for it.
|
||||
#[serde(default, rename = "internet_up_cost_per_tb")]
|
||||
pub inet_up_cost_per_tb: f64,
|
||||
/// Marketplace host that owns the machine. Used for host-level blacklist and
|
||||
/// distinct-host provisioning.
|
||||
#[serde(default)]
|
||||
pub host_id: Option<u64>,
|
||||
/// vast.ai host verification state: `verified`, `unverified`, or
|
||||
/// `deverified`.
|
||||
#[serde(default)]
|
||||
pub verification: Option<String>,
|
||||
}
|
||||
|
||||
/// Connection details for a running instance.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RunningInstance {
|
||||
pub ip: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
/// SSH endpoint + identity of a held instance, discovered by label.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LabeledInstance {
|
||||
pub contract_id: u64,
|
||||
/// vast.ai SSH proxy host (e.g. `ssh5.vast.ai`); empty if not yet assigned.
|
||||
pub ssh_host: String,
|
||||
pub ssh_port: u16,
|
||||
pub public_ipaddr: String,
|
||||
pub actual_status: String,
|
||||
}
|
||||
|
||||
/// Offer selection knobs. Apps choose policy; this crate applies it.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SelectionPolicy {
|
||||
pub gpu_name: Option<String>,
|
||||
pub min_gpu_ram_mb: Option<u64>,
|
||||
pub min_reliability: f64,
|
||||
pub require_verified: bool,
|
||||
pub min_down_mbps: f64,
|
||||
pub min_up_mbps: Option<f64>,
|
||||
pub blacklist_hosts: Vec<u64>,
|
||||
pub drop_cheap_frac: f64,
|
||||
pub image_size_gb: Option<f64>,
|
||||
}
|
||||
|
||||
impl Default for SelectionPolicy {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
gpu_name: None,
|
||||
min_gpu_ram_mb: None,
|
||||
min_reliability: 0.95,
|
||||
require_verified: false,
|
||||
min_down_mbps: 100.0,
|
||||
min_up_mbps: None,
|
||||
blacklist_hosts: vec![59017],
|
||||
drop_cheap_frac: 0.30,
|
||||
image_size_gb: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Provisioning and monitoring retry/timing knobs.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LifecyclePolicy {
|
||||
pub lease_pace: Duration,
|
||||
pub poll_interval: Duration,
|
||||
}
|
||||
|
||||
impl Default for LifecyclePolicy {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
lease_pace: Duration::from_millis(600),
|
||||
poll_interval: Duration::from_secs(10),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Generic vast.ai create-instance payload.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CreateInstanceRequest {
|
||||
pub offer_id: u64,
|
||||
pub image: String,
|
||||
pub disk_gb: u32,
|
||||
pub label: Option<String>,
|
||||
pub env: BTreeMap<String, String>,
|
||||
/// Command passed to vast.ai's `onstart`; `None` omits the field.
|
||||
pub onstart: Option<String>,
|
||||
}
|
||||
|
||||
/// Generic N-instance provision request. Apps decide what env each machine runs.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProvisionRequest {
|
||||
pub count: u32,
|
||||
pub image: String,
|
||||
pub label: Option<String>,
|
||||
pub disk_gb: u32,
|
||||
/// Env applied to every instance.
|
||||
pub env: BTreeMap<String, String>,
|
||||
/// Per-index env overlays, merged after `env`.
|
||||
pub per_instance_env: Vec<BTreeMap<String, String>>,
|
||||
pub onstart: Option<String>,
|
||||
pub selection: SelectionPolicy,
|
||||
pub lifecycle: LifecyclePolicy,
|
||||
/// Whether to print the lease plan and ask on TTY before spending money.
|
||||
pub confirm_lease: bool,
|
||||
}
|
||||
|
||||
/// A successfully-provisioned instance with the offer facts used to pick it.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ProvisionedInstance {
|
||||
pub index: u32,
|
||||
pub contract_id: u64,
|
||||
pub offer_id: u64,
|
||||
pub host_id: Option<u64>,
|
||||
pub gpu_name: String,
|
||||
pub gpu_ram: Option<f64>,
|
||||
pub dph_total: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ProvisionedFleet {
|
||||
pub label: Option<String>,
|
||||
pub instances: Vec<ProvisionedInstance>,
|
||||
}
|
||||
|
||||
/// Generic held-fleet handle file.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct FleetState {
|
||||
pub label: String,
|
||||
pub image: String,
|
||||
pub contracts: Vec<ContractRef>,
|
||||
pub created_at: u64,
|
||||
#[serde(default)]
|
||||
pub metadata: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct ContractRef {
|
||||
pub id: u64,
|
||||
pub index: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct SearchResponse {
|
||||
pub offers: Vec<Offer>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct CreateResponse {
|
||||
pub new_contract: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct InstanceResponse {
|
||||
pub instances: InstanceStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct InstanceStatus {
|
||||
pub actual_status: Option<String>,
|
||||
pub intended_status: Option<String>,
|
||||
#[serde(default)]
|
||||
pub status_msg: Option<String>,
|
||||
#[serde(default)]
|
||||
pub public_ipaddr: Option<String>,
|
||||
#[serde(default)]
|
||||
pub ssh_port: Option<u16>,
|
||||
#[serde(default)]
|
||||
pub disk_usage: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct InstanceListResponse {
|
||||
pub instances: Vec<InstanceListEntry>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct InstanceListEntry {
|
||||
pub id: u64,
|
||||
#[serde(default)]
|
||||
pub label: Option<String>,
|
||||
#[serde(default)]
|
||||
pub actual_status: Option<String>,
|
||||
#[serde(default)]
|
||||
pub ssh_host: Option<String>,
|
||||
#[serde(default)]
|
||||
pub ssh_port: Option<u16>,
|
||||
#[serde(default)]
|
||||
pub public_ipaddr: Option<String>,
|
||||
}
|
||||
|
||||
impl From<InstanceListEntry> for LabeledInstance {
|
||||
fn from(e: InstanceListEntry) -> Self {
|
||||
Self {
|
||||
contract_id: e.id,
|
||||
ssh_host: e.ssh_host.unwrap_or_default(),
|
||||
ssh_port: e.ssh_port.unwrap_or(0),
|
||||
public_ipaddr: e.public_ipaddr.unwrap_or_default(),
|
||||
actual_status: e.actual_status.unwrap_or_else(|| "unknown".to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue