compute/src/config.rs

//! Config types + loaders for the model registry and scenes.
//!
//! The registry (`registry/models.yaml`) resolves stable model *ids* to weights +
//! launch commands, and evolves as models requantize. Scenes (`scenes/*.scene.yaml`)
//! reference those ids and stay stable — the public, shareable contract.
//!
//! Several fields below are deserialized but never read by the binary
//! (`api_version`, `channel`, `entry`, `app`, `app_dir`, `tags`, `author`).
//! That is deliberate: they are the published `lumbridge/v1` manifest surface,
//! and declaring them is what makes a manifest carrying them parse rather than
//! fail. Deleting them to satisfy the lint would silently narrow the contract.
#![allow(dead_code)]

use anyhow::{Context, Result};
use serde::Deserialize;
use std::collections::BTreeMap;
use std::fs;
use std::path::Path;

#[derive(Debug, Deserialize)]
pub struct Registry {
    #[serde(rename = "apiVersion")]
    pub api_version: String,
    pub models: BTreeMap<String, Model>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct Model {
    pub name: String,
    /// Worst-case unified memory held once serving (weights + KV/cache + encoder).
    pub footprint_gb: f64,
    /// `stable` pins the exact weights; `latest` may resolve a newer quant.
    #[serde(default)]
    pub channel: Option<String>,
    #[serde(default)]
    pub health: Option<String>,
    /// Optional marker that must appear in the health response. This
    /// distinguishes mutually-exclusive models that intentionally share a port.
    #[serde(default)]
    pub health_contains: Option<String>,
    pub serve: Serve,
}

/// How to launch a model. Fields are permissive across runtimes (vllm / diffusers /
/// python / uvicorn); only the ones a given `kind` needs are populated.
#[derive(Debug, Clone, Deserialize)]
pub struct Serve {
    pub kind: String,
    /// Optional runtime executable for kind-based builders (for example a
    /// model-specific vLLM virtualenv). Defaults to the kind's command on PATH.
    #[serde(default)]
    pub executable: Option<String>,
    #[serde(default)]
    pub port: Option<u16>,
    /// Explicit launch argv. If set, it is used verbatim (argv[0] = program) and
    /// takes precedence over any `kind`-based builder. This is where launch commands
    /// live — never in a scene — so downloaded scenes can't smuggle code.
    #[serde(default)]
    pub command: Vec<String>,
    #[serde(default)]
    pub weights: Option<String>,
    #[serde(default)]
    pub entry: Option<String>,
    #[serde(default)]
    pub app: Option<String>,
    #[serde(default)]
    pub app_dir: Option<String>,
    #[serde(default)]
    pub served_name: Vec<String>,
    #[serde(default)]
    pub args: BTreeMap<String, serde_yaml::Value>,
    #[serde(default)]
    pub env: BTreeMap<String, String>,
}

#[derive(Debug, Deserialize)]
pub struct Scene {
    #[serde(rename = "apiVersion")]
    pub api_version: String,
    pub metadata: SceneMeta,
    /// Stable model ids resolved via the registry.
    pub models: Vec<String>,
    /// Optional per-scene budget override (GB); defaults to the Governor's global budget.
    #[serde(default)]
    pub budget_gb: Option<f64>,
    #[serde(default)]
    pub activation: Option<Activation>,
}

#[derive(Debug, Deserialize)]
pub struct SceneMeta {
    pub name: String,
    #[serde(default)]
    pub version: u32,
    #[serde(default)]
    pub description: String,
    #[serde(default)]
    pub tags: Vec<String>,
    #[serde(default)]
    pub author: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct Activation {
    /// `footprint-asc` (default) | `listed`.
    #[serde(default)]
    pub order: Option<String>,
    #[serde(default)]
    pub wait_healthy: Option<bool>,
}

impl Registry {
    pub fn load(root: &Path) -> Result<Registry> {
        let p = root.join("registry/models.yaml");
        let s =
            fs::read_to_string(&p).with_context(|| format!("reading registry {}", p.display()))?;
        serde_yaml::from_str(&s).with_context(|| format!("parsing registry {}", p.display()))
    }
}

/// Load every `*.scene.yaml` under `<root>/scenes`, sorted by name.
pub fn load_scenes(root: &Path) -> Result<Vec<Scene>> {
    let dir = root.join("scenes");
    let mut out = Vec::new();
    if !dir.exists() {
        return Ok(out);
    }
    for entry in fs::read_dir(&dir).with_context(|| format!("reading {}", dir.display()))? {
        let path = entry?.path();
        let is_scene = path
            .file_name()
            .and_then(|n| n.to_str())
            .map(|n| n.ends_with(".scene.yaml"))
            .unwrap_or(false);
        if is_scene {
            let s = fs::read_to_string(&path)?;
            let scene: Scene =
                serde_yaml::from_str(&s).with_context(|| format!("parsing {}", path.display()))?;
            out.push(scene);
        }
    }
    out.sort_by(|a, b| a.metadata.name.cmp(&b.metadata.name));
    Ok(out)
}

pub fn find_scene(root: &Path, name: &str) -> Result<Scene> {
    load_scenes(root)?
        .into_iter()
        .find(|s| s.metadata.name == name)
        .with_context(|| format!("no scene named '{name}' in {}/scenes", root.display()))
}