airfRANS-model-exploration/src/airfrans_frontier/remote/cli.py

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())