superfluid_daemon/
nodeagent.rs

1//! The node-agent.
2
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5use std::time::{Duration, Instant};
6
7use superfluid_abi::finish;
8use superfluid_linkf::{AgentAction, AgentSession, Endpoint};
9use superfluid_proto::handshake::{CapacitySummary, HelloAckF};
10use superfluid_proto::linkf::{self, msg_type, LinkFLimits, TokenEventPayload};
11
12use crate::fleet::{decode_params, decode_tokens};
13use crate::wal::GenParams;
14use crate::{Daemon, DaemonError};
15
16pub const FINISH_EXPIRED: u32 = u32::MAX;
17
18pub const FINISH_REFUSED: u32 = u32::MAX - 1;
19
20pub const FINISH_DECLINED: u32 = u32::MAX - 2;
21
22pub const DEFAULT_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
23
24pub const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(120);
25
26struct AgentShared {
27    agent: AgentSession,
28    contexts: HashMap<u64, Vec<u32>>,
29    params: HashMap<u64, GenParams>,
30    last_load_sample: Option<(u64, u64, u64)>,
31    decode_rate_ewma: u32,
32    prefill_rate_ewma: u32,
33    busy: Busy,
34}
35
36/// Time with at least one generation running, which the node's rates are measured over: a
37/// rate over wall time would read an idle node as a slow one.
38#[derive(Default)]
39struct Busy {
40    running: u32,
41    since: Option<Instant>,
42    total: Duration,
43}
44
45impl Busy {
46    fn start(&mut self) {
47        if self.running == 0 {
48            self.since = Some(Instant::now());
49        }
50        self.running += 1;
51    }
52
53    fn stop(&mut self) {
54        self.running = self.running.saturating_sub(1);
55        if self.running == 0 {
56            if let Some(t) = self.since.take() {
57                self.total += t.elapsed();
58            }
59        }
60    }
61
62    fn ms(&self) -> u64 {
63        (self.total + self.since.map(|t| t.elapsed()).unwrap_or_default()).as_millis() as u64
64    }
65}
66
67pub struct NodeAgent {
68    daemon: Arc<Daemon>,
69    shared: Arc<Mutex<AgentShared>>,
70    identity: String,
71    model_name: String,
72    start: Instant,
73    handshake_timeout: Duration,
74    idle_timeout: Duration,
75    auth: Vec<u8>,
76    max_lanes: u32,
77    last_status_ms: u64,
78}
79
80const STATUS_INTERVAL_MS: u64 = 2_000;
81
82const EMIT_COALESCE_TOKENS: usize = 16;
83
84const EMIT_COALESCE_MS: u64 = 200;
85
86impl NodeAgent {
87    pub fn new(daemon: Arc<Daemon>, identity: impl Into<String>, model_name: impl Into<String>) -> NodeAgent {
88        NodeAgent {
89            max_lanes: 0,
90            last_status_ms: 0,
91            daemon,
92            shared: Arc::new(Mutex::new(AgentShared {
93                agent: AgentSession::new(),
94                contexts: HashMap::new(),
95                params: HashMap::new(),
96                last_load_sample: None,
97                decode_rate_ewma: 0,
98                prefill_rate_ewma: 0,
99                busy: Busy::default(),
100            })),
101            identity: identity.into(),
102            model_name: model_name.into(),
103            start: Instant::now(),
104            handshake_timeout: DEFAULT_HANDSHAKE_TIMEOUT,
105            idle_timeout: DEFAULT_IDLE_TIMEOUT,
106            auth: Vec::new(),
107        }
108    }
109
110    pub fn with_max_lanes(mut self, lanes: u32) -> NodeAgent {
111        self.max_lanes = lanes;
112        self
113    }
114
115    pub fn connection(&self) -> NodeAgent {
116        NodeAgent {
117            daemon: Arc::clone(&self.daemon),
118            shared: Arc::clone(&self.shared),
119            identity: self.identity.clone(),
120            model_name: self.model_name.clone(),
121            start: self.start,
122            handshake_timeout: self.handshake_timeout,
123            idle_timeout: self.idle_timeout,
124            auth: self.auth.clone(),
125            max_lanes: self.max_lanes,
126            last_status_ms: 0,
127        }
128    }
129
130    pub fn with_timeouts(mut self, handshake: Duration, idle: Duration) -> NodeAgent {
131        self.handshake_timeout = handshake;
132        self.idle_timeout = idle;
133        self
134    }
135
136    pub fn with_auth(mut self, auth: Vec<u8>) -> NodeAgent {
137        self.auth = auth;
138        self
139    }
140
141    fn sample_load(&mut self) -> (u64, u64, u32, u32, u32, u32, u32, u64) {
142        let mut sh = self.shared.lock().expect("agent shared");
143        let stats = self.daemon.sched_stats();
144        let now = self.start.elapsed().as_millis() as u64;
145        let decoded = stats.decode_tokens.load(std::sync::atomic::Ordering::Relaxed);
146        let prefilled = stats.prefill_tokens.load(std::sync::atomic::Ordering::Relaxed);
147
148        let smooth = |prev: u32, inst: u32| -> u32 {
149            match prev {
150                0 => inst,
151                p => ((p as u64 * 3 + inst as u64) / 4) as u32,
152            }
153        };
154        const MIN_BUSY_MS: u64 = 200;
155        let busy = sh.busy.ms();
156        match sh.last_load_sample {
157            Some((prev_busy, prev_decoded, prev_prefilled)) => {
158                let dt_ms = busy.saturating_sub(prev_busy);
159                if dt_ms >= MIN_BUSY_MS {
160                    let rate = |n: u64, prev: u64| (n.saturating_sub(prev).saturating_mul(1000) / dt_ms) as u32;
161                    let (d, pf) = (rate(decoded, prev_decoded), rate(prefilled, prev_prefilled));
162                    if d > 0 {
163                        sh.decode_rate_ewma = smooth(sh.decode_rate_ewma, d);
164                    }
165                    if pf > 0 {
166                        sh.prefill_rate_ewma = smooth(sh.prefill_rate_ewma, pf);
167                    }
168                    sh.last_load_sample = Some((busy, decoded, prefilled));
169                }
170            }
171            None => sh.last_load_sample = Some((busy, decoded, prefilled)),
172        }
173        let (rate, prefill_rate) = (sh.decode_rate_ewma, sh.prefill_rate_ewma);
174
175        let queue: u64 = stats.queue_depth.iter().map(|q| q.load(std::sync::atomic::Ordering::Relaxed)).sum();
176        (
177            stats.pool_blocks_used.load(std::sync::atomic::Ordering::Relaxed),
178            stats.pool_blocks_total.load(std::sync::atomic::Ordering::Relaxed),
179            stats.lanes_active.load(std::sync::atomic::Ordering::Relaxed) as u32,
180            self.max_lanes,
181            queue as u32,
182            rate,
183            prefill_rate,
184            now,
185        )
186    }
187
188    fn capacity(&mut self) -> CapacitySummary {
189        let (kv_used, kv_total, lanes, max_lanes, queue, rate, prefill_rate, at) =
190            self.sample_load();
191        CapacitySummary {
192            gpu_memory_bytes: 0,
193            gpu_memory_free_bytes: 0,
194            host_memory_bytes: 0,
195            active_sessions: self.shared.lock().expect("agent shared").contexts.len() as u32,
196            max_sessions: u32::MAX,
197            max_context_tokens: self.daemon.max_stream_tokens(),
198            kv_blocks_used: kv_used,
199            kv_blocks_total: kv_total,
200            lanes_active: lanes,
201            max_lanes,
202            queue_depth: queue,
203            decode_tokens_per_s: rate,
204            prefill_tokens_per_s: prefill_rate,
205            sampled_at_ms: at,
206        }
207    }
208
209    fn send_status(&mut self, endpoint: &mut Endpoint) {
210        let capacity = self.capacity();
211        let msg = linkf::NodeStatusMsg {
212            mem: superfluid_proto::linkw::MemCountersMsg {
213                allocated_bytes: 0,
214                host_retained_bytes: 0,
215                pool_blocks_total: capacity.kv_blocks_total,
216                pool_blocks_used: capacity.kv_blocks_used,
217                pool_bytes_evictable: 0,
218                tentative_bytes: 0,
219            },
220            capacity,
221            loadable_models: vec![self.model_name.clone()],
222        };
223        let _ = endpoint.send(msg_type::NODE_STATUS, &msg);
224        self.last_status_ms = self.start.elapsed().as_millis() as u64;
225    }
226
227    fn ackf(&mut self) -> HelloAckF {
228        HelloAckF {
229            chosen_version: superfluid_linkf::driver::PROTO_VERSION,
230            node_identity: self.identity.clone(),
231            capacity: self.capacity(),
232            loadable_models: vec![self.model_name.clone()],
233            feature_flags: 0,
234            limits: LinkFLimits {
235                max_frame: superfluid_proto::MAX_FRAME,
236                transfer_window: 8,
237                max_transfer: 1 << 26,
238                transfer_budget: 1 << 28,
239            },
240        }
241    }
242
243    pub fn serve(&mut self, mut endpoint: Endpoint) -> Result<(), DaemonError> {
244        let hs = if self.handshake_timeout.is_zero() {
245            endpoint.answer_hello_as_agent(&self.ackf(), &self.auth)
246        } else {
247            endpoint.answer_hello_as_agent_within(&self.ackf(), self.handshake_timeout, &self.auth)
248        };
249        match hs {
250            Ok(()) => {}
251            Err(superfluid_linkf::LinkFError::TimedOut) => {
252                tracing::warn!("node-agent dropped a silent connection (no Hello)");
253                return Ok(());
254            }
255            Err(superfluid_linkf::LinkFError::Unauthorized) => {
256                return Err(superfluid_linkf::LinkFError::Unauthorized.into());
257            }
258            Err(e) => return Err(e.into()),
259        }
260        let mut last_activity = Instant::now();
261        loop {
262            let now = self.start.elapsed();
263            for renew in self.shared.lock().expect("agent shared").agent.needs_renew(now) {
264                endpoint.send(msg_type::LEASE_RENEW, &renew)?;
265            }
266            {
267                let mut sh = self.shared.lock().expect("agent shared");
268                for action in sh.agent.poll(now) {
269                    match action {
270                        AgentAction::SelfFence { session_id } => {
271                            tracing::warn!(session_id, "node-agent self-fenced on lease expiry");
272                            sh.contexts.remove(&session_id);
273                            sh.params.remove(&session_id);
274                        }
275                    }
276                }
277            }
278            if now.as_millis() as u64 >= self.last_status_ms + STATUS_INTERVAL_MS {
279                self.send_status(&mut endpoint);
280            }
281            let frames = match endpoint.pump() {
282                Ok(f) => f,
283                Err(superfluid_linkf::LinkFError::Closed) => return Ok(()),
284                Err(e) => return Err(e.into()),
285            };
286            if frames.is_empty() {
287                if !self.idle_timeout.is_zero() && last_activity.elapsed() >= self.idle_timeout {
288                    tracing::info!("node-agent dropped an idle connection");
289                    return Ok(());
290                }
291            } else {
292                last_activity = Instant::now();
293                for frame in frames {
294                    self.handle(&mut endpoint, frame)?;
295                }
296            }
297        }
298    }
299
300    fn handle(
301        &mut self,
302        endpoint: &mut Endpoint,
303        frame: superfluid_proto::Frame,
304    ) -> Result<(), DaemonError> {
305        let now = self.start.elapsed();
306        match frame.msg_type {
307            msg_type::RECONCILE => {
308                let m: linkf::ReconcileMsg = postcard::from_bytes(&frame.payload)?;
309                let (ack, retransmissions) =
310                    self.shared.lock().expect("agent shared").agent.on_reconcile(&m, now);
311                endpoint.send(msg_type::RECONCILE_ACK, &ack)?;
312                for batch in retransmissions {
313                    endpoint.send(msg_type::TOKEN_EVENTS, &batch)?;
314                }
315            }
316            msg_type::SESSION_ASSIGN => {
317                let m: linkf::SessionAssignMsg = postcard::from_bytes(&frame.payload)?;
318                {
319                    let mut sh = self.shared.lock().expect("agent shared");
320                    if sh.agent.on_assign(&m, now) {
321                        sh.params.insert(m.session_id, decode_params(&m.params));
322                    }
323                }
324            }
325            msg_type::SESSION_REVOKE => {
326                let m: linkf::SessionRevokeMsg = postcard::from_bytes(&frame.payload)?;
327                let mut sh = self.shared.lock().expect("agent shared");
328                if sh.agent.on_revoke(&m) {
329                    sh.contexts.remove(&m.session_id);
330                    sh.params.remove(&m.session_id);
331                }
332            }
333            msg_type::APPEND_CONTENT => {
334                let m: linkf::AppendContentMsg = postcard::from_bytes(&frame.payload)?;
335                let mut sh = self.shared.lock().expect("agent shared");
336                if sh.agent.on_append(m.session_id, m.epoch, m.command_id) {
337                    let tokens = decode_tokens(&m.blocks)?;
338                    sh.contexts.insert(m.session_id, tokens);
339                }
340            }
341            msg_type::GENERATE_REQ => {
342                let m: linkf::GenerateReqMsg = postcard::from_bytes(&frame.payload)?;
343                let declined = !self.shared.lock().expect("agent shared").agent.on_generate(&m);
344                if declined {
345                    tracing::warn!(
346                        session = m.session_id,
347                        generation = m.generation_id,
348                        epoch = m.epoch,
349                        "declining a generate (unknown session, superseded epoch, or \
350                         replayed command); reporting it rather than going silent"
351                    );
352                    let msg = linkf::TokenEventsMsg {
353                        events: vec![linkf::TokenEvent {
354                            session_id: m.session_id,
355                            epoch: m.epoch,
356                            generation_id: m.generation_id,
357                            event_seq: 0,
358                            payload: TokenEventPayload::Finish {
359                                reason: FINISH_DECLINED,
360                            },
361                        }],
362                    };
363                    endpoint.send(msg_type::TOKEN_EVENTS, &msg)?;
364                } else {
365                    match self.run_generation(endpoint, &m) {
366                        Ok(()) => {}
367                        Err(e) if is_client_refusal(&e) => {
368                            tracing::warn!(
369                                session = m.session_id,
370                                generation = m.generation_id,
371                                error = %e,
372                                "refusing a generation the node cannot stage; \
373                                 the head should have caught this before dispatch"
374                            );
375                            let batch = self.shared.lock().expect("agent shared").agent.emit(
376                                m.session_id,
377                                m.generation_id,
378                                TokenEventPayload::Finish {
379                                    reason: FINISH_REFUSED,
380                                },
381                            );
382                            if let Some(batch) = batch {
383                                endpoint.send(msg_type::TOKEN_EVENTS, &batch)?;
384                            }
385                        }
386                        Err(e) => return Err(e),
387                    }
388                }
389            }
390            msg_type::STAGE_FORWARD_REQ => {
391                let m: linkf::StageForwardReqMsg = postcard::from_bytes(&frame.payload)?;
392                let resp = match self.daemon.forward_stage(
393                    m.tokens.clone(),
394                    m.start_layer,
395                    m.end_layer,
396                    m.hidden_in.clone(),
397                ) {
398                    Ok(crate::StageOutput::Hidden(h)) => linkf::StageForwardRespMsg {
399                        request_id: m.request_id,
400                        hidden: Some(h),
401                        token: None,
402                    },
403                    Ok(crate::StageOutput::Token(t)) => linkf::StageForwardRespMsg {
404                        request_id: m.request_id,
405                        hidden: None,
406                        token: Some(t),
407                    },
408                    Err(_) => linkf::StageForwardRespMsg {
409                        request_id: m.request_id,
410                        hidden: None,
411                        token: None,
412                    },
413                };
414                endpoint.send(msg_type::STAGE_FORWARD_RESP, &resp)?;
415            }
416            msg_type::WATERMARK_ACK => {
417                let m: linkf::WatermarkAckMsg = postcard::from_bytes(&frame.payload)?;
418                self.shared.lock().expect("agent shared").agent.on_watermark_ack(&m);
419            }
420            msg_type::LEASE_GRANT => {
421                let m: linkf::LeaseGrantMsg = postcard::from_bytes(&frame.payload)?;
422                self.shared.lock().expect("agent shared").agent.on_lease_grant(&m, now);
423            }
424            _ => {  }
425        }
426        Ok(())
427    }
428
429    fn flush_emit(
430        endpoint: &mut Endpoint,
431        shared: &Arc<Mutex<AgentShared>>,
432        session_id: u64,
433        generation_id: u64,
434        pending: &mut Vec<u32>,
435        now: Duration,
436    ) -> Result<(), Stop> {
437        let emitted = shared.lock().expect("agent shared").agent.emit(
438            session_id,
439            generation_id,
440            TokenEventPayload::Tokens(std::mem::take(pending)),
441        );
442        match emitted {
443            Some(b) => {
444                if endpoint.send(msg_type::TOKEN_EVENTS, &b).is_err() {
445                    return Err(Stop::Peer);
446                }
447            }
448            None => return Err(Stop::Fenced),
449        }
450        let frames = endpoint.pump().map_err(|_| Stop::Peer)?;
451        for frame in frames {
452            match frame.msg_type {
453                msg_type::GENERATE_CANCEL => {
454                    if let Ok(m) = postcard::from_bytes::<linkf::GenerateCancelMsg>(&frame.payload) {
455                        if m.session_id == session_id && m.generation_id == generation_id {
456                            return Err(Stop::Cancelled);
457                        }
458                    }
459                }
460                msg_type::WATERMARK_ACK => {
461                    if let Ok(m) = postcard::from_bytes::<linkf::WatermarkAckMsg>(&frame.payload) {
462                        shared.lock().expect("agent shared").agent.on_watermark_ack(&m);
463                    }
464                }
465                msg_type::LEASE_GRANT => {
466                    if let Ok(m) = postcard::from_bytes::<linkf::LeaseGrantMsg>(&frame.payload) {
467                        shared.lock().expect("agent shared").agent.on_lease_grant(&m, now);
468                    }
469                }
470                _ => {  }
471            }
472        }
473        Ok(())
474    }
475
476    fn run_generation(
477        &mut self,
478        endpoint: &mut Endpoint,
479        req: &linkf::GenerateReqMsg,
480    ) -> Result<(), DaemonError> {
481        let session_id = req.session_id;
482        let generation_id = req.generation_id;
483        let (context, params) = {
484            let sh = self.shared.lock().expect("agent shared");
485            (
486                sh.contexts.get(&session_id).cloned().unwrap_or_default(),
487                sh.params.get(&session_id).copied().unwrap_or_default(),
488            )
489        };
490        let max_new = req.budgets.max_new_tokens.min(u32::MAX as u64) as u32;
491
492        let daemon = Arc::clone(&self.daemon);
493        let start = self.start;
494        let shared = Arc::clone(&self.shared);
495        let mut stop: Option<Stop> = None;
496        let mut pending: Vec<u32> = Vec::with_capacity(EMIT_COALESCE_TOKENS);
497        let mut last_flush = start.elapsed();
498        let extras = crate::scheduler::GenExtras::default();
499        self.shared.lock().expect("agent shared").busy.start();
500        let generated = daemon.generate_tokens_streaming(
501            context,
502            params,
503            max_new,
504            req.budgets.max_wall_ms,
505            extras,
506            |batch| {
507                pending.extend_from_slice(batch);
508                let now = start.elapsed();
509                let due = pending.len() >= EMIT_COALESCE_TOKENS
510                    || now.saturating_sub(last_flush).as_millis() as u64 >= EMIT_COALESCE_MS;
511                if !due {
512                    return std::ops::ControlFlow::Continue(());
513                }
514                last_flush = now;
515                match Self::flush_emit(
516                    endpoint,
517                    &shared,
518                    session_id,
519                    generation_id,
520                    &mut pending,
521                    now,
522                ) {
523                    Ok(()) => std::ops::ControlFlow::Continue(()),
524                    Err(s) => {
525                        stop = Some(s);
526                        std::ops::ControlFlow::Break(())
527                    }
528                }
529            },
530        );
531        self.shared.lock().expect("agent shared").busy.stop();
532        let (tokens, _finish, expired) = generated?;
533
534        if stop.is_none() && !pending.is_empty() {
535            let now = start.elapsed();
536            if let Err(s) = Self::flush_emit(
537                endpoint,
538                &self.shared,
539                session_id,
540                generation_id,
541                &mut pending,
542                now,
543            ) {
544                stop = Some(s);
545            }
546        }
547
548        let reason = match stop {
549            Some(Stop::Fenced) | Some(Stop::Peer) => return Ok(()),
550            Some(Stop::Cancelled) => finish::CANCELLED,
551            None if expired => FINISH_EXPIRED,
552            None if tokens.len() as u32 >= max_new => finish::LENGTH,
553            None => finish::EOS,
554        };
555        if let Some(batch) =
556            self.shared
557                .lock()
558                .expect("agent shared")
559                .agent
560                .emit(session_id, generation_id, TokenEventPayload::Finish { reason })
561        {
562            endpoint.send(msg_type::TOKEN_EVENTS, &batch)?;
563        }
564        Ok(())
565    }
566}
567
568fn is_client_refusal(e: &DaemonError) -> bool {
569    matches!(
570        e,
571        DaemonError::StreamTooLong { .. } | DaemonError::Constraint(_)
572    )
573}
574
575enum Stop {
576    Cancelled,
577    Fenced,
578    Peer,
579}
superfluid_daemon/nodeagent.rs

superfluid_daemon/
nodeagent.rs
