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

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