launcher/src/pre_serve.rs

//! Types for the `before_serve` configuration and `exec_aliases` table.
//!
//! Lives outside `config.rs` so the latter stays under the 500-line limit
//! the project enforces (see AGENTS.md).

use serde::Deserialize;

/// Behaviour on a non-zero exit code for a [`PreServeCommand`].
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum OnFailure {
    #[default]
    Abort,
    Warn,
    Continue,
}

/// One entry in a `before_serve` / `after_stop` / `mods` block.
///
/// Three mutually exclusive forms (untagged, so TOML picks by shape):
/// - `Bare(String)` — a raw shell command, run as-is over SSH.
/// - `Configured { command, on_failure }` — a raw shell command with an
///   explicit `on_failure` policy.
/// - `Script { script, on_failure }` — a built-in script referenced by
///   alias (e.g. `stop_dgx_services`). The script body is baked into the
///   binary at compile time (see [`crate::builtin_scripts`]) and uploaded
///   + run via SSH on each node — no file copy from the operator's machine.
///
///   Unknown aliases error at resolution time with the list of known ones.
#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
pub enum PreServeCommand {
    Bare(String),
    Configured {
        command: String,
        #[serde(default)]
        on_failure: OnFailure,
    },
    Script {
        script: String,
        #[serde(default)]
        on_failure: OnFailure,
    },
}

impl PreServeCommand {
    /// For `Bare`/`Configured`, the raw shell command. For `Script`, the
    /// alias (NOT the script body) — use [`resolve`] to get the body.
    pub fn command(&self) -> &str {
        match self {
            PreServeCommand::Bare(c) => c,
            PreServeCommand::Configured { command, .. } => command,
            PreServeCommand::Script { script, .. } => script,
        }
    }

    pub fn on_failure(&self) -> OnFailure {
        match self {
            PreServeCommand::Bare(_) => OnFailure::default(),
            PreServeCommand::Configured { on_failure, .. } => *on_failure,
            PreServeCommand::Script { on_failure, .. } => *on_failure,
        }
    }
}

#[derive(Clone, Debug, Default, Deserialize)]
pub struct BeforeServeTable {
    #[serde(default)]
    pub commands: Vec<PreServeCommand>,
}

#[derive(Clone, Debug, Default, Deserialize)]
pub struct ModsTable {
    #[serde(default)]
    pub commands: Vec<PreServeCommand>,
}

/// Distributed execution backend — `mp` or `ray`.  Absent means single-node.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum DistributedBackend {
    Mp,
    Ray,
}

/// Recipe arguments — either an inline TOML string (split by whitespace) or
/// a `Vec<String>` list.  Accepted in TOML as either:
///
/// ```toml
/// args = """
///   --gpu-memory-utilization 0.90
///   --max-model-len 262144
/// """
/// ```
///
/// ```toml
/// args = ["--gpu-memory-utilization", "0.90"]
/// ```
#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
pub enum RecipeArgs {
    List(Vec<String>),
    Inline(String),
}

impl Default for RecipeArgs {
    fn default() -> Self {
        RecipeArgs::List(Vec::new())
    }
}

impl RecipeArgs {
    pub fn to_vec(&self) -> Vec<String> {
        match self {
            RecipeArgs::Inline(s) => s.split_whitespace().map(String::from).collect(),
            RecipeArgs::List(v) => v.clone(),
        }
    }

    /// Replace all occurrences of `placeholder` with `replacement` in each
    /// arg string. Used to substitute model path references (e.g.
    /// `{{mtp_draft}}`) with resolved HF cache paths per-node.
    pub fn substitute(&self, placeholder: &str, replacement: &str) -> RecipeArgs {
        match self {
            RecipeArgs::Inline(s) => RecipeArgs::Inline(s.replace(placeholder, replacement)),
            RecipeArgs::List(v) => RecipeArgs::List(
                v.iter()
                    .map(|a| a.replace(placeholder, replacement))
                    .collect(),
            ),
        }
    }
}

/// Venv setup commands — either an inline multi-line TOML string (one
/// command per line) or a `Vec<String>` list. Accepted in TOML as either:
///
/// ```toml
/// venv_setup = """
///   uv venv ~/venvs/vllm025 -p 3.12
///   uv pip install vllm==0.25.1 --torch-backend=cu130
/// """
/// ```
///
/// ```toml
/// venv_setup = [
///     "uv venv ~/venvs/vllm025 -p 3.12",
///     "uv pip install vllm==0.25.1 --torch-backend=cu130",
/// ]
/// ```
///
/// Blank lines and lines starting with `#` are skipped (shell comments).
#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
pub enum VenvSetup {
    List(Vec<String>),
    Inline(String),
}

impl Default for VenvSetup {
    fn default() -> Self {
        VenvSetup::List(Vec::new())
    }
}

impl VenvSetup {
    /// Expand to a list of commands. For `Inline`, each non-blank,
    /// non-comment line becomes one command. For `List`, each entry is
    /// a command as-is.
    pub fn to_vec(&self) -> Vec<String> {
        match self {
            VenvSetup::List(v) => v.clone(),
            VenvSetup::Inline(s) => s
                .lines()
                .map(|l| l.trim().to_string())
                .filter(|l| !l.is_empty() && !l.starts_with('#'))
                .collect(),
        }
    }
}

impl std::fmt::Display for DistributedBackend {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DistributedBackend::Mp => write!(f, "mp"),
            DistributedBackend::Ray => write!(f, "ray"),
        }
    }
}