|
|
@@ -0,0 +1,263 @@
|
|
|
+//! 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 => {
|
|
|
+ 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);
|
|
|
+ }
|
|
|
+}
|