File size: 3,035 Bytes
e85df1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | //! 热驻留（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());
}
}
|