superfluid_daemon/
fleet.rs

1//! The head daemon's fleet edge.
2
3use std::collections::BTreeMap;
4
5use superfluid_linkf::{Endpoint, HeadSession};
6use superfluid_proto::handshake::HelloAckF;
7use superfluid_proto::linkf::{
8    self, msg_type, AppendContentMsg, GenerateCancelMsg, GenerationBudgetsMsg, Lease, TokenEventPayload,
9};
10
11use crate::codec::TextCodec;
12use superfluid_engine::StageOutput;
13use crate::wal::GenParams;
14use crate::DaemonError;
15
16pub fn encode_params(params: &GenParams) -> Vec<u8> {
17    postcard::to_stdvec(params).expect("GenParams postcard")
18}
19
20pub fn decode_params(bytes: &[u8]) -> GenParams {
21    postcard::from_bytes(bytes).unwrap_or_default()
22}
23
24pub fn encode_tokens(tokens: &[u32]) -> Vec<u8> {
25    postcard::to_stdvec(&tokens.to_vec()).expect("tokens postcard")
26}
27
28pub fn decode_tokens(bytes: &[u8]) -> Result<Vec<u32>, DaemonError> {
29    postcard::from_bytes(bytes).map_err(DaemonError::Codec)
30}
31
32#[derive(Debug, Clone)]
33pub struct FleetGeneration {
34    pub tokens: Vec<u32>,
35    pub finish: u32,
36    pub expired: bool,
37}
38
39#[derive(Debug, Clone)]
40pub struct FleetSegment {
41    pub channel: u32,
42    pub text: String,
43    pub tool: Option<(String, String)>,
44}
45
46#[derive(Debug, Clone)]
47struct Placement {
48    params: GenParams,
49    qos_class: u8,
50    codec: String,
51    lease: Lease,
52    tokens: Vec<u32>,
53}
54
55pub struct FleetHead {
56    head: HeadSession,
57    endpoint: Endpoint,
58    node: HelloAckF,
59    next_generation: u64,
60    placements: BTreeMap<u64, Placement>,
61    addr: String,
62    auth: Vec<u8>,
63}
64
65/// Epochs and command ids are numbered across every connection this head makes: a node keeps
66/// a session's across connections, and refuses an epoch or command it has already seen.
67static NEXT_EPOCH: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
68static NEXT_COMMANDS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
69
70fn fresh_session_state() -> HeadSession {
71    HeadSession::starting_at(NEXT_COMMANDS.fetch_add(1 << 32, std::sync::atomic::Ordering::Relaxed))
72}
73
74pub const CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
75
76pub const HANDSHAKE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
77
78fn closed_at_handshake() -> String {
79    "closed the connection at the handshake; check that --fleet-auth holds the token in the node's --auth-file".into()
80}
81
82/// Dials `addr` and opens the link: encrypted and keyed by the token when there is one.
83fn handshake(addr: &str, auth: Vec<u8>) -> Result<(Endpoint, HelloAckF), DaemonError> {
84    let node = |why: String| DaemonError::FleetNode { addr: addr.to_string(), why };
85    let stream = Endpoint::tcp_connect_within(addr, CONNECT_TIMEOUT).map_err(|e| node(format!("unreachable ({e})")))?;
86    let opened = if auth.is_empty() {
87        Endpoint::from_tcp(stream)
88    } else {
89        Endpoint::from_tcp_secure(stream, &auth, superfluid_linkf::Role::Initiator, HANDSHAKE_TIMEOUT)
90    };
91    hello(opened, addr, auth, "--fleet-auth and the node's --auth-file differ")
92}
93
94/// The head's side of the handshake on a stream just opened to `peer`, dialed or accepted,
95/// with each way it can fail named for the operator; `mismatch` says what differs when the
96/// token is refused.
97fn hello(opened: Result<Endpoint, superfluid_linkf::LinkFError>, peer: &str, auth: Vec<u8>, mismatch: &str) -> Result<(Endpoint, HelloAckF), DaemonError> {
98    use superfluid_linkf::LinkFError;
99    let node = |why: String| DaemonError::FleetNode { addr: peer.to_string(), why };
100    let mut endpoint = match opened {
101        Ok(ep) => ep,
102        Err(LinkFError::Unauthorized) => return Err(node(format!("refused the fleet token: {mismatch}"))),
103        Err(LinkFError::TimedOut) => {
104            return Err(node(format!("did not finish the handshake within {} s", HANDSHAKE_TIMEOUT.as_secs())))
105        }
106        Err(LinkFError::Closed | LinkFError::Io(_)) if !auth.is_empty() => return Err(node(closed_at_handshake())),
107        Err(e) => return Err(node(e.to_string())),
108    };
109    match endpoint.hello_as_head_within(auth, HANDSHAKE_TIMEOUT) {
110        Ok(ack) => Ok((endpoint, ack)),
111        Err(LinkFError::Unauthorized) => Err(node(format!("refused the fleet token: {mismatch}"))),
112        Err(LinkFError::Closed | LinkFError::Io(_)) => Err(node(closed_at_handshake())),
113        Err(LinkFError::TimedOut) => Err(node(format!(
114            "did not answer the handshake within {} s; is superfluid-noded listening there?",
115            HANDSHAKE_TIMEOUT.as_secs()
116        ))),
117        Err(e) => Err(node(e.to_string())),
118    }
119}
120
121fn is_transient(e: &DaemonError) -> bool {
122    matches!(
123        e,
124        DaemonError::Fleet(superfluid_linkf::LinkFError::Closed)
125            | DaemonError::Fleet(superfluid_linkf::LinkFError::TimedOut)
126            | DaemonError::Fleet(superfluid_linkf::LinkFError::Io(_))
127    )
128}
129
130impl FleetHead {
131    pub fn connect(addr: &str, auth: Vec<u8>) -> Result<FleetHead, DaemonError> {
132        let (endpoint, node) = handshake(addr, auth.clone())?;
133        Ok(FleetHead::from_parts(endpoint, node, addr.to_string(), auth))
134    }
135
136    /// A node that dialed in; the head still opens the handshake, encrypted when it has a token.
137    pub fn accept(sock: std::net::TcpStream, auth: Vec<u8>) -> Result<FleetHead, DaemonError> {
138        let peer = sock.peer_addr().map(|a| a.to_string()).unwrap_or_default();
139        let opened = if auth.is_empty() {
140            Endpoint::from_tcp(sock)
141        } else {
142            Endpoint::from_tcp_secure(sock, &auth, superfluid_linkf::Role::Initiator, HANDSHAKE_TIMEOUT)
143        };
144        let (endpoint, node) = hello(opened, &peer, auth.clone(), "its token differs from this head's")?;
145        Ok(FleetHead::from_parts(endpoint, node, String::new(), auth))
146    }
147
148    fn from_parts(endpoint: Endpoint, node: HelloAckF, addr: String, auth: Vec<u8>) -> FleetHead {
149        FleetHead {
150            head: fresh_session_state(),
151            endpoint,
152            node,
153            next_generation: 1,
154            placements: BTreeMap::new(),
155            addr,
156            auth,
157        }
158    }
159
160    fn reconnect_self(&mut self) -> Result<(), DaemonError> {
161        let (addr, auth) = (self.addr.clone(), self.auth.clone());
162        self.reconnect(&addr, auth)?;
163        Ok(())
164    }
165
166    pub fn node(&self) -> &HelloAckF {
167        &self.node
168    }
169
170    /// Reads what the node sent while the connection sat idle; an error means it is gone.
171    pub fn drain_status(&mut self) -> Result<(), DaemonError> {
172        for frame in self.endpoint.pump()? {
173            if frame.msg_type == msg_type::NODE_STATUS {
174                if let Ok(m) = postcard::from_bytes::<linkf::NodeStatusMsg>(&frame.payload) {
175                    self.node.capacity = m.capacity;
176                    if !m.loadable_models.is_empty() {
177                        self.node.loadable_models = m.loadable_models;
178                    }
179                }
180            }
181        }
182        Ok(())
183    }
184
185    pub fn keepalive(&mut self) -> Result<(), DaemonError> {
186        self.endpoint.send(msg_type::KEEPALIVE, &())?;
187        self.drain_status()
188    }
189
190    pub fn knows(&self, session_id: u64) -> bool {
191        self.head.current_epoch(session_id).is_some()
192    }
193
194
195    pub fn reconnect(&mut self, addr: &str, auth: Vec<u8>) -> Result<Vec<u64>, DaemonError> {
196        let _ = self.endpoint.shutdown();
197        if addr.is_empty() {
198            return Err(DaemonError::FleetNode {
199                addr: self.node.node_identity.clone(),
200                why: "dialed in on its own; waiting for it to join again".into(),
201            });
202        }
203        let (endpoint, node) = handshake(addr, auth)?;
204        self.node = node;
205        self.endpoint = endpoint;
206        let rec = self.head.reconcile();
207        self.endpoint.send(msg_type::RECONCILE, &rec)?;
208
209        let mut received: std::collections::BTreeSet<(u64, u64, u64)> =
210            std::collections::BTreeSet::new();
211        let ack = loop {
212            let mut got = None;
213            for frame in self.endpoint.wait_frames()? {
214                match frame.msg_type {
215                    msg_type::RECONCILE_ACK => {
216                        got = Some(postcard::from_bytes::<linkf::ReconcileAckMsg>(&frame.payload)?);
217                    }
218                    msg_type::TOKEN_EVENTS => {
219                        let m: linkf::TokenEventsMsg = postcard::from_bytes(&frame.payload)?;
220                        for ev in &m.events {
221                            received.insert((ev.session_id, ev.generation_id, ev.event_seq));
222                        }
223                        for a in self.head.on_token_events(&m) {
224                            self.endpoint.send(msg_type::WATERMARK_ACK, &a)?;
225                        }
226                    }
227                    _ => {}
228                }
229            }
230            if let Some(a) = got {
231                break a;
232            }
233        };
234
235        let mut expected: std::collections::BTreeSet<(u64, u64, u64)> =
236            std::collections::BTreeSet::new();
237        for s in &ack.sessions {
238            if self.head.current_epoch(s.session_id).is_some() {
239                for r in &s.retry_ranges {
240                    for seq in r.from_seq..=r.to_seq {
241                        let id = (s.session_id, r.generation_id, seq);
242                        if !received.contains(&id) {
243                            expected.insert(id);
244                        }
245                    }
246                }
247            }
248        }
249        while !expected.is_empty() {
250            for frame in self.endpoint.wait_frames()? {
251                if frame.msg_type != msg_type::TOKEN_EVENTS {
252                    continue;
253                }
254                let m: linkf::TokenEventsMsg = postcard::from_bytes(&frame.payload)?;
255                for ev in &m.events {
256                    expected.remove(&(ev.session_id, ev.generation_id, ev.event_seq));
257                }
258                for a in self.head.on_token_events(&m) {
259                    self.endpoint.send(msg_type::WATERMARK_ACK, &a)?;
260                }
261            }
262        }
263
264        // A session the node still holds may have missed the last append, written into the
265        // connection that dropped; resend each one's latest content.
266        let reassigned = self.head.on_reconcile_ack(&ack);
267        for s in ack.sessions.iter().filter(|s| s.held && !reassigned.contains(&s.session_id)) {
268            if let (Some(p), Some(epoch)) =
269                (self.placements.get(&s.session_id).cloned(), self.head.current_epoch(s.session_id))
270            {
271                if !p.tokens.is_empty() {
272                    self.send_append(s.session_id, epoch, &p.tokens)?;
273                }
274            }
275        }
276        let mut replaced = Vec::new();
277        for session_id in reassigned {
278            if let Some(p) = self.placements.get(&session_id).cloned() {
279                if let Some(epoch) = self.head.current_epoch(session_id) {
280                    let mut msg = self.head.assign_session(
281                        session_id,
282                        epoch,
283                        p.lease,
284                        p.qos_class,
285                        &p.codec,
286                    );
287                    msg.params = encode_params(&p.params);
288                    self.endpoint.send(msg_type::SESSION_ASSIGN, &msg)?;
289                    if !p.tokens.is_empty() {
290                        self.send_append(session_id, epoch, &p.tokens)?;
291                    }
292                    replaced.push(session_id);
293                }
294            }
295        }
296        Ok(replaced)
297    }
298
299    pub fn assign(
300        &mut self,
301        session_id: u64,
302        params: GenParams,
303        qos_class: u8,
304        codec: &str,
305        lease: Lease,
306    ) -> Result<u64, DaemonError> {
307        let entry = self.placements.entry(session_id).or_insert(Placement {
308            params,
309            qos_class,
310            codec: codec.to_string(),
311            lease,
312            tokens: Vec::new(),
313        });
314        entry.params = params;
315        entry.qos_class = qos_class;
316        entry.codec = codec.to_string();
317        entry.lease = lease;
318
319        let epoch = NEXT_EPOCH.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
320        let mut msg = self.head.assign_session(session_id, epoch, lease, qos_class, codec);
321        msg.params = encode_params(&params);
322        match self.endpoint.send(msg_type::SESSION_ASSIGN, &msg) {
323            Ok(()) => Ok(epoch),
324            Err(e) => {
325                let e: DaemonError = e.into();
326                if is_transient(&e) {
327                    self.reconnect_self()?;
328                    Ok(self.head.current_epoch(session_id).unwrap_or(epoch))
329                } else {
330                    Err(e)
331                }
332            }
333        }
334    }
335
336    pub fn park(&mut self, session_id: u64) -> Result<(), DaemonError> {
337        let epoch = self.head.current_epoch(session_id).unwrap_or(0);
338        self.revoke(session_id, epoch)
339    }
340
341    /// Ends the session on the node, which frees the context it kept. `epoch` is the one it was
342    /// last assigned, so this also reaches a session placed over another connection.
343    pub fn revoke(&mut self, session_id: u64, epoch: u64) -> Result<(), DaemonError> {
344        self.placements.remove(&session_id);
345        let msg = match self.head.revoke(session_id) {
346            Some(msg) => msg,
347            None => linkf::SessionRevokeMsg { command_id: self.head.next_command(session_id), session_id, epoch },
348        };
349        self.endpoint.send(msg_type::SESSION_REVOKE, &msg)?;
350        Ok(())
351    }
352
353    pub fn append(&mut self, session_id: u64, tokens: &[u32]) -> Result<(), DaemonError> {
354        let epoch = self
355            .head
356            .current_epoch(session_id)
357            .ok_or(DaemonError::Protocol("append before assign"))?;
358        if let Some(p) = self.placements.get_mut(&session_id) {
359            p.tokens = tokens.to_vec();
360        }
361        match self.send_append(session_id, epoch, tokens) {
362            Ok(()) => Ok(()),
363            Err(e) if is_transient(&e) => {
364                self.reconnect_self()
365            }
366            Err(e) => Err(e),
367        }
368    }
369
370    fn send_append(&mut self, session_id: u64, epoch: u64, tokens: &[u32]) -> Result<(), DaemonError> {
371        let msg = AppendContentMsg {
372            command_id: self.head.next_command(session_id),
373            session_id,
374            epoch,
375            blocks: encode_tokens(tokens),
376            spans: Vec::new(),
377            media_refs: Vec::new(),
378        };
379        self.endpoint.send(msg_type::APPEND_CONTENT, &msg)?;
380        Ok(())
381    }
382
383    pub fn generate(
384        &mut self,
385        session_id: u64,
386        max_new_tokens: u64,
387        max_wall_ms: u64,
388    ) -> Result<FleetGeneration, DaemonError> {
389        let generation_id = self.next_generation;
390        self.next_generation += 1;
391        for attempt in 0..2 {
392            match self.generate_once(session_id, generation_id, max_new_tokens, max_wall_ms) {
393                Ok(g) => return Ok(g),
394                Err(e) if attempt == 0 && is_transient(&e) => {
395                    self.reconnect_self()?;
396                    if let Some(g) = self.recovered_generation(session_id, generation_id) {
397                        self.head.retire_generation(session_id, generation_id);
398                        return Ok(g);
399                    }
400                    continue;
401                }
402                Err(e) => return Err(e),
403            }
404        }
405        unreachable!("generate retried at most once")
406    }
407
408    pub fn generate_streaming(
409        &mut self,
410        session_id: u64,
411        max_new_tokens: u64,
412        max_wall_ms: u64,
413        cancel: &std::sync::atomic::AtomicBool,
414        mut on_tokens: impl FnMut(&[u32]),
415    ) -> Result<FleetGeneration, DaemonError> {
416        let mut delivered = false;
417        let first = self.generate_streaming_once(session_id, max_new_tokens, max_wall_ms, cancel, |t| {
418            delivered = true;
419            on_tokens(t)
420        });
421        match first {
422            // A connection the node dropped while idle fails before anything streams; nothing
423            // reached the caller, so the generation can start over on a fresh one.
424            Err(e) if !delivered && is_transient(&e) => {
425                self.reconnect_self()?;
426                self.generate_streaming_once(session_id, max_new_tokens, max_wall_ms, cancel, on_tokens)
427            }
428            r => r,
429        }
430    }
431
432    fn generate_streaming_once(
433        &mut self,
434        session_id: u64,
435        max_new_tokens: u64,
436        max_wall_ms: u64,
437        cancel: &std::sync::atomic::AtomicBool,
438        mut on_tokens: impl FnMut(&[u32]),
439    ) -> Result<FleetGeneration, DaemonError> {
440        let generation_id = self.next_generation;
441        self.next_generation += 1;
442        let req = self
443            .head
444            .generate(
445                session_id,
446                generation_id,
447                GenerationBudgetsMsg { max_new_tokens, max_wall_ms },
448            )
449            .ok_or(DaemonError::Protocol("generate before assign"))?;
450        self.endpoint.send(msg_type::GENERATE_REQ, &req)?;
451
452        let mut tokens = Vec::new();
453        let mut finish = 0u32;
454        let mut done = false;
455        let mut cancel_sent = false;
456        while !done {
457            if !cancel_sent && cancel.load(std::sync::atomic::Ordering::Relaxed) {
458                self.endpoint.send(
459                    msg_type::GENERATE_CANCEL,
460                    &GenerateCancelMsg { session_id, generation_id },
461                )?;
462                cancel_sent = true;
463            }
464            for frame in self.endpoint.wait_frames()? {
465                match frame.msg_type {
466                    msg_type::TOKEN_EVENTS => {
467                        let m: linkf::TokenEventsMsg = postcard::from_bytes(&frame.payload)?;
468                        for ev in &m.events {
469                            if ev.session_id == session_id && ev.generation_id == generation_id {
470                                match &ev.payload {
471                                    TokenEventPayload::Tokens(t) => {
472                                        tokens.extend_from_slice(t);
473                                        on_tokens(t);
474                                    }
475                                    TokenEventPayload::Finish { reason } => {
476                                        finish = *reason;
477                                        done = true;
478                                    }
479                                    _ => {}
480                                }
481                            }
482                        }
483                        for ack in self.head.on_token_events(&m) {
484                            self.endpoint.send(msg_type::WATERMARK_ACK, &ack)?;
485                        }
486                    }
487                    msg_type::LEASE_RENEW => {
488                        let m: linkf::LeaseRenewMsg = postcard::from_bytes(&frame.payload)?;
489                        if let Some(grant) = self.head.on_lease_renew(&m) {
490                            self.endpoint.send(msg_type::LEASE_GRANT, &grant)?;
491                        }
492                    }
493                    msg_type::NODE_STATUS => {
494                        if let Ok(m) = postcard::from_bytes::<linkf::NodeStatusMsg>(&frame.payload)
495                        {
496                            self.node.capacity = m.capacity;
497                            if !m.loadable_models.is_empty() {
498                                self.node.loadable_models = m.loadable_models;
499                            }
500                        }
501                    }
502                    _ => {}
503                }
504            }
505        }
506        let expired = finish == u32::MAX;
507        self.head.retire_generation(session_id, generation_id);
508        Ok(FleetGeneration { tokens, finish, expired })
509    }
510
511    pub fn forward_stage(
512        &mut self,
513        tokens: &[u32],
514        start_layer: u32,
515        end_layer: u32,
516        hidden_in: Option<Vec<u8>>,
517    ) -> Result<superfluid_engine::StageOutput, DaemonError> {
518        let request_id = self.next_generation;
519        self.next_generation += 1;
520        let req = linkf::StageForwardReqMsg {
521            request_id,
522            tokens: tokens.to_vec(),
523            start_layer,
524            end_layer,
525            hidden_in,
526        };
527        self.endpoint.send(msg_type::STAGE_FORWARD_REQ, &req)?;
528        loop {
529            for frame in self.endpoint.wait_frames()? {
530                if frame.msg_type == msg_type::STAGE_FORWARD_RESP {
531                    let m: linkf::StageForwardRespMsg = postcard::from_bytes(&frame.payload)?;
532                    if m.request_id != request_id {
533                        continue;
534                    }
535                    return match (m.hidden, m.token) {
536                        (Some(h), None) => Ok(superfluid_engine::StageOutput::Hidden(h)),
537                        (None, Some(t)) => Ok(superfluid_engine::StageOutput::Token(t)),
538                        _ => Err(DaemonError::Generation(
539                            "stage node could not run the forward (engine lacks staged forward)".into(),
540                        )),
541                    };
542                }
543            }
544        }
545    }
546
547    pub fn generate_chat(
548        &mut self,
549        session_id: u64,
550        codec: &dyn TextCodec,
551        messages: &[crate::codec::ChatMessage],
552        tools: &[String],
553        max_new_tokens: u64,
554        max_wall_ms: u64,
555    ) -> Result<Vec<FleetSegment>, DaemonError> {
556        let prompt = codec
557            .render_prompt_structured(messages, tools)
558            .ok_or(DaemonError::Protocol("codec has no chat dialect"))?;
559        self.append(session_id, &prompt)?;
560        let gen = self.generate(session_id, max_new_tokens, max_wall_ms)?;
561
562        let mut chan = crate::codec::primed_channelizer(codec, &prompt);
563        let schemas = crate::codec::ToolSchemas::from_jsons(tools);
564        let mut segments = Vec::new();
565        for run in chan.split(&gen.tokens) {
566            let text = codec.decode(&run.text);
567            let tool = if run.channel == crate::wal::channel::TOOL_CALL && run.closes {
568                codec.parse_tool_call_with(&text, Some(&schemas))
569            } else {
570                None
571            };
572            if !text.is_empty() || tool.is_some() {
573                segments.push(FleetSegment {
574                    channel: run.channel,
575                    text,
576                    tool,
577                });
578            }
579        }
580        Ok(segments)
581    }
582
583    fn recovered_generation(
584        &self,
585        session_id: u64,
586        generation_id: u64,
587    ) -> Option<FleetGeneration> {
588        let epoch = self.head.current_epoch(session_id)?;
589        let mut events: Vec<_> = self
590            .head
591            .wal()
592            .session_events(session_id)
593            .iter()
594            .filter(|c| c.event.epoch == epoch && c.event.generation_id == generation_id)
595            .collect();
596        events.sort_by_key(|c| c.event.event_seq);
597        let mut tokens = Vec::new();
598        let mut finish = None;
599        for c in events {
600            match &c.event.payload {
601                TokenEventPayload::Tokens(t) => tokens.extend_from_slice(t),
602                TokenEventPayload::Finish { reason } => finish = Some(*reason),
603                _ => {}
604            }
605        }
606        let finish = finish?;
607        Some(FleetGeneration {
608            tokens,
609            finish,
610            expired: finish == crate::nodeagent::FINISH_EXPIRED,
611        })
612    }
613
614    fn generate_once(
615        &mut self,
616        session_id: u64,
617        generation_id: u64,
618        max_new_tokens: u64,
619        max_wall_ms: u64,
620    ) -> Result<FleetGeneration, DaemonError> {
621        let req = self
622            .head
623            .generate(
624                session_id,
625                generation_id,
626                GenerationBudgetsMsg {
627                    max_new_tokens,
628                    max_wall_ms,
629                },
630            )
631            .ok_or(DaemonError::Protocol("generate before assign"))?;
632        self.endpoint.send(msg_type::GENERATE_REQ, &req)?;
633
634        let mut tokens = Vec::new();
635        let mut finish = 0u32;
636        let mut expired = false;
637        let mut done = false;
638        while !done {
639            for frame in self.endpoint.wait_frames()? {
640                match frame.msg_type {
641                    msg_type::TOKEN_EVENTS => {
642                        let m: linkf::TokenEventsMsg = postcard::from_bytes(&frame.payload)?;
643                        for ev in &m.events {
644                            if ev.session_id == session_id && ev.generation_id == generation_id {
645                                match &ev.payload {
646                                    TokenEventPayload::Tokens(t) => tokens.extend_from_slice(t),
647                                    TokenEventPayload::Finish { reason } => {
648                                        finish = *reason;
649                                        done = true;
650                                    }
651                                    _ => {}
652                                }
653                            }
654                        }
655                        for ack in self.head.on_token_events(&m) {
656                            self.endpoint.send(msg_type::WATERMARK_ACK, &ack)?;
657                        }
658                    }
659                    msg_type::LEASE_RENEW => {
660                        let m: linkf::LeaseRenewMsg = postcard::from_bytes(&frame.payload)?;
661                        if let Some(grant) = self.head.on_lease_renew(&m) {
662                            self.endpoint.send(msg_type::LEASE_GRANT, &grant)?;
663                        }
664                    }
665                    msg_type::NODE_STATUS => {
666                        if let Ok(m) = postcard::from_bytes::<linkf::NodeStatusMsg>(&frame.payload)
667                        {
668                            self.node.capacity = m.capacity;
669                            if !m.loadable_models.is_empty() {
670                                self.node.loadable_models = m.loadable_models;
671                            }
672                        }
673                    }
674                    _ => {}
675                }
676            }
677        }
678        if finish == u32::MAX {
679            expired = true;
680        }
681        self.head.retire_generation(session_id, generation_id);
682        Ok(FleetGeneration {
683            tokens,
684            finish,
685            expired,
686        })
687    }
688
689    pub fn wal(&self) -> &superfluid_linkf::WalStub {
690        self.head.wal()
691    }
692}
693
694pub struct PipelineFleet {
695    stages: Vec<FleetHead>,
696    boundaries: Vec<u32>,
697}
698
699impl PipelineFleet {
700    pub fn connect(nodes: Vec<(String, Vec<u8>)>, n_layers: u32) -> Result<PipelineFleet, DaemonError> {
701        if nodes.is_empty() {
702            return Err(DaemonError::Config("a pipeline needs at least one stage node"));
703        }
704        let p = nodes.len() as u32;
705        let boundaries: Vec<u32> = (0..p).map(|i| ((i + 1) * n_layers) / p).collect();
706        let mut stages = Vec::with_capacity(nodes.len());
707        for (addr, auth) in nodes {
708            stages.push(FleetHead::connect(&addr, auth)?);
709        }
710        Ok(PipelineFleet { stages, boundaries })
711    }
712
713    pub fn boundaries(&self) -> &[u32] {
714        &self.boundaries
715    }
716
717    fn forward_one(&mut self, context: &[u32]) -> Result<u32, DaemonError> {
718        let mut hidden: Option<Vec<u8>> = None;
719        let mut start = 0u32;
720        let n = self.stages.len();
721        for i in 0..n {
722            let end = self.boundaries[i];
723            let empty: Vec<u32> = Vec::new();
724            let tokens: &[u32] = if i == 0 { context } else { &empty };
725            match self.stages[i].forward_stage(tokens, start, end, hidden.take())? {
726                StageOutput::Hidden(h) => hidden = Some(h),
727                StageOutput::Token(t) => return Ok(t),
728            }
729            start = end;
730        }
731        Err(DaemonError::Protocol(
732            "pipeline finished every stage without a token — last stage must reach n_layers",
733        ))
734    }
735
736    pub fn generate(&mut self, context: &[u32], max_tokens: u32) -> Result<Vec<u32>, DaemonError> {
737        let mut ctx = context.to_vec();
738        let mut out = Vec::with_capacity(max_tokens as usize);
739        for _ in 0..max_tokens {
740            let tok = self.forward_one(&ctx)?;
741            out.push(tok);
742            ctx.push(tok);
743        }
744        Ok(out)
745    }
746}
superfluid_daemon/fleet.rs

superfluid_daemon/
fleet.rs
