launcher/src/cli.rs

use crate::runlock::RunlockCommands;
pub(crate) mod hf;
use clap::{Parser, Subcommand, ValueEnum};
pub use hf::HfCommands;

pub use crate::cli_sub::{
    ConfigCommands, PowerAction, PublishArgs, RegistryCommands, ReposArgs, ReposCommands,
    SyncImageArgs, TmuxCommands,
};

/// Output format for human-readable reports (`--format`).
#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
pub enum OutputFormat {
    Plain,
    Md,
    Html,
}

#[derive(Parser)]
#[command(author, version, about = "DGX Spark Launcher for vLLM and SGLang")]
pub struct Cli {
    #[arg(long, short, default_value = "config.toml")]
    pub config: String,

    #[arg(long, short, help = "Show commands without executing")]
    pub dry_run: bool,

    #[command(subcommand)]
    pub command: Commands,
}

#[derive(Subcommand)]
pub enum Commands {
    /// Distribute a recipe's image across the cluster nodes (node-to-node,
    /// one save + parallel fan-out; no serve started).
    SyncImage(SyncImageArgs),

    /// Launch a cluster
    Serve {
        #[arg(long, required = true)]
        recipe: String,

        #[arg(long, default_value = "auto")]
        cluster: String,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,

        /// Keep containers after stop (default: auto-remove on stop)
        #[arg(long)]
        keep: bool,

        /// Skip pre-serve hooks (cluster-level + recipe-level).
        #[arg(long)]
        skip_before_serve: bool,

        /// Don't follow logs after serving (default: follow).
        #[arg(long, action = clap::ArgAction::SetTrue)]
        no_follow: bool,

        /// Skip the engine-readiness wait after launch (no /v1/models poll).
        #[arg(long)]
        no_wait: bool,

        /// Power on cluster if switch is off (requires HA config).
        #[arg(short = 'O', long = "on")]
        auto_power_on: bool,

        /// Sync Docker image across nodes before starting.
        #[arg(long, default_value_t = true)]
        sync_image: bool,

        /// Explicit source node for image sync (default: auto-detect).
        #[arg(long)]
        sync_image_source: Option<String>,

        /// Build image on this node if not found (requires docker-images/<recipe>/).
        #[arg(long)]
        build_node: Option<String>,
        /// Force rebuild even if checksum matches.
        #[arg(long)]
        force_build: bool,

        /// Run even on occupied nodes (dangerous for GPU recipes: OOM,
        /// port conflicts). Skips the occupancy check in cluster selection.
        #[arg(long)]
        force: bool,

        /// Serve single-node on the first free node (overrides recipe `solo`);
        /// `--no-solo` forces multi-node.
        #[arg(long)]
        solo: Option<bool>,

        /// Mount a per-node local HF mirror (bare = ~/.cache/hf-models;
        /// `--mirror <path>` = that path; first serve copies per node).
        /// Mutually exclusive with `--hf-cache` (both pick the HF mount source).
        #[arg(long, num_args = 0..=1, default_missing_value = "")]
        mirror: Option<String>,

        /// Override the host-side HF cache dir for this serve (config
        /// `hf_cache_dir`, else `~/.cache/huggingface`). `~` expands to the
        /// node's SSH-user home. Mutually exclusive with `--mirror`.
        #[arg(long)]
        hf_cache: Option<String>,

        /// How hard to run the recipe's `[[staging]]` blocks: `auto`
        /// (ledger-verified, copy what is new), `skip` (assume staged),
        /// `force` (re-verify everything).
        #[arg(long)]
        stage: Option<String>,

        /// Scale the recipe to N nodes (`--fit=N`; bare = whole cluster, TP
        /// tracks it; incompatible with --solo).
        #[arg(long, num_args = 0..=1, default_missing_value = "")]
        fit: Option<String>,

        /// Apply the recipe's `cold_args` over its `args` (overlay by flag
        /// name) — the first-launch / tight-node safety valve for memory
        /// knobs (GMU, KV/context limits). No-op if the recipe has no
        /// `cold_args`.
        #[arg(long)]
        cold: bool,

        /// `sudo rm -rf` the host-side extra_mounts folders on the launch
        /// nodes before containers start (repo-based mounts untouched).
        #[arg(long)]
        clear: bool,

        /// Skip the post-build registry push (only meaningful when `[registry]`
        /// is configured and the image needs (re)building).
        #[arg(long)]
        no_push: bool,

        /// Where the image comes from: build from `docker-images/` (default)
        /// or pull the prebuilt `{base}/{name}:{checksum}` from `[registry]`
        /// (skips the build; the image must already be pushed).
        #[arg(long, value_enum, default_value_t = crate::docker::ImageSource::Repo)]
        image_source: crate::docker::ImageSource,
    },

    /// Show container logs
    Logs {
        #[arg(long, required = true)]
        recipe: String,

        #[arg(long, action = clap::ArgAction::SetTrue)]
        follow: bool,

        #[arg(long, default_value = "auto")]
        cluster: String,

        /// Target a specific node within the cluster (default: head node if following, all nodes if not).
        #[arg(long)]
        node: Option<String>,
    },

    /// Stop a cluster
    Stop {
        #[arg(long, required = true)]
        recipe: String,

        /// Target specific nodes (repeatable, e.g. --node dgx1 --node dgx2).
        #[arg(long)]
        node: Vec<String>,

        #[arg(long, default_value = "auto")]
        cluster: String,

        /// Power off cluster after stopping containers (requires HA config).
        #[arg(long = "off")]
        power_off: bool,

        /// Don't wait until the containers are actually stopped and removed
        /// (default: wait and verify, up to 60 s per node, so an immediate
        /// re-serve on the same nodes doesn't race the dying containers).
        #[arg(long)]
        no_wait: bool,
    },

    /// Bench a live serve with tool-eval-bench: reports auto-land in
    /// recipes/<lane>/reports/<UTC-date>_<label>.{md,json,log} (seed 42, hardmode, perf).
    Bench {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, required = true)]
        cluster: String,
        /// Bench target node within the cluster (default: rank 0). Use when
        /// the serve was launched with `serve --node <n>` — a pinned/solo
        /// run may live on a non-rank-0 node.
        #[arg(long)]
        node: Option<String>,
        /// Label for the run (campaign record). Optional for `--perf-only`
        /// runs (default `perf`); required for quality runs.
        #[arg(long)]
        label: Option<String>,
        #[arg(long, default_value = "1,2")]
        concurrency: String,
        #[arg(long, default_value_t = 42)]
        seed: u64,
        /// Run scenario evaluations over N concurrent streams (tool
        /// `--parallel`). Speeds up the scenario phase ~Nx at the cost of
        /// mild engine contention; perf numbers are measured separately.
        #[arg(long)]
        parallel: Option<u32>,
        /// Local tokenizer for llama-benchy (tokenizer.json or a dir with
        /// one) — needed when the served id isn't resolvable offline.
        #[arg(long)]
        tokenizer: Option<String>,
        /// Skip the scenario phase — run only the perf micro-benchmarks
        /// (use to re-baseline the perf table after a cold first row).
        #[arg(long)]
        perf_only: bool,
        /// Run the context-depth suite (16K..256K) after the main pass — off
        /// by default (a full sweep takes ~45 min on a 2x lane).
        #[arg(long)]
        depth: bool,
        /// Run tool-eval-bench on a remote host over SSH instead of locally,
        /// so the bench client's own network (a WiFi laptop) stays off the
        /// measurement path. The tokenizer is uploaded to the host; the
        /// report + run package still land locally. The host must have
        /// `tool-eval-bench` on PATH.
        #[arg(long)]
        ssh_host: Option<String>,
        /// Disable thinking/reasoning in the harness (`--no-think`, sets
        /// `enable_thinking=false`). Needed for thinking-by-default models
        /// (Ling, GLM-4.5-style) whose long default reasoning blows the
        /// per-turn timeout and excludes most scenarios.
        #[arg(long)]
        no_think: bool,
        /// Per-turn request timeout seconds forwarded to the harness
        /// (`--timeout`). Raise above the 120 s default for slow models.
        #[arg(long)]
        timeout: Option<u64>,
    },

    /// Check cluster health
    Health {
        #[arg(
            long,
            help = "Cluster name (e.g. 4nodes-switch). If omitted — checks all."
        )]
        cluster: Option<String>,
        #[arg(long, help = "Show node host keys (for adding to config.toml)")]
        fingerprint: bool,
        /// Output format (cluster-health mode only; ignored with --fingerprint).
        #[arg(long, value_enum, default_value_t = OutputFormat::Plain)]
        format: OutputFormat,
    },

    /// Check and repair a recipe (doctor mode)
    Doctor {
        #[arg(long, required = true, help = "Recipe name to check")]
        recipe: String,
        #[arg(
            long,
            default_value = "auto",
            help = "Cluster (if not auto — uses recipe.cluster)"
        )]
        cluster: String,
        #[arg(long, help = "Auto-fix issues (create Dockerfile + build image)")]
        repair: bool,
    },

    /// Publish a recipe (+ image overlay) as a standalone reproducible bundle
    Publish(PublishArgs),

    /// Browse and install recipes from the public registry
    #[command(subcommand)]
    Registry(RegistryCommands),

    /// Manage tmux sessions (Docker alternative)
    #[command(subcommand)]
    Tmux(TmuxCommands),

    /// Auto-discover host keys and write them into config.toml
    InitHostKeys {
        #[arg(long, help = "Cluster name (e.g. 4nodes-switch). If omitted — all.")]
        cluster: Option<String>,
    },

    /// Run an arbitrary shell command on a cluster or single node
    Exec {
        /// Cluster name (mutually exclusive with --recipe).
        #[arg(long, conflicts_with = "recipe")]
        cluster: Option<String>,

        /// Recipe name; resolves to the recipe's pinned cluster.
        #[arg(long, conflicts_with = "cluster")]
        recipe: Option<String>,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,

        /// Name of an alias from `[exec_aliases]` in config.toml.
        /// Mutually exclusive with a trailing command.
        #[arg(long)]
        alias: Option<String>,

        /// Command to run (everything after `--`).
        #[arg(trailing_var_arg = true)]
        command: Vec<String>,
    },

    /// Manage configuration (global config path, init)
    #[command(subcommand)]
    Config(ConfigCommands),

    /// Run-ledger management (per-recipe .lock; offline)
    #[command(subcommand)]
    Runlock(RunlockCommands),

    /// Hugging Face catalogue: search models, read a card, list revisions
    Hf {
        #[command(subcommand)]
        cmd: HfCommands,
    },

    /// Shared Hugging Face cache on the NAS: `status` / `doctor` (read-only
    /// here; update/prune are separate subcommands).
    Cache {
        #[command(subcommand)]
        cmd: CacheCommands,
    },

    /// Manage the configured git recipe repos (clone/pull cache, archive
    /// registry). `sm121 repos` lists state; `sm121 repos sync [--force]`
    /// re-pulls.
    Repos(ReposArgs),

    /// Archive a recipe repo (it is skipped by sync; state in repos.json)
    Archive {
        /// The repo URL (a `[[recipe_repos]]` entry)
        url: String,
    },

    /// Un-archive a recipe repo (it takes part in sync again)
    Unarchive {
        /// The repo URL (a `[[recipe_repos]]` entry)
        url: String,
    },

    /// Move a recipe across recipe repos (git mv + commit + push; the
    /// target is updated first, so a failure duplicates, never deletes).
    /// Default (or `-d`) prints the exact commands; `--apply` executes.
    Promote {
        /// Recipe name (its directory under the repo's `recipes/`)
        recipe: String,

        /// Source repo URL (a `[[recipe_repos]]` entry)
        #[arg(long, required = true)]
        from: String,

        /// Target repo URL (a `[[recipe_repos]]` entry)
        #[arg(long, required = true)]
        to: String,

        /// Execute the move (default: print the plan only)
        #[arg(long)]
        apply: bool,
    },

    /// Launch the web dashboard (axum + SSE) on a local port
    Web {
        #[arg(long, default_value_t = 3100)]
        port: u16,

        #[arg(long, default_value = "127.0.0.1")]
        bind: String,

        /// Bearer token for data endpoints (required when --bind is not loopback,
        /// unless --insecure)
        #[arg(long)]
        token: Option<String>,

        /// Serve WITHOUT authentication (the token check is skipped entirely —
        /// anyone who can reach the port can serve/stop/power the cluster). For
        /// isolated LANs where the operator explicitly accepts that.
        #[arg(long)]
        insecure: bool,

        /// Open the browser once the server is up
        #[arg(long)]
        open: bool,

        /// Re-sync the configured `[[recipe_repos]]` in the background every
        /// N minutes so the dashboard serves fresh recipes without a restart
        /// (0 disables; default 10). The HA add-on relies on this instead of
        /// its own entrypoint loop.
        #[arg(long, default_value_t = 10)]
        repo_sync_interval: u32,
    },

    /// Table of all tested recipes and their benchmark history
    Report {
        /// Interactive HTML table via --serve (default: print Markdown)
        #[arg(long)]
        serve: bool,
        /// Port for --serve
        #[arg(long, default_value_t = 3101)]
        port: u16,
        /// Open the browser (implies --serve)
        #[arg(long)]
        open: bool,
        #[arg(long, help = "Archive a recipe (hidden from report output)")]
        archive: Option<String>,
        #[arg(long, help = "Un-archive a recipe")]
        unarchive: Option<String>,
        #[arg(long, help = "Show the archived recipes instead of the active ones")]
        archived: bool,
    },

    /// Build Docker image for a recipe on a remote node
    Build {
        #[arg(long, required = true, help = "Recipe name")]
        recipe: String,

        #[arg(long, required = true, help = "Node to build on (e.g. dgx1.lan)")]
        build_node: String,

        #[arg(long, default_value = "auto", help = "Cluster name")]
        cluster: String,

        /// Force rebuild even if checksum matches.
        #[arg(long)]
        force_build: bool,

        /// Skip the post-build registry push (only meaningful when `[registry]`
        /// is configured).
        #[arg(long)]
        no_push: bool,

        /// Where the image comes from: build from `docker-images/` (default)
        /// or pull the prebuilt tag from `[registry]` (no build, no push).
        #[arg(long, value_enum, default_value_t = crate::docker::ImageSource::Repo)]
        image_source: crate::docker::ImageSource,
    },

    /// Power on/off nodes via Home Assistant (smart plugs / Wake-on-LAN).
    Power {
        #[command(subcommand)]
        action: PowerAction,
    },

    /// Validate recipes against the launcher contract (build dirs exist,
    /// port convention, no floating :latest, floating model revisions,
    /// node counts). Errors fail the run (exit 1); archived lanes are
    /// grandfathered (their errors downgrade to warnings). Also runs as a
    /// `cargo test` gate over the repo's own `recipes/`.
    Validate,
}

#[cfg(test)]
#[path = "cli_tests.rs"]
mod tests;

/// `sm121 cache …`
#[derive(Subcommand)]
pub enum CacheCommands {
    /// What is cached and how much room is left (read-only)
    Status {
        /// Cluster whose rank-0 node inspects the cache
        #[arg(long)]
        cluster: Option<String>,
        /// Machine-readable output (for the dashboard)
        #[arg(long)]
        json: bool,
    },
    /// Refresh every cached repo to its refs (dry-run unless --apply)
    Update {
        #[arg(long)]
        cluster: Option<String>,
        #[arg(long)]
        json: bool,
        /// Only these repos (repeatable, org/name)
        #[arg(long = "repo")]
        repos: Vec<String>,
        /// Really download (default: preview)
        #[arg(long)]
        apply: bool,
    },
    /// Is the tooling the maintenance job needs present on that node?
    Doctor {
        #[arg(long)]
        cluster: Option<String>,
        #[arg(long)]
        json: bool,
    },
}