awesome-bookmarks/goose
3
0
Fork 0
mirror of https://github.com/block/goose.git synced 2026-07-17 12:56:20 +02:00
Code Issues Packages Projects Releases Wiki Activity

Clean up

Browse Source
This commit is contained in:
Douwe Osinga
2026-03-18 06:35:44 -04:00
parent 12eb0421c7
commit cf733ea272
3 changed files with 40 additions and 276 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/routes/local_inference.rs
-1
View File
@@ -48,7 +48,6 @@ pub struct LocalModelResponse {
pub status: ModelDownloadStatus,
pub recommended: bool,
pub settings: ModelSettings,
/// Estimated inference speed tier (fast/medium/slow)
pub speed_tier: SpeedTier,
}
crates/goose/src/providers/local_inference/local_model_registry.rs
-1
View File
@@ -132,7 +132,6 @@ pub struct LocalModelEntry {
pub settings: ModelSettings,
#[serde(default)]
pub size_bytes: u64,
/// Quality rank from quantization (1-28, higher is better quality)
#[serde(default = "default_quality_rank")]
pub quality_rank: u8,
}
crates/goose/src/providers/local_inference/recommender.rs
+40 -274
View File
@@ -1,10 +1,4 @@
//! Model recommendation engine for local inference.
//!
//! Provides hardware-aware model recommendations by:
//! - Detecting available memory (GPU VRAM or system RAM)
//! - Estimating model memory requirements including KV cache
//! - Scoring models based on quality and fit
//! - Estimating inference speed
use super::hf_models::HfQuantVariant;
use super::local_model_registry::{
@@ -16,8 +10,6 @@ use serde::{Deserialize, Serialize};
use std::fmt;
use utoipa::ToSchema;
/// Bytes per parameter for each quantization level.
/// Used to estimate parameter count from file size and for memory calculations.
const QUANT_BPP: &[(&str, f64)] = &[
("F32", 4.0),
("BF16", 2.0),
@@ -55,7 +47,6 @@ const QUANT_BPP: &[(&str, f64)] = &[
("MXFP4_MOE", 0.58),
];
/// Get bytes-per-parameter for a quantization level.
pub fn quant_bpp(quant: &str) -> f64 {
let upper = quant.to_uppercase();
QUANT_BPP
@@ -65,7 +56,6 @@ pub fn quant_bpp(quant: &str) -> f64 {
.unwrap_or(0.58) // Default to Q4_K_M-ish
}
/// Estimate parameter count in billions from file size and quantization.
pub fn estimate_params_billion(size_bytes: u64, quant: &str) -> f64 {
let bpp = quant_bpp(quant);
if bpp == 0.0 || size_bytes == 0 {
@@ -74,12 +64,6 @@ pub fn estimate_params_billion(size_bytes: u64, quant: &str) -> f64 {
(size_bytes as f64) / bpp / 1e9
}
/// Estimate memory required in GB to run a model with given context length.
///
/// Memory formula: model_weights + KV_cache + overhead
/// - model_weights = params_b * bpp (already in the file size)
/// - KV_cache ≈ 0.000008 * params_b * context_length (in GB)
/// - overhead ≈ 0.5 GB for runtime buffers
pub fn estimate_memory_gb(size_bytes: u64, quant: &str, context_length: u32) -> f64 {
let params_b = estimate_params_billion(size_bytes, quant);
let model_mem_gb = size_bytes as f64 / 1e9;
@@ -88,7 +72,6 @@ pub fn estimate_memory_gb(size_bytes: u64, quant: &str, context_length: u32) ->
model_mem_gb + kv_cache_gb + overhead_gb
}
/// Speed tier for user-friendly display.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ToSchema)]
#[serde(rename_all = "lowercase")]
pub enum SpeedTier {
@@ -97,19 +80,14 @@ pub enum SpeedTier {
Slow,
}
/// How the model will be executed based on available memory.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RunMode {
/// Fully loaded in GPU VRAM - fastest
Gpu,
/// Partial GPU with CPU offload - slower due to memory transfers
CpuOffload,
/// CPU only, no GPU acceleration
CpuOnly,
}
impl RunMode {
/// Speed multiplier for this run mode (relative to full GPU).
pub fn speed_multiplier(&self) -> f64 {
match self {
RunMode::Gpu => 1.0,
@@ -119,20 +97,16 @@ impl RunMode {
}
}
/// Predict how the model will run based on memory requirements and availability.
pub fn predict_run_mode(model_memory_gb: f64, vram_gb: f64, has_gpu: bool) -> RunMode {
if !has_gpu {
return RunMode::CpuOnly;
}
// If model fits in VRAM with some headroom, it'll run fully on GPU
if model_memory_gb <= vram_gb * 0.95 {
RunMode::Gpu
} else if model_memory_gb <= vram_gb * 2.0 {
// Model is larger than VRAM but not massively so - will offload to CPU
RunMode::CpuOffload
} else {
// Model is way too large for GPU - effectively CPU only
RunMode::CpuOnly
}
}
@@ -149,10 +123,6 @@ impl fmt::Display for SpeedTier {
impl SpeedTier {
pub fn from_tps(tps: f64) -> Self {
// Thresholds calibrated for practical usability:
// - Fast: comfortable interactive use
// - Medium: usable but noticeable delay
// - Slow: significant wait between responses
if tps >= 20.0 {
SpeedTier::Fast
} else if tps >= 8.0 {
@@ -163,7 +133,6 @@ impl SpeedTier {
}
}
/// Quantization speed multiplier - lower precision is faster.
fn quant_speed_multiplier(quant: &str) -> f64 {
let upper = quant.to_uppercase();
if upper.starts_with("F32") {
@@ -189,14 +158,12 @@ fn quant_speed_multiplier(quant: &str) -> f64 {
}
}
/// Get the number of CPU cores available.
pub fn cpu_core_count() -> usize {
std::thread::available_parallelism()
.map(|p| p.get())
.unwrap_or(4)
}
/// Get available GPU/accelerator VRAM in GB.
pub fn available_vram_gb() -> f64 {
let devices = list_llama_ggml_backend_devices();
devices
@@ -214,55 +181,25 @@ pub fn available_vram_gb() -> f64 {
.unwrap_or(0.0)
}
/// Estimate inference speed in tokens per second.
///
/// Accounts for model size, quantization, GPU availability, core count, and run mode.
/// A large model on a small GPU will be penalized for CPU offloading.
pub fn estimate_speed_tps(params_billion: f64, quant: &str, has_gpu: bool) -> f64 {
estimate_speed_tps_with_memory(params_billion, quant, has_gpu, available_vram_gb())
}
/// Estimate inference speed with explicit VRAM availability.
/// This allows callers to provide known memory values for more accurate estimates.
pub fn estimate_speed_tps_with_memory(
params_billion: f64,
quant: &str,
has_gpu: bool,
vram_gb: f64,
) -> f64 {
// Base speed constant by backend (empirically calibrated)
// These represent rough tok/s for a 1B parameter model
// Modern GPUs (M3, RTX 4090, etc) can do 200+ tok/s on 1B models
let base_k = if has_gpu {
// GPU backends - Metal/CUDA/etc
// Calibrated for M3/M4 and modern NVIDIA cards
300.0
} else if cfg!(target_arch = "aarch64") {
80.0
} else {
// CPU-only is much slower
if cfg!(target_arch = "aarch64") {
80.0 // ARM CPUs (Apple Silicon CPU-only, etc)
} else {
60.0 // x86 CPUs
}
60.0
};
// Speed scales inversely with parameter count
let params = params_billion.max(0.1);
let base_tps = base_k / params;
// 10% bonus for systems with 8+ cores (matches llmfit)
let core_mult = if cpu_core_count() >= 8 { 1.1 } else { 1.0 };
// Estimate model memory to predict run mode
// Use a rough approximation: params * 0.6 for Q4_K_M-ish models + overhead
let model_mem_gb = params_billion * quant_bpp(quant) + 1.0;
let run_mode = predict_run_mode(model_mem_gb, vram_gb, has_gpu);
// Apply all multipliers
let run_mode = predict_run_mode(model_mem_gb, available_vram_gb(), has_gpu);
base_tps * quant_speed_multiplier(quant) * core_mult * run_mode.speed_multiplier()
}
/// Check if a GPU/accelerator is available for inference.
pub fn has_gpu_accelerator() -> bool {
let devices = list_llama_ggml_backend_devices();
devices.iter().any(|d| {
@@ -275,7 +212,6 @@ pub fn has_gpu_accelerator() -> bool {
})
}
/// Result of model fit analysis.
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct ModelFit {
/// Whether the model fits in available memory
@@ -293,8 +229,6 @@ pub struct ModelFit {
}
impl ModelFit {
/// Compute a combined score for ranking models.
/// Higher is better. Balances fit (headroom) with quality.
pub fn score(&self) -> f64 {
if !self.fits {
return -1.0;
@@ -314,7 +248,6 @@ impl ModelFit {
}
}
/// Analyze how well a model fits the available memory.
pub fn analyze_model_fit(
size_bytes: u64,
quant: &str,
@@ -343,10 +276,6 @@ pub fn analyze_model_fit(
}
}
/// Get available memory for inference in bytes.
///
/// Prefers GPU/accelerator memory if available, falls back to CPU/system RAM.
/// On Apple Silicon with unified memory, this returns the available unified memory.
pub fn available_inference_memory_bytes(runtime: &InferenceRuntime) -> u64 {
let _ = &runtime.backend();
let devices = list_llama_ggml_backend_devices();
@@ -379,15 +308,8 @@ pub fn available_inference_memory_bytes(runtime: &InferenceRuntime) -> u64 {
}
}
/// Default context length used for recommendations when not specified.
const DEFAULT_CONTEXT_LENGTH: u32 = 8192;
/// Recommend the best local model from the registry for the current hardware.
///
/// Returns the model ID of the recommended model. Prefers:
/// 1. Downloaded models that fit in memory
/// 2. Highest quality model that fits
/// 3. Falls back to smallest featured model if nothing fits
pub fn recommend_local_model(runtime: &InferenceRuntime) -> String {
let available_memory = available_inference_memory_bytes(runtime);
@@ -434,10 +356,19 @@ pub fn recommend_local_model(runtime: &InferenceRuntime) -> String {
FEATURED_MODELS[0].to_string()
}
/// Recommend the best quantization variant for a given model based on available memory.
///
/// Returns the index of the recommended variant in the input slice.
/// Memory requirements (including KV cache and overhead) are calculated by analyze_model_fit.
#[cfg(test)]
fn make_variant(quant: &str, size_bytes: u64) -> HfQuantVariant {
use super::hf_models::quant_quality_rank;
HfQuantVariant {
quantization: quant.to_string(),
size_bytes,
filename: format!("model-{}.gguf", quant),
download_url: String::new(),
description: "",
quality_rank: quant_quality_rank(quant),
}
}
pub fn recommend_variant(
variants: &[HfQuantVariant],
available_memory_bytes: u64,
@@ -472,197 +403,32 @@ pub fn recommend_variant(
#[cfg(test)]
mod tests {
use super::*;
use test_case::test_case;
#[test]
fn test_quant_bpp() {
assert!((quant_bpp("Q4_K_M") - 0.58).abs() < 0.01);
assert!((quant_bpp("Q8_0") - 1.05).abs() < 0.01);
assert!((quant_bpp("F16") - 2.0).abs() < 0.01);
assert!((quant_bpp("unknown") - 0.58).abs() < 0.01); // Default
fn standard_variants() -> Vec<HfQuantVariant> {
vec![
make_variant("Q2_K", 3_000_000_000),
make_variant("Q4_K_M", 5_000_000_000),
make_variant("Q6_K", 7_000_000_000),
make_variant("Q8_0", 9_000_000_000),
]
}
#[test]
fn test_estimate_params_billion() {
// A 4GB Q4_K_M file should be ~7B params
let params = estimate_params_billion(4_000_000_000, "Q4_K_M");
assert!(params > 6.0 && params < 8.0);
// A 1GB Q4_K_M file should be ~1.7B params
let params = estimate_params_billion(1_000_000_000, "Q4_K_M");
assert!(params > 1.5 && params < 2.0);
}
#[test]
fn test_estimate_memory_gb() {
// 4GB file with 8K context
let mem = estimate_memory_gb(4_000_000_000, "Q4_K_M", 8192);
// Should be slightly more than 4GB (file + KV cache + overhead)
assert!(mem > 4.0 && mem < 5.0);
}
#[test]
fn test_analyze_model_fit() {
// 4GB model with 8GB available
let fit = analyze_model_fit(4_000_000_000, "Q4_K_M", 19, 8192, 8_000_000_000);
assert!(fit.fits);
assert!(fit.headroom_percent > 0.0);
// 4GB model with 3GB available - shouldn't fit
let fit = analyze_model_fit(4_000_000_000, "Q4_K_M", 19, 8192, 3_000_000_000);
assert!(!fit.fits);
assert!(fit.headroom_percent < 0.0);
}
#[test]
fn test_model_fit_score() {
// High quality, good headroom
let fit1 = ModelFit {
fits: true,
estimated_memory_gb: 4.0,
available_memory_gb: 8.0,
headroom_percent: 50.0,
quality_rank: 25,
params_billion: 7.0,
};
// Lower quality, same headroom
let fit2 = ModelFit {
fits: true,
estimated_memory_gb: 4.0,
available_memory_gb: 8.0,
headroom_percent: 50.0,
quality_rank: 10,
params_billion: 7.0,
};
assert!(fit1.score() > fit2.score());
// Doesn't fit
let fit3 = ModelFit {
fits: false,
estimated_memory_gb: 10.0,
available_memory_gb: 8.0,
headroom_percent: -25.0,
quality_rank: 25,
params_billion: 14.0,
};
assert!(fit3.score() < 0.0);
}
#[test]
fn test_recommend_variant() {
let variants = vec![
HfQuantVariant {
quantization: "Q2_K".into(),
size_bytes: 2_000_000_000,
filename: "m-Q2_K.gguf".into(),
download_url: String::new(),
description: "Small",
quality_rank: 7,
},
HfQuantVariant {
quantization: "Q4_K_M".into(),
size_bytes: 4_000_000_000,
filename: "m-Q4_K_M.gguf".into(),
download_url: String::new(),
description: "Medium",
quality_rank: 19,
},
HfQuantVariant {
quantization: "Q8_0".into(),
size_bytes: 8_000_000_000,
filename: "m-Q8_0.gguf".into(),
download_url: String::new(),
description: "Large",
quality_rank: 25,
},
];
// Memory estimation includes KV cache + overhead:
// Q2_K (2GB file): ~2GB + 0.28GB KV + 0.5GB overhead = ~2.78GB needed
// Q4_K_M (4GB file): ~4GB + 0.45GB KV + 0.5GB overhead = ~4.95GB needed
// Q8_0 (8GB file): ~8GB + 0.50GB KV + 0.5GB overhead = ~9.0GB needed
//
// Scoring balances quality (70%) with headroom (30%), so a model with
// comfortable headroom may beat a higher-quality model that barely fits.
// With 6GB available: Q4_K_M fits (4.95 < 6), Q8_0 doesn't (9.0 > 6)
assert_eq!(recommend_variant(&variants, 6_000_000_000), Some(1));
// With 10GB available: Q8_0 fits with ~10% headroom (ideal range ≤20%)
// Q8_0: headroom = 10%, headroom_norm = 1.0, quality = 25/28 = 0.89
// score = 0.7*0.89 + 0.3*1.0 = 0.92
// Q4_K_M: headroom = 50%, headroom_norm = 1.0 - 0.75*0.5 = 0.62, quality = 19/28 = 0.68
// score = 0.7*0.68 + 0.3*0.62 = 0.66
// Q8_0 wins - higher quality with good headroom
assert_eq!(recommend_variant(&variants, 10_000_000_000), Some(2));
// With 16GB available: Q8_0 still wins (higher quality, acceptable headroom)
// Q8_0: headroom = 44%, headroom_norm = 1.0 - 0.75*0.4 = 0.7
// score = 0.7*0.89 + 0.3*0.7 = 0.83
assert_eq!(recommend_variant(&variants, 16_000_000_000), Some(2));
// With 2GB available: nothing fits (Q2_K needs ~2.78GB)
assert_eq!(recommend_variant(&variants, 2_000_000_000), None);
// With 4GB available: Q2_K fits (2.78 < 4)
assert_eq!(recommend_variant(&variants, 4_000_000_000), Some(0));
}
#[test]
fn test_predict_run_mode() {
// Model fits in VRAM
assert_eq!(predict_run_mode(4.0, 8.0, true), RunMode::Gpu);
// Model slightly larger than VRAM - will offload
assert_eq!(predict_run_mode(10.0, 8.0, true), RunMode::CpuOffload);
// Model way too large for VRAM
assert_eq!(predict_run_mode(20.0, 8.0, true), RunMode::CpuOnly);
// No GPU
assert_eq!(predict_run_mode(4.0, 8.0, false), RunMode::CpuOnly);
assert_eq!(predict_run_mode(4.0, 0.0, true), RunMode::CpuOnly);
}
#[test]
fn test_estimate_speed_tps_with_memory() {
// 1B model fitting in 16GB VRAM should be fast (full GPU)
let tps = estimate_speed_tps_with_memory(1.0, "Q4_K_M", true, 16.0);
assert!(tps > 100.0, "1B on GPU should be >100 tps, got {}", tps);
// 7B model fitting in 16GB VRAM
let tps = estimate_speed_tps_with_memory(7.0, "Q4_K_M", true, 16.0);
assert!(tps > 20.0, "7B on GPU should be >20 tps, got {}", tps);
// 7B model on small 4GB VRAM - will need offload, slower
let tps_offload = estimate_speed_tps_with_memory(7.0, "Q4_K_M", true, 4.0);
let tps_full = estimate_speed_tps_with_memory(7.0, "Q4_K_M", true, 16.0);
assert!(
tps_offload < tps_full,
"Offload should be slower: {} vs {}",
tps_offload,
tps_full
// available memory → expected quant
#[test_case(4_500_000_000, Some("Q2_K") ; "tight memory only fits smallest")]
#[test_case(7_000_000_000, Some("Q4_K_M") ; "moderate memory picks mid quality")]
#[test_case(9_000_000_000, Some("Q6_K") ; "good memory picks high quality")]
#[test_case(24_000_000_000, Some("Q8_0") ; "plenty of memory picks best quality")]
#[test_case(2_000_000_000, None ; "insufficient memory recommends nothing")]
fn recommend_variant_picks_best_quality_that_fits(
available_bytes: u64,
expected_quant: Option<&str>,
) {
let variants = standard_variants();
let result = recommend_variant(&variants, available_bytes);
assert_eq!(
result.map(|i| variants[i].quantization.as_str()),
expected_quant,
);
// CPU is slower than GPU
let cpu_tps = estimate_speed_tps_with_memory(7.0, "Q4_K_M", false, 0.0);
assert!(tps_full > cpu_tps, "GPU should be faster than CPU");
// Lower quant is faster
let q4_tps = estimate_speed_tps_with_memory(7.0, "Q4_K_M", true, 16.0);
let q8_tps = estimate_speed_tps_with_memory(7.0, "Q8_0", true, 16.0);
assert!(q4_tps > q8_tps, "Q4 should be faster than Q8");
}
#[test]
fn test_speed_tier() {
// New thresholds: Fast >= 20, Medium >= 8, Slow < 8
assert_eq!(SpeedTier::from_tps(100.0), SpeedTier::Fast);
assert_eq!(SpeedTier::from_tps(20.0), SpeedTier::Fast);
assert_eq!(SpeedTier::from_tps(19.9), SpeedTier::Medium);
assert_eq!(SpeedTier::from_tps(8.0), SpeedTier::Medium);
assert_eq!(SpeedTier::from_tps(7.9), SpeedTier::Slow);
assert_eq!(SpeedTier::from_tps(3.0), SpeedTier::Slow);
}
}