launcher/src/cold.rs

//! Cold-start overrides.
//!
//! A recipe may declare `cold_args` — a flat list of engine flags in exactly
//! the same shape as the main `args` (so it is engine-agnostic: a vLLM recipe
//! writes `--gpu-memory-utilization 0.75 --max-model-len 32768`, an SGLang
//! recipe writes `--mem-fraction-static 0.6 …`). `serve --cold` overlays those
//! onto the main args: any flag that appears in `cold_args` **replaces** its
//! value in the main args (or is appended if the recipe omitted it). Flags not
//! mentioned in `cold_args` are left untouched.
//!
//! This is the "boot it safely first, then serve full" mode: the first launch
//! (or a launch onto a tight node) uses the conservative `cold_args` to avoid a
//! KV-cache OOM or a GB10 unified-memory *stall*, without editing the recipe
//! per launch. The values are taken as written — the recipe author is expected
//! to write cold values lower than the serving values.

/// Overlay `cold` onto `main`: for each flag in `cold`, replace that flag's
/// value in `main` (or append `flag value` when `main` lacks it). Handles the
/// space form (`--max-len 32768`), the `=` form (`--max-len=32768`), and
/// boolean flags (a flag with no following value is just ensured present).
/// Returns a new vec; `main` is unchanged.
pub fn apply_cold(main: &[String], cold: &[String]) -> Vec<String> {
    let mut out = main.to_vec();
    let mut i = 0;
    while i < cold.len() {
        let tok = cold[i].as_str();
        if !tok.starts_with('-') {
            // A stray value token with no preceding flag — skip defensively.
            i += 1;
            continue;
        }
        if let Some(eq) = tok.find('=') {
            // `--flag=value` form.
            set_flag(&mut out, &tok[..eq], tok[eq + 1..].to_string());
            i += 1;
            continue;
        }
        // Space form: the value is the next token iff it exists and is not
        // itself a flag.
        let has_value = i + 1 < cold.len() && !cold[i + 1].starts_with('-');
        if has_value {
            set_flag(&mut out, tok, cold[i + 1].clone());
            i += 2;
        } else {
            // Boolean flag: ensure it is present in `main`.
            if !out.iter().any(|a| a == tok) {
                out.push(tok.to_string());
            }
            i += 1;
        }
    }
    out
}

/// Replace the value of `flag` in `args`, or append `flag value` when absent.
/// Handles both the space form (`--max-len 32768`) and the `=` form
/// (`--max-len=32768`).
fn set_flag(args: &mut Vec<String>, flag: &str, value: String) {
    let eq = format!("{flag}=");
    if let Some(i) = args.iter().position(|a| a == flag || a.starts_with(&eq)) {
        if args[i] == flag {
            // Space form: the value is the next token.
            if i + 1 < args.len() {
                args[i + 1] = value;
            } else {
                args.push(value);
            }
        } else {
            // `--flag=value` form: rewrite the whole token.
            args[i] = format!("{flag}={value}");
        }
    } else {
        args.push(flag.to_string());
        args.push(value);
    }
}

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

    fn main_args() -> Vec<String> {
        [
            "--served-model-name",
            "m",
            "--max-model-len",
            "500000",
            "--max-num-seqs",
            "4",
            "--gpu-memory-utilization",
            "0.87",
        ]
        .iter()
        .map(|s| s.to_string())
        .collect()
    }

    fn get(args: &[String], flag: &str) -> String {
        args.iter()
            .position(|a| a == flag)
            .map(|i| args[i + 1].clone())
            .unwrap_or_default()
    }

    #[test]
    fn cold_replaces_present_flags_and_appends_absent() {
        let cold = [
            "--gpu-memory-utilization",
            "0.75",
            "--max-model-len",
            "32768",
            "--max-num-seqs",
            "2",
            "--max-num-batched-tokens",
            "2048",
        ]
        .iter()
        .map(|s| s.to_string())
        .collect::<Vec<_>>();
        let out = apply_cold(&main_args(), &cold);
        // Each flag appears exactly once (replaced, not duplicated).
        for f in [
            "--max-model-len",
            "--max-num-seqs",
            "--gpu-memory-utilization",
        ] {
            assert_eq!(
                out.iter().filter(|a| a.as_str() == f).count(),
                1,
                "{f} duplicated"
            );
        }
        assert_eq!(get(&out, "--max-model-len"), "32768");
        assert_eq!(get(&out, "--max-num-seqs"), "2");
        assert_eq!(get(&out, "--gpu-memory-utilization"), "0.75");
        // Absent flag appended.
        assert_eq!(get(&out, "--max-num-batched-tokens"), "2048");
        // Untouched flags stay.
        assert_eq!(get(&out, "--served-model-name"), "m");
    }

    #[test]
    fn cold_handles_eq_form_and_appends() {
        let base = vec!["--max-model-len=500000".to_string()];
        let cold = ["--max-model-len=16384", "--max-num-seqs", "1"]
            .iter()
            .map(|s| s.to_string())
            .collect::<Vec<_>>();
        let out = apply_cold(&base, &cold);
        assert_eq!(out, vec!["--max-model-len=16384", "--max-num-seqs", "1"]);
    }

    #[test]
    fn cold_boolean_flag_ensured_present_once() {
        let base = vec!["--enable-prefix-caching".to_string(), "--foo".to_string()];
        let cold = ["--enable-prefix-caching", "--bar"]
            .iter()
            .map(|s| s.to_string())
            .collect::<Vec<_>>();
        let out = apply_cold(&base, &cold);
        assert_eq!(
            out.iter()
                .filter(|a| a.as_str() == "--enable-prefix-caching")
                .count(),
            1
        );
        assert!(out.iter().any(|a| a.as_str() == "--bar"));
    }

    #[test]
    fn empty_cold_is_noop() {
        assert_eq!(apply_cold(&main_args(), &[]), main_args());
    }

    #[test]
    fn engine_agnostic_sglang_flags() {
        // A cold block is not tied to vLLM flag names — any flag overlays.
        let base = vec!["--mem-fraction-static".to_string(), "0.9".to_string()];
        let cold = ["--mem-fraction-static", "0.6"]
            .iter()
            .map(|s| s.to_string())
            .collect::<Vec<_>>();
        let out = apply_cold(&base, &cold);
        assert_eq!(get(&out, "--mem-fraction-static"), "0.6");
    }
}