mesh_llm_routing/
lib.rs

1use std::collections::HashMap;
2use std::path::Path;
3use std::sync::Arc;
4use std::sync::atomic::{AtomicU64, Ordering};
5
6pub mod prefix_affinity;
7
8/// Calculate total model size, summing all split files if present.
9/// Split files follow the pattern: name-00001-of-00004.gguf.
10pub fn total_model_bytes(model: &Path) -> u64 {
11    let name = model.to_string_lossy();
12    if let Some(pos) = name.find("-00001-of-") {
13        let of_pos = pos + 10;
14        if let Some(ext_pos) = name[of_pos..].find(".gguf")
15            && let Ok(n_split) = name[of_pos..of_pos + ext_pos].parse::<u32>()
16        {
17            let prefix = &name[..pos + 1];
18            let suffix = &name[of_pos + ext_pos..];
19            let mut total: u64 = 0;
20            for i in 1..=n_split {
21                let split_name = format!("{}{:05}-of-{:05}{}", prefix, i, n_split, suffix);
22                total += std::fs::metadata(&split_name).map(|m| m.len()).unwrap_or(0);
23            }
24            return total;
25        }
26    }
27    std::fs::metadata(model).map(|m| m.len()).unwrap_or(0)
28}
29
30/// The current inference target selected by runtime planning.
31#[derive(Clone, Debug, PartialEq, Eq, Hash)]
32pub enum InferenceTarget {
33    /// No backend running anywhere.
34    None,
35    /// This node serves the model on the given local HTTP port.
36    Local(u16),
37    /// Another node serves the model; proxy via QUIC to this peer.
38    Remote(iroh::EndpointId),
39}
40
41/// Per-model routing table.
42#[derive(Clone, Debug, Default)]
43pub struct ModelTargets {
44    /// model_name -> list of inference targets.
45    pub targets: HashMap<String, Vec<InferenceTarget>>,
46    /// Shared round-robin counter across clones.
47    counter: Arc<AtomicU64>,
48}
49
50impl ModelTargets {
51    /// Get target for a specific model. Round-robins across multiple hosts.
52    pub fn get(&self, model: &str) -> InferenceTarget {
53        match self.targets.get(model) {
54            Some(targets) if !targets.is_empty() => {
55                let idx = self.counter.fetch_add(1, Ordering::Relaxed) as usize % targets.len();
56                targets[idx].clone()
57            }
58            _ => InferenceTarget::None,
59        }
60    }
61
62    /// All candidate targets for a model, preserving their current order.
63    pub fn candidates(&self, model: &str) -> Vec<InferenceTarget> {
64        self.targets.get(model).cloned().unwrap_or_default()
65    }
66
67    /// Round-robin pick from a caller-supplied candidate slice.
68    pub fn pick_from(&self, candidates: &[InferenceTarget]) -> InferenceTarget {
69        if candidates.is_empty() {
70            InferenceTarget::None
71        } else {
72            let idx = self.counter.fetch_add(1, Ordering::Relaxed) as usize % candidates.len();
73            candidates[idx].clone()
74        }
75    }
76
77    /// Sticky pick from a caller-supplied candidate slice.
78    pub fn pick_sticky_from(candidates: &[InferenceTarget], sticky_key: u64) -> InferenceTarget {
79        if candidates.is_empty() {
80            InferenceTarget::None
81        } else {
82            let idx = sticky_key as usize % candidates.len();
83            candidates[idx].clone()
84        }
85    }
86}
mesh_llm_routing/lib.rs

mesh_llm_routing/
lib.rs
