mesh_llm_routing/
lib.rs

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

mesh_llm_routing/
lib.rs
