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.add_argument("--resume", help="path to checkpoint_latest.pt to resume from") 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, resume_path=args.resume or os.environ.get("AIRFRANS_RESUME_CHECKPOINT"), ) 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") (local_run_dir / "config.toml").write_text(config.raw_text) write_skyignore(config) if dry_run: sky_yaml_path = _write_attempt_sky_yaml( config=config, selection=selection, run_id=run_id, local_run_dir=local_run_dir, resume_checkpoint=None, attempt=1, ) 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() last_error: str | None = None for attempt in range(1, config.run.max_attempts + 1): resume_checkpoint = _stage_resume_checkpoint(local_run_dir, run_id) sky_yaml_path = _write_attempt_sky_yaml( config=config, selection=selection, run_id=run_id, local_run_dir=local_run_dir, resume_checkpoint=resume_checkpoint, attempt=attempt, ) state( "PROVISIONING", selected_offer_id=selection.selected_offer_id, attempt=attempt, resume_checkpoint=str(resume_checkpoint) if resume_checkpoint is not None else None, ) return_code = _run_sky_with_periodic_collection( cluster=run_id, sky_yaml_path=sky_yaml_path, config=config, local_run_dir=local_run_dir, env=env, ) state("COLLECTING", selected_offer_id=selection.selected_offer_id, attempt=attempt, return_code=return_code) _collect_best_effort(cluster=run_id, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, env=env) status = _classify_artifacts(local_run_dir) if status == "success": state("VERIFYING_ARTIFACTS", selected_offer_id=selection.selected_offer_id, attempt=attempt) verify_artifacts(local_run_dir, required=config.artifacts.required) if config.cleanup.on_success == "sky_down" and not skip_down: state("CLEANING_UP", selected_offer_id=selection.selected_offer_id, attempt=attempt) _run_checked(["sky", "down", run_id, "-y"], env=env, timeout=300) state("SUCCEEDED", selected_offer_id=selection.selected_offer_id, attempt=attempt) print(f"run_id: {run_id}") print(f"artifacts: {local_run_dir}") return 0 if status == "failure_report": last_error = "training wrote failure_report.json" state("FAILED_TRAINING", selected_offer_id=selection.selected_offer_id, attempt=attempt, error=last_error) if config.cleanup.on_failure == "sky_down" and not skip_down: _run_best_effort(["sky", "down", run_id, "-y"], env=env) raise RuntimeError(last_error) last_error = f"SkyPilot job ended without terminal artifacts, return_code={return_code}" state("RETRYING", selected_offer_id=selection.selected_offer_id, attempt=attempt, error=last_error) if config.cleanup.on_failure == "sky_down" and not skip_down: _run_best_effort(["sky", "down", run_id, "-y"], env=env) state("FAILED", selected_offer_id=selection.selected_offer_id, error=last_error or "max attempts exhausted") raise RuntimeError(last_error or "max attempts exhausted") def _write_attempt_sky_yaml( *, config: RemoteRunConfig, selection: SelectionResult, run_id: str, local_run_dir: Path, resume_checkpoint: Path | None, attempt: int, ) -> Path: sky_yaml = render_skypilot_yaml( config, selection, run_id=run_id, resume_checkpoint=resume_checkpoint, ) sky_yaml_path = local_run_dir / f"sky_attempt_{attempt}.yaml" sky_yaml_path.write_text(sky_yaml) if attempt == 1: (local_run_dir / "sky.yaml").write_text(sky_yaml) return sky_yaml_path def _stage_resume_checkpoint(local_run_dir: Path, run_id: str) -> Path | None: latest = local_run_dir / "checkpoint_latest.pt" if not latest.is_file(): return None resume_dir = Path(".airfrans_resume") / run_id resume_dir.mkdir(parents=True, exist_ok=True) destination = resume_dir / "checkpoint_latest.pt" shutil.copy2(latest, destination) return destination def _run_sky_with_periodic_collection( *, cluster: str, sky_yaml_path: Path, config: RemoteRunConfig, local_run_dir: Path, env: dict[str, str], ) -> int: process = subprocess.Popen(["sky", "launch", "-c", cluster, str(sky_yaml_path), "-y"], env=env) deadline = time.monotonic() + config.run.timeout_minutes * 60 next_collect = time.monotonic() + config.run.artifact_sync_interval_seconds while True: return_code = process.poll() if return_code is not None: return int(return_code) now = time.monotonic() if now >= deadline: process.terminate() try: process.wait(timeout=30) except subprocess.TimeoutExpired: process.kill() return int(process.returncode or 124) if config.run.artifact_sync_interval_seconds == 0 or now >= next_collect: _collect_best_effort(cluster=cluster, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, env=env) if _classify_artifacts(local_run_dir) in {"success", "failure_report"}: return_code = process.poll() if return_code is not None: return int(return_code) next_collect = now + max(1, config.run.artifact_sync_interval_seconds) time.sleep(min(5.0, max(0.1, next_collect - now))) 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 _collect_best_effort(*, cluster: str, remote_dir: Path, local_dir: Path, env: dict[str, str]) -> None: try: _collect_with_rsync(cluster=cluster, remote_dir=remote_dir, local_dir=local_dir, env=env) except Exception: pass def _classify_artifacts(local_run_dir: Path) -> str: if (local_run_dir / "final_metrics.json").is_file() and (local_run_dir / "checkpoint_final.pt").is_file(): return "success" if (local_run_dir / "failure_report.json").is_file(): return "failure_report" if (local_run_dir / "checkpoint_latest.pt").is_file(): return "restartable" return "incomplete" 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())