1use std::collections::BTreeMap;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::sync::{Arc, Mutex};
6
7use superfluid_linkf::Liveness;
8use superfluid_proto::linkf::Lease;
9
10use crate::fleet::{FleetGeneration, FleetHead};
11use crate::wal::GenParams;
12use crate::DaemonError;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum PlacementPolicy {
16 LoadAware,
17 LeastLoaded,
18 RoundRobin,
19}
20
21#[derive(Clone, Copy, Default)]
22struct NodeLoad {
23 kv_used: u64,
24 kv_total: u64,
25 lanes_active: u32,
26 queue_depth: u32,
27 decode_tokens_per_s: u32,
28 prefill_tokens_per_s: u32,
29 at: Option<std::time::Instant>,
30}
31
32impl NodeLoad {
33 const TTL: std::time::Duration = std::time::Duration::from_secs(10);
34
35 fn fresh(&self) -> bool {
36 self.at.is_some_and(|t| t.elapsed() < Self::TTL)
37 }
38
39 fn pool_pct(&self) -> Option<u64> {
40 (self.fresh() && self.kv_total > 0).then(|| self.kv_used * 100 / self.kv_total)
41 }
42
43 fn eta_ms(&self, inflight: u32, prompt_tokens: u64, expected_out: u64, typical: (u64, u64)) -> u64 {
46 let ahead = if self.fresh() {
47 self.lanes_active as u64 + self.queue_depth as u64 + inflight as u64
48 } else {
49 inflight as u64
50 };
51 const NOMINAL_DECODE: u64 = 20;
52 const NOMINAL_PREFILL: u64 = 200;
53 let or = |reported: u32, typical: u64, nominal: u64| match (reported, typical) {
56 (0, 0) => nominal,
57 (0, t) => t,
58 (r, _) => r as u64,
59 };
60 let decode = or(self.decode_tokens_per_s, typical.0, NOMINAL_DECODE);
61 let prefill = or(self.prefill_tokens_per_s, typical.1, NOMINAL_PREFILL);
62 let queued = ahead.saturating_mul(expected_out).saturating_mul(1_000) / decode;
63 let mine = prompt_tokens.saturating_mul(1_000) / prefill
64 + expected_out.saturating_mul(1_000) / decode;
65 queued.saturating_add(mine)
66 }
67}
68
69fn watch(node: &Arc<NodeConn>, head: &mut FleetHead) {
71 let node = Arc::downgrade(node);
72 head.on_status(Arc::new(move |cap| {
73 if let Some(node) = node.upgrade() {
74 *node.live.lock().expect("live") = FleetManager::load_from(cap);
75 }
76 }));
77}
78
79#[derive(Debug, Clone)]
80pub struct NodeSpec {
81 pub addr: String,
82 pub auth: Vec<u8>,
83}
84
85struct NodeConn {
86 spec: NodeSpec,
87 conns: Vec<Mutex<Option<FleetHead>>>,
88 control: Mutex<Option<FleetHead>>,
92 open: std::sync::atomic::AtomicUsize,
94 cursor: std::sync::atomic::AtomicUsize,
95 decoded: std::sync::atomic::AtomicU64,
97 live: Mutex<NodeLoad>,
100}
101
102impl NodeMeta {
103 fn new() -> NodeMeta {
104 NodeMeta {
105 models: Vec::new(),
106 identity: String::new(),
107 peer: None,
108 down_since: None,
109 last_retry: None,
110 learned: false,
111 max_context_tokens: 0,
112 load: NodeLoad::default(),
113 tokenizer_refused: Vec::new(),
114 }
115 }
116
117 fn up(&self) -> bool {
118 self.down_since.is_none()
119 }
120}
121
122impl NodeConn {
123 fn checkout_next(&self) -> (usize, std::sync::MutexGuard<'_, Option<FleetHead>>) {
124 let open = self.open.load(std::sync::atomic::Ordering::Relaxed).clamp(1, self.conns.len());
125 let i = self.cursor.fetch_add(1, std::sync::atomic::Ordering::Relaxed) % open;
126 (i, self.conns[i].lock().expect("node conn"))
127 }
128
129 fn checkout_free(&self, deadline: Option<std::time::Instant>) -> Option<(usize, std::sync::MutexGuard<'_, Option<FleetHead>>)> {
134 let open = self.open.load(std::sync::atomic::Ordering::Relaxed).clamp(1, self.conns.len());
135 let start = self.cursor.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
136 let dials = !self.spec.addr.starts_with(JOINED);
137 let mut pause = std::time::Duration::from_millis(1);
138 loop {
139 let mut busy = false;
140 for k in 0..open {
141 let i = (start + k) % open;
142 match self.conns[i].try_lock() {
143 Ok(guard) if guard.is_some() || dials => return Some((i, guard)),
144 Ok(_) | Err(std::sync::TryLockError::Poisoned(_)) => {}
145 Err(std::sync::TryLockError::WouldBlock) => busy = true,
146 }
147 }
148 if !busy {
150 let i = start % open;
151 return Some((i, self.conns[i].lock().expect("node conn")));
152 }
153 if deadline.is_some_and(|at| std::time::Instant::now() >= at) {
154 return None;
155 }
156 std::thread::sleep(pause);
157 pause = (pause * 2).min(std::time::Duration::from_millis(10));
158 }
159 }
160
161 fn checkout_idle(&self) -> std::sync::MutexGuard<'_, Option<FleetHead>> {
165 let open = self.open.load(std::sync::atomic::Ordering::Relaxed).clamp(1, self.conns.len());
166 for slot in &self.conns[..open] {
167 if let Ok(guard) = slot.try_lock() {
168 return guard;
169 }
170 }
171 let control = self.control.lock().expect("node conn");
172 if control.is_some() || !self.spec.addr.starts_with(JOINED) {
173 return control;
174 }
175 drop(control);
176 self.checkout_next().1
177 }
178
179 fn checkout_at(&self, i: usize) -> std::sync::MutexGuard<'_, Option<FleetHead>> {
180 self.conns[i.min(self.conns.len() - 1)].lock().expect("node conn")
181 }
182
183 fn checkout_at_by(
185 &self,
186 i: usize,
187 deadline: Option<std::time::Instant>,
188 ) -> Option<std::sync::MutexGuard<'_, Option<FleetHead>>> {
189 let Some(at) = deadline else { return Some(self.checkout_at(i)) };
190 let slot = &self.conns[i.min(self.conns.len() - 1)];
191 let mut pause = std::time::Duration::from_millis(1);
192 loop {
193 match slot.try_lock() {
194 Ok(guard) => return Some(guard),
195 Err(std::sync::TryLockError::Poisoned(_)) => return Some(self.checkout_at(i)),
196 Err(std::sync::TryLockError::WouldBlock) => {}
197 }
198 let left = at.saturating_duration_since(std::time::Instant::now());
199 if left.is_zero() {
200 return None;
201 }
202 std::thread::sleep(pause.min(left));
203 pause = (pause * 2).min(std::time::Duration::from_millis(10));
204 }
205 }
206
207 fn reset(&self) {
208 for slot in self.conns.iter().chain([&self.control]) {
209 if let Ok(mut g) = slot.try_lock() {
210 *g = None;
211 }
212 }
213 }
214}
215
216#[derive(Clone)]
217struct NodeMeta {
218 models: Vec<String>,
219 identity: String,
220 peer: Option<std::net::IpAddr>,
222 down_since: Option<std::time::Instant>,
224 last_retry: Option<std::time::Instant>,
225 learned: bool,
226 max_context_tokens: u64,
227 load: NodeLoad,
228 tokenizer_refused: Vec<String>,
230}
231
232#[derive(Debug, Clone)]
233struct Placed {
234 node: usize,
235 params: GenParams,
236 qos_class: u8,
237 codec: String,
238 model: String,
239 tokens: Vec<u32>,
240 slot: usize,
241 epoch: u64,
242 parked: bool,
243}
244
245struct Shared {
246 meta: Vec<NodeMeta>,
247 forgets: Vec<(usize, usize, u64, u64)>,
249 sessions: BTreeMap<u64, Placed>,
250 rr_cursor: usize,
251 cancels: BTreeMap<u64, Arc<AtomicBool>>,
252 inflight: Vec<u32>,
253 claims: BTreeMap<u64, usize>,
254}
255
256pub struct FleetManager {
257 nodes: std::sync::RwLock<Vec<Arc<NodeConn>>>,
258 conns_by_lanes: bool,
259 retry_every: std::time::Duration,
260 liveness: Liveness,
261 tokenizers: std::sync::RwLock<std::collections::HashMap<String, u64>>,
264 shared: Mutex<Shared>,
265 warm: Mutex<crate::fleet_warm::Warm>,
267 policy: PlacementPolicy,
268 lease: Lease,
269 pool_high_pct: u32,
270}
271
272struct InFlight<'a> {
273 shared: &'a Mutex<Shared>,
274 node: usize,
275 session: u64,
276}
277
278impl<'a> InFlight<'a> {
279 fn claim(shared: &'a Mutex<Shared>, node: usize, session: u64) {
280 let mut sh = shared.lock().expect("shared");
281 if let Some(prev) = sh.claims.insert(session, node) {
282 if prev != node {
283 if let Some(n) = sh.inflight.get_mut(prev) {
284 *n = n.saturating_sub(1);
285 }
286 } else {
287 return;
288 }
289 }
290 if let Some(n) = sh.inflight.get_mut(node) {
291 *n = n.saturating_add(1);
292 }
293 }
294
295 fn release(shared: &Mutex<Shared>, session: u64) {
296 let mut sh = shared.lock().expect("shared");
297 if let Some(node) = sh.claims.remove(&session) {
298 if let Some(n) = sh.inflight.get_mut(node) {
299 *n = n.saturating_sub(1);
300 }
301 }
302 }
303
304 fn guard(shared: &'a Mutex<Shared>, node: usize, session: u64) -> InFlight<'a> {
305 Self::claim(shared, node, session);
306 InFlight { shared, node, session }
307 }
308}
309
310impl Drop for InFlight<'_> {
311 fn drop(&mut self) {
312 let _ = self.node;
313 Self::release(self.shared, self.session);
314 }
315}
316
317#[derive(Clone, Copy, Default)]
318pub struct RequestShape {
319 pub prompt_tokens: u64,
320 pub expected_out: u64,
321 pub warm: Option<usize>,
324}
325
326pub const DEFAULT_POOL_HIGH_PCT: u32 = 90;
327
328pub const DEFAULT_CONNS_PER_NODE: usize = 0;
330
331const MAX_CONNS_PER_NODE: usize = 64;
332
333const JOINED: &str = "joined:";
335
336pub const RETRY_DOWN_AFTER: std::time::Duration = std::time::Duration::from_secs(10);
338
339#[derive(Debug, Clone)]
341pub struct NodeStatus {
342 pub name: String,
344 pub addr: String,
346 pub up: bool,
347 pub lanes: usize,
349 pub lanes_active: u32,
351 pub queue_depth: u32,
352 pub kv_blocks_used: u64,
353 pub kv_blocks_total: u64,
354 pub decode_tokens: u64,
355}
356
357#[derive(Debug, Clone)]
358pub struct NodeReport {
359 pub addr: String,
360 pub identity: String,
361 pub models: Vec<String>,
362 pub max_context_tokens: u64,
363 pub max_lanes: u32,
364}
365
366impl FleetManager {
367 pub fn new(nodes: Vec<NodeSpec>, policy: PlacementPolicy, lease: Lease) -> FleetManager {
368 Self::with_pool_high(nodes, policy, lease, DEFAULT_POOL_HIGH_PCT)
369 }
370
371 pub fn with_pool_high(
372 nodes: Vec<NodeSpec>,
373 policy: PlacementPolicy,
374 lease: Lease,
375 pool_high_pct: u32,
376 ) -> FleetManager {
377 Self::with_options(nodes, policy, lease, pool_high_pct, DEFAULT_CONNS_PER_NODE)
378 }
379
380 pub fn with_options(
381 nodes: Vec<NodeSpec>,
382 policy: PlacementPolicy,
383 lease: Lease,
384 pool_high_pct: u32,
385 conns_per_node: usize,
386 ) -> FleetManager {
387 let meta = vec![NodeMeta::new(); nodes.len()];
388 let node_count = nodes.len();
389 let nodes = nodes
390 .into_iter()
391 .map(|spec| {
392 Arc::new(NodeConn {
393 spec,
394 conns: (0..if conns_per_node == 0 { MAX_CONNS_PER_NODE } else { conns_per_node })
395 .map(|_| Mutex::new(None))
396 .collect(),
397 control: Mutex::new(None),
398 open: std::sync::atomic::AtomicUsize::new(conns_per_node.max(1)),
399 cursor: std::sync::atomic::AtomicUsize::new(0),
400 decoded: std::sync::atomic::AtomicU64::new(0),
401 live: Mutex::new(NodeLoad::default()),
402 })
403 })
404 .collect();
405 FleetManager {
406 nodes: std::sync::RwLock::new(nodes),
407 conns_by_lanes: conns_per_node == 0,
408 retry_every: RETRY_DOWN_AFTER,
409 liveness: Liveness::DEFAULT,
410 tokenizers: Default::default(),
411 warm: Default::default(),
412 shared: Mutex::new(Shared {
413 meta,
414 sessions: BTreeMap::new(),
415 forgets: Vec::new(),
416 rr_cursor: 0,
417 inflight: vec![0; node_count],
418 claims: BTreeMap::new(),
419 cancels: BTreeMap::new(),
420 }),
421 policy,
422 lease,
423 pool_high_pct: pool_high_pct.clamp(1, 100),
424 }
425 }
426
427 fn node(&self, i: usize) -> Arc<NodeConn> {
428 Arc::clone(&self.nodes.read().expect("nodes")[i])
429 }
430
431 fn node_count(&self) -> usize {
432 self.nodes.read().expect("nodes").len()
433 }
434
435 pub fn adopt(&self, head: FleetHead, peer: std::net::IpAddr) -> Result<(usize, bool, usize), String> {
443 let identity = head.node().node_identity.clone();
444 let lanes = (head.node().capacity.max_lanes as usize).max(1);
445 let (i, new) = {
446 let mut sh = self.shared.lock().expect("shared");
447 match sh.meta.iter().position(|m| m.identity == identity) {
448 Some(i) => {
449 let node = self.node(i);
450 let connected = node.conns.iter().any(in_use);
451 if connected && sh.meta[i].peer.is_some_and(|p| p != peer) {
452 return Err(format!(
453 "another node named '{identity}' is joined from {}; give this one --identity",
454 sh.meta[i].peer.expect("peer")
455 ));
456 }
457 sh.meta[i].peer = Some(peer);
458 (i, false)
459 }
460 None => {
461 let mut nodes = self.nodes.write().expect("nodes");
462 nodes.push(Arc::new(NodeConn {
463 spec: NodeSpec { addr: format!("{JOINED}{identity}"), auth: Vec::new() },
464 conns: (0..lanes).map(|_| Mutex::new(None)).collect(),
465 control: Mutex::new(None),
466 open: std::sync::atomic::AtomicUsize::new(1),
467 cursor: std::sync::atomic::AtomicUsize::new(0),
468 decoded: std::sync::atomic::AtomicU64::new(0),
469 live: Mutex::new(NodeLoad::default()),
470 }));
471 let mut meta = NodeMeta::new();
474 meta.identity = identity.clone();
475 meta.peer = Some(peer);
476 meta.down_since = Some(std::time::Instant::now());
477 sh.meta.push(meta);
478 sh.inflight.push(0);
479 (nodes.len() - 1, true)
480 }
481 }
482 };
483 self.sync_meta(i, &head);
484 let node = self.node(i);
485 let mut head = head;
486 head.keep_alive(self.liveness);
487 watch(&node, &mut head);
488 let slot = node.conns.iter().position(|c| !in_use(c));
489 let has_control = node.control.lock().expect("node conn").is_some();
490 match slot {
491 Some(slot) => {
492 if !has_control {
493 let _ = head.ask_control();
494 }
495 *node.checkout_at(slot) = Some(head);
496 }
497 None if !has_control => *node.control.lock().expect("node conn") = Some(head),
499 None => *node.checkout_at(0) = Some(head),
500 }
501 let open = node.conns.iter().filter(|c| in_use(c)).count();
502 node.open.store(open.max(1), std::sync::atomic::Ordering::Relaxed);
503 if new {
504 self.shared.lock().expect("shared").meta[i].down_since = None;
506 } else {
507 self.mark_up(i);
508 }
509 Ok((i, new, open))
510 }
511
512 pub fn connections(&self, i: usize) -> usize {
514 self.node(i).open.load(std::sync::atomic::Ordering::Relaxed)
515 }
516
517 pub fn load(&self) -> Vec<usize> {
518 let sh = self.shared.lock().expect("shared");
519 let mut counts = vec![0usize; self.node_count()];
520 for p in sh.sessions.values() {
521 if !p.parked {
522 counts[p.node] += 1;
523 }
524 }
525 counts
526 }
527
528 pub fn status(&self) -> Vec<NodeStatus> {
530 let nodes: Vec<Arc<NodeConn>> = self.nodes.read().expect("nodes").clone();
531 let sh = self.shared.lock().expect("shared");
532 nodes
533 .iter()
534 .zip(&sh.meta)
535 .zip(&sh.inflight)
536 .map(|((node, meta), &inflight)| {
537 let live = *node.live.lock().expect("live");
538 let load = if live.at > meta.load.at { live } else { meta.load };
539 let addr = match (&meta.peer, node.spec.addr.strip_prefix(JOINED)) {
540 (Some(peer), Some(_)) => peer.to_string(),
541 _ => node.spec.addr.clone(),
542 };
543 NodeStatus {
544 name: if meta.identity.is_empty() { addr.clone() } else { meta.identity.clone() },
545 addr,
546 up: meta.up() && meta.learned,
547 lanes: node.open.load(Ordering::Relaxed),
548 lanes_active: inflight,
549 queue_depth: load.queue_depth,
550 kv_blocks_used: load.kv_used,
551 kv_blocks_total: load.kv_total,
552 decode_tokens: node.decoded.load(Ordering::Relaxed),
553 }
554 })
555 .collect()
556 }
557
558 pub fn nodes_up(&self) -> (usize, usize) {
561 let sh = self.shared.lock().expect("shared");
562 (sh.meta.iter().filter(|m| m.up() && m.learned).count(), sh.meta.len())
563 }
564
565 pub fn serving(&self, model: &str) -> (usize, usize) {
566 let sh = self.shared.lock().expect("shared");
567 let up = (0..sh.meta.len())
568 .filter(|&i| sh.meta[i].up() && sh.meta[i].learned && Self::serves(&sh, i, model))
569 .count();
570 (up, sh.meta.len())
571 }
572
573 pub fn node_of(&self, session: u64) -> Option<String> {
574 let sh = self.shared.lock().expect("shared");
575 let node = sh.sessions.get(&session)?.node;
576 Some(sh.meta[node].identity.clone())
577 }
578
579 pub fn with_retry_every(mut self, every: std::time::Duration) -> FleetManager {
582 self.retry_every = every;
583 self
584 }
585
586 pub fn expect_tokenizer(&self, model: &str, digest: u64) {
589 self.tokenizers.write().expect("tokenizers").insert(model.to_string(), digest);
590 }
591
592 pub fn with_liveness(mut self, liveness: Liveness) -> FleetManager {
595 self.liveness = liveness;
596 self
597 }
598
599 pub fn with_warm_log(mut self, path: std::path::PathBuf) -> FleetManager {
602 self.warm = Mutex::new(crate::fleet_warm::Warm::kept_in(path));
603 self
604 }
605
606 fn mark_unhealthy(&self, i: usize, why: &dyn std::fmt::Display) {
607 let mut sh = self.shared.lock().expect("shared");
608 if sh.meta[i].down_since.is_none() {
609 let now = std::time::Instant::now();
610 sh.meta[i].down_since = Some(now);
611 sh.meta[i].last_retry = Some(now);
612 tracing::warn!("fleet: node {} is down ({why}); retrying it in the background", self.label(&sh, i));
613 }
614 }
615
616 fn mark_up(&self, i: usize) {
617 let mut sh = self.shared.lock().expect("shared");
618 if let Some(since) = sh.meta[i].down_since.take() {
619 tracing::info!("fleet: node {} is back after {} s", self.label(&sh, i), since.elapsed().as_secs());
620 }
621 }
622
623 fn label(&self, sh: &Shared, i: usize) -> String {
624 let node = self.node(i);
625 let addr = &node.spec.addr;
626 match sh.meta[i].identity.as_str() {
627 "" => addr.clone(),
628 id => format!("{id} ({addr})"),
629 }
630 }
631
632 fn identity(&self, i: usize) -> String {
633 let sh = self.shared.lock().expect("shared");
634 self.label(&sh, i)
635 }
636
637 fn active_node(&self, session: u64) -> Result<usize, DaemonError> {
638 let sh = self.shared.lock().expect("shared");
639 let p = sh
640 .sessions
641 .get(&session)
642 .ok_or(DaemonError::UnknownSession(session))?;
643 if p.parked {
644 return Err(DaemonError::Config("session is parked; resume it first"));
645 }
646 Ok(p.node)
647 }
648
649 fn ensure_connected(&self, i: usize, guard: &mut Option<FleetHead>) -> Result<(), DaemonError> {
650 if guard.is_some() {
651 return Ok(());
652 }
653 let node = self.node(i);
654 let spec = &node.spec;
655 if spec.addr.starts_with(JOINED) {
656 let e = DaemonError::FleetNode { addr: spec.addr.clone(), why: "left; waiting for it to join again".into() };
657 self.mark_unhealthy(i, &e);
658 return Err(e);
659 }
660 match FleetHead::connect(&spec.addr, spec.auth.clone()) {
661 Ok(mut head) => {
662 head.keep_alive(self.liveness);
663 watch(&node, &mut head);
664 self.sync_meta(i, &head);
665 *guard = Some(head);
666 Ok(())
667 }
668 Err(e) => {
669 self.mark_unhealthy(i, &e);
670 Err(e)
671 }
672 }
673 }
674
675 fn rebind(&self, session: u64, head: &mut FleetHead) -> Result<(), DaemonError> {
678 if head.knows(session) {
679 return Ok(());
680 }
681 let p = self
682 .shared
683 .lock()
684 .expect("shared")
685 .sessions
686 .get(&session)
687 .cloned()
688 .ok_or(DaemonError::UnknownSession(session))?;
689 let epoch = head.assign(session, p.params, p.qos_class, &p.codec, self.lease)?;
690 if !p.tokens.is_empty() {
691 head.append(session, &p.tokens)?;
692 }
693 if let Some(p) = self.shared.lock().expect("shared").sessions.get_mut(&session) {
694 p.epoch = epoch;
695 }
696 Ok(())
697 }
698
699 fn rebind_on(&self, session: u64, guard: &mut Option<FleetHead>) -> Result<(), DaemonError> {
702 let r = self.rebind(session, guard.as_mut().expect("connected"));
703 if let Err(e) = &r {
704 if is_fatal_node(e) {
705 *guard = None;
706 }
707 }
708 r
709 }
710
711 pub fn advertised_context(&self, model: &str) -> Option<u64> {
712 let sh = self.shared.lock().expect("shared");
713 (0..self.node_count())
714 .filter(|&i| sh.meta[i].up() && Self::serves(&sh, i, model))
715 .map(|i| sh.meta[i].max_context_tokens)
716 .filter(|&c| c > 0)
717 .max()
718 }
719
720 fn serves(sh: &Shared, i: usize, model: &str) -> bool {
721 !sh.meta[i].learned || sh.meta[i].models.iter().any(|m| m == model)
722 }
723
724 fn holds(sh: &Shared, i: usize, len: u64) -> bool {
725 let max = sh.meta[i].max_context_tokens;
726 max == 0 || len <= max
727 }
728
729 fn length_refusal(&self, cand: &[usize], model: &str, len: u64) -> Option<DaemonError> {
730 let sh = self.shared.lock().expect("shared");
731 let eligible: Vec<usize> =
732 cand.iter().copied().filter(|&i| Self::serves(&sh, i, model)).collect();
733 if eligible.is_empty() || eligible.iter().any(|&i| Self::holds(&sh, i, len)) {
734 return None;
735 }
736 Some(DaemonError::StreamTooLong {
737 len,
738 max: eligible.iter().map(|&i| sh.meta[i].max_context_tokens).max().unwrap_or(0),
739 })
740 }
741
742 fn sync_meta(&self, i: usize, head: &FleetHead) {
743 let mut sh = self.shared.lock().expect("shared");
744 let cap = &head.node().capacity;
745 let theirs = head.tokenizer();
746 let (models, refused): (Vec<String>, Vec<String>) = {
747 let ours = self.tokenizers.read().expect("tokenizers");
748 head.node().loadable_models.iter().cloned().partition(|m| {
749 ours.get(m).is_none_or(|&d| d == 0 || theirs == 0 || d == theirs)
750 })
751 };
752 if !refused.is_empty() && sh.meta[i].tokenizer_refused != refused {
753 tracing::warn!(
754 "fleet: node {} tokenizes {} otherwise than this head does, so nothing for it is placed there; \
755 serve the same build of the model on both",
756 self.label(&sh, i),
757 refused.join(", ")
758 );
759 }
760 sh.meta[i].tokenizer_refused = refused;
761 sh.meta[i].models = models;
762 sh.meta[i].max_context_tokens = cap.max_context_tokens;
763 sh.meta[i].identity = head.node().node_identity.clone();
764 sh.meta[i].learned = true;
765 sh.meta[i].load = Self::load_from(cap);
766 let node = self.node(i);
767 if self.conns_by_lanes && !node.spec.addr.starts_with(JOINED) {
768 let (lanes, pool) = (cap.max_lanes.max(1) as usize, node.conns.len());
769 let was = node.open.swap(lanes.min(pool), std::sync::atomic::Ordering::Relaxed);
770 if lanes > pool && was != pool {
771 tracing::warn!(
772 "fleet: node {} has {lanes} lanes; the head opens at most {pool} connections to it",
773 self.label(&sh, i)
774 );
775 }
776 }
777 }
778
779 fn load_from(cap: &superfluid_proto::handshake::CapacitySummary) -> NodeLoad {
780 NodeLoad {
781 kv_used: cap.kv_blocks_used,
782 kv_total: cap.kv_blocks_total,
783 lanes_active: cap.lanes_active,
784 queue_depth: cap.queue_depth,
785 decode_tokens_per_s: cap.decode_tokens_per_s,
786 prefill_tokens_per_s: cap.prefill_tokens_per_s,
787 at: Some(std::time::Instant::now()),
788 }
789 }
790
791 fn replacement_shape(tokens: &[u32], expected_out: u64) -> RequestShape {
792 RequestShape {
793 prompt_tokens: tokens.len() as u64,
794 expected_out,
795 warm: None,
796 }
797 }
798
799 fn candidates_claiming(
800 &self,
801 model: &str,
802 claim_for: Option<u64>,
803 shape: RequestShape,
804 ) -> Vec<usize> {
805 let mut sh = self.shared.lock().expect("shared");
806 let mut load = vec![0usize; self.node_count()];
807 for p in sh.sessions.values() {
808 if !p.parked {
809 load[p.node] += 1;
810 }
811 }
812 let mut cand: Vec<usize> = (0..self.node_count())
813 .filter(|&i| sh.meta[i].up() && Self::serves(&sh, i, model))
814 .collect();
815 match self.policy {
816 PlacementPolicy::LoadAware => {
817 let fits: Vec<usize> = cand
818 .iter()
819 .copied()
820 .filter(|&i| Self::holds(&sh, i, shape.prompt_tokens))
821 .collect();
822 let judged: &[usize] = if fits.is_empty() { &cand } else { &fits };
823 let hot: Vec<usize> = judged
824 .iter()
825 .copied()
826 .filter(|&i| {
827 sh.meta[i].load.pool_pct().is_some_and(|p| p >= self.pool_high_pct as u64)
828 })
829 .collect();
830 if hot.len() < judged.len() {
831 cand.retain(|i| !hot.contains(i));
832 }
833 let inflight = sh.inflight.clone();
834 let mean = |rate: fn(&NodeLoad) -> u32| {
835 let known: Vec<u64> =
836 cand.iter().map(|&i| rate(&sh.meta[i].load) as u64).filter(|&r| r > 0).collect();
837 known.iter().sum::<u64>() / (known.len() as u64).max(1)
838 };
839 let typical = (mean(|l| l.decode_tokens_per_s), mean(|l| l.prefill_tokens_per_s));
840 cand.sort_by_key(|&i| {
841 (
842 sh.meta[i].load.eta_ms(inflight[i], shape.prompt_tokens, shape.expected_out, typical),
843 load[i],
844 )
845 });
846 }
847 PlacementPolicy::LeastLoaded => {
848 let inflight = sh.inflight.clone();
849 cand.sort_by_key(|&i| load[i] + inflight[i] as usize);
850 }
851 PlacementPolicy::RoundRobin => {
852 let n = self.node_count().max(1);
853 let c = sh.rr_cursor % n;
854 cand.sort_by_key(|&i| (i + n - c) % n);
855 sh.rr_cursor = sh.rr_cursor.wrapping_add(1);
856 }
857 }
858 if let Some(w) = shape.warm.filter(|_| self.policy != PlacementPolicy::RoundRobin) {
859 if let Some(at) = cand.iter().position(|&i| i == w) {
860 let lanes = self.node(w).open.load(Ordering::Relaxed).max(1);
861 if (sh.inflight[w] as usize) < lanes {
862 let w = cand.remove(at);
863 cand.insert(0, w);
864 }
865 }
866 }
867 if let (Some(session), Some(&winner)) = (claim_for, cand.first()) {
868 Self::claim_locked(&mut sh, winner, session);
869 }
870 cand
871 }
872
873 fn checkout_live(node_conn: &NodeConn, slot: usize) -> std::sync::MutexGuard<'_, Option<FleetHead>> {
876 Self::checkout_live_by(node_conn, slot, None).expect("no deadline to pass")
877 }
878
879 fn checkout_live_by(
881 node_conn: &NodeConn,
882 slot: usize,
883 deadline: Option<std::time::Instant>,
884 ) -> Option<std::sync::MutexGuard<'_, Option<FleetHead>>> {
885 let mut guard = node_conn.checkout_at_by(slot, deadline)?;
886 if guard.as_mut().is_some_and(|h| h.drain_status().is_err()) {
887 *guard = None;
888 }
889 Some(guard)
890 }
891
892 fn slot_of(&self, session: u64) -> usize {
893 self.shared
894 .lock()
895 .expect("shared")
896 .sessions
897 .get(&session)
898 .map(|p| p.slot)
899 .unwrap_or(0)
900 }
901
902 fn claim_locked(sh: &mut Shared, node: usize, session: u64) {
903 if let Some(prev) = sh.claims.insert(session, node) {
904 if prev == node {
905 return;
906 }
907 if let Some(n) = sh.inflight.get_mut(prev) {
908 *n = n.saturating_sub(1);
909 }
910 }
911 if let Some(n) = sh.inflight.get_mut(node) {
912 *n = n.saturating_add(1);
913 }
914 }
915
916 #[allow(clippy::too_many_arguments)]
917 pub fn place(
918 &self,
919 session: u64,
920 params: GenParams,
921 qos_class: u8,
922 codec: &str,
923 model: &str,
924 tokens: &[u32],
925 expected_out: u64,
926 ) -> Result<usize, DaemonError> {
927 let placed = self.place_by(session, params, qos_class, codec, model, tokens, expected_out, None)?;
928 Ok(placed.expect("no deadline to pass"))
929 }
930
931 #[allow(clippy::too_many_arguments)]
935 pub fn place_by(
936 &self,
937 session: u64,
938 params: GenParams,
939 qos_class: u8,
940 codec: &str,
941 model: &str,
942 tokens: &[u32],
943 expected_out: u64,
944 deadline: Option<std::time::Instant>,
945 ) -> Result<Option<usize>, DaemonError> {
946 let warm = self.warm_node(model, tokens);
947 let cand = self.candidates_claiming(
948 model,
949 Some(session),
950 RequestShape {
951 prompt_tokens: tokens.len() as u64,
952 expected_out,
953 warm,
954 },
955 );
956 let release_on_error = |e: DaemonError| {
957 InFlight::release(&self.shared, session);
958 e
959 };
960 if cand.is_empty() {
961 return Err(release_on_error(DaemonError::NoNode(
962 "no healthy node can serve the model",
963 )));
964 }
965 let len = tokens.len() as u64;
966 if let Some(e) = self.length_refusal(&cand, model, len) {
967 return Err(release_on_error(e));
968 }
969 let late = || deadline.is_some_and(|at| std::time::Instant::now() >= at);
970 let mut last_err = None;
971 for &i in &cand {
972 if late() {
973 break;
974 }
975 match self.try_place_by(i, session, params, qos_class, codec, model, tokens, deadline) {
976 Ok(true) => {
977 tracing::info!(
978 "fleet: session {session} ({} prompt tokens) placed on {}",
979 tokens.len(),
980 self.identity(i)
981 );
982 return Ok(Some(i));
983 }
984 Ok(false) => continue,
985 Err(e) => {
986 self.mark_unhealthy(i, &e);
987 self.node(i).reset();
988 last_err = Some(e);
989 }
990 }
991 }
992 if late() {
993 InFlight::release(&self.shared, session);
994 return Ok(None);
995 }
996 if let Some(e) = self.length_refusal(&cand, model, len) {
997 return Err(release_on_error(e));
998 }
999 Err(release_on_error(last_err.unwrap_or(DaemonError::NoNode(
1000 "no reachable node serves the model",
1001 ))))
1002 }
1003
1004 #[allow(clippy::too_many_arguments)]
1005 fn try_place_on(
1006 &self,
1007 i: usize,
1008 session: u64,
1009 params: GenParams,
1010 qos_class: u8,
1011 codec: &str,
1012 model: &str,
1013 tokens: &[u32],
1014 ) -> Result<bool, DaemonError> {
1015 self.try_place_by(i, session, params, qos_class, codec, model, tokens, None)
1016 }
1017
1018 #[allow(clippy::too_many_arguments)]
1021 fn try_place_by(
1022 &self,
1023 i: usize,
1024 session: u64,
1025 params: GenParams,
1026 qos_class: u8,
1027 codec: &str,
1028 model: &str,
1029 tokens: &[u32],
1030 deadline: Option<std::time::Instant>,
1031 ) -> Result<bool, DaemonError> {
1032 let node = self.node(i);
1033 let Some((slot, mut guard)) = node.checkout_free(deadline) else {
1034 return Ok(false);
1035 };
1036 if let Some(h) = guard.as_mut() {
1037 if h.drain_status().is_err() {
1038 *guard = None;
1039 }
1040 }
1041 self.ensure_connected(i, &mut guard)?;
1042 self.flush_forgets(i, slot, guard.as_mut().expect("connected"));
1043 self.sync_meta(i, guard.as_ref().expect("connected"));
1044 {
1045 let sh = self.shared.lock().expect("shared");
1046 if !Self::serves(&sh, i, model) {
1047 return Ok(false);
1048 }
1049 if !Self::holds(&sh, i, tokens.len() as u64) {
1050 return Ok(false);
1051 }
1052 }
1053 let head = guard.as_mut().expect("connected");
1054 let epoch = match head.assign(session, params, qos_class, codec, self.lease) {
1055 Ok(epoch) => epoch,
1056 Err(e) => {
1057 *guard = None;
1058 return Err(e);
1059 }
1060 };
1061 if !tokens.is_empty() {
1062 if let Err(e) = head.append(session, tokens) {
1063 *guard = None;
1064 return Err(e);
1065 }
1066 }
1067 InFlight::claim(&self.shared, i, session);
1068 self.shared.lock().expect("shared").sessions.insert(
1069 session,
1070 Placed {
1071 node: i,
1072 params,
1073 qos_class,
1074 codec: codec.to_string(),
1075 model: model.to_string(),
1076 tokens: tokens.to_vec(),
1077 slot,
1078 epoch,
1079 parked: false,
1080 },
1081 );
1082 Ok(true)
1083 }
1084
1085 pub fn append(&self, session: u64, tokens: &[u32]) -> Result<(), DaemonError> {
1086 let node = self.active_node(session)?;
1087 let len = tokens.len() as u64;
1088 let fits = {
1089 let sh = self.shared.lock().expect("shared");
1090 Self::holds(&sh, node, len)
1091 };
1092 if !fits {
1093 return self.re_place_grown(session, node, tokens);
1094 }
1095 {
1096 let node_conn = self.node(node);
1097 let mut guard = Self::checkout_live(&node_conn, self.slot_of(session));
1098 self.ensure_connected(node, &mut guard)?;
1099 self.rebind_on(session, &mut guard)?;
1100 let head = guard.as_mut().expect("connected");
1101 if let Err(e) = head.append(session, tokens) {
1102 *guard = None;
1103 return Err(e);
1104 }
1105 self.sync_meta(node, head);
1106 }
1107 if let Some(p) = self.shared.lock().expect("shared").sessions.get_mut(&session) {
1108 p.tokens = tokens.to_vec();
1109 }
1110 Ok(())
1111 }
1112
1113 fn re_place_grown(
1114 &self,
1115 session: u64,
1116 from: usize,
1117 tokens: &[u32],
1118 ) -> Result<(), DaemonError> {
1119 let placed = {
1120 let sh = self.shared.lock().expect("shared");
1121 sh.sessions.get(&session).cloned().ok_or(DaemonError::UnknownSession(session))?
1122 };
1123 let cand: Vec<usize> = self
1124 .candidates_claiming(
1125 &placed.model,
1126 Some(session),
1127 Self::replacement_shape(tokens, 0),
1128 )
1129 .into_iter()
1130 .filter(|&i| i != from)
1131 .collect();
1132 let len = tokens.len() as u64;
1133 if let Some(e) = self.length_refusal(&cand, &placed.model, len) {
1134 return Err(e);
1135 }
1136 for &i in &cand {
1137 match self.try_place_on(
1138 i,
1139 session,
1140 placed.params,
1141 placed.qos_class,
1142 &placed.codec,
1143 &placed.model,
1144 tokens,
1145 ) {
1146 Ok(true) => return Ok(()),
1147 Ok(false) => continue,
1148 Err(e) => {
1149 self.mark_unhealthy(i, &e);
1150 self.node(i).reset();
1151 }
1152 }
1153 }
1154 if let Some(e) = self.length_refusal(&cand, &placed.model, len) {
1155 return Err(e);
1156 }
1157 Err(DaemonError::Config(
1158 "no reachable node can hold the session's grown context",
1159 ))
1160 }
1161
1162 pub fn generate(
1163 &self,
1164 session: u64,
1165 max_new_tokens: u64,
1166 max_wall_ms: u64,
1167 ) -> Result<FleetGeneration, DaemonError> {
1168 let node = self.active_node(session)?;
1169
1170 match self.generate_on(node, session, max_new_tokens, max_wall_ms) {
1171 Ok(g) => return Ok(g),
1172 Err(e) if !is_fatal_node(&e) => return Err(e),
1173 Err(e) => self.mark_unhealthy(node, &e),
1174 }
1175 self.node(node).reset();
1176 let placed = self
1177 .shared
1178 .lock()
1179 .expect("shared")
1180 .sessions
1181 .get(&session)
1182 .cloned()
1183 .expect("placed session");
1184 let cand: Vec<usize> = self
1185 .candidates_claiming(
1186 &placed.model,
1187 Some(session),
1188 Self::replacement_shape(&placed.tokens, max_new_tokens),
1189 )
1190 .into_iter()
1191 .filter(|&i| i != node)
1192 .collect();
1193 for i in cand {
1194 match self.try_place_on(
1195 i,
1196 session,
1197 placed.params,
1198 placed.qos_class,
1199 &placed.codec,
1200 &placed.model,
1201 &placed.tokens,
1202 ) {
1203 Ok(true) => {}
1204 Ok(false) => continue,
1205 Err(e) => {
1206 self.mark_unhealthy(i, &e);
1207 self.node(i).reset();
1208 continue;
1209 }
1210 }
1211 tracing::warn!("fleet: session {session} moved to {}", self.identity(i));
1212 match self.generate_on(i, session, max_new_tokens, max_wall_ms) {
1213 Ok(g) => return Ok(g),
1214 Err(e) => self.mark_unhealthy(i, &e),
1215 }
1216 self.node(i).reset();
1217 }
1218 Err(DaemonError::Generation(
1219 "generation failed and no healthy node could take over the session".into(),
1220 ))
1221 }
1222
1223 fn map_refusal(
1224 &self,
1225 node: usize,
1226 session: u64,
1227 r: Result<FleetGeneration, DaemonError>,
1228 ) -> Result<FleetGeneration, DaemonError> {
1229 let g = r?;
1230 if g.finish == crate::nodeagent::FINISH_DECLINED {
1231 return Err(DaemonError::Generation(format!(
1232 "node {} declined session {session} (stale epoch or replayed command)",
1233 self.shared.lock().expect("shared").meta[node].identity
1234 )));
1235 }
1236 if g.finish != crate::nodeagent::FINISH_REFUSED {
1237 return Ok(g);
1238 }
1239 let sh = self.shared.lock().expect("shared");
1240 let len = sh.sessions.get(&session).map(|p| p.tokens.len() as u64).unwrap_or(0);
1241 Err(DaemonError::StreamTooLong {
1242 len,
1243 max: sh.meta[node].max_context_tokens,
1244 })
1245 }
1246
1247 fn generate_on(
1248 &self,
1249 node: usize,
1250 session: u64,
1251 max_new_tokens: u64,
1252 max_wall_ms: u64,
1253 ) -> Result<FleetGeneration, DaemonError> {
1254 let _inflight = InFlight::guard(&self.shared, node, session);
1255 let node_conn = self.node(node);
1256 let mut guard = Self::checkout_live(&node_conn, self.slot_of(session));
1257 self.ensure_connected(node, &mut guard)?;
1258 self.rebind_on(session, &mut guard)?;
1259 let head = guard.as_mut().expect("connected");
1260 let r = head.generate(session, max_new_tokens, max_wall_ms);
1261 if let Ok(g) = &r {
1262 node_conn.decoded.fetch_add(g.tokens.len() as u64, Ordering::Relaxed);
1263 }
1264 if let Err(e) = &r {
1265 if is_fatal_node(e) {
1266 *guard = None;
1267 }
1268 }
1269 if let Some(head) = guard.as_ref() {
1270 self.sync_meta(node, head);
1271 }
1272 drop(guard);
1273 self.map_refusal(node, session, r)
1274 }
1275
1276 pub fn generate_streaming(
1277 &self,
1278 session: u64,
1279 max_new_tokens: u64,
1280 max_wall_ms: u64,
1281 on_tokens: impl FnMut(&[u32]),
1282 ) -> Result<FleetGeneration, DaemonError> {
1283 self.generate_streaming_with(session, max_new_tokens, max_wall_ms, None, &mut |_| {}, on_tokens)
1284 }
1285
1286 pub fn generate_streaming_with(
1289 &self,
1290 session: u64,
1291 max_new_tokens: u64,
1292 max_wall_ms: u64,
1293 opts: Option<&crate::fleet::Ask<'_>>,
1294 on_logprobs: &mut dyn FnMut(&[superfluid_proto::linkf::TokenLogprobMsg]),
1295 on_tokens: impl FnMut(&[u32]),
1296 ) -> Result<FleetGeneration, DaemonError> {
1297 self.generate_streaming_by(session, max_new_tokens, max_wall_ms, None, opts, on_logprobs, on_tokens)
1298 }
1299
1300 #[allow(clippy::too_many_arguments)]
1304 pub fn generate_streaming_by(
1305 &self,
1306 session: u64,
1307 max_new_tokens: u64,
1308 max_wall_ms: u64,
1309 deadline: Option<std::time::Instant>,
1310 opts: Option<&crate::fleet::Ask<'_>>,
1311 on_logprobs: &mut dyn FnMut(&[superfluid_proto::linkf::TokenLogprobMsg]),
1312 mut on_tokens: impl FnMut(&[u32]),
1313 ) -> Result<FleetGeneration, DaemonError> {
1314 let (node, flag) = {
1315 let mut sh = self.shared.lock().expect("shared");
1316 let p = sh
1317 .sessions
1318 .get(&session)
1319 .ok_or(DaemonError::UnknownSession(session))?;
1320 if p.parked {
1321 return Err(DaemonError::Config("session is parked; resume it first"));
1322 }
1323 let node = p.node;
1324 let flag = sh
1325 .cancels
1326 .entry(session)
1327 .or_insert_with(|| Arc::new(AtomicBool::new(false)))
1328 .clone();
1329 flag.store(false, Ordering::Relaxed);
1330 (node, flag)
1331 };
1332
1333 let mut all_tokens: Vec<u32> = Vec::new();
1334
1335 match self.generate_streaming_on(node, session, max_new_tokens, max_wall_ms, deadline, &flag, opts, on_logprobs, |b| {
1336 all_tokens.extend_from_slice(b);
1337 on_tokens(b);
1338 }) {
1339 Ok(g) => return Ok(g),
1340 Err(e) if !is_fatal_node(&e) => return Err(e),
1341 Err(e) => self.mark_unhealthy(node, &e),
1342 }
1343
1344 let placed = self
1345 .shared
1346 .lock()
1347 .expect("shared")
1348 .sessions
1349 .get(&session)
1350 .cloned()
1351 .expect("placed session");
1352 let input = placed.tokens.clone();
1353 let mut tried = vec![node];
1354 loop {
1355 let remaining = max_new_tokens.saturating_sub(all_tokens.len() as u64);
1356 if remaining == 0 {
1357 return Ok(FleetGeneration {
1358 tokens: all_tokens,
1359 finish: superfluid_abi::finish::LENGTH,
1360 expired: false,
1361 });
1362 }
1363 if deadline.is_some_and(|at| std::time::Instant::now() >= at) {
1364 return Ok(FleetGeneration {
1365 tokens: all_tokens,
1366 finish: crate::nodeagent::FINISH_EXPIRED,
1367 expired: true,
1368 });
1369 }
1370 let cand: Vec<usize> = self
1371 .candidates_claiming(
1372 &placed.model,
1373 Some(session),
1374 Self::replacement_shape(&input, remaining),
1375 )
1376 .into_iter()
1377 .filter(|i| !tried.contains(i))
1378 .collect();
1379 if cand.is_empty() {
1380 return Err(DaemonError::Generation(
1381 "streaming generation failed and no healthy node could take over".into(),
1382 ));
1383 }
1384 let i = cand[0];
1385 tried.push(i);
1386 let mut ctx = input.clone();
1387 ctx.extend_from_slice(&all_tokens);
1388 match self.try_place_by(
1389 i,
1390 session,
1391 placed.params,
1392 placed.qos_class,
1393 &placed.codec,
1394 &placed.model,
1395 &ctx,
1396 deadline,
1397 ) {
1398 Ok(true) => {}
1399 Ok(false) => continue,
1400 Err(e) => {
1401 self.mark_unhealthy(i, &e);
1402 self.node(i).reset();
1403 continue;
1404 }
1405 }
1406 tracing::warn!(
1407 "fleet: session {session} moved to {} after {} streamed tokens",
1408 self.identity(i),
1409 all_tokens.len()
1410 );
1411 match self.generate_streaming_on(i, session, remaining, max_wall_ms, deadline, &flag, opts, on_logprobs, |b| {
1412 all_tokens.extend_from_slice(b);
1413 on_tokens(b);
1414 }) {
1415 Ok(g) => {
1416 return Ok(FleetGeneration {
1417 tokens: all_tokens,
1418 finish: g.finish,
1419 expired: g.expired,
1420 })
1421 }
1422 Err(e) if !is_fatal_node(&e) => return Err(e),
1423 Err(e) => self.mark_unhealthy(i, &e),
1424 }
1425 }
1426 }
1427
1428 #[allow(clippy::too_many_arguments)]
1429 fn generate_streaming_on(
1430 &self,
1431 node: usize,
1432 session: u64,
1433 max_new_tokens: u64,
1434 max_wall_ms: u64,
1435 deadline: Option<std::time::Instant>,
1436 flag: &AtomicBool,
1437 opts: Option<&crate::fleet::Ask<'_>>,
1438 on_logprobs: &mut dyn FnMut(&[superfluid_proto::linkf::TokenLogprobMsg]),
1439 mut on_tokens: impl FnMut(&[u32]),
1440 ) -> Result<FleetGeneration, DaemonError> {
1441 let _inflight = InFlight::guard(&self.shared, node, session);
1442 let node_conn = self.node(node);
1443 let unstarted = || FleetGeneration { tokens: Vec::new(), finish: crate::nodeagent::FINISH_EXPIRED, expired: true };
1446 let Some(mut guard) = Self::checkout_live_by(&node_conn, self.slot_of(session), deadline) else {
1447 return Ok(unstarted());
1448 };
1449 let max_wall_ms = match deadline.map(|at| at.saturating_duration_since(std::time::Instant::now())) {
1450 Some(left) if left.is_zero() => return Ok(unstarted()),
1451 Some(left) => max_wall_ms.min(left.as_millis().max(1) as u64),
1452 None => max_wall_ms,
1453 };
1454 self.ensure_connected(node, &mut guard)?;
1455 self.rebind_on(session, &mut guard)?;
1456 let head = guard.as_mut().expect("connected");
1457 let decoded = &node_conn.decoded;
1458 let r = head.generate_streaming_with(session, max_new_tokens, max_wall_ms, flag, opts, on_logprobs, |b: &[u32]| {
1459 decoded.fetch_add(b.len() as u64, Ordering::Relaxed);
1460 on_tokens(b);
1461 });
1462 if let Err(e) = &r {
1463 if is_fatal_node(e) {
1464 *guard = None;
1465 }
1466 }
1467 if let Some(head) = guard.as_ref() {
1468 self.sync_meta(node, head);
1469 }
1470 drop(guard);
1471 self.map_refusal(node, session, r)
1472 }
1473
1474 pub fn state_sync(
1477 &self,
1478 model: &str,
1479 req: superfluid_proto::linkw::StateSyncReqMsg,
1480 blobs: &[crate::fleet::Blob],
1481 ) -> Result<superfluid_proto::linkw::StateSyncOkMsg, DaemonError> {
1482 self.on_a_node(model, |head| head.state_sync(req.clone(), blobs))
1483 }
1484
1485 pub fn state_sync_streaming(
1489 &self,
1490 model: &str,
1491 req: superfluid_proto::linkw::StateSyncReqMsg,
1492 blobs: &[crate::fleet::Blob],
1493 on_segment: &mut dyn FnMut(superfluid_proto::linkw::TranscribeSegmentMsg),
1494 ) -> Result<superfluid_proto::linkw::StateSyncOkMsg, DaemonError> {
1495 let mut told = false;
1496 self.on_a_node(model, |head| {
1497 let answer = head.state_sync_streaming(req.clone(), blobs, &mut |seg| {
1498 told = true;
1499 on_segment(seg);
1500 });
1501 match answer {
1502 Err(e) if told && is_fatal_node(&e) => Err(DaemonError::Generation(format!("the node transcribing was lost: {e}"))),
1503 answer => answer,
1504 }
1505 })
1506 }
1507
1508 pub fn state_sync_all(
1511 &self,
1512 model: &str,
1513 req: superfluid_proto::linkw::StateSyncReqMsg,
1514 ) -> Result<superfluid_proto::linkw::StateSyncOkMsg, DaemonError> {
1515 let cand = self.candidates_claiming(model, None, RequestShape::default());
1516 let mut answer = Err(DaemonError::NoNode("no healthy node can serve the model"));
1517 for i in cand {
1518 answer = Ok(self.state_sync_on(i, req.clone(), &[])?);
1519 }
1520 answer
1521 }
1522
1523 fn state_sync_on(
1524 &self,
1525 i: usize,
1526 req: superfluid_proto::linkw::StateSyncReqMsg,
1527 blobs: &[crate::fleet::Blob],
1528 ) -> Result<superfluid_proto::linkw::StateSyncOkMsg, DaemonError> {
1529 self.on_node(i, |head| head.state_sync(req, blobs))
1530 }
1531
1532 fn on_node<R>(&self, i: usize, call: impl FnOnce(&mut FleetHead) -> Result<R, DaemonError>) -> Result<R, DaemonError> {
1535 let node = self.node(i);
1536 let mut guard = node.checkout_idle();
1537 if guard.as_mut().is_some_and(|h| h.drain_status().is_err()) {
1538 *guard = None;
1539 }
1540 self.ensure_connected(i, &mut guard)?;
1541 let r = call(guard.as_mut().expect("connected"));
1542 if let Err(e) = &r {
1543 if is_fatal_node(e) {
1544 *guard = None;
1545 self.mark_unhealthy(i, e);
1546 }
1547 }
1548 r
1549 }
1550
1551 fn on_a_node<R>(&self, model: &str, mut call: impl FnMut(&mut FleetHead) -> Result<R, DaemonError>) -> Result<R, DaemonError> {
1553 let cand = self.candidates_claiming(model, None, RequestShape::default());
1554 let mut last = None;
1555 for i in cand {
1556 match self.on_node(i, &mut call) {
1557 Ok(r) => return Ok(r),
1558 Err(e) if is_fatal_node(&e) => last = Some(e),
1559 Err(e) => return Err(e),
1560 }
1561 }
1562 Err(last.unwrap_or(DaemonError::NoNode("no healthy node can serve the model")))
1563 }
1564
1565 pub fn remote_tokenizer(self: &Arc<Self>, model: &str) -> Result<crate::fleet::RemoteTokenizer, DaemonError> {
1568 let export = self.on_a_node(model, |head| head.tokenizer_export())?;
1569 if export.vocab.is_empty() {
1570 return Err(DaemonError::Generation(format!(
1571 "no node serving {model} shares a tokenizer; give the head --model a path to the model's files"
1572 )));
1573 }
1574 Ok(crate::fleet::RemoteTokenizer::new(Arc::clone(self), model.to_string(), export))
1575 }
1576
1577 pub fn tokenize(&self, model: &str, text: &superfluid_proto::linkf::TokenizeMsg) -> Result<Option<Vec<u32>>, DaemonError> {
1579 self.on_a_node(model, |head| head.tokenize(text.clone()))
1580 }
1581
1582 pub fn cancel(&self, session: u64) -> Result<(), DaemonError> {
1583 let sh = self.shared.lock().expect("shared");
1584 if !sh.sessions.contains_key(&session) {
1585 return Err(DaemonError::UnknownSession(session));
1586 }
1587 if let Some(flag) = sh.cancels.get(&session) {
1588 flag.store(true, Ordering::Relaxed);
1589 }
1590 Ok(())
1591 }
1592
1593 pub fn park(&self, session: u64) -> Result<(), DaemonError> {
1594 let node = {
1595 let mut sh = self.shared.lock().expect("shared");
1596 let p = sh
1597 .sessions
1598 .get_mut(&session)
1599 .ok_or(DaemonError::UnknownSession(session))?;
1600 if p.parked {
1601 return Ok(());
1602 }
1603 p.parked = true;
1604 (p.node, p.slot, p.epoch)
1605 };
1606 let (node, slot, epoch) = node;
1607 let node_conn = self.node(node);
1608 let mut guard = Self::checkout_live(&node_conn, slot);
1609 if self.ensure_connected(node, &mut guard).is_ok() {
1610 if let Err(e) = guard.as_mut().expect("connected").revoke(session, epoch) {
1611 *guard = None;
1612 self.mark_unhealthy(node, &e);
1613 }
1614 }
1615 Ok(())
1616 }
1617
1618 pub fn resume(&self, session: u64) -> Result<usize, DaemonError> {
1619 let placed = {
1620 let sh = self.shared.lock().expect("shared");
1621 let p = sh
1622 .sessions
1623 .get(&session)
1624 .ok_or(DaemonError::UnknownSession(session))?;
1625 if !p.parked {
1626 return Err(DaemonError::Config("session is not parked"));
1627 }
1628 p.clone()
1629 };
1630 let cand = self.candidates_claiming(
1631 &placed.model,
1632 Some(session),
1633 Self::replacement_shape(&placed.tokens, 0),
1634 );
1635 if cand.is_empty() {
1636 InFlight::release(&self.shared, session);
1637 return Err(DaemonError::NoNode("no healthy node can serve the model"));
1638 }
1639 let mut last_err = None;
1640 for i in cand {
1641 match self.try_place_on(
1642 i,
1643 session,
1644 placed.params,
1645 placed.qos_class,
1646 &placed.codec,
1647 &placed.model,
1648 &placed.tokens,
1649 ) {
1650 Ok(true) => return Ok(i),
1651 Ok(false) => continue,
1652 Err(e) => {
1653 self.mark_unhealthy(i, &e);
1654 self.node(i).reset();
1655 last_err = Some(e);
1656 }
1657 }
1658 }
1659 Err(last_err.unwrap_or(DaemonError::NoNode("no reachable node serves the model")))
1660 }
1661
1662 pub fn finish(&self, session: u64) {
1665 let placed = {
1666 let mut sh = self.shared.lock().expect("shared");
1667 sh.cancels.remove(&session);
1668 sh.sessions.remove(&session)
1669 };
1670 InFlight::release(&self.shared, session);
1671 let Some(p) = placed else { return };
1672 let identity = self.shared.lock().expect("shared").meta[p.node].identity.clone();
1673 if !identity.is_empty() {
1674 self.warm.lock().expect("warm").ran(&p.model, &p.tokens, &identity);
1675 }
1676 self.shared.lock().expect("shared").forgets.push((p.node, p.slot, session, p.epoch));
1677 let node = self.node(p.node);
1678 let lock = node.conns[p.slot].try_lock();
1679 if let Ok(mut guard) = lock {
1680 if let Some(head) = guard.as_mut() {
1681 self.flush_forgets(p.node, p.slot, head);
1682 }
1683 }
1684 }
1685
1686 fn flush_forgets(&self, node: usize, slot: usize, head: &mut FleetHead) {
1687 let due: Vec<(u64, u64)> = {
1688 let mut sh = self.shared.lock().expect("shared");
1689 if sh.forgets.is_empty() {
1690 return;
1691 }
1692 let (due, keep) = sh.forgets.drain(..).partition(|&(n, s, _, _)| n == node && s == slot);
1693 sh.forgets = keep;
1694 due.into_iter().map(|(_, _, session, epoch)| (session, epoch)).collect()
1695 };
1696 for (session, epoch) in due {
1697 let _ = head.revoke(session, epoch);
1698 }
1699 }
1700
1701 pub fn probe(&self) -> Vec<Result<NodeReport, DaemonError>> {
1703 std::thread::scope(|scope| {
1704 let jobs: Vec<_> = (0..self.node_count())
1705 .map(|i| {
1706 scope.spawn(move || {
1707 let node = self.node(i);
1708 let mut guard = node.checkout_at(0);
1709 self.ensure_connected(i, &mut guard)?;
1710 let node = guard.as_ref().expect("connected").node();
1711 Ok(NodeReport {
1712 addr: self.node(i).spec.addr.clone(),
1713 identity: node.node_identity.clone(),
1714 models: node.loadable_models.clone(),
1715 max_context_tokens: node.capacity.max_context_tokens,
1716 max_lanes: node.capacity.max_lanes,
1717 })
1718 })
1719 })
1720 .collect();
1721 jobs.into_iter().map(|j| j.join().expect("probe thread")).collect()
1722 })
1723 }
1724
1725 pub fn keepalive(&self) {
1729 let nodes: Vec<Arc<NodeConn>> = self.nodes.read().expect("nodes").clone();
1730 for (i, node) in nodes.iter().enumerate() {
1731 let mut dropped = None;
1732 for (slot, conn) in node.conns.iter().enumerate() {
1733 let Ok(mut guard) = conn.try_lock() else { continue };
1734 let Some(head) = guard.as_mut() else { continue };
1735 self.flush_forgets(i, slot, head);
1736 match head.drain_status() {
1737 Ok(()) => self.sync_meta(i, head),
1738 Err(e) => {
1739 *guard = None;
1740 dropped = Some(e);
1741 }
1742 }
1743 }
1744 if let Ok(mut guard) = node.control.try_lock() {
1745 if guard.as_mut().is_some_and(|h| h.drain_status().is_err()) {
1746 *guard = None;
1747 }
1748 }
1749 if let Some(e) = dropped {
1753 if node.conns.iter().all(|c| c.try_lock().is_ok_and(|g| g.is_none())) {
1754 self.mark_unhealthy(i, &e);
1755 }
1756 }
1757 }
1758 }
1759
1760 pub fn retry_down(&self) {
1763 let due: Vec<usize> = {
1764 let mut sh = self.shared.lock().expect("shared");
1765 let now = std::time::Instant::now();
1766 let mut due = Vec::new();
1767 for (i, m) in sh.meta.iter_mut().enumerate() {
1768 if m.down_since.is_some() && m.last_retry.is_none_or(|t| now.duration_since(t) >= self.retry_every) {
1769 m.last_retry = Some(now);
1770 due.push(i);
1771 }
1772 }
1773 due
1774 };
1775 std::thread::scope(|scope| {
1776 for i in due {
1777 scope.spawn(move || {
1778 let node = self.node(i);
1779 let spec = &node.spec;
1780 if spec.addr.starts_with(JOINED) {
1782 return;
1783 }
1784 match FleetHead::connect(&spec.addr, spec.auth.clone()) {
1785 Ok(mut head) => {
1786 head.keep_alive(self.liveness);
1787 watch(&node, &mut head);
1788 self.sync_meta(i, &head);
1789 self.mark_up(i);
1790 if let Ok(mut slot) = node.conns[0].try_lock() {
1791 if slot.is_none() {
1792 *slot = Some(head);
1793 }
1794 }
1795 }
1796 Err(e) => tracing::debug!("fleet: {e}"),
1797 }
1798 });
1799 }
1800 });
1801 }
1802
1803 pub fn tend(&self) {
1804 self.keepalive();
1805 self.retry_down();
1806 self.warm.lock().expect("warm").save();
1807 }
1808
1809 fn warm_node(&self, model: &str, tokens: &[u32]) -> Option<usize> {
1811 let warm = self.warm.lock().expect("warm");
1812 let (identity, _) = warm.warmest(model, tokens)?;
1813 let sh = self.shared.lock().expect("shared");
1814 (0..sh.meta.len()).find(|&i| sh.meta[i].identity == identity && sh.meta[i].up() && Self::serves(&sh, i, model))
1815 }
1816
1817 pub fn spawn_keepalive(self: &Arc<Self>) {
1818 let mgr = Arc::downgrade(self);
1819 let tick = self.retry_every.min(self.liveness.heartbeat);
1820 std::thread::Builder::new()
1821 .name("fleet-keepalive".into())
1822 .spawn(move || loop {
1823 std::thread::sleep(tick);
1824 let Some(mgr) = mgr.upgrade() else { return };
1825 mgr.tend();
1826 })
1827 .expect("spawn the fleet keepalive");
1828 }
1829}
1830
1831fn in_use(slot: &Mutex<Option<FleetHead>>) -> bool {
1835 slot.try_lock().map_or(true, |g| g.is_some())
1836}
1837
1838fn is_fatal_node(e: &DaemonError) -> bool {
1839 matches!(
1840 e,
1841 DaemonError::FleetNode { .. }
1842 | DaemonError::Fleet(superfluid_linkf::LinkFError::Closed)
1843 | DaemonError::Fleet(superfluid_linkf::LinkFError::TimedOut)
1844 | DaemonError::Fleet(superfluid_linkf::LinkFError::Io(_))
1845 )
1846}