launcher/src/tmux/mod.rs

pub mod commands;
pub mod discover;
pub mod ops;
pub mod serve;
pub mod session;

use anyhow::Result;

use crate::cli::TmuxCommands;
use crate::config::{Config, NcclConfig, Node, Recipe};

/// Manages tmux sessions for model inference across a DGX Spark cluster.
#[derive(Clone)]
pub struct TmuxManager;

impl TmuxManager {
    /// Generates a namespaced session name for a recipe and node rank.
    /// Format: `{recipe_name}-{engine}-r{rank}`
    pub fn session_name(&self, recipe: &Recipe, rank: u32) -> String {
        // Replace shell-unsafe characters (parentheses, quotes, etc.)
        // with underscores so the session name doesn't trigger glob
        // expansion or shell injection when used in tmux commands.
        let name = recipe
            .name
            .chars()
            .map(|c| {
                if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
                    c
                } else {
                    '_'
                }
            })
            .collect::<String>();
        format!("{}-{}-r{}", name, recipe.engine, rank)
    }

    /// Builds the list of `export` commands for NCCL environment variables.
    pub fn build_nccl_setup(&self, nccl: Option<&NcclConfig>) -> Vec<String> {
        match nccl {
            Some(nccl) => vec![
                format!("export NCCL_IB_HCA=\"{}\"", nccl.ib_hca),
                format!("export NCCL_SOCKET_IFNAME=\"{}\"", nccl.socket_ifname),
                format!(
                    "export NCCL_IB_GID_INDEX={}",
                    // Unresolved (dry-run) renders `auto`; real launches
                    // always have the per-node detected/pinned value.
                    nccl.ib_gid_index
                        .map_or_else(|| "auto".to_string(), |g| g.to_string())
                ),
                format!("export NCCL_IB_DISABLE={}", nccl.ib_disable),
            ],
            None => Vec::new(),
        }
    }

    /// Builds engine-specific host IP environment variable.
    pub fn build_host_ip_export(&self, engine: &str, node: &Node) -> Vec<String> {
        let ip = node.effective_ib_ip();
        match engine {
            "vllm" => vec![format!("export VLLM_HOST_IP={}", ip)],
            "sglang" => vec![format!("export SGLANG_WORKER_IP={}", ip)],
            _ => vec![
                format!("export VLLM_HOST_IP={}", ip),
                format!("export SGLANG_WORKER_IP={}", ip),
            ],
        }
    }
}

/// Detailed health status of a tmux session.
#[derive(Debug, Clone)]
pub struct SessionHealth {
    pub session_exists: bool,
    pub process_alive: bool,
    pub gpu_ok: bool,
    pub api_ok: bool,
}

impl SessionHealth {
    pub fn is_healthy(&self) -> bool {
        self.session_exists && self.process_alive && self.gpu_ok && self.api_ok
    }

    pub fn is_degraded(&self) -> bool {
        self.session_exists && !self.process_alive
    }
}

/// Top-level handler for all `cargo run tmux <subcommand>` commands.
pub async fn handle_tmux(config: &Config, cmd: TmuxCommands, dry_run: bool) -> Result<()> {
    match cmd {
        TmuxCommands::Serve {
            recipe,
            cluster,
            skip_before_serve,
        } => {
            serve::tmux_serve(
                config,
                &recipe,
                &cluster,
                skip_before_serve,
                dry_run,
                false,
                false,
            )
            .await?;
        }
        TmuxCommands::Stop { recipe, cluster } => {
            ops::tmux_stop(config, &recipe, &cluster, dry_run).await?;
        }
        TmuxCommands::Status { recipe, cluster } => {
            ops::tmux_status(config, &recipe, &cluster, dry_run).await?;
        }
        TmuxCommands::Logs {
            recipe,
            follow,
            cluster,
        } => {
            ops::tmux_logs(config, &recipe, &cluster, follow, dry_run).await?;
        }
        TmuxCommands::Attach { recipe, cluster } => {
            ops::tmux_attach(config, &recipe, &cluster, dry_run).await?;
        }
    }
    Ok(())
}