mesh_client/mesh/
types.rs

1use iroh::{EndpointAddr, EndpointId};
2pub use mesh_llm_types::mesh::{
3    DEMAND_TTL_SECS, MAX_SPLIT_RTT_MS, ModelDemand, ModelRuntimeDescriptor, ModelSourceKind,
4    ServedModelDescriptor, ServedModelIdentity, infer_available_model_descriptors,
5    infer_local_served_model_descriptor, infer_served_model_descriptors, merge_demand,
6};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
11pub enum NodeRole {
12    #[default]
13    Worker,
14    Host {
15        http_port: u16,
16    },
17    Client,
18}
19
20#[derive(Debug, Clone)]
21pub struct PeerInfo {
22    pub id: EndpointId,
23    pub addr: EndpointAddr,
24    pub tunnel_port: Option<u16>,
25    pub role: NodeRole,
26    pub models: Vec<String>,
27    pub vram_bytes: u64,
28    pub rtt_ms: Option<u32>,
29    pub model_source: Option<String>,
30    pub serving_models: Vec<String>,
31    pub hosted_models: Vec<String>,
32    pub hosted_models_known: bool,
33    pub available_models: Vec<String>,
34    pub requested_models: Vec<String>,
35    pub last_seen: std::time::Instant,
36    pub version: Option<String>,
37    pub gpu_name: Option<String>,
38    pub hostname: Option<String>,
39    pub is_soc: Option<bool>,
40    pub gpu_vram: Option<String>,
41    pub gpu_bandwidth_gbps: Option<String>,
42    pub available_model_metadata: Vec<crate::proto::node::CompactModelMetadata>,
43    pub experts_summary: Option<crate::proto::node::ExpertsSummary>,
44    pub available_model_sizes: HashMap<String, u64>,
45    pub served_model_descriptors: Vec<ServedModelDescriptor>,
46    pub served_model_runtime: Vec<ModelRuntimeDescriptor>,
47    pub owner_id: Option<String>,
48}
49
50impl PeerInfo {
51    pub fn is_assigned_model(&self, model: &str) -> bool {
52        self.serving_models.iter().any(|m| m == model)
53    }
54
55    pub fn routable_models(&self) -> Vec<String> {
56        if self.hosted_models_known {
57            self.hosted_models.clone()
58        } else {
59            self.serving_models.clone()
60        }
61    }
62
63    pub fn routes_model(&self, model: &str) -> bool {
64        if self.hosted_models_known {
65            self.hosted_models.iter().any(|m| m == model)
66        } else {
67            self.is_assigned_model(model)
68        }
69    }
70
71    pub fn advertised_context_length(&self, model: &str) -> Option<u32> {
72        self.served_model_runtime
73            .iter()
74            .find(|r| r.model_name == model)
75            .and_then(ModelRuntimeDescriptor::advertised_context_length)
76    }
77}
78
79#[derive(Debug, Clone)]
80pub struct PeerAnnouncement {
81    pub addr: EndpointAddr,
82    pub role: NodeRole,
83    pub models: Vec<String>,
84    pub vram_bytes: u64,
85    pub model_source: Option<String>,
86    pub serving_models: Vec<String>,
87    pub hosted_models: Option<Vec<String>>,
88    pub available_models: Vec<String>,
89    pub requested_models: Vec<String>,
90    pub version: Option<String>,
91    pub model_demand: HashMap<String, ModelDemand>,
92    pub mesh_id: Option<String>,
93    pub gpu_name: Option<String>,
94    pub hostname: Option<String>,
95    pub is_soc: Option<bool>,
96    pub gpu_vram: Option<String>,
97    pub gpu_bandwidth_gbps: Option<String>,
98    pub available_model_metadata: Vec<crate::proto::node::CompactModelMetadata>,
99    pub experts_summary: Option<crate::proto::node::ExpertsSummary>,
100    pub available_model_sizes: HashMap<String, u64>,
101    pub served_model_descriptors: Vec<ServedModelDescriptor>,
102    pub served_model_runtime: Vec<ModelRuntimeDescriptor>,
103    pub owner_id: Option<String>,
104}
105
106pub fn should_be_host_for_model(my_id: EndpointId, my_vram: u64, model_peers: &[PeerInfo]) -> bool {
107    for peer in model_peers {
108        if matches!(peer.role, NodeRole::Client) {
109            continue;
110        }
111        if peer.vram_bytes > my_vram {
112            return false;
113        }
114        if peer.vram_bytes == my_vram && peer.id > my_id {
115            return false;
116        }
117    }
118    true
119}
mesh_client/mesh/types.rs

mesh_client/mesh/
types.rs
