launcher/src/builtin_scripts.rs

//! Built-in host scripts baked into the binary at compile time.
//!
//! Each script lives as a real `.sh` file under `host-scripts/` (so it stays
//! diffable, shellcheck-able, and editor-friendly) and is embedded here via
//! `include_str!`. Recipes reference a script by **alias** in a
//! `before_serve` / `after_stop` list (e.g. `{ script = "stop_dgx_services" }`);
//! the launcher resolves the alias to the script body and runs it over SSH on
//! each node. No file copy from the operator's machine is needed.
//!
//! Only **generic** helpers belong here (cluster-wide service toggles, display
//! memory unlocks). Anything that knows about a specific recipe — NAS paths,
//! model revisions, per-rank staging — lives NEXT TO THAT RECIPE in the shared
//! recipe repo and is referenced as a relative path:
//!
//! ```toml
//! before_serve = [{ script = "scripts/prepare.sh" }]   # <recipe dir>/scripts/prepare.sh
//! ```
//!
//! The launcher reads it from the recipe directory, uploads it, and runs it
//! with `NODE_RANK` exported — same contract as a built-in alias. That keeps
//! recipe-specific staging versioned with the recipe (and updatable by the
//! add-on's recipe-repo sync) instead of requiring a launcher release.
//!
//! To add a new built-in:
//! 1. Drop the `.sh` file into `host-scripts/`.
//! 2. Add an `include_str!` line + a `(alias, body)` entry below.
//! 3. Document the alias in AGENTS.md.

/// A built-in script: its alias and the embedded body.
pub struct BuiltinScript {
    pub alias: &'static str,
    pub body: &'static str,
}

/// The registry of all baked-in scripts. Order is the display order for
/// error hints.
pub static BUILTIN_SCRIPTS: &[BuiltinScript] = &[
    BuiltinScript {
        alias: "stop_dgx_services",
        body: include_str!("../host-scripts/stop-dgx-services.sh"),
    },
    BuiltinScript {
        alias: "restore_dgx_services",
        body: include_str!("../host-scripts/restore-dgx-services.sh"),
    },
    // Temporary (per-boot) 2 GiB display-reserve unlock for the EXL3 2x
    // runtime (nvidia_drm modeset=1 fbdev=0 remod; no UEFI/boot changes).
    BuiltinScript {
        alias: "unlock_display_ram",
        body: include_str!("../host-scripts/unlock-display-ram.sh"),
    },
    BuiltinScript {
        alias: "restore_display_ram",
        body: include_str!("../host-scripts/restore-display-ram.sh"),
    },
];

/// Resolve an alias to its script body. Returns `None` (and the caller
/// formats the list of known aliases) when the alias is unknown.
pub fn resolve(alias: &str) -> Option<&'static str> {
    BUILTIN_SCRIPTS
        .iter()
        .find(|s| s.alias == alias)
        .map(|s| s.body)
}

/// Comma-separated list of known aliases, for error messages.
pub fn known_aliases() -> String {
    BUILTIN_SCRIPTS
        .iter()
        .map(|s| s.alias)
        .collect::<Vec<_>>()
        .join(", ")
}

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

    #[test]
    fn resolves_known_aliases() {
        assert!(resolve("stop_dgx_services").is_some());
        assert!(resolve("restore_dgx_services").is_some());
        let body = resolve("stop_dgx_services").unwrap();
        // Sanity: the embedded script is non-empty and starts with a shebang.
        assert!(body.starts_with("#!/usr/bin/env bash"));
        assert!(body.contains("systemctl"));
    }

    #[test]
    fn unknown_alias_returns_none() {
        assert!(resolve("nope").is_none());
    }

    #[test]
    fn known_aliases_lists_all() {
        let list = known_aliases();
        assert!(list.contains("stop_dgx_services"));
        assert!(list.contains("restore_dgx_services"));
    }

    #[test]
    fn scripts_reference_each_other_correctly() {
        // The restore script's docstring mentions the stop script by its
        // filename (`stop-dgx-services.sh`, not the old `glmnf3` name) —
        // regression guard for the rename.
        let restore = resolve("restore_dgx_services").unwrap();
        assert!(restore.contains("stop-dgx-services.sh"));
        assert!(!restore.contains("glmnf3"));
        let stop = resolve("stop_dgx_services").unwrap();
        assert!(!stop.contains("glmnf3"));
    }
}