launcher/src/tmux/mod.rs

pub mod ops;
pub mod serve;

use anyhow::Result;
use std::time::Duration;
use tokio::time::sleep;
use tracing::{info, warn};

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

/// Manages tmux sessions for model inference across a DGX Spark cluster.
#[derive(Clone)]
pub struct TmuxManager {
    #[allow(dead_code)]
    pub session_prefix: String,
}

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 {
        let name = recipe.name.replace(' ', "_");
        format!("{}-{}-r{}", name, recipe.engine, rank)
    }

    /// Builds the list of `export` commands for NCCL environment variables.
    pub fn build_nccl_setup(&self, nccl: &NcclConfig) -> Vec<String> {
        vec![
            format!("export NCCL_IB_HCA=\"{}\"", nccl.ib_hca),
            format!("export NCCL_SOCKET_IFNAME=\"{}\"", nccl.socket_ifname),
            format!("export NCCL_IB_GID_INDEX={}", nccl.ib_gid_index),
            format!("export NCCL_IB_DISABLE={}", nccl.ib_disable),
        ]
    }

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

    /// Builds the full model launch command for a given node.
    pub fn build_launch_command(&self, recipe: &Recipe, node: &Node, cluster: &Cluster) -> String {
        let nnodes = cluster.nodes.len() as u32;
        let pp_size = recipe.pipeline_parallel_size.unwrap_or(nnodes);
        let is_sglang = recipe.engine == "sglang";

        let mut parts: Vec<String> = Vec::new();

        // Engine entrypoint
        match recipe.engine.as_str() {
            "vllm" => {
                parts.push("python -m vllm.entrypoints.openai.api_server".to_string());
            }
            "sglang" => {
                parts.push("python -m sglang.launch_server".to_string());
            }
            other => {
                parts.push(format!("python -m {}", other));
            }
        }

        // Model — engine-specific
        if is_sglang {
            parts.push("--model-path".to_string());
        }
        parts.push(recipe.model.clone());

        // Recipe-specific args
        parts.extend(recipe.args.clone());

        // Distributed flags — engine-specific
        if is_sglang {
            // SGLang uses --tp (not --tensor-parallel-size),
            // and has no --pipeline-parallel-size / --distributed-executor-backend
            parts.extend(vec![
                "--tp".to_string(),
                recipe.tensor_parallel_size.to_string(),
                "--nnodes".to_string(),
                nnodes.to_string(),
                "--node-rank".to_string(),
                node.rank.to_string(),
                "--master-addr".to_string(),
                cluster.master_ib_ip.clone(),
            ]);
        } else {
            match recipe.distributed_backend.as_str() {
                "ray" => {
                    parts.extend(vec![
                        "--tensor-parallel-size".to_string(),
                        "1".to_string(),
                        "--pipeline-parallel-size".to_string(),
                        pp_size.to_string(),
                        "--nnodes".to_string(),
                        nnodes.to_string(),
                        "--node-rank".to_string(),
                        node.rank.to_string(),
                        "--master-addr".to_string(),
                        cluster.master_ib_ip.clone(),
                        "--distributed-executor-backend".to_string(),
                        "ray".to_string(),
                    ]);
                }
                _ => {
                    // MP mode (default)
                    parts.extend(vec![
                        "--tensor-parallel-size".to_string(),
                        "1".to_string(),
                        "--pipeline-parallel-size".to_string(),
                        pp_size.to_string(),
                        "--nnodes".to_string(),
                        nnodes.to_string(),
                        "--node-rank".to_string(),
                        node.rank.to_string(),
                        "--master-addr".to_string(),
                        cluster.master_ib_ip.clone(),
                        "--distributed-executor-backend".to_string(),
                        recipe.distributed_backend.clone(),
                    ]);
                }
            }
        }

        // Port — SGLang needs it on all nodes, vLLM only on head node
        if is_sglang || node.rank == 0 {
            parts.extend(vec!["--port".to_string(), recipe.port.to_string()]);
        }

        parts.join(" ")
    }
    /// Builds the full sequence of tmux commands for a session, with multi-window layout.
    pub fn build_session_commands(
        &self,
        node: &Node,
        recipe: &Recipe,
        nccl: &NcclConfig,
        cluster: &Cluster,
    ) -> Vec<String> {
        let sess = self.session_name(recipe, node.rank);
        let mut cmds: Vec<String> = Vec::new();

        // Kill existing session if it exists
        cmds.push(format!("tmux kill-session -t {} 2>/dev/null || true", sess));

        // Create new detached session (window 0)
        cmds.push(format!("tmux new-session -d -s {} -n model", sess));

        // --- Window 0: model process ---
        for export in self.build_nccl_setup(nccl) {
            cmds.push(format!("tmux send-keys -t {}:0 '{}' Enter", sess, export));
        }
        for export in self.build_host_ip_export(&recipe.engine, node) {
            cmds.push(format!("tmux send-keys -t {}:0 '{}' Enter", sess, export));
        }
        if recipe.distributed_backend == "ray" && node.rank > 0 {
            cmds.push(format!(
                "tmux send-keys -t {}:0 'export RAY_ADDRESS={}:6379' Enter",
                sess, cluster.master_ib_ip
            ));
        }
        for (k, v) in &recipe.env {
            cmds.push(format!(
                "tmux send-keys -t {}:0 'export {}=\"{}\"' Enter",
                sess, k, v
            ));
        }

        let launch_cmd = self.build_launch_command(recipe, node, cluster);
        cmds.push(format!(
            "tmux send-keys -t {}:0 '{}' Enter",
            sess, launch_cmd
        ));

        // --- Window 1: GPU monitor ---
        cmds.push(format!("tmux new-window -t {}:1 -n gpu", sess));
        cmds.push(format!(
            "tmux send-keys -t {}:1 'nvidia-smi -l 5' Enter",
            sess
        ));

        // --- Window 2: shell (debug) ---
        cmds.push(format!("tmux new-window -t {}:2 -n shell", sess));
        cmds.push(format!(
            "tmux send-keys -t {}:2 'echo \"ready for debugging\"' Enter",
            sess
        ));
        cmds.push(format!("tmux select-window -t {}:0", sess));

        cmds
    }
    /// Executes a list of commands via SSH on a node to create a tmux session.
    pub async fn create_session(
        &self,
        config: &Config,
        node: &Node,
        commands: &[String],
        dry_run: bool,
    ) -> Result<()> {
        if dry_run {
            info!("  ssh {}@{}:", config.ssh_user, node.ssh_ip);
            for cmd in commands {
                info!("    {}", cmd);
            }
            return Ok(());
        }

        let script = commands.join(" && ");
        let output = run_ssh(config, node, script, dry_run).await?;

        if output.trim().is_empty() {
            info!("  ✅ tmux session created on {}", node.name);
        } else {
            info!("  tmux session output on {}: {}", node.name, output.trim());
        }

        Ok(())
    }

    /// Checks whether a tmux session exists on the given node.
    pub async fn check_session(
        &self,
        config: &Config,
        node: &Node,
        session_name: &str,
        dry_run: bool,
    ) -> Result<bool> {
        if dry_run {
            info!(
                "  ssh {}@{}: tmux list-sessions 2>/dev/null | grep '{}'",
                config.ssh_user, node.ssh_ip, session_name
            );
            return Ok(true);
        }

        let cmd = format!(
            "tmux list-sessions 2>/dev/null | grep -c '{}'",
            session_name
        );
        let output = run_ssh(config, node, cmd, dry_run).await?;

        let found = output.trim().parse::<usize>().unwrap_or(0) > 0;
        if found {
            info!("  Session '{}' found on {}", session_name, node.name);
        } else {
            warn!("  Session '{}' NOT found on {}", session_name, node.name);
        }

        Ok(found)
    }
    /// Performs a deep health check on a tmux session.
    pub async fn deep_check_session(
        &self,
        config: &Config,
        node: &Node,
        _cluster: &Cluster,
        session_name: &str,
        recipe: &Recipe,
        dry_run: bool,
    ) -> Result<SessionHealth> {
        if dry_run {
            return Ok(SessionHealth {
                session_exists: true,
                process_alive: true,
                gpu_ok: true,
                api_ok: true,
            });
        }

        // 1. Check session exists
        let exists = self
            .check_session(config, node, session_name, dry_run)
            .await?;

        // 2. Check process is alive inside tmux
        let process_alive = if exists {
            let pid_cmd = format!(
                "tmux list-panes -t {}:0 -F '#{{pane_pid}}' | head -1",
                session_name
            );
            match run_ssh(config, node, pid_cmd, dry_run).await {
                Ok(pid_str) => {
                    let pid = pid_str.trim();
                    if !pid.is_empty() && pid != "0" {
                        let ps_cmd = format!("ps -p {} -o pid= 2>/dev/null", pid);
                        run_ssh(config, node, ps_cmd, dry_run)
                            .await
                            .map(|o| !o.trim().is_empty())
                            .unwrap_or(false)
                    } else {
                        false
                    }
                }
                Err(_) => false,
            }
        } else {
            false
        };

        // 3. Check GPU is available
        let gpu_ok = match run_ssh(
            config,
            node,
            "nvidia-smi --query-gpu=index,utilization.gpu,memory.used --format=csv,noheader 2>&1 | head -5"
                .to_string(),
            dry_run,
        )
        .await
        {
            Ok(out) => !out.contains("failed") && !out.contains("error") && !out.trim().is_empty(),
            Err(_) => false,
        };

        // 4. Check API endpoint (head node only)
        let api_ok = if node.rank == 0 && process_alive {
            let api_check = format!(
                "curl -s -o /dev/null -w '%{{http_code}}' http://127.0.0.1:{}/v1/models 2>/dev/null || echo 'fail'",
                recipe.port
            );
            match run_ssh(config, node, api_check, dry_run).await {
                Ok(out) => {
                    let code = out.trim();
                    code == "200" || !code.contains("fail")
                }
                Err(_) => false,
            }
        } else {
            node.rank != 0 || process_alive
        };

        Ok(SessionHealth {
            session_exists: exists,
            process_alive,
            gpu_ok,
            api_ok,
        })
    }

    /// Gracefully stops a tmux session with timeout escalation.
    pub async fn graceful_stop(
        &self,
        config: &Config,
        node: &Node,
        session_name: &str,
        dry_run: bool,
    ) -> Result<()> {
        if dry_run {
            info!(
                "  ssh {}@{}: graceful stop of session '{}'",
                config.ssh_user, node.ssh_ip, session_name
            );
            return Ok(());
        }

        info!(
            "  Sending Ctrl+C to session '{}' on {}",
            session_name, node.name
        );
        let _ = run_ssh(
            config,
            node,
            format!("tmux send-keys -t {}:0 C-c", session_name),
            dry_run,
        )
        .await;

        let mut exited = false;
        for i in 0..15 {
            sleep(Duration::from_secs(2)).await;
            let check_cmd = format!(
                "tmux list-panes -t {}:0 -F '#{{pane_pid}}' 2>/dev/null | head -1",
                session_name
            );
            match run_ssh(config, node, check_cmd, dry_run).await {
                Ok(pid_str) => {
                    let pid = pid_str.trim();
                    if pid.is_empty() || pid == "0" || pid == "No" {
                        exited = true;
                        info!(
                            "  Session '{}' exited gracefully (attempt {})",
                            session_name,
                            i + 1
                        );
                        break;
                    }
                    let ps_cmd = format!("ps -p {} -o pid= 2>/dev/null", pid);
                    match run_ssh(config, node, ps_cmd, dry_run).await {
                        Ok(ps_out) if ps_out.trim().is_empty() => {
                            exited = true;
                            info!(
                                "  Session '{}' process exited (attempt {})",
                                session_name,
                                i + 1
                            );
                            break;
                        }
                        _ => continue,
                    }
                }
                Err(_) => {
                    exited = true;
                    break;
                }
            }
        }

        if !exited {
            warn!(
                "  Session '{}' did not exit gracefully, force killing...",
                session_name
            );
            let _ = run_ssh(
                config,
                node,
                format!("tmux kill-session -t {}", session_name),
                dry_run,
            )
            .await;
        }

        Ok(())
    }
    /// Streams the content of a tmux pane, optionally following new output.
    pub async fn stream_pane(
        &self,
        config: &Config,
        node: &Node,
        session_name: &str,
        window: u32,
        follow: bool,
        dry_run: bool,
    ) -> Result<String> {
        if dry_run {
            if follow {
                info!(
                    "  ssh {}@{}: tail -f (streaming from session '{}' window {})",
                    config.ssh_user, node.ssh_ip, session_name, window
                );
            } else {
                info!(
                    "  ssh {}@{}: tmux capture-pane -t {}:{} -p -S -100",
                    config.ssh_user, node.ssh_ip, session_name, window
                );
            }
            return Ok(String::new());
        }

        let cmd = if follow {
            format!(
                "tmux pipe-pane -o -t {}:{} 'cat > /tmp/sm121_pipe_{}_{}.log' && tail -f /tmp/sm121_pipe_{}_{}.log",
                session_name,
                window,
                session_name.replace('-', "_"),
                window,
                session_name.replace('-', "_"),
                window
            )
        } else {
            format!(
                "tmux capture-pane -t {}:{} -p -S -200 2>/dev/null || echo 'no output'",
                session_name, window
            )
        };

        run_ssh_with_retry(config, node, cmd, dry_run, 1).await
    }

    /// Attaches to a tmux session on a node via SSH (interactive).
    pub async fn attach_session(
        &self,
        config: &Config,
        node: &Node,
        session_name: &str,
        dry_run: bool,
    ) -> Result<()> {
        if dry_run {
            info!(
                "  Would attach: ssh -t {}@{} tmux attach-session -t {}",
                config.ssh_user, node.ssh_ip, session_name
            );
            return Ok(());
        }

        info!(
            "  Attaching to '{}' on {}... Press Ctrl+B D to detach.",
            session_name, node.name
        );
        info!("  (russh does not support interactive PTY. Use the following command directly:)");
        println!(
            "  ssh -t {}@{} tmux attach-session -t {}",
            config.ssh_user, node.ssh_ip, session_name
        );
        println!();
        info!("  Streaming pane content instead:");
        let output = self
            .stream_pane(config, node, session_name, 0, true, dry_run)
            .await?;
        println!("{}", output);

        Ok(())
    }
}

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

    #[allow(dead_code)]
    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 } => {
            serve::tmux_serve(config, &recipe, &cluster, dry_run).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(())
}