from __future__ import annotations import argparse import json import os import signal 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_hf_upload_smoke, run_smoke_training, run_wandb_smoke 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") hf_smoke = subparsers.add_parser("hf-smoke", help="upload tiny smoke artifacts to Hugging Face") hf_smoke.add_argument("--artifact-dir", required=True) hf_smoke.add_argument("--run-id", required=True) hf_smoke.add_argument("--repo-id", help="HF repo id or repo slug; default is /airfrans-hf-smoke") hf_smoke.set_defaults(command="hf-smoke") wandb_smoke = subparsers.add_parser("wandb-smoke", help="log tiny smoke metrics to Weights & Biases") wandb_smoke.add_argument("--artifact-dir", required=True) wandb_smoke.add_argument("--run-id", required=True) wandb_smoke.add_argument("--entity", default="zacheryasc-personal") wandb_smoke.add_argument("--project", default="airfRANS-model-sweep") wandb_smoke.add_argument("--hf-repo-url") wandb_smoke.set_defaults(command="wandb-smoke") 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 == "hf-smoke": output = run_hf_upload_smoke( artifact_dir=args.artifact_dir, run_id=args.run_id, repo_id=args.repo_id, ) print(f"artifact_dir: {output}") return 0 if args.command == "wandb-smoke": output = run_wandb_smoke( artifact_dir=args.artifact_dir, run_id=args.run_id, entity=args.entity, project=args.project, hf_repo_url=args.hf_repo_url, ) print(f"artifact_dir: {output}") 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("REMOTE_FINISHED", selected_offer_id=selection.selected_offer_id, attempt=attempt, return_code=return_code) _collect_terminal_best_effort(cluster=run_id, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, required=config.artifacts.required, env=env) status = _classify_artifacts(local_run_dir) if status == "success": try: state("COLLECTING_REQUIRED", selected_offer_id=selection.selected_offer_id, attempt=attempt) _collect_required_artifacts(cluster=run_id, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, required=config.artifacts.required, env=env) state("VERIFYING_ARTIFACTS", selected_offer_id=selection.selected_offer_id, attempt=attempt) verify_artifacts(local_run_dir, required=config.artifacts.required) except Exception as exc: _cleanup_partial_artifacts(local_run_dir) last_error = f"remote job succeeded but artifact collection failed: {exc}" state("FAILED_COLLECTION", selected_offer_id=selection.selected_offer_id, attempt=attempt, return_code=return_code, 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) from exc 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, return_code=return_code, 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) if return_code == 0: last_error = "remote job succeeded but terminal artifacts were not collected" state("FAILED_COLLECTION", selected_offer_id=selection.selected_offer_id, attempt=attempt, return_code=return_code, 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) _collect_restart_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 == "failure_report": last_error = "training wrote failure_report.json" state("FAILED_TRAINING", selected_offer_id=selection.selected_offer_id, attempt=attempt, return_code=return_code, 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"remote job failed without terminal success, return_code={return_code}, local_artifacts={status}" if attempt < config.run.max_attempts: state("RETRYING", selected_offer_id=selection.selected_offer_id, attempt=attempt, return_code=return_code, error=last_error) if config.cleanup.on_failure == "sky_down" and not skip_down: _run_best_effort(["sky", "down", run_id, "-y"], env=env) else: state("FAILED", selected_offer_id=selection.selected_offer_id, attempt=attempt, return_code=return_code, error=last_error) 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: argv = ["sky", "launch", "-c", cluster, str(sky_yaml_path), "-y"] if config.artifacts.mode == "object_store_upload": _ensure_hf_secret_env(env) argv.extend(["--secret", "HF_TOKEN"]) if _load_secret_env(env, "WANDB_API_KEY", required=False, purpose="W&B remote runs"): argv.extend(["--secret", "WANDB_API_KEY"]) process = subprocess.Popen(argv, env=env, start_new_session=True) 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: _terminate_process_group(process) return 124 if config.run.artifact_sync_interval_seconds == 0 or now >= next_collect: _collect_terminal_best_effort(cluster=cluster, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, required=config.artifacts.required, env=env) next_collect = now + max(1, config.run.artifact_sync_interval_seconds) time.sleep(min(5.0, max(0.1, next_collect - now))) def _collect_paths_with_rsync( *, cluster: str, remote_dir: Path, local_dir: Path, paths: tuple[str, ...], env: dict[str, str], timeout: int, ) -> None: local_dir.mkdir(parents=True, exist_ok=True) for relative_path in paths: source = f"{cluster}:~/sky_workdir/{remote_dir}/./{relative_path}" _run_checked( [ "rsync", "-Pavz", "--relative", "--ignore-missing-args", "--delay-updates", "--timeout=30", source, f"{local_dir}/", ], env=env, timeout=timeout, ) def _collect_required_artifacts(*, cluster: str, remote_dir: Path, local_dir: Path, required: tuple[str, ...], env: dict[str, str]) -> None: _collect_paths_with_rsync( cluster=cluster, remote_dir=remote_dir, local_dir=local_dir, paths=_large_artifact_names(required), env=env, timeout=3600, ) def _collect_terminal_best_effort(*, cluster: str, remote_dir: Path, local_dir: Path, required: tuple[str, ...], env: dict[str, str]) -> None: try: _collect_paths_with_rsync( cluster=cluster, remote_dir=remote_dir, local_dir=local_dir, paths=_terminal_artifact_names(required), env=env, timeout=120, ) except Exception: _cleanup_partial_artifacts(local_dir) def _collect_restart_best_effort(*, cluster: str, remote_dir: Path, local_dir: Path, env: dict[str, str]) -> None: try: _collect_paths_with_rsync( cluster=cluster, remote_dir=remote_dir, local_dir=local_dir, paths=("checkpoint_latest.pt",), env=env, timeout=3600, ) except Exception: _cleanup_partial_artifacts(local_dir) _LARGE_ARTIFACT_SUFFIXES = (".pt", ".pth", ".ckpt", ".safetensors") _TERMINAL_ARTIFACT_NAMES = ( "artifact_manifest.json", "checksums.txt", "config.toml", "data_manifest.json", "environment_manifest.json", "failure_report.json", "final_metrics.json", "heartbeat.json", "hf_upload_manifest.json", "latest_metrics.json", "metrics.jsonl", "normalization.json", "run_manifest.json", "split_manifest.json", "wandb_smoke_manifest.json", ) def _terminal_artifact_names(required: tuple[str, ...]) -> tuple[str, ...]: names = set(_TERMINAL_ARTIFACT_NAMES) names.update(name for name in required if not _is_large_artifact(name)) return tuple(sorted(names)) def _large_artifact_names(required: tuple[str, ...]) -> tuple[str, ...]: return tuple(name for name in required if _is_large_artifact(name)) def _is_large_artifact(name: str) -> bool: return name.endswith(_LARGE_ARTIFACT_SUFFIXES) def _cleanup_partial_artifacts(local_run_dir: Path) -> None: for partial_dir in (".rsync-partial", ".~tmp~"): for path in local_run_dir.rglob(partial_dir): if path.is_dir() and not path.is_symlink(): shutil.rmtree(path, ignore_errors=True) elif path.exists(): try: path.unlink() except OSError: pass def _terminate_process_group(process: subprocess.Popen[Any]) -> None: if process.poll() is not None: return try: os.killpg(process.pid, signal.SIGTERM) except ProcessLookupError: return try: process.wait(timeout=30) return except subprocess.TimeoutExpired: pass try: os.killpg(process.pid, signal.SIGKILL) except ProcessLookupError: return process.wait() def _classify_artifacts(local_run_dir: Path) -> str: if (local_run_dir / "final_metrics.json").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: process = subprocess.Popen(argv, env=env, start_new_session=True) try: return_code = process.wait(timeout=timeout) except subprocess.TimeoutExpired: _terminate_process_group(process) raise if return_code != 0: raise subprocess.CalledProcessError(return_code, argv) 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 _ensure_hf_secret_env(env: dict[str, str]) -> None: _load_secret_env(env, "HF_TOKEN", required=True, purpose="object_store_upload runs") def _load_secret_env(env: dict[str, str], name: str, *, required: bool, purpose: str) -> bool: if env.get(name): return True for path in (Path(name), Path(".env") / name): if path.is_file(): value = path.read_text().strip() if value: env[name] = value return True env_file = Path(".env") if env_file.is_file(): for raw_line in env_file.read_text().splitlines(): line = raw_line.strip() if not line or line.startswith("#") or "=" not in line: continue key, value = line.split("=", 1) if key.strip() == name: value = value.strip().strip("\"'") if value: env[name] = value return True if required: raise RuntimeError(f"{name} env var or local secret file is required for {purpose}") return False 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())