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

superfluid_daemon/
fleet.rs
