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