launcher/src/cli.rs

use clap::{Args, Parser, Subcommand, ValueEnum};

/// 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 {
    /// 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).
        #[arg(long, num_args = 0..=1, default_missing_value = "")]
        mirror: 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>,

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

    /// 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,
        #[arg(long, required = true)]
        label: String,
        #[arg(long, default_value = "1,2")]
        concurrency: String,
        #[arg(long, default_value_t = 42)]
        seed: 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),

    /// 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)
        #[arg(long)]
        token: Option<String>,

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

    /// 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,
    },

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

#[derive(Subcommand)]
pub enum PowerAction {
    /// Power off: SSH `sudo shutdown now` + wait + HA `turn_off`.
    Off {
        #[arg(long)]
        cluster: String,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,
    },
    /// Power on: HA `turn_on` only (no SSH; the smart plug wakes the machine).
    On {
        #[arg(long)]
        cluster: String,

        /// Target a specific node within the cluster (default: all).
        #[arg(long)]
        node: Option<String>,
    },
    /// Switch clusters: disable nodes from source, enable nodes in target.
    Switch {
        /// Current cluster (nodes + entities to shut down).
        #[arg(long)]
        from: String,
        /// Target cluster (nodes + entities to wake up).
        #[arg(long)]
        to: String,
    },
}

/// `sm121 config` subcommands — manage the global config file.
#[derive(Subcommand)]
pub enum ConfigCommands {
    /// Show the active config path and where the global config lives
    Path,
    /// Create the global config file (template) if it does not exist
    InitGlobal,
}

#[derive(Args)]
pub struct PublishArgs {
    #[arg(long, required = true)]
    pub recipe: String,

    /// Cluster whose fabric env + node count fill the generated run script
    #[arg(long)]
    pub cluster: Option<String>,

    /// Output dir (default: ./publish/<recipe>)
    #[arg(long)]
    pub out: Option<std::path::PathBuf>,

    /// Also create a public GitHub gist (requires `gh` CLI, authenticated)
    #[arg(long)]
    pub gist: bool,

    /// Upload the bundle to a recipe registry (see `sm121 registry --help`)
    #[arg(long, value_name = "URL")]
    pub registry: Option<String>,

    /// Registry publish token (or $SM121RECIPES_TOKEN)
    #[arg(long)]
    pub token: Option<String>,
}

#[derive(Subcommand)]
pub enum RegistryCommands {
    /// List recipes on the registry
    List {
        #[arg(long, value_name = "URL")]
        registry: Option<String>,
    },
    /// Print a registry recipe (recipe.toml)
    Show {
        name: String,
        #[arg(long, value_name = "URL")]
        registry: Option<String>,
    },
    /// Download a recipe into ./recipes/<name>/ (+ ./docker-images/<name>/ if it has an image)
    Install {
        name: String,
        #[arg(long, value_name = "URL")]
        registry: Option<String>,
        /// Overwrite an existing ./recipes/<name>/
        #[arg(long)]
        force: bool,
    },
}

#[derive(Subcommand)]
pub enum TmuxCommands {
    /// Launch cluster via tmux
    Serve {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
        /// Skip pre-serve hooks (cluster-level + recipe-level).
        #[arg(long)]
        skip_before_serve: bool,
    },
    /// Stop tmux sessions
    Stop {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
    /// Detailed tmux session status
    Status {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
    /// Show tmux session logs
    Logs {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, action = clap::ArgAction::SetTrue)]
        follow: bool,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
    /// Attach to head node tmux session
    Attach {
        #[arg(long, required = true)]
        recipe: String,
        #[arg(long, default_value = "auto")]
        cluster: String,
    },
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Parse `sm121 health <args>` and return the parsed `--format`.
    fn health_format(args: &[&str]) -> Result<(OutputFormat, bool), clap::Error> {
        let cli = Cli::try_parse_from(
            std::iter::once("sm121")
                .chain(std::iter::once("health"))
                .chain(args.iter().copied()),
        )?;
        let Commands::Health {
            fingerprint,
            format,
            ..
        } = cli.command
        else {
            panic!("expected Health")
        };
        Ok((format, fingerprint))
    }

    #[test]
    fn health_format_defaults_to_plain() {
        assert_eq!(health_format(&[]).unwrap(), (OutputFormat::Plain, false));
    }

    #[test]
    fn health_format_parses_each_variant() {
        assert_eq!(
            health_format(&["--format", "plain"]).unwrap().0,
            OutputFormat::Plain
        );
        assert_eq!(
            health_format(&["--format", "md"]).unwrap().0,
            OutputFormat::Md
        );
        assert_eq!(
            health_format(&["--format", "html"]).unwrap().0,
            OutputFormat::Html
        );
    }

    #[test]
    fn health_format_combines_with_cluster_and_fingerprint() {
        let (fmt, fp) = health_format(&["--cluster", "c1", "--format", "md"]).unwrap();
        assert_eq!((fmt, fp), (OutputFormat::Md, false));
        let (fmt, fp) = health_format(&["--fingerprint", "--format", "html"]).unwrap();
        assert_eq!((fmt, fp), (OutputFormat::Html, true));
    }

    #[test]
    fn health_format_rejects_invalid_value() {
        assert!(health_format(&["--format", "bogus"]).is_err());
    }

    #[test]
    fn report_flag_defaults_and_overrides() {
        let cli = Cli::try_parse_from(["sm121", "report"]).unwrap();
        let Commands::Report {
            serve,
            port,
            open,
            archive: _,
            unarchive: _,
            archived: _,
        } = cli.command
        else {
            unreachable!()
        };
        assert!(!serve && !open);
        assert_eq!(port, 3101);
        let cli = Cli::try_parse_from(["sm121", "report", "--open", "--port", "4000"]).unwrap();
        let Commands::Report { open, port, .. } = cli.command else {
            unreachable!()
        };
        assert!(open);
        assert_eq!(port, 4000);
    }
}