clawft_kernel/infer_proxy/
table.rs

1//! The placement table: which node serves `infer.<role>` right now.
2//!
3//! Local instances come from the card-18 adapter ([`InferRuntime`]:
4//! `reconcile`, `health`, `endpoint`); remote ones from `infer.<role>`
5//! [`ServiceAdvertisement`]s of admitted peers. Resolution prefers the
6//! local instance (sticky: warm KV), then the first admitted peer in node-id
7//! order. Every change bumps a generation so caches (the `ProviderRouter`
8//! TTL cache) can be invalidated.
9
10use std::collections::{BTreeMap, HashMap, HashSet};
11use std::sync::{Arc, RwLock};
12use std::time::{Duration, Instant};
13
14use tokio::sync::watch;
15
16use super::types::{MeshDialer, ProxyAudit, ProxyError, Target};
17use super::upstream::parse_loopback_base;
18use super::wire::MAX_ROLE;
19use crate::ipc::GlobalPid;
20use crate::mesh_service_adv::ServiceAdvertisement;
21use crate::workload_pkg::manifest::valid_token;
22use crate::workload_runtime::infer::{Health, InferRuntime, Reconcile};
23use crate::workload_runtime::types::InstanceHandle;
24
25/// Prefix of the advertised service name.
26pub const SERVICE_PREFIX: &str = "infer.";
27
28#[derive(Debug, Clone)]
29struct LocalEntry {
30    base: String,
31    model: Option<String>,
32    api: String,
33    runtime: String,
34}
35
36/// How long a received advertisement counts without being refreshed.
37pub const DEFAULT_ADVERT_TTL: Duration = Duration::from_secs(60);
38
39struct RemoteEntry {
40    ad: ServiceAdvertisement,
41    /// When *we* received it. The peer's own timestamp is never trusted
42    /// for freshness.
43    received: Instant,
44}
45
46/// A local instance a peer's request may reach.
47#[derive(Debug, Clone, PartialEq, Eq)]
48pub struct MeshLocal {
49    /// Verified loopback base URL.
50    pub base: String,
51    /// The model the server knows the instance by; requests are pinned to it.
52    pub model: Option<String>,
53    /// Server software (`LlamaCpp`, `MlxLm`, `Ollama`).
54    pub runtime: String,
55}
56
57#[derive(Default)]
58struct Inner {
59    local: HashMap<String, LocalEntry>,
60    remote: HashMap<String, BTreeMap<String, RemoteEntry>>,
61    /// Consumer-side: nodes allowed to serve each role. Default deny.
62    allow: HashMap<String, HashSet<String>>,
63    /// Serving-side: peers this node serves each role to. Default deny.
64    serve_allow: HashMap<String, HashSet<String>>,
65    exposed: HashSet<String>,
66    proxy_ports: HashMap<String, u16>,
67}
68
69/// What [`PlacementTable::sync_local`] found.
70#[derive(Debug, Clone, PartialEq)]
71pub struct SyncOutcome {
72    /// Whether the role is now served locally.
73    pub registered: bool,
74    /// What reconcile did (restart, give up, ...), for the caller to chain.
75    pub reconcile: Option<Reconcile>,
76    /// Health seen.
77    pub health: Option<Health>,
78}
79
80/// Role to serving node.
81pub struct PlacementTable {
82    node_id: String,
83    dialer: Option<Arc<dyn MeshDialer>>,
84    audit: Option<Arc<dyn ProxyAudit>>,
85    inner: RwLock<Inner>,
86    gen_tx: watch::Sender<u64>,
87    advert_ttl: Duration,
88}
89
90fn valid_role(role: &str) -> bool {
91    valid_token(role, MAX_ROLE) && !role.contains('/')
92}
93
94impl PlacementTable {
95    /// Table for node `node_id`. Without a dialer there is no mesh: only
96    /// local instances resolve.
97    pub fn new(
98        node_id: impl Into<String>,
99        dialer: Option<Arc<dyn MeshDialer>>,
100        audit: Option<Arc<dyn ProxyAudit>>,
101    ) -> Self {
102        Self {
103            node_id: node_id.into(),
104            dialer,
105            audit,
106            inner: RwLock::new(Inner::default()),
107            gen_tx: watch::channel(0).0,
108            advert_ttl: DEFAULT_ADVERT_TTL,
109        }
110    }
111
112    /// Lifetime of a received advertisement (refreshed by each repeat).
113    pub fn with_advert_ttl(mut self, ttl: Duration) -> Self {
114        self.advert_ttl = ttl;
115        self
116    }
117
118    /// Allow or stop allowing `node` to serve `role` to this node. This is
119    /// the consumer-side counterpart of [`expose_to_mesh`](Self::expose_to_mesh)
120    /// and the governed decision to send this node's prompts to that node:
121    /// without it no remote instance resolves, however it advertises.
122    pub fn allow_remote_node(&self, role: &str, node: &str, on: bool) {
123        let changed = {
124            let mut g = self.inner.write().unwrap();
125            let set = g.allow.entry(role.to_string()).or_default();
126            if on {
127                set.insert(node.to_string())
128            } else {
129                set.remove(node)
130            }
131        };
132        if changed {
133            self.audit(
134                "infer.remote.allow",
135                serde_json::json!({"role": role, "node": node, "allowed": on, "consumer": self.node_id}),
136            );
137            self.bump();
138        }
139    }
140
141    /// This node's id.
142    pub fn node_id(&self) -> &str {
143        &self.node_id
144    }
145
146    /// The mesh dialer, when there is one.
147    pub fn dialer(&self) -> Option<&Arc<dyn MeshDialer>> {
148        self.dialer.as_ref()
149    }
150
151    /// Changes so far. A router cache keyed on this is never stale.
152    pub fn generation(&self) -> u64 {
153        *self.gen_tx.borrow()
154    }
155
156    /// Wakes on every change (`invalidate_all_placement` on it).
157    pub fn subscribe(&self) -> watch::Receiver<u64> {
158        self.gen_tx.subscribe()
159    }
160
161    fn bump(&self) {
162        self.gen_tx.send_modify(|g| *g += 1);
163    }
164
165    fn audit(&self, kind: &str, payload: serde_json::Value) {
166        if let Some(a) = &self.audit {
167            a.record(kind, payload);
168        }
169    }
170
171    /// Serve `role` from the loopback server at `base`. Refused when the
172    /// endpoint is not `http://<loopback ip>:<port>`.
173    pub fn register_local(
174        &self,
175        role: &str,
176        base: &str,
177        model: Option<String>,
178        api: &str,
179        runtime: &str,
180    ) -> Result<(), ProxyError> {
181        if !valid_role(role) {
182            return Err(ProxyError::BadRequest(format!("bad role '{role}'")));
183        }
184        parse_loopback_base(base)?;
185        let entry = LocalEntry {
186            base: base.trim_end_matches('/').to_string(),
187            model,
188            api: api.to_string(),
189            runtime: runtime.to_string(),
190        };
191        let changed = {
192            let mut g = self.inner.write().unwrap();
193            let prev = g.local.insert(role.to_string(), entry.clone());
194            prev.is_none_or(|p| p.base != entry.base)
195        };
196        if changed {
197            self.bump();
198        }
199        Ok(())
200    }
201
202    /// Stop serving `role` locally.
203    pub fn deregister_local(&self, role: &str) {
204        let had = self.inner.write().unwrap().local.remove(role).is_some();
205        if had {
206            self.bump();
207        }
208    }
209
210    /// Re-evaluate `role` against its adapter instance: reconcile (a dead
211    /// managed server is restarted within its policy), probe, and serve it
212    /// only while it answers.
213    pub async fn sync_local(
214        &self,
215        role: &str,
216        rt: &InferRuntime,
217        h: &InstanceHandle,
218    ) -> SyncOutcome {
219        let reconcile = rt.reconcile(h).await.ok();
220        let health = rt.health(h).await.ok().map(|r| r.health);
221        let endpoint = rt.endpoint(h).await;
222        let spec = rt.spec_of(h).await;
223        let served = rt.served_model(h).await;
224        let up = matches!(health, Some(Health::Up));
225        let mut registered = false;
226        if let (true, Some(base), Some(spec)) = (up, endpoint, spec) {
227            registered = self
228                .register_local(
229                    role,
230                    &base,
231                    served,
232                    &spec.api,
233                    &format!("{:?}", spec.runtime),
234                )
235                .is_ok();
236        }
237        if !registered {
238            self.deregister_local(role);
239        }
240        SyncOutcome {
241            registered,
242            reconcile,
243            health,
244        }
245    }
246
247    /// Allow or stop serving `role` to mesh peers. Beyond loopback this is
248    /// a governed decision (ADR-101 section 7): the caller decides, this
249    /// records it.
250    pub fn expose_to_mesh(&self, role: &str, on: bool) {
251        let changed = {
252            let mut g = self.inner.write().unwrap();
253            if on {
254                g.exposed.insert(role.to_string())
255            } else {
256                g.exposed.remove(role)
257            }
258        };
259        if changed {
260            self.audit(
261                "infer.mesh.expose",
262                serde_json::json!({"role": role, "node_id": self.node_id, "exposed": on}),
263            );
264            self.bump();
265        }
266    }
267
268    /// The loopback proxy port registered for `role` (tests, status).
269    pub fn proxy_port(&self, role: &str) -> Option<u16> {
270        self.inner.read().unwrap().proxy_ports.get(role).copied()
271    }
272
273    /// The loopback proxy port that fronts `role` on this node.
274    pub fn set_proxy_port(&self, role: &str, port: u16) {
275        self.inner
276            .write()
277            .unwrap()
278            .proxy_ports
279            .insert(role.to_string(), port);
280    }
281
282    /// Allow or stop allowing `peer` to use this node's instance of `role`
283    /// (the serving-side counterpart of
284    /// [`allow_remote_node`](Self::allow_remote_node)). Exposure
285    /// ([`expose_to_mesh`](Self::expose_to_mesh)) says the role may leave
286    /// loopback at all; this says to whom. Default deny, audited.
287    pub fn allow_mesh_peer(&self, role: &str, peer: &str, on: bool) {
288        let changed = {
289            let mut g = self.inner.write().unwrap();
290            let set = g.serve_allow.entry(role.to_string()).or_default();
291            if on {
292                set.insert(peer.to_string())
293            } else {
294                set.remove(peer)
295            }
296        };
297        if changed {
298            self.audit(
299                "infer.mesh.allow",
300                serde_json::json!({"role": role, "peer": peer, "allowed": on, "server": self.node_id}),
301            );
302        }
303    }
304
305    /// Take an `infer.<role>` advertisement received from `sender`, the
306    /// node id the connection was verified as. Heard only when the
307    /// advertisement is the sender's own (`ad.node_id == sender`), the sender
308    /// is an enforced-admission full node with a trusted scope, and this
309    /// node allowlisted `sender` for the role. A repeat refreshes the TTL.
310    pub fn ingest_advertisement(&self, sender: &str, ad: &ServiceAdvertisement) -> bool {
311        let Some(role) = ad.name.strip_prefix(SERVICE_PREFIX) else {
312            return false;
313        };
314        if !valid_role(role) || sender == self.node_id {
315            return false;
316        }
317        let refuse = |why: &str| {
318            self.audit(
319                "infer.advert.refused",
320                serde_json::json!({"role": role, "sender": sender, "advertised_node": ad.node_id, "why": why}),
321            );
322            false
323        };
324        if ad.node_id != sender {
325            return refuse("advertisement is for another node");
326        }
327        if !self.dialer.as_ref().is_some_and(|d| d.is_admitted(sender)) {
328            return refuse("peer not admitted as a trusted node");
329        }
330        let mut g = self.inner.write().unwrap();
331        if !g.allow.get(role).is_some_and(|s| s.contains(sender)) {
332            drop(g);
333            return refuse("node is not allowed to serve this role");
334        }
335        let per = g.remote.entry(role.to_string()).or_default();
336        let changed = per.get(sender).is_none_or(|e| {
337            e.ad.metadata != ad.metadata || e.ad.methods != ad.methods || e.received.elapsed() > self.advert_ttl
338        });
339        per.insert(
340            sender.to_string(),
341            RemoteEntry {
342                ad: ad.clone(),
343                received: Instant::now(),
344            },
345        );
346        drop(g);
347        if changed {
348            self.bump();
349        }
350        true
351    }
352
353    /// A peer left or was revoked: forget what it advertised.
354    pub fn remove_node(&self, node_id: &str) {
355        let changed = {
356            let mut g = self.inner.write().unwrap();
357            let mut any = false;
358            for per in g.remote.values_mut() {
359                any |= per.remove(node_id).is_some();
360            }
361            g.remote.retain(|_, per| !per.is_empty());
362            any
363        };
364        if changed {
365            self.bump();
366        }
367    }
368
369    /// Where `role` is served: this node first, then an admitted peer.
370    /// Admission is re-checked here, so a revoked peer stops resolving at
371    /// once even before its entry is removed.
372    pub fn resolve(&self, role: &str) -> Option<Target> {
373        let g = self.inner.read().unwrap();
374        if let Some(l) = g.local.get(role) {
375            return Some(Target::Local {
376                base: l.base.clone(),
377            });
378        }
379        let dialer = self.dialer.as_ref()?;
380        let allowed = g.allow.get(role)?;
381        g.remote
382            .get(role)?
383            .iter()
384            .find(|(n, e)| {
385                allowed.contains(*n)
386                    && e.received.elapsed() <= self.advert_ttl
387                    && dialer.is_admitted(n)
388            })
389            .map(|(n, _)| Target::Remote { node_id: n.clone() })
390    }
391
392    /// The local instance of `role`, only when it is exposed to the mesh
393    /// and `peer` is on the role's serve allowlist: what that peer's
394    /// forwarded request may reach. Never a remote target, so a request
395    /// cannot be bounced from node to node.
396    pub fn local_for_peer(&self, role: &str, peer: &str) -> Option<MeshLocal> {
397        let g = self.inner.read().unwrap();
398        (g.exposed.contains(role) && g.serve_allow.get(role).is_some_and(|s| s.contains(peer)))
399            .then(|| {
400                g.local.get(role).map(|l| MeshLocal {
401                    base: l.base.clone(),
402                    model: l.model.clone(),
403                    runtime: l.runtime.clone(),
404                })
405            })
406            .flatten()
407    }
408
409    /// Whether `peer` is on `role`'s serve allowlist.
410    pub fn mesh_peer_allowed(&self, role: &str, peer: &str) -> bool {
411        self.inner
412            .read()
413            .unwrap()
414            .serve_allow
415            .get(role)
416            .is_some_and(|s| s.contains(peer))
417    }
418
419    /// Roles currently exposed to the mesh.
420    pub fn exposed_roles(&self) -> Vec<String> {
421        let mut v: Vec<String> = self.inner.read().unwrap().exposed.iter().cloned().collect();
422        v.sort();
423        v
424    }
425
426    /// Whether `peer` is on any role's serve allowlist.
427    pub fn peer_listed_any(&self, peer: &str) -> bool {
428        self.inner
429            .read()
430            .unwrap()
431            .serve_allow
432            .values()
433            .any(|s| s.contains(peer))
434    }
435
436    /// Base URL for in-process consumers (`PlacementResolver`): the
437    /// instance itself when it is local, this node's loopback proxy when it
438    /// is remote (the proxy speaks the mesh). `None` falls back.
439    pub fn base_url_for_role(&self, role: &str) -> Option<String> {
440        match self.resolve(role)? {
441            Target::Local { base } => Some(format!("{base}/v1")),
442            Target::Remote { .. } => {
443                let port = *self.inner.read().unwrap().proxy_ports.get(role)?;
444                Some(format!("http://127.0.0.1:{port}/v1"))
445            }
446        }
447    }
448
449    /// Base URL of a server for `role` on this node only (`.../v1`), for
450    /// consumers that cannot go through the mesh.
451    pub fn local_base_for_role(&self, role: &str) -> Option<String> {
452        match self.resolve(role)? {
453            Target::Local { base } => Some(format!("{base}/v1")),
454            Target::Remote { .. } => None,
455        }
456    }
457
458    /// `infer.<role>` advertisement for a locally served role that is
459    /// exposed to the mesh.
460    pub fn advertisement(&self, role: &str, now_secs: u64) -> Option<ServiceAdvertisement> {
461        let g = self.inner.read().unwrap();
462        let l = g.local.get(role).filter(|_| g.exposed.contains(role))?;
463        let mut metadata = HashMap::new();
464        metadata.insert("role".to_string(), role.to_string());
465        metadata.insert("node_id".to_string(), self.node_id.clone());
466        metadata.insert("api".to_string(), l.api.clone());
467        metadata.insert("runtime".to_string(), l.runtime.clone());
468        metadata.insert("load".to_string(), "0".to_string());
469        if let Some(m) = &l.model {
470            metadata.insert("model".to_string(), m.clone());
471        }
472        Some(ServiceAdvertisement {
473            name: format!("{SERVICE_PREFIX}{role}"),
474            methods: ["chat.completions", "completions", "embeddings", "models"]
475                .map(String::from)
476                .to_vec(),
477            node_id: self.node_id.clone(),
478            global_pid: GlobalPid::local(0, &self.node_id),
479            version: "1".to_string(),
480            contract_hash: None,
481            metadata,
482            last_updated: now_secs,
483        })
484    }
485
486    /// Advertisements of every exposed local role.
487    pub fn advertisements(&self, now_secs: u64) -> Vec<ServiceAdvertisement> {
488        let roles: Vec<String> = self.inner.read().unwrap().local.keys().cloned().collect();
489        roles
490            .iter()
491            .filter_map(|r| self.advertisement(r, now_secs))
492            .collect()
493    }
494}
clawft_kernel/infer_proxy/table.rs

clawft_kernel/infer_proxy/
table.rs
