launcher/src/docker/mods.rs

//! Mods: commands that run *inside* the container via entrypoint wrapping,
//! before the engine process starts.  Each command runs inside `sh -c` on
//! container start; the original entrypoint is `exec`'d at the end so the
//! engine becomes PID 1 and signals pass through correctly.

use anyhow::{Context, Result};
use tracing::info;

use crate::config::{Cluster, Config, Node, OnFailure, PreServeCommand, Recipe, resolve_mods_dir};
use crate::ssh::upload_dir;

use super::shq;

/// Build a shell script that runs `mods` then `exec`s the original tail
/// (entrypoint + model + args + auto-generated flags).
///
/// Mod commands are shell code (piped directly into the script) — only
/// embedded single quotes are escaped so the outer [`shq`] wrapper doesn't
/// break.  Each tail element is individually [`shq`]'d so multi-word or
/// special-containing parts survive word-splitting inside `sh -c`.
///
/// The returned string is already wrapped in single quotes (via `shq`),
/// ready to use as the `-c` argument to `sh -c`.
pub fn build_mod_script(mods: &[PreServeCommand], tail: &[String]) -> String {
    let mut script = String::new();

    for m in mods {
        // Mod commands are shell code — do NOT wrap in shq (which would
        // prevent word splitting).  Embed directly; the outer shq wrapper
        // handles any embedded single quotes.
        let cmd = m.command();
        match m.on_failure() {
            OnFailure::Abort => {
                script.push_str(&format!("{cmd} || exit 1; "));
            }
            OnFailure::Warn | OnFailure::Continue => {
                // NOTE: `|| true` swallows the exit code and stderr is
                // lost to the container log.  If mod debugging is needed,
                // consider redirecting stderr to a temp file, e.g.:
                //   cmd 2>/tmp/sm121-mod-err.log; true
                script.push_str(&format!("{cmd} || true; "));
            }
        }
    }

    script.push_str("exec ");
    for part in tail {
        script.push_str(&shq(part));
        script.push(' ');
    }

    let trimmed = script.trim_end().to_string();
    shq(&trimmed)
}

/// Upload the mods directory to every node and return the remote path, or
/// empty string when no mods directory exists.  Also validates that
/// `entrypoint` is set when mods are present.
pub async fn upload_mods_to_cluster(
    config: &Config,
    cluster: &Cluster,
    recipe: &Recipe,
    recipe_name: &str,
    node_name: Option<&str>,
    dry_run: bool,
) -> Result<String> {
    if recipe.mods.commands.is_empty() {
        return Ok(String::new());
    }
    if recipe.entrypoint.is_empty() {
        anyhow::bail!(
            "{}",
            fluent_i18n::t!("mods-entrypoint-required", { "name" => &recipe.name })
        );
    }
    let nodes: Vec<&Node> = if let Some(ref nn) = node_name {
        vec![
            cluster
                .nodes
                .iter()
                .find(|n| n.name == *nn)
                .context(format!(
                    "Node '{}' not found in cluster '{}'",
                    nn, cluster.name
                ))?,
        ]
    } else {
        cluster.nodes.iter().collect()
    };
    let remote = format!("/tmp/sm121-mods-{}", recipe_name);
    // Upload global mods first (shared scripts), then recipe mods on top
    // (recipe-specific overrides). Both are merged into the same /mods/ dir.
    if let Some(ref global_dir) = config.global_mods_dir
        && global_dir.is_dir()
    {
        for node in &nodes {
            upload_dir(config, node, global_dir, &remote, dry_run).await?;
        }
    }
    if let Some(local_mods_dir) = resolve_mods_dir(recipe_name).await {
        for node in &nodes {
            info!(
                "{}",
                fluent_i18n::t!("mods-uploading", { "node" => &node.name })
            );
            upload_dir(config, node, &local_mods_dir, &remote, dry_run).await?;
            info!(
                "{}",
                fluent_i18n::t!("mods-uploaded", { "node" => &node.name })
            );
        }
    }
    Ok(remote)
}

#[cfg(test)]
fn unshq(s: &str) -> String {
    let inner = s
        .strip_prefix('\'')
        .and_then(|s| s.strip_suffix('\''))
        .unwrap_or(s);
    inner.replace(r"'\''", "'")
}

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

    #[test]
    fn simple_mods_wraps_tail() {
        let mods = vec![
            PreServeCommand::Bare("pip install flash-attn".into()),
            PreServeCommand::Bare("python patch.py".into()),
        ];
        let tail = vec![
            "vllm".to_string(),
            "serve".to_string(),
            "org/model".to_string(),
            "--host".to_string(),
            "0.0.0.0".to_string(),
        ];

        let script = build_mod_script(&mods, &tail);
        assert!(script.starts_with('\''), "shq-wrapped: {script}");
        assert!(script.ends_with('\''), "shq-wrapped: {script}");

        let inner = unshq(&script);
        // Mod commands are NOT shq-wrapped — they're raw shell code
        assert!(inner.starts_with("pip install flash-attn || exit 1; "));
        assert!(inner.contains("python patch.py || exit 1; "));
        assert!(
            inner.contains("exec 'vllm' 'serve' 'org/model' '--host' '0.0.0.0'"),
            "inner: {inner}"
        );
    }

    #[test]
    fn on_failure_abort_exits_on_error() {
        let mods = vec![PreServeCommand::Bare("false".into())];
        let tail = vec!["echo".to_string(), "ok".to_string()];
        let script = build_mod_script(&mods, &tail);
        let inner = unshq(&script);
        assert!(inner.contains("|| exit 1"));
    }

    #[test]
    fn on_failure_warn_continues() {
        let mods = vec![PreServeCommand::Configured {
            command: "false".into(),
            on_failure: OnFailure::Warn,
        }];
        let tail = vec!["echo".to_string(), "ok".to_string()];
        let script = build_mod_script(&mods, &tail);
        let inner = unshq(&script);
        assert!(inner.contains("|| true"), "warn continues: {inner}");
        assert!(inner.contains("exec"));
    }

    #[test]
    fn on_failure_continue_is_silent() {
        let mods = vec![PreServeCommand::Configured {
            command: "false".into(),
            on_failure: OnFailure::Continue,
        }];
        let tail = vec!["true".to_string()];
        let script = build_mod_script(&mods, &tail);
        let inner = unshq(&script);
        assert!(inner.contains("|| true"));
        assert!(inner.contains("exec 'true'"));
    }

    #[test]
    fn empty_mods_preserves_tail() {
        let mods = vec![];
        let tail = vec!["vllm".to_string(), "serve".to_string(), "noop".to_string()];
        let script = build_mod_script(&mods, &tail);
        let inner = unshq(&script);
        assert_eq!(inner, "exec 'vllm' 'serve' 'noop'");
    }

    #[test]
    fn escaping_special_chars_in_commands() {
        let mods = vec![PreServeCommand::Bare("echo 'hello world'".into())];
        let tail = vec!["true".to_string()];
        let script = build_mod_script(&mods, &tail);
        // The inner single quote in `echo 'hello world'` gets escaped as
        // '\'' so the outer shq doesn't break.  After unshq, the inner
        // script should have `echo '\''hello world'\''` (the outer shq
        // wrapper's quotes are stripped, but the escaped inner quotes
        // are restored to literal single quotes).
        let inner = unshq(&script);
        assert!(
            inner.contains("echo 'hello world'"),
            "inner single quotes should survive: {inner}"
        );
    }

    #[test]
    fn mods_with_entrypoint_sh_c() {
        let mods = vec![PreServeCommand::Bare("echo mod".into())];
        let tail = vec![
            "sh".to_string(),
            "-c".to_string(),
            "i=0; while true; do echo $i; i=$((i+1)); done".to_string(),
            "noop".to_string(),
        ];
        let script = build_mod_script(&mods, &tail);
        let inner = unshq(&script);
        assert!(inner.starts_with("echo mod || exit 1;"));
        assert!(inner.contains("exec 'sh' '-c' '"));
        assert!(inner.contains("'noop'"));
    }

    #[test]
    fn mods_wrapping_uses_entrypoint_sh() {
        let mods = vec![PreServeCommand::Bare("pip list".into())];
        let tail = vec!["vllm".to_string(), "serve".to_string()];
        let script = build_mod_script(&mods, &tail);
        let inner = unshq(&script);
        assert!(
            inner.contains("exec 'vllm' 'serve'"),
            "tail should be exec'd: {inner}"
        );
    }

    #[test]
    fn deserialize_mods_toml() {
        let toml_str = r#"
name = "Modded"
model = "test/model"
entrypoint = ["vllm", "serve"]

[mods]
commands = [
    "pip install flash-attn",
    { command = "./patch.sh", on_failure = "warn" },
]
"#;
        let recipe: crate::config::Recipe = toml::from_str(toml_str).unwrap();
        assert_eq!(recipe.mods.commands.len(), 2);
        assert_eq!(recipe.mods.commands[0].command(), "pip install flash-attn");
        assert_eq!(recipe.mods.commands[0].on_failure(), OnFailure::Abort);
        assert_eq!(recipe.mods.commands[1].command(), "./patch.sh");
        assert_eq!(recipe.mods.commands[1].on_failure(), OnFailure::Warn);
    }
}