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

e85df1f verified 1 day ago

5.68 kB

//! 本机 GPU 状态诊断：读 Ollama `/api/ps` 的真实 VRAM 分配。
//!
//! # 为什么要读 `/api/ps` 而不是猜
//!
//! 大多数应用会通过 `num_gpu` 参数"控制"GPU 层数——**这是错的**。
//! Ollama 自己会按模型结构和显存选择 Metal/CUDA 后端与层数；
//! 强行传 `num_gpu` 会触发 runner 重载，反而更慢。
//!
//! 正确的做法：**让 Ollama 自己决定，然后如实报告它实际分配了多少显存。**

use serde::{Deserialize, Serialize};

/// 本机 Ollama 的默认地址。
pub const DEFAULT_OLLAMA_BASE_URL: &str = "http://127.0.0.1:11434";

/// 单个模型的 GPU 驻留状态。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GpuModelStatus {
    /// 模型名（如 `qwen3-06b:latest`）。
    pub name: String,
    /// 模型总大小（字节）。
    pub model_bytes: u64,
    /// 实际分配在 VRAM 的字节数（`size_vram`）。
    pub gpu_memory_bytes: u64,
    /// 上下文长度（Ollama 新版才返回）。
    pub context_length: Option<u64>,
    /// 驻留过期时间（keep_alive 到期）。
    pub expires_at: Option<String>,
    /// GPU 驻留百分比（0~100，近似值）。
    /// 这是**驻留估算**，不是 GPU 利用率，也不是 tokens/s。
    pub gpu_residency_percent: u8,
}

/// 本机 GPU 加速总状态。
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GpuAccelerationStatus {
    /// Ollama 是否在线。
    pub online: bool,
    /// 查询时间（毫秒时间戳）。
    pub checked_at_ms: i64,
    /// 已加载的模型列表。
    pub models: Vec<GpuModelStatus>,
}

#[derive(Debug, Default, Deserialize)]
struct PsResponse {
    #[serde(default)]
    models: Vec<PsModel>,
}

#[derive(Debug, Default, Deserialize)]
struct PsModel {
    #[serde(default)]
    name: String,
    #[serde(default)]
    size: u64,
    #[serde(default, rename = "size_vram")]
    size_vram: u64,
    #[serde(default, rename = "context_length")]
    context_length: Option<u64>,
    #[serde(default, rename = "expires_at")]
    expires_at: Option<String>,
}

/// 计算 GPU 驻留百分比（纯函数，可单测）。
pub fn residency_percent(size: u64, gpu: u64) -> u8 {
    if size == 0 {
        return 0;
    }
    ((gpu.min(size) as u128 * 100 / size as u128).min(100)) as u8
}

/// 把 `/api/ps` 的原始响应转成结构化状态（纯函数，可单测）。
pub fn from_ps_json(value: &str, now_ms: i64) -> Result<GpuAccelerationStatus, String> {
    let parsed: PsResponse =
        serde_json::from_str(value).map_err(|e| format!("无法解析 /api/ps 响应：{e}"))?;
    let models = parsed
        .models
        .into_iter()
        .map(|model| GpuModelStatus {
            gpu_residency_percent: residency_percent(model.size, model.size_vram),
            name: model.name,
            model_bytes: model.size,
            gpu_memory_bytes: model.size_vram,
            context_length: model.context_length,
            expires_at: model.expires_at,
        })
        .collect();
    Ok(GpuAccelerationStatus {
        online: true,
        checked_at_ms: now_ms,
        models,
    })
}

/// 查询本机 Ollama 的 GPU 状态（HTTP 请求）。
pub async fn query_status(base_url: Option<&str>) -> Result<GpuAccelerationStatus, 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(1))
        .build()
        .map_err(|e| format!("初始化状态查询失败：{e}"))?;
    let response = client
        .get(format!("{base}/api/ps"))
        .timeout(std::time::Duration::from_secs(3))
        .send()
        .await
        .map_err(|e| format!("无法读取本机 Ollama GPU 状态：{e}"))?;
    if !response.status().is_success() {
        return Err(format!("本机 Ollama 状态接口返回 {}", response.status()));
    }
    let text = response
        .text()
        .await
        .map_err(|e| format!("无法读取响应体：{e}"))?;
    from_ps_json(&text, now_ms())
}

fn now_ms() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_millis() as i64)
        .unwrap_or(0)
}

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

    #[test]
    fn reports_real_residency_and_handles_cpu_fallback() {
        let status = from_ps_json(
            r#"{"models":[
                {"name":"gpu-model","size":1000,"size_vram":1000,"context_length":4096,"expires_at":"soon"},
                {"name":"partial-model","size":1000,"size_vram":500},
                {"name":"cpu-model","size":1000,"size_vram":0},
                {"name":"unknown-size","size":0,"size_vram":0}
            ]}"#,
            0,
        )
        .unwrap();
        assert!(status.online);
        assert_eq!(status.models[0].gpu_residency_percent, 100);
        assert_eq!(status.models[1].gpu_residency_percent, 50);
        assert_eq!(status.models[2].gpu_residency_percent, 0);
        assert_eq!(status.models[3].gpu_residency_percent, 0);
        assert_eq!(status.models[0].context_length, Some(4096));
        // 序列化字段名（camelCase）
        let wire = serde_json::to_value(&status).unwrap();
        assert_eq!(wire["models"][0]["gpuResidencyPercent"], 100);
        assert_eq!(wire["models"][0]["gpuMemoryBytes"], 1000);
    }

    #[test]
    fn accepts_no_loaded_models() {
        let status = from_ps_json("{}", 0).unwrap();
        assert!(status.online);
        assert!(status.models.is_empty());
    }
}

	//! 本机 GPU 状态诊断：读 Ollama `/api/ps` 的真实 VRAM 分配。
	//!
	//! # 为什么要读 `/api/ps` 而不是猜
	//!
	//! 大多数应用会通过 `num_gpu` 参数"控制"GPU 层数——这是错的。
	//! Ollama 自己会按模型结构和显存选择 Metal/CUDA 后端与层数；
	//! 强行传 `num_gpu` 会触发 runner 重载，反而更慢。
	//!
	//! 正确的做法：让 Ollama 自己决定，然后如实报告它实际分配了多少显存。

	use serde::{Deserialize, Serialize};

	/// 本机 Ollama 的默认地址。
	pub const DEFAULT_OLLAMA_BASE_URL: &str = "http://127.0.0.1:11434";

	/// 单个模型的 GPU 驻留状态。
	#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
	#[serde(rename_all = "camelCase")]
	pub struct GpuModelStatus {
	/// 模型名（如 `qwen3-06b:latest`）。
	pub name: String,
	/// 模型总大小（字节）。
	pub model_bytes: u64,
	/// 实际分配在 VRAM 的字节数（`size_vram`）。
	pub gpu_memory_bytes: u64,
	/// 上下文长度（Ollama 新版才返回）。
	pub context_length: Option<u64>,
	/// 驻留过期时间（keep_alive 到期）。
	pub expires_at: Option<String>,
	/// GPU 驻留百分比（0~100，近似值）。
	/// 这是驻留估算，不是 GPU 利用率，也不是 tokens/s。
	pub gpu_residency_percent: u8,
	}

	/// 本机 GPU 加速总状态。
	#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
	#[serde(rename_all = "camelCase")]
	pub struct GpuAccelerationStatus {
	/// Ollama 是否在线。
	pub online: bool,
	/// 查询时间（毫秒时间戳）。
	pub checked_at_ms: i64,
	/// 已加载的模型列表。
	pub models: Vec<GpuModelStatus>,
	}

	#[derive(Debug, Default, Deserialize)]
	struct PsResponse {
	#[serde(default)]
	models: Vec<PsModel>,
	}

	#[derive(Debug, Default, Deserialize)]
	struct PsModel {
	#[serde(default)]
	name: String,
	#[serde(default)]
	size: u64,
	#[serde(default, rename = "size_vram")]
	size_vram: u64,
	#[serde(default, rename = "context_length")]
	context_length: Option<u64>,
	#[serde(default, rename = "expires_at")]
	expires_at: Option<String>,
	}

	/// 计算 GPU 驻留百分比（纯函数，可单测）。
	pub fn residency_percent(size: u64, gpu: u64) -> u8 {
	if size == 0 {
	return 0;
	}
	((gpu.min(size) as u128 * 100 / size as u128).min(100)) as u8
	}

	/// 把 `/api/ps` 的原始响应转成结构化状态（纯函数，可单测）。
	pub fn from_ps_json(value: &str, now_ms: i64) -> Result<GpuAccelerationStatus, String> {
	let parsed: PsResponse =
	serde_json::from_str(value).map_err(\|e\| format!("无法解析 /api/ps 响应：{e}"))?;
	let models = parsed
	.models
	.into_iter()
	.map(\|model\| GpuModelStatus {
	gpu_residency_percent: residency_percent(model.size, model.size_vram),
	name: model.name,
	model_bytes: model.size,
	gpu_memory_bytes: model.size_vram,
	context_length: model.context_length,
	expires_at: model.expires_at,
	})
	.collect();
	Ok(GpuAccelerationStatus {
	online: true,
	checked_at_ms: now_ms,
	models,
	})
	}

	/// 查询本机 Ollama 的 GPU 状态（HTTP 请求）。
	pub async fn query_status(base_url: Option<&str>) -> Result<GpuAccelerationStatus, 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(1))
	.build()
	.map_err(\|e\| format!("初始化状态查询失败：{e}"))?;
	let response = client
	.get(format!("{base}/api/ps"))
	.timeout(std::time::Duration::from_secs(3))
	.send()
	.await
	.map_err(\|e\| format!("无法读取本机 Ollama GPU 状态：{e}"))?;
	if !response.status().is_success() {
	return Err(format!("本机 Ollama 状态接口返回 {}", response.status()));
	}
	let text = response
	.text()
	.await
	.map_err(\|e\| format!("无法读取响应体：{e}"))?;
	from_ps_json(&text, now_ms())
	}

	fn now_ms() -> i64 {
	std::time::SystemTime::now()
	.duration_since(std::time::UNIX_EPOCH)
	.map(\|d\| d.as_millis() as i64)
	.unwrap_or(0)
	}

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

	#[test]
	fn reports_real_residency_and_handles_cpu_fallback() {
	let status = from_ps_json(
	r#"{"models":[
	{"name":"gpu-model","size":1000,"size_vram":1000,"context_length":4096,"expires_at":"soon"},
	{"name":"partial-model","size":1000,"size_vram":500},
	{"name":"cpu-model","size":1000,"size_vram":0},
	{"name":"unknown-size","size":0,"size_vram":0}
	]}"#,
	0,
	)
	.unwrap();
	assert!(status.online);
	assert_eq!(status.models[0].gpu_residency_percent, 100);
	assert_eq!(status.models[1].gpu_residency_percent, 50);
	assert_eq!(status.models[2].gpu_residency_percent, 0);
	assert_eq!(status.models[3].gpu_residency_percent, 0);
	assert_eq!(status.models[0].context_length, Some(4096));
	// 序列化字段名（camelCase）
	let wire = serde_json::to_value(&status).unwrap();
	assert_eq!(wire["models"][0]["gpuResidencyPercent"], 100);
	assert_eq!(wire["models"][0]["gpuMemoryBytes"], 1000);
	}

	#[test]
	fn accepts_no_loaded_models() {
	let status = from_ps_json("{}", 0).unwrap();
	assert!(status.online);
	assert!(status.models.is_empty());
	}
	}