|
|
@@ -0,0 +1,186 @@
|
|
|
+//! 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");
|
|
|
+ }
|
|
|
+}
|
