feat: 开源本机 GPU 加速工具箱（诊断/热驻留/资源调度）

e85df1f verified 1 day ago

3.04 kB

//! 热驻留（warm residency）：把模型常驻在显存里，避免每次请求冷加载。
//!
//! # 为什么重要
//!
//! Ollama 默认会在空闲后卸载模型（free VRAM）。下次请求需要**冷加载 3~5 秒**——
//! 比云端还慢，混合流的"本机秒回"就名存实亡。
//!
//! 解法：一次轻量预热 + `keep_alive` 保活。之后每次请求首 token 只要 0.04~0.18 秒。
//!
//! # 关键：**不要**传 `num_gpu`
//!
//! 强行指定 GPU 层数会触发 runner 重载，反而破坏热驻留。
//! 让 Ollama 自己选择后端与层数，我们只负责"保持它醒着"。

use crate::status::DEFAULT_OLLAMA_BASE_URL;
use serde_json::json;

/// 预热请求体（纯函数，可单测）。
///
/// - `prompt: "hi"` + `num_predict: 1`——只生成 1 个 token，几乎零开销
/// - `keep_alive`——保活时长（Ollama 语法：`"30m"` / `"-1"` 永久 / 数字秒）
/// - **不含 `num_gpu`**——见模块文档
pub fn warmup_body(model: &str, keep_alive: &str) -> serde_json::Value {
    json!({
        "model": model,
        "prompt": "hi",
        "stream": false,
        "keep_alive": keep_alive,
        "options": {
            "num_predict": 1
        }
    })
}

/// 预热一个模型（把请求发到本机 Ollama）。
///
/// 这个提示**不会**进入对话历史，也不出本机 loopback。
pub async fn warmup_model(
    model: &str,
    keep_alive: &str,
    base_url: Option<&str>,
) -> Result<(), String> {
    let model = model.trim();
    if model.is_empty() {
        return Err("模型名不能为空".to_string());
    }
    if model.len() > 256 {
        return Err("模型名过长".to_string());
    }
    let base = base_url
        .unwrap_or(DEFAULT_OLLAMA_BASE_URL)
        .trim_end_matches('/');
    let client = reqwest::Client::builder()
        .connect_timeout(std::time::Duration::from_secs(3))
        .build()
        .map_err(|e| format!("初始化预热请求失败：{e}"))?;
    let response = client
        .post(format!("{base}/api/generate"))
        .json(&warmup_body(model, keep_alive))
        .timeout(std::time::Duration::from_secs(120))
        .send()
        .await
        .map_err(|e| format!("预热请求失败：{e}"))?;
    if !response.status().is_success() {
        return Err(format!("预热返回 {}", response.status()));
    }
    Ok(())
}

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

    #[test]
    fn warmup_body_is_local_bounded_and_never_forces_gpu_layers() {
        let body = warmup_body("qwen3-06b:latest", "30m");
        assert_eq!(body["model"], "qwen3-06b:latest");
        assert_eq!(body["stream"], false);
        assert_eq!(body["keep_alive"], "30m");
        assert_eq!(body["options"]["num_predict"], 1);
        // 红线：绝不强行指定 GPU 层数（会触发 runner 重载）
        assert!(body["options"].get("num_gpu").is_none());
    }

    #[tokio::test]
    async fn rejects_empty_model_name() {
        assert!(warmup_model("", "30m", None).await.is_err());
    }
}

	//! 热驻留（warm residency）：把模型常驻在显存里，避免每次请求冷加载。
	//!
	//! # 为什么重要
	//!
	//! Ollama 默认会在空闲后卸载模型（free VRAM）。下次请求需要冷加载 3~5 秒——
	//! 比云端还慢，混合流的"本机秒回"就名存实亡。
	//!
	//! 解法：一次轻量预热 + `keep_alive` 保活。之后每次请求首 token 只要 0.04~0.18 秒。
	//!
	//! # 关键：不要传 `num_gpu`
	//!
	//! 强行指定 GPU 层数会触发 runner 重载，反而破坏热驻留。
	//! 让 Ollama 自己选择后端与层数，我们只负责"保持它醒着"。

	use crate::status::DEFAULT_OLLAMA_BASE_URL;
	use serde_json::json;

	/// 预热请求体（纯函数，可单测）。
	///
	/// - `prompt: "hi"` + `num_predict: 1`——只生成 1 个 token，几乎零开销
	/// - `keep_alive`——保活时长（Ollama 语法：`"30m"` / `"-1"` 永久 / 数字秒）
	/// - 不含 `num_gpu`——见模块文档
	pub fn warmup_body(model: &str, keep_alive: &str) -> serde_json::Value {
	json!({
	"model": model,
	"prompt": "hi",
	"stream": false,
	"keep_alive": keep_alive,
	"options": {
	"num_predict": 1
	}
	})
	}

	/// 预热一个模型（把请求发到本机 Ollama）。
	///
	/// 这个提示不会进入对话历史，也不出本机 loopback。
	pub async fn warmup_model(
	model: &str,
	keep_alive: &str,
	base_url: Option<&str>,
	) -> Result<(), String> {
	let model = model.trim();
	if model.is_empty() {
	return Err("模型名不能为空".to_string());
	}
	if model.len() > 256 {
	return Err("模型名过长".to_string());
	}
	let base = base_url
	.unwrap_or(DEFAULT_OLLAMA_BASE_URL)
	.trim_end_matches('/');
	let client = reqwest::Client::builder()
	.connect_timeout(std::time::Duration::from_secs(3))
	.build()
	.map_err(\|e\| format!("初始化预热请求失败：{e}"))?;
	let response = client
	.post(format!("{base}/api/generate"))
	.json(&warmup_body(model, keep_alive))
	.timeout(std::time::Duration::from_secs(120))
	.send()
	.await
	.map_err(\|e\| format!("预热请求失败：{e}"))?;
	if !response.status().is_success() {
	return Err(format!("预热返回 {}", response.status()));
	}
	Ok(())
	}

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

	#[test]
	fn warmup_body_is_local_bounded_and_never_forces_gpu_layers() {
	let body = warmup_body("qwen3-06b:latest", "30m");
	assert_eq!(body["model"], "qwen3-06b:latest");
	assert_eq!(body["stream"], false);
	assert_eq!(body["keep_alive"], "30m");
	assert_eq!(body["options"]["num_predict"], 1);
	// 红线：绝不强行指定 GPU 层数（会触发 runner 重载）
	assert!(body["options"].get("num_gpu").is_none());
	}

	#[tokio::test]
	async fn rejects_empty_model_name() {
	assert!(warmup_model("", "30m", None).await.is_err());
	}
	}