Skip to main content

superfluid_daemon/
fleet_manager.rs

1//! Head-side multi-node orchestration.
2
3use 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    /// `typical` stands in for a rate this node has not reported yet: the mean of the nodes
44    /// that have, so a node that has served nothing is neither shunned nor flooded.
45    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        // A rate is what the node does when busy, so it stays true while the node idles; only the
54        // load that piles up ahead of a request goes stale.
55        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
69/// Keeps `node`'s live load from each status `head` reads; weak, as the node owns the head.
70fn 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    /// A connection no generation takes, for short calls (a prompt's encoding, a state
89    /// operation) while every lane's connection is held by a generation. The head dials it for
90    /// a node it lists; a joined node dials it when asked, and one that predates the ask has none.
91    control: Mutex<Option<FleetHead>>,
92    /// How many of `conns` are in use: the count asked for, or the node's lanes once learned.
93    open: std::sync::atomic::AtomicUsize,
94    cursor: std::sync::atomic::AtomicUsize,
95    /// Tokens this node has generated for the head, for `/metrics`.
96    decoded: std::sync::atomic::AtomicU64,
97    /// The last status any of its connections read, mid-generation included, for `/metrics`;
98    /// the load placement reads is synced from a connection between calls.
99    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    /// A lane's connection for a new session: the next in turn that no call holds (connected,
130    /// or for a node the head lists, one it can dial), else the first of those to come free.
131    /// Taken strictly in turn, a session would wait out a generation still running while
132    /// another lane is idle. None once `deadline` passes with every one still held.
133    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            // None taken and none busy: a joined node's lanes not adopted yet.
149            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    /// A connection no call is using, so a short operation does not wait out a generation:
162    /// a lane's when one is idle, else the control connection, else (a joined node without
163    /// one) the next lane's in turn.
164    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    /// As [`Self::checkout_at`], none once `deadline` passes with the slot still held.
184    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    /// Where a joined node dialed from, so a second machine with its name is told apart.
221    peer: Option<std::net::IpAddr>,
222    /// Set when the node failed; requests pass it over until the background retry connects.
223    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    /// The models the node serves with a tokenizer other than the head's.
229    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    /// Finished sessions to revoke on a node: (node, slot, session, epoch).
248    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    /// The tokenizer digest the head serves each model with: a node whose differs is not used
262    /// for that model.
263    tokenizers: std::sync::RwLock<std::collections::HashMap<String, u64>>,
264    shared: Mutex<Shared>,
265    /// Which node last ran each prompt prefix, to place the next turn where it is warm.
266    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    /// The node holding the longest prefix of the prompt warm, tried first while it has a
322    /// lane free.
323    pub warm: Option<usize>,
324}
325
326pub const DEFAULT_POOL_HIGH_PCT: u32 = 90;
327
328/// 0: as many connections to a node as it has lanes, so each lane can be busy at once.
329pub const DEFAULT_CONNS_PER_NODE: usize = 0;
330
331const MAX_CONNS_PER_NODE: usize = 64;
332
333/// The address a joined node is listed under: it dialed in, so the head cannot dial it.
334const JOINED: &str = "joined:";
335
336/// How often the background retry tries a node that failed.
337pub const RETRY_DOWN_AFTER: std::time::Duration = std::time::Duration::from_secs(10);
338
339/// One node as the head sees it now, for `/metrics`.
340#[derive(Debug, Clone)]
341pub struct NodeStatus {
342    /// The node's identity, or the address the head dials while it has not answered.
343    pub name: String,
344    /// The address it joined from or is dialed at.
345    pub addr: String,
346    pub up: bool,
347    /// Connections open to it: its lanes, once it answered.
348    pub lanes: usize,
349    /// Generations the head has in flight there.
350    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    /// Takes in a node that dialed the head: one seen before, by identity, gets its slot
436    /// back; a new one gets a slot of its own. `true` for a new one.
437    /// Takes in a connection from a node that dialed the head. The node has a slot per lane it
438    /// advertised and fills one per connection it opens, so it runs as many generations at
439    /// once as it dialed; a node seen before, by identity, takes its slots back, unless it is
440    /// still connected from another address, which is a second machine with the same name.
441    /// Returns the node's index, whether it is new, and its open connections.
442    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                    // Listed as down until its connection is in place, so no request lands
472                    // on an empty slot meanwhile.
473                    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            // One past its lanes is the connection it dialed when asked.
498            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            // Listed as down only until this connection was in place: not a return.
505            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    /// The connections a node has open.
513    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    /// Every node the head knows, in the order it learned them.
529    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    /// The nodes up now that serve `model`, and every node the head knows.
559    /// Nodes up of all the head knows, whatever each serves.
560    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    /// How often a node that failed is tried again in the background (default
580    /// [`RETRY_DOWN_AFTER`]).
581    pub fn with_retry_every(mut self, every: std::time::Duration) -> FleetManager {
582        self.retry_every = every;
583        self
584    }
585
586    /// Serves `model` only on nodes whose tokenizer has `digest` ([`crate::fleet::tokenizer_digest`]);
587    /// a node that does not say is taken at its word.
588    pub fn expect_tokenizer(&self, model: &str, digest: u64) {
589        self.tokenizers.write().expect("tokenizers").insert(model.to_string(), digest);
590    }
591
592    /// How the head's connections are kept alive (default [`Liveness::DEFAULT`]); the
593    /// background tend reads each idle one every `heartbeat`.
594    pub fn with_liveness(mut self, liveness: Liveness) -> FleetManager {
595        self.liveness = liveness;
596        self
597    }
598
599    /// Remembers in `path` which node ran each prompt prefix, so a restarted head still places
600    /// a conversation's next turn on the node that has it warm.
601    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    /// A fresh connection knows none of the sessions placed over the one it replaced; this
676    /// places the session on it again from what the head kept.
677    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    /// Rebinds the session on the slot's connection, dropping the connection only when the link
700    /// itself failed.
701    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    /// The slot's connection, dropped first if the node closed it while it sat idle: a write
874    /// into a closed socket can look like it succeeded.
875    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    /// As [`Self::checkout_live`], none once `deadline` passes with the slot still held.
880    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    /// As [`Self::place`], for a generation that must start by `deadline`: none once it passes
932    /// before the session is placed (every connection to a node serving the model held by
933    /// another generation meanwhile), rather than placed for a generation already too late.
934    #[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    /// As [`Self::try_place_on`], not placed once `deadline` passes while every connection to
1019    /// the node is held.
1020    #[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    /// As [`Self::generate_streaming`], with what the generation asks beyond its sampling, sent
1287    /// to whichever node runs it, and `on_logprobs` told the log-probabilities it asked for.
1288    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    /// As [`Self::generate_streaming_with`], for a generation that must start by `deadline`: the
1301    /// wait for the session's connection counts against it, one that passes meanwhile ends the
1302    /// generation unstarted (expired, no tokens), and the node is given only the time left.
1303    #[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        // Another session placed on the same connection may hold it now: the wait for it counts
1444        // against the deadline, and the node is given what is left once it is had.
1445        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    /// Runs a state operation on one node serving `model`, the next one tried when a node
1475    /// fails.
1476    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    /// [`FleetManager::state_sync`] for a streamed transcription, each segment told as the node
1486    /// has it. Once one has been told, a node lost after it fails the call rather than starting
1487    /// over on another, which would tell those segments twice.
1488    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    /// Runs a state operation on every node up that serves `model`, as a LoRA must be on each
1509    /// node a request can land on; the answer is the last node's.
1510    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    /// Runs `call` on an idle connection to node `i`, connecting it first if it has to; a node
1533    /// whose link fails is down.
1534    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    /// Runs `call` on one node serving `model`, the next tried when a node fails.
1552    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    /// The tokenizer of a node serving `model`, for a head with no copy of it; encoding asks a
1566    /// node serving the model each time.
1567    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    /// Text a node serving `model` encodes.
1578    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    /// Ends a session the caller is done with: the head drops its record, and the node the
1663    /// context it kept (now if its connection is free, else the next time it is).
1664    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    /// Connects to every node once, so the head learns each one (or why it cannot) at start.
1702    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    /// Reads every idle connection, which sends the node a heartbeat when one is due and
1726    /// answers its lease renewals; drops the ones a node closed or went silent on, and finishes
1727    /// deferred session ends. A connection in use is read by the call using it.
1728    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            // With its last connection gone the node is down until the background retry reaches
1750            // it again, which also learns what it serves now: a node that restarted under another
1751            // model name would otherwise be passed over for the name it had.
1752            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    /// Tries every node that failed and is due a retry, in parallel and outside the connection
1761    /// slots; one that answers is placed on again. Requests never wait on a down node.
1762    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                    // A joined node comes back when it dials in again.
1781                    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    /// The node, up and serving `model`, that last ran the longest prefix of `tokens`.
1810    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
1831/// Whether a slot holds a connection, for [`FleetManager::adopt`], which must not wait on one:
1832/// a call holds a slot's lock while it waits for the shared state adopt holds. A slot whose lock
1833/// is held is in use, so connected.
1834fn 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}