1use std::collections::{HashMap, HashSet, VecDeque};
4use std::path::PathBuf;
5use std::sync::atomic::Ordering;
6use std::sync::mpsc::{Receiver, Sender, TryRecvError};
7use std::sync::{Arc, Mutex};
8
9use superfluid_abi::{encoding, finish, op_state};
10use superfluid_proto::linkw::{self, ALL_SPACES};
11
12use crate::bus::SessionBus;
13use crate::codec::{Channelizer, TextCodec, Utf8Stream};
14use crate::park;
15use crate::runtime::EngineHost;
16use crate::store::{CommittedEvent, SessionStore};
17use crate::wal::GenParams;
18use crate::DaemonError;
19
20const KV_SPACE: u32 = 1;
21const MAX_WORKER_RESPAWNS: u32 = 3;
22const STAGES_PER_ADMISSION: u32 = 2;
24
25fn lane_cap(configured: usize, max_seqs: Option<u64>, emit_slots: Option<u32>) -> usize {
26 let lanes = match max_seqs {
27 Some(n) if (n as usize) < configured => {
28 eprintln!("[superfluid] {configured} lanes asked for; the runtime holds {n} sequences, so {n} run");
29 (n as usize).max(1)
30 }
31 _ => configured,
32 };
33 match emit_slots {
36 Some(n) if (n as usize) < lanes => {
37 eprintln!("[superfluid] {lanes} lanes asked for; the worker's token-out ring holds {n} emits a tick, so {n} run");
38 (n as usize).max(1)
39 }
40 _ => lanes,
41 }
42}
43
44fn ring_slots(host: &mut EngineHost, ring_id: u32) -> Option<u32> {
45 if !host.is_alive() {
46 return None;
47 }
48 host.client()
49 .hello
50 .limits
51 .ring_specs
52 .iter()
53 .find(|r| r.ring_id == ring_id)
54 .map(|r| r.slots)
55}
56
57fn worker_is_gone(e: &DaemonError) -> bool {
58 matches!(
59 e,
60 DaemonError::Agent(
61 superfluid_agent::AgentError::Io(_)
62 | superfluid_agent::AgentError::Closed
63 | superfluid_agent::AgentError::Envelope(_)
64 )
65 )
66}
67
68pub const PREFILL_BUDGET: u32 = 4096;
69pub const SOLO_PREFILL_BUDGET: u32 = 16384;
70const TICK_DECODE: u16 = 32;
71const TICK_DECODE_MAX: u16 = 256;
72const DECODE_SHARE_BESIDE_PREFILL: f64 = 0.25;
73const DECODE_SHARE_BESIDE_PREFILL_INTERACTIVE: f64 = 2.0;
76const TICK_DECODE_BESIDE_PREFILL: u16 = 8;
77const TARGET_TICK_TOTAL_SECS: f64 = 2.0;
78const OPEN_TICK_SECS: f64 = 0.25;
79const MAX_PREFILL_TICK_SECS: f64 = 8.0;
80const MAX_LONE_PREFILL_TICK_SECS: f64 = 16.0;
81const OPEN_PREFILL_TICK_SECS: f64 = 0.5;
84const OPEN_PREFILL_TICK_MAX: u32 = 1024;
87
88const OPEN_PREFILL_TICK_SECS_YIELDING: f64 = 2.0;
91
92const LATENCY_HOLD: std::time::Duration = std::time::Duration::from_secs(300);
95
96const SHARED_PREFIX_MIN_TOKENS: u64 = 1024;
99const SHARED_PREFIX_TAIL_TOKENS: u32 = 64;
101const SHARED_PREFIX_MAX_HOLD_TICKS: u64 = 64;
103const PREFIX_KEEP_MIN_TOKENS: u64 = 1024;
108const PREFIX_KEPT_QUARTERS: u64 = 3;
111const PREFIX_INDEX_MAX: usize = 256;
113const TURN_CHECKPOINT_MIN_TOKENS: u64 = 128;
117const TURN_CHECKPOINTS_TRACKED: usize = 64;
120
121const PREFILL_RATE_MIN_TOKENS: u32 = 256;
122
123
124const SHARED_PREFIX_CLASS: u64 = 2;
127const HOST_EXPORT_CLASS: u64 = 4;
131
132#[derive(Clone, Copy, Debug, PartialEq)]
133struct TickShape {
134 open_secs: f64,
135 full_secs: f64,
136 decode_floor: u16,
137 beside_prefill: u16,
138}
139
140const NATIVE_TICKS: TickShape =
141 TickShape { open_secs: OPEN_TICK_SECS, full_secs: f64::INFINITY, decode_floor: TICK_DECODE, beside_prefill: TICK_DECODE_BESIDE_PREFILL };
142const ROUND_TICKS: TickShape = TickShape { open_secs: 0.05, full_secs: 0.25, decode_floor: 4, beside_prefill: 1 };
143
144fn tick_target(shape: &TickShape, configured: f64, lanes: usize, max_lanes: usize, finishing: bool) -> f64 {
145 if lanes < max_lanes || finishing {
146 configured.min(shape.open_secs)
147 } else {
148 configured.min(shape.full_secs)
149 }
150}
151
152pub struct GenerateJob {
153 pub session: u64,
154 pub stream: Vec<u32>,
155 pub input_end: u64,
158 pub params: GenParams,
159 pub max_tokens: u32,
160 pub events: Sender<LaneEvent>,
161 pub class: u8,
162 pub batch_invariant: bool,
163 pub carry: Option<LaneCarry>,
164 pub queued_at: u64,
165 pub completion: bool,
166 pub deadline: Option<std::time::Instant>,
167 pub media: Vec<(u64, std::path::PathBuf, u32)>,
168 pub grammar_handle: u32,
169 pub extras: GenExtras,
170 pub trace: tracing::Span,
171}
172
173#[derive(Debug, Clone, Default)]
174pub struct GenExtras {
175 pub presence_penalty: f32,
176 pub frequency_penalty: f32,
177 pub repeat_penalty: f32,
178 pub grammar_handle: u32,
179 pub logit_bias_handle: u32,
180 pub want_logprobs: bool,
181 pub top_logprobs: u8,
182 pub wants_deltas: bool,
183 pub ephemeral: bool,
184 pub keep_prefix: bool,
188 pub open_channel: u32,
189 pub thinking: Option<bool>,
190 pub ignore_eos: bool,
191 pub seed_drawn: bool,
192 pub tool_schemas: Option<std::sync::Arc<crate::codec::ToolSchemas>>,
193}
194
195pub struct LaneCarry {
196 out: LaneOutput,
197 produced_before: u32,
198 spec_before: (u64, u64),
199 decode_before: (u64, f64),
200 warm_prefix_first: Option<u64>,
201}
202
203#[derive(Debug, Clone)]
204pub struct TokenLogprob {
205 pub token: u32,
206 pub logprob: f32,
207 pub top: Vec<(u32, f32)>,
208}
209
210pub enum LaneEvent {
211 Committed(CommittedEvent),
212 Provisional { channel: u32, text: String, produced: u32, closes: bool },
215 TokenLogprobs(Vec<TokenLogprob>),
216 CompletionTokens(Vec<u32>),
217 Expired,
218 Done {
219 tokens_generated: u32,
220 finish: u32,
221 warm_prefix: u64,
222 spec: (u64, u64),
223 decode: (u64, f64),
226 },
227 Failed(String),
228 Error(DaemonError),
231}
232
233pub struct LaneOutput {
238 seg: Box<dyn Channelizer>,
239 utf8: Utf8Stream,
240 disp_seg: Box<dyn Channelizer>,
241 disp_utf8: Utf8Stream,
242 tool_buf: String,
243 disp_tokens: u32,
244}
245
246pub struct OutputSink<'a> {
248 pub codec: &'a Arc<dyn TextCodec + Send + Sync>,
249 pub store: &'a Arc<Mutex<SessionStore>>,
250 pub bus: &'a SessionBus,
251 pub tool_lease_ms: u64,
252}
253
254impl LaneOutput {
255 pub fn new(codec: &Arc<dyn TextCodec + Send + Sync>, open_channel: u32, produced_before: u32) -> LaneOutput {
258 let mk = || {
259 let mut seg = codec.channelizer();
260 if open_channel != 0 {
261 seg.prime(open_channel);
262 }
263 seg
264 };
265 LaneOutput {
266 seg: mk(),
267 utf8: Utf8Stream::default(),
268 disp_seg: mk(),
269 disp_utf8: Utf8Stream::default(),
270 tool_buf: String::new(),
271 disp_tokens: produced_before,
272 }
273 }
274
275 pub fn provisional(
277 &mut self,
278 codec: &Arc<dyn TextCodec + Send + Sync>,
279 bus: &SessionBus,
280 session: u64,
281 events: &Sender<LaneEvent>,
282 tokens: &[u32],
283 ) {
284 if tokens.is_empty() {
285 return;
286 }
287 self.disp_tokens += tokens.len() as u32;
288 let produced = self.disp_tokens;
289 for run in self.disp_seg.split(tokens) {
290 let mut text = String::new();
291 for &t in &run.text {
292 text.push_str(&self.disp_utf8.push(&codec.token_bytes(t)));
293 }
294 let channel = run.channel;
295 let tool = channel == crate::wal::channel::TOOL_CALL;
296 let closes = tool && run.closes;
299 if text.is_empty() && !closes {
300 continue;
301 }
302 let text_like = channel == crate::wal::channel::TEXT || channel == crate::wal::channel::REASONING;
303 if text_like {
304 bus.publish_provisional(session, channel, &text);
305 }
306 if text_like || tool {
309 let _ = events.send(LaneEvent::Provisional { channel, text, produced, closes });
310 }
311 }
312 }
313
314 pub fn provisional_flush(&mut self, bus: &SessionBus, session: u64, events: &Sender<LaneEvent>) {
315 let text = self.disp_utf8.flush();
316 if text.is_empty() {
317 return;
318 }
319 let channel = self.disp_seg.current_channel();
320 let text_like = channel == crate::wal::channel::TEXT || channel == crate::wal::channel::REASONING;
321 if text_like {
322 bus.publish_provisional(session, channel, &text);
323 }
324 if text_like || channel == crate::wal::channel::TOOL_CALL {
325 let produced = self.disp_tokens;
326 let _ = events.send(LaneEvent::Provisional { channel, text, produced, closes: false });
327 }
328 }
329
330 pub fn close(
333 &mut self,
334 sink: &OutputSink<'_>,
335 session: u64,
336 events: &Sender<LaneEvent>,
337 fin: u32,
338 quiet: bool,
339 ) -> Result<(), DaemonError> {
340 self.provisional_flush(sink.bus, session, events);
341 let text = self.utf8.flush();
342 let e = sink.store.lock().expect("store").commit_generated(
343 session,
344 Vec::new(),
345 text,
346 self.seg.current_channel(),
347 fin,
348 )?;
349 if !quiet {
350 let _ = events.send(LaneEvent::Committed(e));
351 }
352 Ok(())
353 }
354
355 #[allow(clippy::too_many_arguments)]
359 pub fn commit(
360 &mut self,
361 sink: &OutputSink<'_>,
362 session: u64,
363 events: &Sender<LaneEvent>,
364 tokens: &[u32],
365 fin: u32,
366 terminal: bool,
367 schemas: Option<&crate::codec::ToolSchemas>,
368 ) -> Result<bool, DaemonError> {
369 if terminal {
370 self.provisional_flush(sink.bus, session, events);
371 }
372 let runs = self.seg.split(tokens);
373 let n_runs = runs.len().max(1);
374 for (i, run) in runs.into_iter().enumerate() {
375 let last = i + 1 == n_runs;
376 let mut text = String::new();
377 for &t in &run.text {
378 text.push_str(&self.utf8.push(&sink.codec.token_bytes(t)));
379 }
380 if last && terminal {
381 text.push_str(&self.utf8.flush());
382 }
383 let channel = run.channel;
384 let tool_closed = run.closes && channel == crate::wal::channel::TOOL_CALL;
385 if channel == crate::wal::channel::TOOL_CALL {
386 self.tool_buf.push_str(&text);
387 }
388 let e = sink.store.lock().expect("store").commit_generated(
389 session,
390 run.span,
391 text,
392 channel,
393 if last { fin } else { finish::NONE },
394 )?;
395 let tool_event = if tool_closed {
396 let raw = std::mem::take(&mut self.tool_buf);
397 Scheduler::commit_tool_block(sink.codec, sink.store, sink.tool_lease_ms, session, raw, schemas)?
398 } else {
399 None
400 };
401 if events.send(LaneEvent::Committed(e)).is_err() {
402 return Ok(false);
403 }
404 if let Some(te) = tool_event {
405 if events.send(LaneEvent::Committed(te)).is_err() {
406 return Ok(false);
407 }
408 }
409 }
410 if terminal && !self.tool_buf.is_empty() {
411 let raw = std::mem::take(&mut self.tool_buf);
412 let te = Scheduler::commit_tool_block(sink.codec, sink.store, sink.tool_lease_ms, session, raw, schemas)?;
413 if let Some(te) = te {
414 let _ = events.send(LaneEvent::Committed(te));
415 }
416 }
417 Ok(true)
418 }
419}
420
421pub enum SchedCmd {
422 Submit(Box<GenerateJob>),
423 Sync(linkw::StateSyncReqMsg, Sender<Result<linkw::StateSyncOkMsg, String>>),
424 SyncAs(u8, Box<SchedCmd>),
425 SyncStream(
426 linkw::StateSyncReqMsg,
427 Sender<linkw::TranscribeSegmentMsg>,
428 Sender<Result<linkw::StateSyncOkMsg, String>>,
429 ),
430 Pin(crate::pins::PinReq),
431 PinHeld(Vec<crate::pins::PinKey>, Sender<u64>),
432 PinUnanchor(u64),
433 Shutdown,
434}
435
436pub struct SchedulerHandle {
437 pub tx: Sender<SchedCmd>,
438 join: std::sync::Mutex<Option<std::thread::JoinHandle<()>>>,
439}
440
441impl SchedulerHandle {
442 pub(crate) fn new(tx: Sender<SchedCmd>, join: std::thread::JoinHandle<()>) -> SchedulerHandle {
444 SchedulerHandle { tx, join: std::sync::Mutex::new(Some(join)) }
445 }
446
447 pub fn signal_shutdown(&self) {
448 let _ = self.tx.send(SchedCmd::Shutdown);
449 }
450
451 pub fn shutdown(&self) {
452 let mut guard = self.join.lock().expect("join");
453 if let Some(j) = guard.take() {
454 self.signal_shutdown();
455 let _ = j.join();
456 }
457 }
458}
459
460impl Drop for SchedulerHandle {
461 fn drop(&mut self) {
462 self.shutdown();
463 }
464}
465
466enum LanePhase {
467 Prefilling,
468 Decoding,
469 Finishing,
470}
471
472struct Lane {
473 tag: u64,
474 session: u64,
475 job: GenerateJob,
476 phase: LanePhase,
477 prefilled: u64,
478 produced: u32,
479 fin: u32,
480 warm_prefix: u64,
481 out: LaneOutput,
482 frame_fed: usize,
483 emit_dead: bool,
484 admitted: bool,
485 seed_handle: u64,
486 span: linkw::TokenRefMsg,
487 all_tokens: Vec<u32>,
488 parked: bool,
489 parked_lossy: bool,
490 preempted: bool,
491 starved_ticks: u32,
492 prefill_starved_ticks: u32,
493 checkpoint_at: u64,
497 checkpoint_due: bool,
498 checkpointed: u64,
499 turn_kept: bool,
501 produced_before: u32,
502 spec: (u64, u64),
503 decode: (u64, f64),
509 decode_ticks: Option<(std::time::Instant, u32, std::time::Instant, u32)>,
512 warm_prefix_first: Option<u64>,
513 admitted_at: std::time::Instant,
514 ttft_observed: bool,
515 ttft_frame_ms: Option<u64>,
516 fault_code: Option<u32>,
517}
518
519impl Lane {
520 fn decode_measured(&self) -> (u64, f64) {
522 let mut d = self.decode;
523 if let Some((at, by, last_at, last_by)) = self.decode_ticks {
524 if last_by > by {
525 d.0 += u64::from(last_by - by);
526 d.1 += last_at.saturating_duration_since(at).as_secs_f64();
527 }
528 }
529 d
530 }
531}
532
533#[derive(Default)]
534struct ClassQueue {
535 jobs: VecDeque<GenerateJob>,
536}
537
538pub fn kv_bits_label(bits: u64) -> Option<&'static str> {
539 match bits {
540 16 => Some("f16"),
541 8 => Some("q8_0"),
542 4 => Some("q4_0"),
543 84 => Some("K q8_0 / V q4_0"),
544 _ => None,
545 }
546}
547
548#[derive(Debug, Default)]
549pub struct SchedStats {
550 pub pressure_evictions: std::sync::atomic::AtomicU64,
551 pub pressure_bytes_evicted: std::sync::atomic::AtomicU64,
552 pub unservable_cold_admissions: std::sync::atomic::AtomicU64,
553 pub unservable_park_refusals: std::sync::atomic::AtomicU64,
554 pub preemptions: std::sync::atomic::AtomicU64,
555 pub admits_deferred: std::sync::atomic::AtomicU64,
556 pub starvation_grants: std::sync::atomic::AtomicU64,
557 pub os_pressure_events: std::sync::atomic::AtomicU64,
558 pub ticks: std::sync::atomic::AtomicU64,
559 pub lanes_active: std::sync::atomic::AtomicU64,
560 pub queue_depth: [std::sync::atomic::AtomicU64; 4],
561 pub prefill_tokens: std::sync::atomic::AtomicU64,
562 pub decode_tokens: std::sync::atomic::AtomicU64,
563 pub warm_prefix_tokens: std::sync::atomic::AtomicU64,
564 pub cold_admissions: std::sync::atomic::AtomicU64,
565 pub parks_lossless: std::sync::atomic::AtomicU64,
566 pub parks_lossy: std::sync::atomic::AtomicU64,
567 pub resumes: std::sync::atomic::AtomicU64,
568 pub worker_respawns: std::sync::atomic::AtomicU64,
569 pub spec_proposed: std::sync::atomic::AtomicU64,
570 pub spec_accepted: std::sync::atomic::AtomicU64,
571 pub ttft_buckets: [std::sync::atomic::AtomicU64; TTFT_BUCKETS_MS.len() + 1],
572 pub ttft_sum_ms: std::sync::atomic::AtomicU64,
573 pub ttft_count: std::sync::atomic::AtomicU64,
574 pub ttft_last_ms: std::sync::atomic::AtomicU64,
575 pub pool_blocks_used: std::sync::atomic::AtomicU64,
576 pub prefill_budget_live: std::sync::atomic::AtomicU64,
577 pub decode_grant_live: std::sync::atomic::AtomicU64,
578 pub tick_target_live_ms: std::sync::atomic::AtomicU64,
579 pub pool_blocks_total: std::sync::atomic::AtomicU64,
580 pub kv_bits: std::sync::atomic::AtomicU64,
581 pub pins_held: std::sync::atomic::AtomicU64,
582 pub pinned_bytes: std::sync::atomic::AtomicU64,
583 pub pins_yielded: std::sync::atomic::AtomicU64,
584 pub pins_expired: std::sync::atomic::AtomicU64,
585 pub tick_wall_ms_live: std::sync::atomic::AtomicU64,
586 pub prefill_rate_live: std::sync::atomic::AtomicU64,
587 pub decode_rate_live: std::sync::atomic::AtomicU64,
588 pub work_ticks: std::sync::atomic::AtomicU64,
589 tick_seq: std::sync::atomic::AtomicU64,
590 tick_since_ms: std::sync::atomic::AtomicU64,
591 tick_prefill_tokens: std::sync::atomic::AtomicU64,
592 tick_prefill_lanes: std::sync::atomic::AtomicU64,
593 tick_decode_lanes: std::sync::atomic::AtomicU64,
594 tick_admits: std::sync::atomic::AtomicU64,
595 tick_retires: std::sync::atomic::AtomicU64,
596 tick_target_ms: std::sync::atomic::AtomicU64,
597}
598
599#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
600pub struct TickInFlight {
601 pub since_ms: u64,
602 pub prefill_tokens: u64,
603 pub prefill_lanes: u64,
604 pub decode_lanes: u64,
605 pub admits: u64,
606 pub retires: u64,
607 pub target_ms: u64,
608}
609
610impl SchedStats {
611 pub fn set_tick_in_flight(&self, t: TickInFlight) {
612 use std::sync::atomic::Ordering::SeqCst;
613 self.tick_seq.fetch_add(1, SeqCst);
614 self.tick_since_ms.store(t.since_ms, SeqCst);
615 self.tick_prefill_tokens.store(t.prefill_tokens, SeqCst);
616 self.tick_prefill_lanes.store(t.prefill_lanes, SeqCst);
617 self.tick_decode_lanes.store(t.decode_lanes, SeqCst);
618 self.tick_admits.store(t.admits, SeqCst);
619 self.tick_retires.store(t.retires, SeqCst);
620 self.tick_target_ms.store(t.target_ms, SeqCst);
621 self.tick_seq.fetch_add(1, SeqCst);
622 }
623
624 pub fn clear_tick_in_flight(&self) {
625 self.set_tick_in_flight(TickInFlight::default());
626 }
627
628 pub fn tick_in_flight(&self) -> Option<TickInFlight> {
629 use std::sync::atomic::Ordering::SeqCst;
630 loop {
631 let s1 = self.tick_seq.load(SeqCst);
632 if s1 & 1 == 1 {
633 std::hint::spin_loop();
634 continue;
635 }
636 let t = TickInFlight {
637 since_ms: self.tick_since_ms.load(SeqCst),
638 prefill_tokens: self.tick_prefill_tokens.load(SeqCst),
639 prefill_lanes: self.tick_prefill_lanes.load(SeqCst),
640 decode_lanes: self.tick_decode_lanes.load(SeqCst),
641 admits: self.tick_admits.load(SeqCst),
642 retires: self.tick_retires.load(SeqCst),
643 target_ms: self.tick_target_ms.load(SeqCst),
644 };
645 if self.tick_seq.load(SeqCst) == s1 {
646 return (t.since_ms != 0).then_some(t);
647 }
648 }
649 }
650}
651
652pub fn uptime_ms() -> u64 {
653 static START: std::sync::OnceLock<std::time::Instant> = std::sync::OnceLock::new();
654 START.get_or_init(std::time::Instant::now).elapsed().as_millis() as u64
655}
656
657fn per_sec(tokens: u64, secs: f64) -> u64 {
658 if tokens == 0 || secs <= 0.0 {
659 0
660 } else {
661 (tokens as f64 / secs) as u64
662 }
663}
664
665pub const TTFT_BUCKETS_MS: [u64; 13] =
666 [5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10_000, 20_000, 40_000];
667
668impl SchedStats {
669 pub fn observe_ttft(&self, ms: u64) {
670 let idx = TTFT_BUCKETS_MS
671 .iter()
672 .position(|&b| ms <= b)
673 .unwrap_or(TTFT_BUCKETS_MS.len());
674 self.ttft_buckets[idx].fetch_add(1, Ordering::Relaxed);
675 self.ttft_sum_ms.fetch_add(ms, Ordering::Relaxed);
676 self.ttft_count.fetch_add(1, Ordering::Relaxed);
677 self.ttft_last_ms.store(ms, Ordering::Relaxed);
678 }
679}
680
681#[derive(Debug, Clone, PartialEq, Eq)]
682pub struct SpeculationGrant {
683 pub strategy_id: String,
684 pub slot: u32,
685 pub class: u8,
686 pub cert_id: u32,
687 pub identity: u64,
688 pub modes: u8,
689}
690
691fn proposal_config(id: &str, opts: &crate::SpeculationOptions) -> String {
692 fn env_depth(key: &str) -> Option<u32> {
693 std::env::var(key)
694 .ok()
695 .and_then(|v| v.trim().parse::<u32>().ok())
696 .filter(|n| (1..=15).contains(n))
697 }
698 fn env_bool(key: &str, default: bool) -> bool {
699 std::env::var(key)
700 .ok()
701 .map(|v| v.trim().parse::<i32>().unwrap_or(i32::from(default)) != 0)
702 .unwrap_or(default)
703 }
704 let strategy_depth = || match id {
705 "mtp-head" => env_depth("SUPERFLUID_MTP_DRAFTS").unwrap_or(0),
706 "eagle3" => env_depth("SUPERFLUID_EAGLE_DRAFTS").unwrap_or(2),
707 "dflash" | "dspark" => env_depth("SUPERFLUID_DFLASH_DRAFTS").unwrap_or(0),
708 _ => 4,
709 };
710 let max_draft = opts.max_draft.unwrap_or_else(strategy_depth);
711 let adaptive = opts
712 .adaptive
713 .unwrap_or_else(|| env_bool("SUPERFLUID_SPEC_ADAPTIVE", true));
714 let min_yield = opts.min_yield.unwrap_or_else(|| {
715 std::env::var("SUPERFLUID_SPEC_MIN_YIELD")
716 .ok()
717 .and_then(|v| v.trim().parse::<f64>().ok())
718 .filter(|v| v.is_finite() && *v >= 0.0)
719 .unwrap_or(0.75)
720 });
721 let yield_rounds = opts.yield_rounds.unwrap_or_else(|| {
722 std::env::var("SUPERFLUID_SPEC_LANE_YIELD_ROUNDS")
723 .ok()
724 .and_then(|v| v.trim().parse::<u32>().ok())
725 .filter(|n| *n > 0)
726 .unwrap_or(24)
727 });
728 let throughput_gate = opts
729 .throughput_gate
730 .unwrap_or_else(|| env_bool("SUPERFLUID_SPEC_GATE", true));
731 let gate_probe_tokens = opts.gate_probe_tokens.unwrap_or_else(|| {
732 std::env::var("SUPERFLUID_SPEC_GATE_PROBE")
733 .ok()
734 .and_then(|v| v.trim().parse::<u32>().ok())
735 .unwrap_or(16)
736 });
737 let min_speedup = opts.min_speedup.unwrap_or_else(|| {
738 std::env::var("SUPERFLUID_SPEC_MIN_SPEEDUP")
739 .ok()
740 .and_then(|v| v.trim().parse::<f64>().ok())
741 .filter(|v| v.is_finite() && *v >= 1.0)
742 .unwrap_or(1.08)
743 });
744 let env_u32 = |key: &str, default: u32| {
745 std::env::var(key)
746 .ok()
747 .and_then(|v| v.trim().parse::<u32>().ok())
748 .unwrap_or(default)
749 };
750 let gate_reprobe = opts
751 .gate_reprobe
752 .unwrap_or_else(|| env_u32("SUPERFLUID_SPEC_GATE_REPROBE", 32));
753 let gate_reprobe_max = opts
754 .gate_reprobe_max
755 .unwrap_or_else(|| env_u32("SUPERFLUID_SPEC_GATE_REPROBE_MAX", 1024));
756 let bitexact = env_bool("BASERT_SPEC_BITEXACT", false);
757 let mut cfg = format!(
758 "max_draft={max_draft},adaptive={},min_yield={min_yield:.6},yield_rounds={yield_rounds},throughput_gate={},gate_probe_tokens={gate_probe_tokens},min_speedup={min_speedup:.6},gate_reprobe={gate_reprobe},gate_reprobe_max={gate_reprobe_max},bitexact={}",
759 u8::from(adaptive),
760 u8::from(throughput_gate),
761 u8::from(bitexact),
762 );
763 if id == "dspark" {
764 let confidence = opts.dspark_confidence.or_else(|| {
765 std::env::var("SUPERFLUID_DSPARK_CONFIDENCE")
766 .ok()
767 .and_then(|v| v.trim().parse::<f32>().ok())
768 .filter(|t| t.is_finite() && *t > 0.0)
769 });
770 if let Some(t) = confidence {
771 cfg.push_str(&format!(",confidence={t:.6}"));
772 }
773 }
774 cfg
775}
776
777fn prefill_tick_cap(secs_per_token: f64, lone: bool) -> u32 {
778 let secs = if lone { MAX_LONE_PREFILL_TICK_SECS } else { MAX_PREFILL_TICK_SECS };
779 (secs / secs_per_token).clamp(512.0, PREFILL_BUDGET as f64) as u32
780}
781
782fn open_prefill_tick_target(yields: bool, latency: bool) -> f64 {
787 if yields && !latency {
788 OPEN_PREFILL_TICK_SECS_YIELDING
789 } else {
790 OPEN_PREFILL_TICK_SECS
791 }
792}
793
794fn open_prefill_tick_cap(secs_per_token: f64, target_secs: f64) -> u32 {
795 let tokens = (target_secs / secs_per_token).clamp(128.0, OPEN_PREFILL_TICK_MAX as f64) as u32;
796 let step = if tokens < 512 { 128 } else { 512 };
797 tokens / step * step
798}
799
800enum CheckpointError {
803 Refused(String),
804 Failed(String),
805}
806
807impl CheckpointError {
808 fn of(e: superfluid_agent::AgentError) -> CheckpointError {
811 if matches!(e, superfluid_agent::AgentError::Rejected(s) if s == superfluid_abi::Status::Unsupported as i32) {
812 CheckpointError::Refused(e.to_string())
813 } else {
814 CheckpointError::Failed(e.to_string())
815 }
816 }
817}
818
819struct PrefillBudget {
820 left: u64,
821 tail_capped: bool,
822 held_below: Option<u8>,
823}
824
825impl PrefillBudget {
826 const TAIL_FLOOR: u64 = 256;
827
828 fn new(budget: u64) -> Self {
829 PrefillBudget { left: budget, tail_capped: false, held_below: None }
830 }
831
832 fn offer(&self, want: u64) -> u64 {
833 want.min(self.left)
834 }
835
836 fn offer_to(&self, want: u64, class: u8) -> u64 {
837 if self.held_below.is_some_and(|c| class > c) {
838 return 0;
839 }
840 self.offer(want)
841 }
842
843 fn hold_below(&mut self, class: u8) {
847 self.held_below = Some(self.held_below.map_or(class, |c| c.min(class)));
848 }
849
850 fn take(&mut self, chunk: u64, completes: bool) {
851 self.left -= chunk;
852 if completes && !self.tail_capped {
853 self.tail_capped = true;
854 self.left = self.left.min(chunk.max(Self::TAIL_FLOOR));
855 }
856 }
857}
858
859fn with_first_admits(plan: &linkw::TickPlanMsg, keep: usize) -> linkw::TickPlanMsg {
862 let mut p = plan.clone();
863 let dropped: std::collections::HashSet<u64> = p.admits.split_off(keep.min(p.admits.len())).iter().map(|a| a.lane_tag).collect();
864 p.prefills.retain(|x| !dropped.contains(&x.lane_tag));
865 p.decodes.retain(|x| !dropped.contains(&x.lane_tag));
866 p.commits.retain(|x| !dropped.contains(&x.lane_tag));
867 p
868}
869
870fn starved_after(prev: u32, waiting: bool, youngest_served: Option<u64>, tag: u64) -> u32 {
874 if waiting && youngest_served.is_some_and(|y| y > tag) {
875 prev.saturating_add(1)
876 } else {
877 prev
878 }
879}
880
881fn prefill_rank(class: u8, remaining: u64, tag: u64, starved_ticks: u32, bound: u32) -> (bool, u8, u64, u64) {
882 let starving = remaining > 0 && starved_ticks >= bound;
883 (!starving, class, remaining, tag)
884}
885
886#[cfg(test)]
887mod prefill_order_tests {
888 use super::{prefill_rank, prefill_tick_cap, starved_after, PrefillBudget, PREFILL_BUDGET};
889
890 #[test]
891 fn prompts_that_arrived_together_prefill_one_after_another() {
892 const BOUND: u32 = 8;
896 let mut lanes: Vec<(u64, u64, u32)> = (0..8).map(|t| (t, 7680, 0)).collect();
897 let mut finished_at = Vec::new();
898 for tick in 0..400u32 {
899 let mut order: Vec<usize> = (0..lanes.len()).filter(|&i| lanes[i].1 > 0).collect();
900 if order.is_empty() {
901 break;
902 }
903 order.sort_by_key(|&i| prefill_rank(3, lanes[i].1, lanes[i].0, lanes[i].2, BOUND));
904 let served = order[0];
905 lanes[served].1 -= 512;
906 lanes[served].2 = 0;
907 if lanes[served].1 == 0 {
908 finished_at.push(tick + 1);
909 }
910 let youngest = Some(lanes[served].0);
911 for (i, l) in lanes.iter_mut().enumerate() {
912 l.2 = starved_after(l.2, i != served && l.1 > 0, youngest, l.0);
913 }
914 }
915 assert_eq!(finished_at, (1..=8).map(|k| 15 * k).collect::<Vec<_>>(), "fifteen ticks each, in turn");
916 }
917
918 #[test]
919 fn only_an_engine_that_cannot_fork_refuses_checkpoints_for_good() {
920 use super::CheckpointError;
921 use superfluid_abi::Status;
922 use superfluid_agent::AgentError;
923 let refused = |e| matches!(CheckpointError::of(e), CheckpointError::Refused(_));
924 assert!(refused(AgentError::Rejected(Status::Unsupported as i32)));
925 assert!(!refused(AgentError::Rejected(Status::UnknownHandle as i32)), "a lane that faulted in its tick");
926 assert!(!refused(AgentError::Rejected(Status::NeedsReplan as i32)), "a full sequence table");
927 assert!(!refused(AgentError::Closed), "a worker that went away");
928 }
929
930 #[test]
931 fn an_open_tick_prefills_whole_chunks_for_about_half_a_second() {
932 use super::open_prefill_tick_cap;
933 assert_eq!(open_prefill_tick_cap(1.0 / 1200.0, super::OPEN_PREFILL_TICK_SECS), 512, "600 tokens round down to one chunk");
934 assert_eq!(open_prefill_tick_cap(1.0 / 5000.0, super::OPEN_PREFILL_TICK_SECS), 1024, "2500 tokens, at most two chunks");
935 assert_eq!(open_prefill_tick_cap(1.0, super::OPEN_PREFILL_TICK_SECS), 128, "never below one step");
936 assert_eq!(open_prefill_tick_cap(1.0 / 330.0, super::OPEN_PREFILL_TICK_SECS), 128, "a 27B: 165 tokens round down to one step");
937 assert_eq!(open_prefill_tick_cap(1.0 / 700.0, super::OPEN_PREFILL_TICK_SECS), 256, "350 tokens, two steps");
938 assert_eq!(open_prefill_tick_cap(1e-6, super::OPEN_PREFILL_TICK_SECS), 1024, "never above two chunks, whatever an early tick suggests");
939 assert_eq!(
940 open_prefill_tick_cap(1.0 / 330.0, super::OPEN_PREFILL_TICK_SECS_YIELDING),
941 512,
942 "an engine that yields mid-tick: a 27B's open tick is a whole chunk"
943 );
944 use super::open_prefill_tick_target as target;
945 assert_eq!(target(true, false), super::OPEN_PREFILL_TICK_SECS_YIELDING, "agents alone, on an engine that yields");
946 assert_eq!(target(true, true), super::OPEN_PREFILL_TICK_SECS, "a chat around: its next tokens wait out each tick");
947 assert_eq!(target(false, false), super::OPEN_PREFILL_TICK_SECS, "an engine that runs a tick whole");
948 }
949
950 #[test]
951 fn lone_lane_plans_longer_prefill_ticks() {
952 let spt = 0.0078125;
953 assert_eq!(prefill_tick_cap(spt, false), 1024);
954 assert_eq!(prefill_tick_cap(spt, true), 2048);
955 assert_eq!(prefill_tick_cap(1.0, true), 512);
956 assert_eq!(prefill_tick_cap(1e-6, false), PREFILL_BUDGET);
957 assert_eq!(prefill_tick_cap(1e-6, true), PREFILL_BUDGET);
958 }
959
960 fn grants(budget: u64, remaining: &[u64]) -> Vec<u64> {
961 let mut b = PrefillBudget::new(budget);
962 remaining
963 .iter()
964 .map(|&r| {
965 let c = b.offer(r);
966 if c > 0 {
967 b.take(c, c == r);
968 }
969 c
970 })
971 .collect()
972 }
973
974 #[test]
975 fn the_tail_after_the_first_completion_is_capped_in_aggregate() {
976 let g = grants(4096, &[500, 4000, 4000, 4000, 4000, 4000, 4000, 4000]);
977 assert_eq!(g, vec![500, 500, 0, 0, 0, 0, 0, 0]);
978 let g = grants(4096, &[40, 4000, 4000]);
979 assert_eq!(g, vec![40, 256, 0]);
980 let g = grants(4096, &[1000, 300, 200, 4000]);
981 assert_eq!(g, vec![1000, 300, 200, 500]);
982 assert_eq!(grants(4096, &[3000, 3000]), vec![3000, 1096]);
983 }
984
985 #[test]
986 fn nothing_of_a_lower_class_prefills_beside_a_chat_that_finishes_its_prompt() {
987 let mut b = PrefillBudget::new(4096);
988 assert_eq!(b.offer_to(22, 0), 22);
989 b.take(22, true);
990 b.hold_below(0);
991 assert_eq!(b.offer_to(4000, 3), 0, "an agent waits a tick");
992 assert_eq!(b.offer_to(100, 0), 100, "another chat does not");
993 let mut b = PrefillBudget::new(4096);
994 b.take(500, false);
995 assert_eq!(b.offer_to(4000, 3), 3596, "a chat still prefilling holds nothing back");
996 }
997
998 #[test]
999 fn a_long_prefill_progresses_under_a_stream_of_short_prompts() {
1000 const BUDGET: u64 = 512;
1001 const BOUND: u32 = 3;
1002 let mut lanes: Vec<(u64, u64, u32)> = vec![(0, 8 * BUDGET, 0)];
1003 let mut next_tag = 1;
1004 let mut long_ticks = Vec::new();
1005 for tick in 0..64u32 {
1006 lanes.retain(|l| l.0 == 0 || l.1 > 0);
1007 while lanes.len() < 3 {
1008 lanes.push((next_tag, BUDGET, 0));
1009 next_tag += 1;
1010 }
1011 let mut order: Vec<usize> = (0..lanes.len()).collect();
1012 order.sort_by_key(|&i| prefill_rank(0, lanes[i].1, lanes[i].0, lanes[i].2, BOUND));
1013 let mut budget = PrefillBudget::new(BUDGET);
1014 let mut served = Vec::new();
1015 for i in order {
1016 let l = &mut lanes[i];
1017 if l.1 == 0 {
1018 continue;
1019 }
1020 let chunk = budget.offer(l.1);
1021 l.1 -= chunk;
1022 if chunk > 0 {
1023 l.2 = 0;
1024 served.push(l.0);
1025 budget.take(chunk, l.1 == 0);
1026 }
1027 if l.0 == 0 && chunk > 0 {
1028 long_ticks.push(tick);
1029 }
1030 }
1031 let youngest = served.iter().copied().max();
1032 for l in lanes.iter_mut() {
1033 l.2 = starved_after(l.2, l.1 > 0 && !served.contains(&l.0), youngest, l.0);
1034 }
1035 assert!(lanes[0].2 <= BOUND, "tick {tick}: the long lane waited past the bound");
1036 if lanes[0].1 == 0 {
1037 break;
1038 }
1039 }
1040 assert_eq!(lanes[0], (0, 0, 0), "the long prompt finished prefilling: {long_ticks:?}");
1041 assert!(long_ticks.windows(2).all(|w| w[1] - w[0] <= BOUND + 1), "bounded gaps: {long_ticks:?}");
1042 assert!(prefill_rank(0, 8 * BUDGET, 0, 0, BOUND) > prefill_rank(0, BUDGET, 9, 0, BOUND));
1043 assert!(prefill_rank(5, 8 * BUDGET, 0, BOUND, BOUND) < prefill_rank(0, 1, 9, 0, BOUND));
1044 assert!(prefill_rank(5, 0, 0, BOUND, BOUND) > prefill_rank(0, 1, 9, 0, BOUND));
1045 }
1046}
1047
1048#[cfg(test)]
1049mod tick_target_tests {
1050 use super::{tick_target, NATIVE_TICKS, OPEN_TICK_SECS, ROUND_TICKS};
1051
1052 #[test]
1053 fn a_tick_is_short_while_a_joiner_could_be_admitted_next() {
1054 assert_eq!(tick_target(&NATIVE_TICKS, 2.0, 1, 2, false), OPEN_TICK_SECS, "a slot is free");
1055 assert_eq!(tick_target(&NATIVE_TICKS, 2.0, 2, 2, false), 2.0, "full and staying full: nothing could admit a joiner");
1056 assert_eq!(tick_target(&ROUND_TICKS, 2.0, 1, 2, false), 0.05, "a round-granular runtime's open tick is a few rounds");
1057 assert_eq!(tick_target(&ROUND_TICKS, 2.0, 2, 2, false), 0.25, "and its full one is never long");
1058 assert_eq!(
1059 tick_target(&NATIVE_TICKS, 2.0, 2, 2, true),
1060 OPEN_TICK_SECS,
1061 "a finishing lane frees a slot, and its client waits on the retire"
1062 );
1063 assert_eq!(tick_target(&NATIVE_TICKS, 0.1, 2, 2, true), 0.1, "a configured target under the open one stands");
1064 }
1065}
1066
1067#[cfg(test)]
1068mod speculation_config_tests {
1069 use super::proposal_config;
1070 use crate::SpeculationOptions;
1071
1072 #[test]
1073 fn cli_policy_is_serialized_for_the_engine_and_fingerprint() {
1074 let cfg = proposal_config(
1075 "dspark",
1076 &SpeculationOptions {
1077 max_draft: Some(2),
1078 adaptive: Some(false),
1079 min_yield: Some(0.5),
1080 yield_rounds: Some(12),
1081 throughput_gate: Some(false),
1082 gate_probe_tokens: Some(8),
1083 min_speedup: Some(1.12),
1084 gate_reprobe: Some(0),
1085 gate_reprobe_max: Some(64),
1086 dspark_confidence: Some(0.42),
1087 },
1088 );
1089 assert!(cfg.starts_with(
1090 "max_draft=2,adaptive=0,min_yield=0.500000,yield_rounds=12,throughput_gate=0,gate_probe_tokens=8,min_speedup=1.120000,gate_reprobe=0,gate_reprobe_max=64,"
1091 ));
1092 assert!(cfg.ends_with(",confidence=0.420000"));
1093 }
1094
1095 #[test]
1096 fn one_draft_override_applies_to_mtp() {
1097 let cfg = proposal_config(
1098 "mtp-head",
1099 &SpeculationOptions {
1100 max_draft: Some(1),
1101 ..SpeculationOptions::default()
1102 },
1103 );
1104 assert!(cfg.starts_with("max_draft=1,"));
1105 assert!(!cfg.contains("confidence="));
1106 }
1107}
1108
1109#[derive(Default)]
1110pub struct Speculation {
1111 reg: Mutex<Option<linkw::StrategyRegisterMsg>>,
1112 grant: Mutex<Option<SpeculationGrant>>,
1113}
1114
1115impl Speculation {
1116 pub fn register(
1117 this: &Arc<Speculation>,
1118 host: &mut EngineHost,
1119 strategy: &str,
1120 opts: &crate::SpeculationOptions,
1121 ) -> Result<(), DaemonError> {
1122 let (id, draft_path) = match strategy.split_once(':') {
1123 Some(("draft-model", path)) => ("draft-model", Some(path.to_string())),
1124 Some(("dflash", path)) => ("dflash", Some(path.to_string())),
1125 Some(("dspark", path)) => ("dspark", Some(path.to_string())),
1126 Some(("eagle3", path)) => ("eagle3", Some(path.to_string())),
1127 _ if strategy == "prompt-lookup" => ("prompt-lookup", None),
1128 _ if strategy == "mtp-head" => ("mtp-head", None),
1129 _ => {
1130 return Err(DaemonError::Config(
1131 "unknown speculation strategy (prompt-lookup | mtp-head | dflash / dspark / eagle3 / draft-model followed by ':' and a model id or .base path)",
1132 ))
1133 }
1134 };
1135 let draft_identity = match draft_path.as_ref() {
1136 Some(p) => Some(draft_weights_identity(p)?),
1137 None => None,
1138 };
1139 let impl_hash = {
1140 let mut h = [0u8; 32];
1141 let tag = if id == "dflash" {
1142 "dflash/1"
1143 } else if id == "dspark" {
1144 "dspark/1"
1145 } else if id == "eagle3" {
1146 "eagle3/1"
1147 } else if draft_path.is_some() {
1148 "draft-model/1"
1149 } else if id == "mtp-head" {
1150 "mtp-head/1"
1151 } else {
1152 "prompt-lookup/1"
1153 };
1154 h[..8].copy_from_slice(&superfluid_fingerprint::fnv1a64(tag.as_bytes(), 0xCBF2_9CE4_8422_2325).to_le_bytes());
1155 h
1156 };
1157 let proposal_cfg = proposal_config(id, opts);
1158 let config_hash = {
1159 let mut h = [0u8; 32];
1160 let cfg = match (&draft_path, &draft_identity) {
1161 (Some(p), Some((fnv, size))) => {
1162 format!("draft={p},weights={fnv:016x},bytes={size},{proposal_cfg}")
1163 }
1164 _ => format!("n_max=4,{proposal_cfg}"),
1165 };
1166 h[..8].copy_from_slice(&superfluid_fingerprint::fnv1a64(cfg.as_bytes(), 0xCBF2_9CE4_8422_2325).to_le_bytes());
1167 h
1168 };
1169 let artifacts = match (&draft_path, &draft_identity) {
1170 (Some(path), Some((fnv, size))) => {
1171 let mut content_hash = [0u8; 32];
1172 content_hash[..8].copy_from_slice(&fnv.to_le_bytes());
1173 vec![linkw::ArtifactMsg {
1174 role: superfluid_abi::artifact_role::DRAFT_WEIGHTS,
1175 content_hash,
1176 byte_size: *size,
1177 load_path: path.clone(),
1178 }]
1179 }
1180 _ => vec![],
1181 };
1182 let reg = linkw::StrategyRegisterMsg {
1183 strategy_id: id.to_string(),
1184 impl_version: "1.0.0".into(),
1185 impl_hash,
1186 config_hash,
1187 artifacts,
1188 taps: vec![],
1189 capabilities: vec![linkw::CapabilityReqMsg {
1190 kind_id: superfluid_abi::strategy_cap::PROPOSAL_LINEAR,
1191 params: proposal_cfg.into_bytes(),
1192 }],
1193 target_archs: vec![],
1194 kernel_caps_required: 0,
1195 est_state_bytes: 0,
1196 claimed_exactness: superfluid_abi::exactness::SEED_PATH_INVARIANT,
1197 rng_contract_version: 1,
1198 };
1199 *this.reg.lock().expect("spec reg") = Some(reg);
1200 this.reregister(host)
1201 }
1202
1203 pub fn reregister(&self, host: &mut EngineHost) -> Result<(), DaemonError> {
1204 let Some(reg) = self.reg.lock().expect("spec reg").clone() else {
1205 return Ok(());
1206 };
1207 let ok = host.client().register_strategy(reg.clone()).map_err(|e| {
1208 DaemonError::Generation(format!("speculation strategy {} refused: {e}", reg.strategy_id))
1209 })?;
1210 let best = ok
1211 .certificates
1212 .iter()
1213 .filter(|c| c.sampling_modes & (superfluid_abi::cert_mode::GREEDY | superfluid_abi::cert_mode::GUMBEL) != 0)
1214 .max_by_key(|c| c.exactness);
1215 let identity = {
1216 let mut h = superfluid_fingerprint::RecordHasher::new();
1217 h.field(1, reg.strategy_id.as_bytes())
1218 .field(2, ®.impl_hash)
1219 .field(3, ®.config_hash);
1220 h.finish()
1221 };
1222 let modes = ok.certificates.iter().fold(0u8, |m, c| m | c.sampling_modes);
1223 *self.grant.lock().expect("spec grant") = Some(SpeculationGrant {
1224 strategy_id: reg.strategy_id.clone(),
1225 slot: ok.strategy_slot,
1226 class: best.map(|c| c.exactness).unwrap_or(superfluid_abi::exactness::APPROXIMATE),
1227 cert_id: best.map(|c| c.cert_id).unwrap_or(0),
1228 identity,
1229 modes,
1230 });
1231 Ok(())
1232 }
1233
1234 pub fn grant(&self) -> Option<SpeculationGrant> {
1235 self.grant.lock().expect("spec grant").clone()
1236 }
1237
1238 pub fn slot_for(&self, sampled: bool) -> u32 {
1239 let mode = if sampled {
1240 superfluid_abi::cert_mode::GUMBEL
1241 } else {
1242 superfluid_abi::cert_mode::GREEDY
1243 };
1244 self.grant
1245 .lock()
1246 .expect("spec grant")
1247 .as_ref()
1248 .filter(|g| g.modes & mode != 0)
1249 .map(|g| g.slot)
1250 .unwrap_or(0)
1251 }
1252
1253 pub fn identity_for_fingerprint(&self) -> Option<u64> {
1254 self.grant
1255 .lock()
1256 .expect("spec grant")
1257 .as_ref()
1258 .filter(|g| g.class < superfluid_abi::exactness::SEED_PATH_INVARIANT)
1259 .map(|g| g.identity)
1260 }
1261}
1262
1263#[derive(Debug, Clone, PartialEq, Eq)]
1264pub enum ParkRefused {
1265 SpaceUnservable,
1266 Failed(String),
1267}
1268
1269impl std::fmt::Display for ParkRefused {
1270 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1271 match self {
1272 ParkRefused::SpaceUnservable => {
1273 write!(f, "a state space has no ops (hybrid GDN before P2.5); park refused")
1274 }
1275 ParkRefused::Failed(m) => f.write_str(m),
1276 }
1277 }
1278}
1279
1280impl From<String> for ParkRefused {
1281 fn from(m: String) -> Self {
1282 ParkRefused::Failed(m)
1283 }
1284}
1285
1286impl From<&str> for ParkRefused {
1287 fn from(m: &str) -> Self {
1288 ParkRefused::Failed(m.to_string())
1289 }
1290}
1291
1292enum ParkTask {
1293 Write(Box<ParkJob>),
1294 Remove { dir: PathBuf, session: u64 },
1295 WritePrefix { dir: PathBuf, tokens: Vec<u32>, encoding: u8, sealed: Vec<u8> },
1297}
1298
1299struct ParkJob {
1300 dir: PathBuf,
1301 session: u64,
1302 covered: u64,
1303 digest: u64,
1304 encoding: u8,
1305 sealed: Vec<u8>,
1306 spaces: Option<Vec<park::ParkedSpace>>,
1307}
1308
1309pub struct Scheduler {
1310 host: EngineHost,
1311 store: Arc<Mutex<SessionStore>>,
1312 codec: Arc<dyn TextCodec + Send + Sync>,
1313 cancels: Arc<crate::CancelSet>,
1314 active: Arc<Mutex<HashSet<u64>>>,
1315 stats: Arc<SchedStats>,
1316 rx: Receiver<SchedCmd>,
1317 max_lanes: usize,
1318 configured_lanes: usize,
1319 retired: Arc<Mutex<Option<String>>>,
1320 park_dir: Option<PathBuf>,
1321 pressure_high_pct: u8,
1322 pressure_low_pct: u8,
1323 park_lossy: bool,
1324 trace_scope: crate::otlp::TraceScope,
1325 park_tx: Option<std::sync::mpsc::SyncSender<ParkTask>>,
1326 ephemeral_parked: std::collections::HashSet<u64>,
1327 prefixes: Vec<(PathBuf, Vec<u32>)>,
1329 park_removals_pending: Vec<u64>,
1330 park_writer: Option<std::thread::JoinHandle<()>>,
1331 bus: Arc<SessionBus>,
1332
1333 plan_seq: u64,
1334 next_lane: u64,
1335 lanes: Vec<Lane>,
1336 queues: [ClassQueue; 4],
1337 class_lanes: [usize; 4],
1338 tick_decode_budget: u32,
1339 tick_target_secs: f64,
1340 tick_shape: TickShape,
1341 prefill_budget: u32,
1342 prefill_budget_fixed: bool,
1343 decode_round_secs: Option<f64>,
1344 prefill_secs_per_token: Option<f64>,
1345 just_retired: bool,
1346 prefill_secs_per_token_bulk: Option<f64>,
1347 prefill_secs_per_token_last: Option<f64>,
1348 prefill_rate_samples: u32,
1349 prefix_checkpoints_off: bool,
1352 turn_checkpoints: std::collections::VecDeque<(u64, [u8; 32])>,
1355 tick_decode_budget_fixed: bool,
1356 starvation_ticks: u32,
1357 background_tick_divisor: Arc<std::sync::atomic::AtomicU32>,
1358 pressure_source: Option<Box<dyn crate::pressure::PressureSource + Send>>,
1359 tick_count: u64,
1360 speculation: Arc<Speculation>,
1361 tool_lease_ms: u64,
1362 respawn_streak: u32,
1363 staged_unread: u32,
1364 deferred_sync: VecDeque<(u8, u64, SchedCmd)>,
1365 admit_retry: bool,
1366 pins: crate::pins::Pins,
1367}
1368
1369impl Scheduler {
1370 fn respawn_worker(&mut self) -> Result<(), DaemonError> {
1371 self.host.respawn()?;
1372 self.stats.kv_bits.store(self.host.kv_bits() as u64, Ordering::Relaxed);
1373 Ok(())
1374 }
1375
1376 #[allow(clippy::too_many_arguments)]
1377 pub fn spawn(
1378 host: EngineHost,
1379 store: Arc<Mutex<SessionStore>>,
1380 codec: Arc<dyn TextCodec + Send + Sync>,
1381 cancels: Arc<crate::CancelSet>,
1382 active: Arc<Mutex<HashSet<u64>>>,
1383 stats: Arc<SchedStats>,
1384 background_tick_divisor: Arc<std::sync::atomic::AtomicU32>,
1385 speculation: Arc<Speculation>,
1386 retired: Arc<Mutex<Option<String>>>,
1387 mut opts: crate::DaemonOptions,
1388 ) -> SchedulerHandle {
1389 let (tx, rx) = std::sync::mpsc::channel();
1390 let join = std::thread::Builder::new()
1391 .name("superfluid-scheduler".into())
1392 .spawn(move || {
1393 let bus = store.lock().expect("store").bus();
1394 let host_max_stream = host.max_stream_tokens();
1395 let mut host = host;
1396 let max_lanes = lane_cap(
1397 opts.max_lanes,
1398 host.max_seqs(),
1399 ring_slots(&mut host, superfluid_shm::TOKEN_RING_OUT),
1400 );
1401 let park_refused = host.park_refused();
1402 let prefixes = match (&park_refused, &opts.park_dir) {
1403 (None, Some(dir)) => park::prefix_index(&park::prefix_dir(dir)),
1404 _ => Vec::new(),
1405 };
1406 let lossy_allowed = host.lossy_park_allowed();
1407 let tick_shape = if host.round_granular() { ROUND_TICKS } else { NATIVE_TICKS };
1408 let (park_tx, park_writer) = if opts.park_dir.is_some() && park_refused.is_none() {
1409 let (tx, rx) = std::sync::mpsc::sync_channel::<ParkTask>(4);
1410 let park_budget = opts.park_budget_bytes;
1411 let writer = std::thread::Builder::new()
1412 .name("superfluid-park-writer".into())
1413 .spawn(move || {
1414 while let Ok(task) = rx.recv() {
1415 let job = match task {
1416 ParkTask::Write(job) => job,
1417 ParkTask::Remove { dir, session } => {
1418 park::remove(&dir, session);
1419 continue;
1420 }
1421 ParkTask::WritePrefix { dir, tokens, encoding, sealed } => {
1422 if let Err(e) = park::write_prefix(&dir, &tokens, encoding, &sealed) {
1423 eprintln!("[superfluid] prefix write failed: {e}");
1424 }
1425 if park_budget > 0 {
1426 park::sweep(&dir, park_budget);
1427 }
1428 continue;
1429 }
1430 };
1431 let res = match &job.spaces {
1432 Some(spaces) => park::write_v3(
1433 &job.dir,
1434 job.session,
1435 job.covered,
1436 job.digest,
1437 spaces,
1438 ),
1439 None => park::write(
1440 &job.dir,
1441 job.session,
1442 job.covered,
1443 job.digest,
1444 job.encoding,
1445 &job.sealed,
1446 ),
1447 };
1448 if let Err(e) = res {
1449 eprintln!(
1450 "[superfluid] park write failed (session {}): {e}",
1451 job.session
1452 );
1453 }
1454 if park_budget > 0 {
1455 let freed = park::sweep(&job.dir, park_budget);
1456 if freed > 0 {
1457 tracing::debug!(
1458 freed_bytes = freed,
1459 budget_bytes = park_budget,
1460 "park retention sweep"
1461 );
1462 }
1463 }
1464 }
1465 })
1466 .expect("spawn park writer");
1467 (Some(tx), Some(writer))
1468 } else {
1469 (None, None)
1470 };
1471 Scheduler {
1472 host,
1473 store,
1474 codec,
1475 cancels,
1476 active,
1477 stats,
1478 rx,
1479 max_lanes,
1480 configured_lanes: opts.max_lanes,
1481 retired,
1482 park_dir: match park_refused {
1483 Some(why) if opts.park_dir.is_some() => {
1484 eprintln!("[superfluid] parking is off for this model: {why}");
1485 None
1486 }
1487 _ => opts.park_dir,
1488 },
1489 pressure_high_pct: opts.pressure_high_pct,
1490 pressure_low_pct: opts.pressure_low_pct,
1491 park_lossy: opts.park_lossy && lossy_allowed,
1492 trace_scope: opts.trace_scope.clone(),
1493 park_tx,
1494 park_writer,
1495 ephemeral_parked: std::collections::HashSet::new(),
1496 prefixes,
1497 park_removals_pending: Vec::new(),
1498 bus,
1499 plan_seq: 0,
1500 next_lane: 1,
1501 lanes: Vec::new(),
1502 queues: Default::default(),
1503 class_lanes: opts.class_lanes,
1504 tick_target_secs: if opts.tick_target_ms == 0 {
1505 TARGET_TICK_TOTAL_SECS
1506 } else {
1507 opts.tick_target_ms as f64 / 1000.0
1508 },
1509 tick_shape,
1510 tick_decode_budget: if opts.tick_decode_budget == 0 {
1511 (max_lanes as u32).saturating_mul(TICK_DECODE as u32)
1512 } else {
1513 opts.tick_decode_budget
1514 },
1515 prefill_budget: if opts.prefill_budget == 0 {
1516 1024
1517 } else {
1518 if opts.prefill_budget > PREFILL_BUDGET {
1519 eprintln!(
1520 "superfluid: --prefill-budget {} is above the {PREFILL_BUDGET}-token per-tick \
1521 ceiling (the media run-size limit); using {PREFILL_BUDGET}",
1522 opts.prefill_budget
1523 );
1524 }
1525 opts.prefill_budget.clamp(64, PREFILL_BUDGET)
1526 },
1527 prefill_budget_fixed: opts.prefill_budget != 0,
1528 decode_round_secs: None,
1529 prefill_secs_per_token: None,
1530 just_retired: false,
1531 prefill_secs_per_token_bulk: None,
1532 prefill_secs_per_token_last: None,
1533 prefill_rate_samples: 0,
1534 prefix_checkpoints_off: false,
1535 turn_checkpoints: std::collections::VecDeque::new(),
1536 tick_decode_budget_fixed: opts.tick_decode_budget != 0,
1537 starvation_ticks: opts.agent_starvation_ticks.max(1),
1538 background_tick_divisor,
1539 pressure_source: opts.pressure_source.take_source(),
1540 tick_count: 0,
1541 speculation,
1542 tool_lease_ms: opts.tool_lease_ms,
1543 respawn_streak: 0,
1544 staged_unread: 0,
1545 deferred_sync: VecDeque::new(),
1546 admit_retry: false,
1547 pins: crate::pins::Pins::new(
1548 opts.pin_budget_pct,
1549 host_max_stream * max_lanes as u64,
1550 ),
1551 }
1552 .run();
1553 })
1554 .expect("spawn scheduler");
1555 SchedulerHandle {
1556 tx,
1557 join: std::sync::Mutex::new(Some(join)),
1558 }
1559 }
1560
1561 fn run(mut self) {
1562 self.run_loop();
1563 drop(self.park_tx.take());
1564 if let Some(w) = self.park_writer.take() {
1565 let _ = w.join();
1566 }
1567 if let Some(dir) = self.park_dir.clone() {
1568 let pending = std::mem::take(&mut self.park_removals_pending);
1569 let tracked = std::mem::take(&mut self.ephemeral_parked);
1570 for session in pending.into_iter().chain(tracked) {
1571 park::remove(&dir, session);
1572 }
1573 }
1574 }
1575
1576 fn run_loop(&mut self) {
1577 loop {
1578 if self.lanes.is_empty() && self.queued() == 0 && self.deferred_sync.is_empty() {
1579 self.retry_park_removals();
1580 let retry = (!self.park_removals_pending.is_empty())
1581 .then(|| std::time::Duration::from_millis(50));
1582 let wait = match (self.pins.until_next_deadline(), retry) {
1583 (Some(a), Some(b)) => Some(a.min(b)),
1584 (a, b) => a.or(b),
1585 };
1586 let cmd = match wait {
1587 None => self.rx.recv().map_err(|_| ()),
1588 Some(wait) => match self.rx.recv_timeout(wait) {
1589 Ok(cmd) => Ok(cmd),
1590 Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
1591 self.pins.sweep(&mut self.host, &self.stats);
1592 continue;
1593 }
1594 Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => Err(()),
1595 },
1596 };
1597 match cmd {
1598 Ok(SchedCmd::Shutdown) | Err(()) => return,
1599 Ok(cmd) => self.dispatch(cmd),
1600 }
1601 }
1602 self.host.clear_yield();
1605 let slot = self.host.yield_slot();
1608 let latency = self.latency_sensitive();
1609 crate::runtime::set_latency(&slot, latency);
1610 if !latency && self.cancels.latency_recent(LATENCY_HOLD) {
1611 crate::runtime::set_latency(&slot, true);
1612 }
1613 loop {
1614 match self.rx.try_recv() {
1615 Ok(SchedCmd::Shutdown) => return,
1616 Ok(cmd) => self.dispatch(cmd),
1617 Err(TryRecvError::Empty) => break,
1618 Err(TryRecvError::Disconnected) => return,
1619 }
1620 }
1621 if !self.host.is_alive() {
1622 if self.respawn_streak < MAX_WORKER_RESPAWNS {
1623 self.respawn_streak += 1;
1624 if self.respawn_worker().is_ok() {
1625 let _ = self.speculation.reregister(&mut self.host);
1626 self.after_respawn();
1627 }
1628 }
1629 if !self.host.is_alive() {
1630 if self.respawn_streak >= MAX_WORKER_RESPAWNS {
1631 self.give_up_on_worker();
1632 }
1633 let why = self
1634 .retired
1635 .lock()
1636 .expect("retired")
1637 .clone()
1638 .unwrap_or_else(|| "engine worker unavailable (respawn failed)".into());
1639 let orphaned: Vec<GenerateJob> =
1640 self.queues.iter_mut().flat_map(|q| q.jobs.drain(..)).collect();
1641 for job in orphaned {
1642 self.finish_failed(&job, &why);
1643 }
1644 for (_, _, cmd) in std::mem::take(&mut self.deferred_sync) {
1645 self.dispatch(cmd);
1646 }
1647 continue;
1648 }
1649 }
1650 self.pins.sweep(&mut self.host, &self.stats);
1651 let tick_span = if self.lanes.is_empty() && self.queued() == 0 {
1652 tracing::Span::none()
1653 } else {
1654 tracing::debug_span!(
1655 parent: None,
1656 "tick",
1657 tick_seq = tracing::field::Empty,
1658 no_tick = tracing::field::Empty,
1659 lanes = tracing::field::Empty,
1660 queued = self.queued(),
1661 model = self.trace_scope.model.as_str(),
1662 store_ns = self.trace_scope.ns
1663 )
1664 };
1665 let _tick = tick_span.enter();
1666 self.poll_os_pressure();
1667 self.preempt_for_queued();
1668 self.admit_from_queue();
1669 self.run_deferred_syncs();
1670 let ticks_before = self.tick_count;
1671 let tick_result = self.tick_once(&tick_span);
1672 if !tick_span.is_disabled() {
1673 if self.tick_count != ticks_before {
1674 tick_span.record("tick_seq", self.tick_count);
1675 } else {
1676 tick_span.record("no_tick", true);
1677 }
1678 }
1679 if tick_result.is_ok() {
1680 self.respawn_streak = 0;
1681 }
1682 if let Err(e) = tick_result {
1683 let admitted: Vec<u64> = self
1684 .lanes
1685 .iter()
1686 .filter(|l| l.admitted)
1687 .map(|l| l.tag)
1688 .collect();
1689 if !admitted.is_empty() {
1690 let mut p = linkw::TickPlanMsg {
1691 plan_seq: {
1692 self.plan_seq += 1;
1693 self.plan_seq
1694 },
1695 flags: 0,
1696 prefill_token_budget: PREFILL_BUDGET,
1697 max_decode_lanes: self.max_lanes as u32,
1698 admits: vec![],
1699 commits: vec![],
1700 prefills: vec![],
1701 decodes: vec![],
1702 retires: vec![],
1703 shed_policy: linkw::ShedPolicyMsg {
1704 victim_lanes: vec![],
1705 evictable_cache_classes: u64::MAX,
1706 protected_quota_bytes: 0,
1707 max_evict_bytes: u64::MAX,
1708 },
1709 };
1710 for tag in admitted {
1711 p.retires.push(linkw::LaneRetireMsg {
1712 lane_tag: tag,
1713 publish_to_cache: true,
1714 });
1715 }
1716 self.stats.set_tick_in_flight(TickInFlight {
1717 since_ms: uptime_ms().max(1),
1718 retires: p.retires.len() as u64,
1719 target_ms: self.stats.tick_target_live_ms.load(Ordering::Relaxed),
1720 ..Default::default()
1721 });
1722 let _ = self.host.client().tick(p);
1723 self.stats.clear_tick_in_flight();
1724 }
1725 let msg = format!("scheduler tick failed: {e}");
1726 let drained: Vec<u64> = self.lanes.iter().map(|l| l.session).collect();
1727 for lane in self.lanes.drain(..) {
1728 self.active.lock().expect("active set").remove(&lane.session);
1729 let _ = lane.job.events.send(LaneEvent::Failed(msg.clone()));
1730 }
1731 for session in drained {
1732 self.drop_ephemeral_artifact(session);
1733 }
1734 self.publish_occupancy();
1735 if worker_is_gone(&e) && self.host.is_process() {
1736 let mut fresh = false;
1737 if self.respawn_streak < MAX_WORKER_RESPAWNS {
1738 self.respawn_streak += 1;
1739 self.stats.worker_respawns.fetch_add(1, Ordering::Relaxed);
1740 match self.respawn_worker().and_then(|()| self.speculation.reregister(&mut self.host)) {
1741 Ok(()) => {
1742 eprintln!("[superfluid] engine worker died ({e}); respawned a fresh one");
1743 self.after_respawn();
1744 fresh = true;
1745 }
1746 Err(err) => eprintln!(
1747 "[superfluid] engine worker died ({e}); respawn FAILED: {err}"
1748 ),
1749 }
1750 }
1751 if !fresh && self.respawn_streak >= MAX_WORKER_RESPAWNS {
1752 self.give_up_on_worker();
1753 }
1754 }
1755 }
1756 self.pins.sweep(&mut self.host, &self.stats);
1757 }
1758 }
1759
1760 fn dispatch(&mut self, cmd: SchedCmd) {
1761 match cmd {
1762 SchedCmd::Submit(job) => self.enqueue(*job),
1763 SchedCmd::Sync(req, reply) => self.run_sync(req, reply),
1764 SchedCmd::SyncStream(req, segs, reply) => self.run_sync_stream(req, segs, reply),
1765 SchedCmd::Pin(req) => self.pins.handle(&mut self.host, &self.stats, self.tick_count, req),
1766 SchedCmd::PinHeld(keys, reply) => {
1767 let _ = reply.send(self.pins.held_max(&self.host, &keys));
1768 }
1769 SchedCmd::PinUnanchor(session) => self.pins.unanchor(session),
1770 SchedCmd::SyncAs(class, inner) => {
1771 self.deferred_sync.push_back((class, self.tick_count, *inner));
1772 }
1773 SchedCmd::Shutdown => {}
1774 }
1775 }
1776
1777 fn sync_must_wait(&self, class: u8) -> bool {
1778 self.lanes.iter().any(|l| l.job.class < class)
1779 || self.queues.iter().take(class as usize).any(|q| !q.jobs.is_empty())
1780 }
1781
1782 fn run_deferred_syncs(&mut self) {
1783 let period = self.starvation_ticks as u64;
1784 let stuck = self.lanes.is_empty() && !self.admit_retry;
1785 let starved = |queued: u64| self.tick_count.saturating_sub(queued) >= period;
1786 let pick = self
1787 .deferred_sync
1788 .iter()
1789 .enumerate()
1790 .filter(|(_, (class, queued, _))| stuck || !self.sync_must_wait(*class) || starved(*queued))
1791 .min_by_key(|&(i, (class, queued, _))| {
1792 let aged = starved(*queued);
1793 (!aged, if aged { 0 } else { self.sync_rank(*class) }, i)
1794 })
1795 .map(|(i, _)| i);
1796 if let Some((_, _, cmd)) = pick.and_then(|i| self.deferred_sync.remove(i)) {
1797 self.dispatch(cmd);
1798 }
1799 }
1800
1801 fn sync_rank(&self, class: u8) -> usize {
1802 class as usize
1803 }
1804
1805 fn run_sync(&mut self, req: linkw::StateSyncReqMsg, reply: Sender<Result<linkw::StateSyncOkMsg, String>>) {
1806 let res = if self.host.is_alive() {
1807 self.host.client().state_sync(req).map_err(|e| e.to_string())
1808 } else {
1809 Err("engine worker unavailable".into())
1810 };
1811 let _ = reply.send(res);
1812 }
1813
1814 fn run_sync_stream(
1815 &mut self,
1816 req: linkw::StateSyncReqMsg,
1817 segs: Sender<linkw::TranscribeSegmentMsg>,
1818 reply: Sender<Result<linkw::StateSyncOkMsg, String>>,
1819 ) {
1820 let res = if self.host.is_alive() {
1821 self.host
1822 .client()
1823 .state_sync_streaming(req, &mut |seg| {
1824 segs.send(seg).is_ok()
1825 })
1826 .map_err(|e| e.to_string())
1827 } else {
1828 Err("engine worker unavailable".into())
1829 };
1830 let _ = reply.send(res);
1831 }
1832
1833 fn queued(&self) -> usize {
1834 self.queues.iter().map(|q| q.jobs.len()).sum()
1835 }
1836
1837 fn enqueue(&mut self, mut job: GenerateJob) {
1838 let c = (job.class as usize).min(3);
1839 job.queued_at = self.tick_count;
1840 self.queues[c].jobs.push_back(job);
1841 }
1842
1843 fn engine_exclusive(&self, job: &GenerateJob) -> bool {
1844 job.batch_invariant || (!job.media.is_empty() && self.host.has_recurrent_space())
1845 }
1846
1847 fn strategy_slot_for(speculation: &Speculation, job: &GenerateJob) -> u32 {
1848 if job.grammar_handle != 0 || job.extras.logit_bias_handle != 0 || job.extras.want_logprobs {
1849 0
1850 } else {
1851 speculation.slot_for(job.params.temperature != 0.0)
1852 }
1853 }
1854
1855 fn decode_replay_for(strategy_slot: u32, produced_before: u32) -> u32 {
1856 if strategy_slot == 0 { produced_before } else { 0 }
1857 }
1858
1859 fn prefill_end_of(speculation: &Speculation, lane: &Lane) -> u64 {
1860 let replay = Self::decode_replay_for(Self::strategy_slot_for(speculation, &lane.job), lane.produced_before);
1861 (lane.job.stream.len() as u64)
1862 .saturating_sub(replay as u64)
1863 .max(lane.prefilled)
1864 }
1865
1866 fn pop_admissible(&mut self) -> Option<GenerateJob> {
1867 if self.lanes.iter().any(|l| self.engine_exclusive(&l.job)) {
1868 return None;
1869 }
1870 let mut per_class = [0usize; 4];
1871 for l in &self.lanes {
1872 per_class[(l.job.class as usize).min(3)] += 1;
1873 }
1874 let now = self.tick_count;
1875 let period = self.starvation_ticks as u64;
1876 let mut best: Option<(u8, usize)> = None;
1877 let mut excl_waiting: Option<u8> = None;
1878 for (c, q) in self.queues.iter().enumerate() {
1879 let Some(head) = q.jobs.front() else {
1880 continue;
1881 };
1882 let cap = if self.class_lanes[c] == 0 {
1883 self.max_lanes
1884 } else {
1885 self.class_lanes[c]
1886 };
1887 if per_class[c] >= cap {
1888 continue;
1889 }
1890 let waited = now.saturating_sub(head.queued_at);
1891 let aged = (waited / period).min(c as u64) as u8;
1892 let eff = c as u8 - aged;
1893 if self.engine_exclusive(head) && !self.lanes.is_empty() {
1894 if excl_waiting.is_none_or(|b| eff < b) {
1895 excl_waiting = Some(eff);
1896 }
1897 continue;
1898 }
1899 if best.is_none_or(|(b, _)| eff < b) {
1900 best = Some((eff, c));
1901 }
1902 }
1903 if let (Some(e), Some((b, _))) = (excl_waiting, best) {
1904 if e <= b {
1905 return None;
1906 }
1907 }
1908 let (_, c) = best?;
1909 self.queues[c].jobs.pop_front()
1910 }
1911
1912 fn awaits_a_shared_prefix(&self, job: &GenerateJob) -> bool {
1915 if job.carry.is_some()
1916 || job.class <= crate::qos::INLINE_COMPLETION
1917 || !job.media.is_empty()
1918 || self.tick_count.saturating_sub(job.queued_at) >= SHARED_PREFIX_MAX_HOLD_TICKS
1919 || self.host.has_unservable_space()
1920 || !self.host.cache_serves_all_spaces()
1921 {
1922 return false;
1923 }
1924 let Some(page) = self.host.page_size().filter(|&p| p > 0) else {
1925 return false;
1926 };
1927 if self.host.needs_prefix_checkpoints() {
1928 if self.prefix_checkpoints_off {
1929 return false;
1930 }
1931 return self.lanes.iter().any(|l| {
1935 if l.preempted || l.job.completion || !l.job.media.is_empty() {
1936 return false;
1937 }
1938 let shared = job.stream.iter().zip(&l.all_tokens).take_while(|(a, b)| a == b).count() as u64;
1939 let total = Self::prefill_end_of(&self.speculation, l);
1940 shared >= SHARED_PREFIX_MIN_TOKENS && shared + 1 < total && !Self::kept_for(l, shared) && l.prefilled < shared
1941 });
1942 }
1943 self.lanes.iter().any(|l| {
1944 if l.preempted
1945 || l.parked_lossy
1946 || l.job.completion
1947 || !l.job.media.is_empty()
1948 || l.job.max_tokens.saturating_sub(l.produced) > SHARED_PREFIX_TAIL_TOKENS
1949 {
1950 return false;
1951 }
1952 let shared = job.stream.iter().zip(&l.all_tokens).take_while(|(a, b)| a == b).count() as u64;
1953 let shared = shared / page * page;
1954 shared >= SHARED_PREFIX_MIN_TOKENS && shared > l.warm_prefix
1955 })
1956 }
1957
1958 fn plan_prefix_checkpoints(&mut self) {
1969 if !self.host.needs_prefix_checkpoints() || self.prefix_checkpoints_off {
1970 for l in self.lanes.iter_mut() {
1971 l.checkpoint_at = 0;
1972 }
1973 return;
1974 }
1975 let mut landed = false;
1979 for l in self.lanes.iter_mut() {
1980 if Self::turn_point(&self.speculation, l) == Some(l.prefilled) {
1981 (l.checkpoint_at, l.checkpoint_due, landed) = (l.prefilled, true, true);
1982 }
1983 }
1984 if landed {
1985 self.publish_prefix_checkpoints();
1986 }
1987 for i in 0..self.lanes.len() {
1988 let l = &self.lanes[i];
1989 let total = Self::prefill_end_of(&self.speculation, l);
1990 let mut at = 0u64;
1991 if Self::checkpoint_eligible(l, total) {
1992 for job in self.queues.iter().flat_map(|q| q.jobs.iter()) {
1993 let shared = job.stream.iter().zip(&l.all_tokens).take_while(|(a, b)| a == b).count() as u64;
1994 let useful = shared >= SHARED_PREFIX_MIN_TOKENS
1995 && shared + 1 < total
1996 && shared > l.prefilled
1997 && !Self::kept_for(l, shared);
1998 if useful && (at == 0 || shared < at) {
1999 at = shared;
2000 }
2001 }
2002 if let Some(end) = Self::turn_point(&self.speculation, l).filter(|&end| end > l.prefilled) {
2003 if at == 0 || end < at {
2004 at = end;
2005 }
2006 }
2007 }
2008 self.lanes[i].checkpoint_at = at;
2009 }
2010 }
2011
2012 fn checkpoint_eligible(l: &Lane, total: u64) -> bool {
2013 !l.preempted
2014 && !l.job.completion
2015 && l.job.media.is_empty()
2016 && !matches!(l.phase, LanePhase::Finishing)
2017 && l.prefilled < total
2018 }
2019
2020 fn turn_point(speculation: &Speculation, l: &Lane) -> Option<u64> {
2023 let end = l.job.input_end;
2024 let total = Self::prefill_end_of(speculation, l);
2025 (Self::checkpoint_eligible(l, total)
2026 && !l.turn_kept
2027 && end >= TURN_CHECKPOINT_MIN_TOKENS
2028 && end + 1 < total
2029 && end >= l.prefilled
2030 && end != l.warm_prefix)
2031 .then_some(end)
2032 }
2033
2034 fn kept_for(l: &Lane, shared: u64) -> bool {
2038 let serves = |at: u64| at >= SHARED_PREFIX_MIN_TOKENS && at <= shared;
2039 serves(l.checkpointed) || serves(l.warm_prefix)
2040 }
2041
2042 fn publish_prefix_checkpoints(&mut self) {
2045 for i in 0..self.lanes.len() {
2046 if !std::mem::take(&mut self.lanes[i].checkpoint_due) {
2047 continue;
2048 }
2049 let (tag, at) = (self.lanes[i].tag, self.lanes[i].checkpoint_at);
2050 let tokens = self.lanes[i].all_tokens[..at as usize].to_vec();
2051 let turn = at == self.lanes[i].job.input_end;
2052 self.lanes[i].turn_kept |= turn;
2053 match self.fork_into_cache(tag, &tokens, turn) {
2056 Ok(()) => {
2057 self.lanes[i].checkpointed = at;
2058 if turn {
2059 self.supersede_turn_checkpoints(&tokens);
2060 }
2061 }
2062 Err(CheckpointError::Refused(e)) => {
2063 tracing::info!(lane = tag, "the engine forks no lane's sequence ({e}): no shared-prefix checkpoints");
2064 self.prefix_checkpoints_off = true;
2065 }
2066 Err(CheckpointError::Failed(e)) => tracing::debug!(lane = tag, at, "prefix checkpoint not kept: {e}"),
2067 }
2068 }
2069 }
2070
2071 fn supersede_turn_checkpoints(&mut self, tokens: &[u32]) {
2074 let digest = superfluid_fingerprint::content_digest(tokens);
2075 let mut kept = std::collections::VecDeque::with_capacity(self.turn_checkpoints.len() + 1);
2076 for (len, d) in std::mem::take(&mut self.turn_checkpoints) {
2077 let n = len as usize;
2078 if n == tokens.len() && d == digest {
2079 continue;
2080 }
2081 if n < tokens.len() && superfluid_fingerprint::content_digest(&tokens[..n]) == d {
2082 self.unpublish(&tokens[..n]);
2083 continue;
2084 }
2085 kept.push_back((len, d));
2086 }
2087 kept.push_back((tokens.len() as u64, digest));
2088 while kept.len() > TURN_CHECKPOINTS_TRACKED {
2089 kept.pop_front();
2090 }
2091 self.turn_checkpoints = kept;
2092 }
2093
2094 fn unpublish(&mut self, tokens: &[u32]) {
2095 let res = match self.stage(tokens) {
2096 Ok(span) => self.host.client().state_sync(linkw::StateSyncReqMsg::Unpublish { span }).map_err(|e| e.to_string()),
2097 Err(e) => Err(e.to_string()),
2098 };
2099 if let Err(e) = res {
2100 tracing::debug!(len = tokens.len(), "superseded turn checkpoint not dropped: {e}");
2101 }
2102 }
2103
2104 fn fork_into_cache(&mut self, tag: u64, tokens: &[u32], private: bool) -> Result<(), CheckpointError> {
2105 use CheckpointError::Failed;
2106 let seq = match self.host.client().state_sync(linkw::StateSyncReqMsg::LaneSequence { lane_tag: tag }) {
2107 Ok(linkw::StateSyncOkMsg::LaneSequence { sequence }) => sequence,
2108 Ok(_) => return Err(Failed("mismatched LaneSequence reply".into())),
2109 Err(e) => return Err(CheckpointError::of(e)),
2110 };
2111 let child = match self.host.client().state_sync(linkw::StateSyncReqMsg::Fork {
2112 parent_sequence: seq,
2113 flags: if private { superfluid_abi::fork_flags::PRIVATE } else { 0 },
2114 }) {
2115 Ok(linkw::StateSyncOkMsg::Fork { child_sequence }) => child_sequence,
2116 Ok(_) => return Err(Failed("mismatched Fork reply".into())),
2117 Err(e) => return Err(CheckpointError::of(e)),
2118 };
2119 let published = match self.stage(tokens) {
2120 Ok(span) => match self.host.client().state_sync(linkw::StateSyncReqMsg::PublishSequence { sequence: child, span }) {
2121 Ok(linkw::StateSyncOkMsg::PublishSequence) => Ok(()),
2122 Ok(_) => Err("mismatched PublishSequence reply".into()),
2123 Err(e) => Err(e.to_string()),
2124 },
2125 Err(e) => Err(e.to_string()),
2126 };
2127 if published.is_err() {
2128 let _ = self.host.client().state_sync(linkw::StateSyncReqMsg::FreeSequence { sequence: child });
2129 }
2130 published.map_err(Failed)
2131 }
2132
2133 fn latency_sensitive(&self) -> bool {
2136 self.cancels.latency_recent(LATENCY_HOLD)
2137 || self.queues[..=crate::qos::INLINE_COMPLETION as usize].iter().any(|q| !q.jobs.is_empty())
2138 || self.lanes.iter().any(|l| l.job.class <= crate::qos::INLINE_COMPLETION && !l.preempted)
2139 }
2140
2141 fn joiner_waits_on_this_tick(&self) -> bool {
2142 if self.max_lanes <= 1 {
2143 return false;
2144 }
2145 self.live_lanes() < self.max_lanes
2146 || self.queues[..=crate::qos::INLINE_COMPLETION as usize]
2147 .iter()
2148 .any(|q| !q.jobs.is_empty())
2149 || self.lanes.iter().any(|l| {
2150 l.job.class > crate::qos::INLINE_COMPLETION
2151 && !l.preempted
2152 && l.job.media.is_empty()
2153 && !self.engine_exclusive(&l.job)
2154 })
2155 }
2156
2157 fn live_lanes(&self) -> usize {
2160 self.lanes
2161 .iter()
2162 .filter(|l| !(l.preempted && matches!(l.phase, LanePhase::Finishing)))
2163 .count()
2164 }
2165
2166 fn preempt_for_queued(&mut self) {
2167 if self.live_lanes() < self.max_lanes {
2168 return;
2169 }
2170 let Some(want) = self
2171 .queues
2172 .iter()
2173 .enumerate()
2174 .find(|(_, q)| !q.jobs.is_empty())
2175 .map(|(c, _)| c as u8)
2176 else {
2177 return;
2178 };
2179 if want > crate::qos::INLINE_COMPLETION {
2180 return;
2181 }
2182 let victim_in = |phase_decoding: bool| {
2185 self.lanes
2186 .iter()
2187 .enumerate()
2188 .filter(|(_, l)| {
2189 l.admitted
2190 && !l.preempted
2191 && !self.engine_exclusive(&l.job)
2192 && l.job.class > want
2193 && l.job.media.is_empty()
2194 && if phase_decoding {
2195 matches!(l.phase, LanePhase::Decoding)
2196 } else {
2197 matches!(l.phase, LanePhase::Prefilling)
2198 }
2199 })
2200 .max_by_key(|(_, l)| (l.job.class, l.tag))
2201 .map(|(i, _)| i)
2202 };
2203 let victim = victim_in(true).or_else(|| victim_in(false));
2204 if let Some(i) = victim {
2205 self.preempt_lane(i);
2206 }
2207 }
2208
2209 fn preempt_lane(&mut self, i: usize) {
2210 let lane = &mut self.lanes[i];
2211 lane.preempted = true;
2212 lane.phase = LanePhase::Finishing;
2213 lane.fin = finish::NONE;
2214 tracing::info!(
2215 session = lane.session,
2216 lane = lane.tag,
2217 class = lane.job.class,
2218 produced = lane.produced,
2219 "lane preempted"
2220 );
2221 self.stats.preemptions.fetch_add(1, Ordering::Relaxed);
2222 }
2223
2224 fn poll_os_pressure(&mut self) {
2225 let Some(src) = self.pressure_source.as_mut() else {
2226 return;
2227 };
2228 match src.level() {
2229 crate::pressure::PressureLevel::Normal => {}
2230 crate::pressure::PressureLevel::Warning => {
2231 self.stats.os_pressure_events.fetch_add(1, Ordering::Relaxed);
2232 self.relief_round_now();
2233 }
2234 crate::pressure::PressureLevel::Critical => {
2235 self.stats.os_pressure_events.fetch_add(1, Ordering::Relaxed);
2236 self.relief_round_now();
2237 let victims: Vec<usize> = self
2238 .lanes
2239 .iter()
2240 .enumerate()
2241 .filter(|(_, l)| {
2242 l.admitted
2243 && !l.preempted
2244 && l.job.class == crate::qos::BACKGROUND_AGENT
2245 && matches!(l.phase, LanePhase::Decoding)
2246 })
2247 .map(|(i, _)| i)
2248 .collect();
2249 for i in victims {
2250 self.preempt_lane(i);
2251 }
2252 }
2253 }
2254 }
2255
2256 fn relief_round_now(&mut self) {
2257 self.evict_cache(0, HOST_EXPORT_CLASS);
2260 if let Ok(pong) = self.host.client().ping() {
2261 let mem = pong.mem;
2262 let low = self.pressure_low_pct as u64;
2263 let (used, total) = (mem.pool_blocks_used, mem.pool_blocks_total);
2264 if total == 0 || used * 100 <= total * low {
2265 return;
2266 }
2267 let saved = self.pressure_high_pct;
2268 self.pressure_high_pct = 1;
2269 self.relieve_pressure(&mem, true);
2270 self.pressure_high_pct = saved;
2271 }
2272 }
2273
2274 fn after_respawn(&mut self) {
2275 self.max_lanes = lane_cap(
2276 self.configured_lanes,
2277 self.host.max_seqs(),
2278 ring_slots(&mut self.host, superfluid_shm::TOKEN_RING_OUT),
2279 );
2280 self.pins.set_fallback_pool_tokens(self.host.max_stream_tokens() * self.max_lanes as u64);
2281 if let Some(why) = self.host.take_record_change() {
2282 self.retire(&why);
2283 }
2284 }
2285
2286 fn give_up_on_worker(&mut self) {
2289 if self.retired.lock().expect("retired").is_none() {
2290 self.retire(&format!(
2291 "the engine worker died and {MAX_WORKER_RESPAWNS} respawns in a row did not keep one up"
2292 ));
2293 }
2294 }
2295
2296 fn retire(&mut self, why: &str) {
2297 let msg = format!("{why}; the model loads again on its next request, so retry");
2298 eprintln!("[superfluid] {msg}");
2299 *self.retired.lock().expect("retired") = Some(msg.clone());
2300 let drained: Vec<u64> = self.lanes.iter().map(|l| l.session).collect();
2301 for lane in self.lanes.drain(..) {
2302 self.active.lock().expect("active set").remove(&lane.session);
2303 let _ = lane.job.events.send(LaneEvent::Failed(msg.clone()));
2304 }
2305 for session in drained {
2306 self.drop_ephemeral_artifact(session);
2307 }
2308 let queued: Vec<GenerateJob> = self.queues.iter_mut().flat_map(|q| q.jobs.drain(..)).collect();
2309 for job in queued {
2310 self.finish_failed(&job, &msg);
2311 }
2312 }
2313
2314 fn stage(&mut self, tokens: &[u32]) -> Result<linkw::TokenRefMsg, superfluid_agent::AgentError> {
2315 self.staged_unread += 1;
2316 self.host.client().stage_prompt(tokens)
2317 }
2318
2319 fn admit_from_queue(&mut self) {
2320 self.admit_retry = false;
2321 let retired = self.retired.lock().expect("retired").clone();
2322 if let Some(msg) = retired {
2323 let queued: Vec<GenerateJob> = self.queues.iter_mut().flat_map(|q| q.jobs.drain(..)).collect();
2324 for job in queued {
2325 self.finish_failed(&job, &msg);
2326 }
2327 return;
2328 }
2329 let mut held: Vec<GenerateJob> = Vec::new();
2330 if self.lanes.iter().all(|l| l.admitted) {
2331 self.staged_unread = 0;
2332 }
2333 let stage_slots = ring_slots(&mut self.host, superfluid_shm::TOKEN_RING_IN).unwrap_or(u32::MAX);
2334 while self.live_lanes() < self.max_lanes {
2335 if self.staged_unread > 0 && self.staged_unread.saturating_add(STAGES_PER_ADMISSION) > stage_slots {
2339 break;
2340 }
2341 let Some(job) = self.pop_admissible() else {
2342 break;
2343 };
2344 if job.stream.is_empty() {
2345 self.finish_failed(&job, "session has no content to generate from");
2346 continue;
2347 }
2348 if self.awaits_a_shared_prefix(&job) {
2349 held.push(job);
2350 continue;
2351 }
2352 if let Some(deadline) = job.deadline {
2353 if std::time::Instant::now() > deadline {
2354 let _ = job.events.send(LaneEvent::Expired);
2355 self.active.lock().expect("active set").remove(&job.session);
2356 self.drop_ephemeral_artifact(job.session);
2357 continue;
2358 }
2359 }
2360 let max_stream = self.host.max_stream_tokens();
2361 if job.stream.len() as u64 > max_stream {
2362 let msg = DaemonError::StreamTooLong {
2363 len: job.stream.len() as u64,
2364 max: max_stream,
2365 }
2366 .to_string();
2367 self.finish_failed(&job, &msg);
2368 continue;
2369 }
2370 let span = match self.stage(&job.stream) {
2371 Ok(s) => s,
2372 Err(e) => {
2373 let transport = matches!(
2374 e,
2375 superfluid_agent::AgentError::Io(_)
2376 | superfluid_agent::AgentError::Closed
2377 | superfluid_agent::AgentError::Shm(_)
2378 );
2379 if transport && self.host.is_process() && self.respawn_streak < MAX_WORKER_RESPAWNS {
2380 self.respawn_streak += 1;
2381 eprintln!(
2382 "[superfluid] engine worker gone at admission ({e}); respawning"
2383 );
2384 self.stats.worker_respawns.fetch_add(1, Ordering::Relaxed);
2385 if self
2386 .respawn_worker()
2387 .and_then(|()| self.speculation.reregister(&mut self.host))
2388 .is_ok()
2389 {
2390 let c = (job.class as usize).min(3);
2391 self.queues[c].jobs.push_front(job);
2392 self.after_respawn();
2393 self.admit_retry = true;
2394 break;
2395 }
2396 }
2397 self.finish_failed(&job, &format!("stage failed: {e}"));
2398 continue;
2399 }
2400 };
2401 let mut seed_handle = 0u64;
2402 let mut warm_prefix = 0u64;
2403 if self.host.has_unservable_space() {
2404 self.stats
2405 .unservable_cold_admissions
2406 .fetch_add(1, Ordering::Relaxed);
2407 } else if let Some(page) = self.host.page_size() {
2408 if self.host.cache_serves_all_spaces() && job.media.is_empty() {
2409 if let Ok(m) = self.host.client().match_prefix(ALL_SPACES, span) {
2410 let align_down = |l: u64| l.checked_div(page).unwrap_or(0) * page;
2411 let replay = Self::decode_replay_for(
2412 Self::strategy_slot_for(&self.speculation, &job),
2413 job.carry.as_ref().map(|c| c.produced_before).unwrap_or(0),
2414 );
2415 let stream_len = job.stream.len() as u64;
2416 let aligned_cap =
2417 align_down(stream_len.saturating_sub(1).min(stream_len.saturating_sub(replay as u64)));
2418 let best = crate::pins::common_seed_lengths(&self.host, &m, aligned_cap)
2419 .first()
2420 .copied()
2421 .unwrap_or(0);
2422 if best >= page {
2423 if let Ok(h) = self.host.client().seed_acquire(span, best, 0) {
2424 seed_handle = h;
2425 warm_prefix = best;
2426 }
2427 }
2428 }
2429 }
2430 (seed_handle, warm_prefix) =
2431 self.resume_from_park(&job, span, page, seed_handle, warm_prefix);
2432 }
2433 if warm_prefix > 0 {
2434 self.stats
2435 .warm_prefix_tokens
2436 .fetch_add(warm_prefix, Ordering::Relaxed);
2437 } else {
2438 self.stats.cold_admissions.fetch_add(1, Ordering::Relaxed);
2439 }
2440 let tag = self.next_lane;
2441 self.next_lane += 1;
2442 let mut job = job;
2443 let mut media_ok = true;
2444 let mut bound: Vec<u64> = Vec::new();
2445 for (offset, path, _n) in job.media.iter().filter(|(o, _, _)| *o >= warm_prefix) {
2446 let path_s = path.to_string_lossy().into_owned();
2447 let enc = self.host.client().state_sync(linkw::StateSyncReqMsg::MediaEncode {
2448 image_path: path_s.clone(),
2449 });
2450 let handle = match enc {
2451 Ok(linkw::StateSyncOkMsg::MediaEncode { media_handle, .. }) => media_handle,
2452 other => {
2453 eprintln!("[superfluid] media encode failed for {path_s}: {other:?}");
2454 media_ok = false;
2455 break;
2456 }
2457 };
2458 let bind = self.host.client().state_sync(linkw::StateSyncReqMsg::MediaBind {
2459 lane_tag: tag,
2460 media_handle: handle,
2461 token_offset: *offset as u32,
2462 });
2463 if matches!(bind, Ok(linkw::StateSyncOkMsg::MediaBind)) {
2464 bound.push(handle);
2465 } else {
2466 let _ = self
2467 .host
2468 .client()
2469 .state_sync(linkw::StateSyncReqMsg::MediaRelease { media_handle: handle });
2470 media_ok = false;
2471 break;
2472 }
2473 }
2474 if !media_ok {
2475 for handle in &bound {
2476 let _ = self.host.client().state_sync(
2477 linkw::StateSyncReqMsg::MediaRelease { media_handle: *handle },
2478 );
2479 }
2480 if seed_handle != 0 {
2481 let _ = self.host.client().seed_release(seed_handle);
2482 }
2483 self.finish_failed(&job, "media encode/bind failed (see daemon log)");
2484 continue;
2485 }
2486 let spec_before = job.carry.as_ref().map(|c| c.spec_before).unwrap_or_default();
2487 let decode_before = job.carry.as_ref().map(|c| c.decode_before).unwrap_or_default();
2488 let (out, produced_before, warm_prefix_first) = match job.carry.take() {
2489 Some(mut c) => {
2490 c.out.disp_tokens = c.produced_before;
2491 (c.out, c.produced_before, c.warm_prefix_first)
2492 }
2493 None => (LaneOutput::new(&self.codec, job.extras.open_channel, 0), 0, None),
2494 };
2495 let all_tokens = job.stream.clone();
2496 self.lanes.push(Lane {
2497 tag,
2498 session: job.session,
2499 job,
2500 phase: LanePhase::Prefilling,
2501 prefilled: warm_prefix,
2502 produced: 0,
2503 fin: finish::NONE,
2504 warm_prefix,
2505 out,
2506 frame_fed: 0,
2507 emit_dead: false,
2508 admitted: false,
2509 seed_handle,
2510 span,
2511 all_tokens,
2512 parked: false,
2513 parked_lossy: false,
2514 preempted: false,
2515 starved_ticks: 0,
2516 prefill_starved_ticks: 0,
2517 checkpoint_at: 0,
2518 checkpoint_due: false,
2519 checkpointed: 0,
2520 turn_kept: false,
2521 produced_before,
2522 spec: spec_before,
2523 decode: decode_before,
2524 decode_ticks: None,
2525 warm_prefix_first,
2526 admitted_at: std::time::Instant::now(),
2527 ttft_observed: false,
2528 ttft_frame_ms: None,
2529 fault_code: None,
2530 });
2531 }
2532 for job in held.into_iter().rev() {
2533 let c = (job.class as usize).min(3);
2534 self.queues[c].jobs.push_front(job);
2535 }
2536 }
2537
2538 fn resume_from_park(
2539 &mut self,
2540 job: &GenerateJob,
2541 span: linkw::TokenRefMsg,
2542 page: u64,
2543 seed_handle: u64,
2544 warm_prefix: u64,
2545 ) -> (u64, u64) {
2546 let Some(dir) = self.park_dir.clone() else {
2547 return (seed_handle, warm_prefix);
2548 };
2549 let Some(p) = park::read(&dir, job.session) else {
2550 return self.resume_from_prefix(job, span, page, seed_handle, warm_prefix);
2551 };
2552 if p.multi_space {
2553 if !self.host.multi_space() && self.host.parkable_spaces().len() != p.spaces.len() {
2554 return (seed_handle, warm_prefix);
2555 }
2556 return match self.resume_adopt(job, &p) {
2557 Ok((h, covered)) => {
2558 if seed_handle != 0 {
2559 let _ = self.host.client().seed_release(seed_handle);
2560 }
2561 self.stats.resumes.fetch_add(1, Ordering::Relaxed);
2562 (h, covered)
2563 }
2564 Err(msg) => {
2565 eprintln!("[superfluid] park resume fell back cold (session {}): {msg}", job.session);
2566 (seed_handle, warm_prefix)
2567 }
2568 };
2569 }
2570 if self.host.multi_space() || !job.media.is_empty() {
2571 return (seed_handle, warm_prefix);
2572 }
2573 let align_down = |l: u64| l.checked_div(page).unwrap_or(0) * page;
2574 let aligned_cap = align_down(job.stream.len() as u64 - 1);
2575 let usable = p.covered.min(aligned_cap);
2576 if p.covered % page != 0
2577 || usable <= warm_prefix
2578 || p.covered as usize > job.stream.len()
2579 || park::stream_digest(&job.stream[..p.covered as usize]) != p.stream_digest
2580 {
2581 return (seed_handle, warm_prefix);
2582 }
2583 let covered_span = match self.stage(&job.stream[..p.covered as usize]) {
2584 Ok(s) => s,
2585 Err(e) => {
2586 eprintln!(
2587 "[superfluid] park resume fell back cold (session {}): stage: {e}",
2588 job.session
2589 );
2590 return (seed_handle, warm_prefix);
2591 }
2592 };
2593 if let Err(msg) = self.restore_to_cache(job.session, &p, covered_span) {
2594 eprintln!(
2595 "[superfluid] park resume fell back cold (session {}): {msg}",
2596 job.session
2597 );
2598 return (seed_handle, warm_prefix);
2599 }
2600 match self.host.client().seed_acquire(span, usable, 0) {
2601 Ok(h) => {
2602 if seed_handle != 0 {
2603 let _ = self.host.client().seed_release(seed_handle);
2604 }
2605 self.stats.resumes.fetch_add(1, Ordering::Relaxed);
2606 (h, usable)
2607 }
2608 Err(e) => {
2609 eprintln!(
2610 "[superfluid] park resume restored but could not seed (session {}): {e}",
2611 job.session
2612 );
2613 (seed_handle, warm_prefix)
2614 }
2615 }
2616 }
2617
2618 fn persist_prefix(&mut self, i: usize) -> Result<(), ParkRefused> {
2622 let (Some(dir), Some(page)) = (self.park_dir.clone(), self.host.page_size()) else {
2623 return Ok(());
2624 };
2625 let lane = &self.lanes[i];
2626 if (lane.job.completion && !lane.job.extras.keep_prefix)
2627 || !lane.job.media.is_empty()
2628 || self.host.multi_space()
2629 || self.host.has_unservable_space()
2630 {
2631 return Ok(());
2632 }
2633 let n = lane.all_tokens.len() as u64;
2634 let aligned = n.saturating_sub(1) / page * page;
2635 if aligned < PREFIX_KEEP_MIN_TOKENS {
2636 return Ok(());
2637 }
2638 let covered = self.prefixes.iter().map(|(_, t)| park::common_prefix(t, &lane.all_tokens) as u64).max().unwrap_or(0);
2639 if covered * 4 >= aligned * PREFIX_KEPT_QUARTERS {
2640 return Ok(());
2641 }
2642 let Some((aligned, encoding, sealed)) = self.export_lane_kv(i)? else {
2643 return Ok(());
2644 };
2645 let tokens = self.lanes[i].all_tokens[..aligned as usize].to_vec();
2646 let dir = park::prefix_dir(&dir);
2647 let path = dir.join(format!("{:016x}.park", park::stream_digest(&tokens)));
2648 let task = ParkTask::WritePrefix { dir, tokens: tokens.clone(), encoding, sealed };
2649 match self.park_tx.as_ref().expect("writer exists with park_dir").try_send(task) {
2650 Ok(()) => {
2651 self.prefixes.retain(|(p, _)| *p != path);
2652 self.prefixes.insert(0, (path, tokens));
2653 self.prefixes.truncate(PREFIX_INDEX_MAX);
2654 Ok(())
2655 }
2656 Err(std::sync::mpsc::TrySendError::Full(_)) => Err("park writer backlogged; prefix dropped".into()),
2657 Err(std::sync::mpsc::TrySendError::Disconnected(_)) => Err("park writer gone".into()),
2658 }
2659 }
2660
2661 fn resume_from_prefix(
2665 &mut self,
2666 job: &GenerateJob,
2667 span: linkw::TokenRefMsg,
2668 page: u64,
2669 seed_handle: u64,
2670 warm_prefix: u64,
2671 ) -> (u64, u64) {
2672 if self.prefixes.is_empty() || self.host.multi_space() || !job.media.is_empty() || job.stream.is_empty() {
2673 return (seed_handle, warm_prefix);
2674 }
2675 let align_down = |l: u64| l.checked_div(page).unwrap_or(0) * page;
2676 let aligned_cap = align_down(job.stream.len() as u64 - 1);
2677 let Some((at, usable)) = self
2678 .prefixes
2679 .iter()
2680 .enumerate()
2681 .map(|(at, (_, t))| (at, align_down((park::common_prefix(t, &job.stream) as u64).min(aligned_cap))))
2682 .max_by_key(|&(at, usable)| (usable, std::cmp::Reverse(at)))
2683 else {
2684 return (seed_handle, warm_prefix);
2685 };
2686 if usable < PREFIX_KEEP_MIN_TOKENS.min(aligned_cap) || usable <= warm_prefix + page {
2687 return (seed_handle, warm_prefix);
2688 }
2689 let (path, tokens) = self.prefixes[at].clone();
2690 let Some(p) = park::read_prefix(&path) else {
2691 self.prefixes.remove(at);
2692 return (seed_handle, warm_prefix);
2693 };
2694 let kept = match self.stage(&tokens) {
2695 Ok(s) => s,
2696 Err(e) => {
2697 eprintln!("[superfluid] prefix restore fell back cold (session {}): stage: {e}", job.session);
2698 return (seed_handle, warm_prefix);
2699 }
2700 };
2701 if let Err(msg) = self.restore_to_cache(job.session, &p, kept) {
2702 eprintln!("[superfluid] prefix restore fell back cold (session {}): {msg}", job.session);
2703 self.prefixes.remove(at);
2704 return (seed_handle, warm_prefix);
2705 }
2706 match self.host.client().seed_acquire(span, usable, 0) {
2707 Ok(h) => {
2708 if seed_handle != 0 {
2709 let _ = self.host.client().seed_release(seed_handle);
2710 }
2711 self.stats.resumes.fetch_add(1, Ordering::Relaxed);
2712 (h, usable)
2713 }
2714 Err(e) => {
2715 eprintln!("[superfluid] prefix restored but could not seed (session {}): {e}", job.session);
2716 (seed_handle, warm_prefix)
2717 }
2718 }
2719 }
2720
2721 fn restore_to_cache(
2722 &mut self,
2723 session: u64,
2724 p: &park::Parked,
2725 span: linkw::TokenRefMsg,
2726 ) -> Result<(), String> {
2727 let _op_span = tracing::info_span!(
2728 "restore",
2729 session,
2730 model = self.trace_scope.model.as_str(),
2731 store_ns = self.trace_scope.ns,
2732 covered = p.covered,
2733 bytes = p.sealed.len(),
2734 encoding = p.encoding
2735 )
2736 .entered();
2737 let client = self.host.client();
2738 let seq = match client.state_sync(linkw::StateSyncReqMsg::CreateSequence) {
2739 Ok(linkw::StateSyncOkMsg::CreateSequence { sequence }) => sequence,
2740 Ok(_) => return Err("mismatched CreateSequence reply".into()),
2741 Err(e) => return Err(e.to_string()),
2742 };
2743 let free_seq = |client: &mut superfluid_agent::WorkerClient| {
2744 let _ = client.state_sync(linkw::StateSyncReqMsg::FreeSequence { sequence: seq });
2745 };
2746 let buf = match client.buf_create(p.sealed.len() as u64) {
2747 Ok(b) => b,
2748 Err(e) => {
2749 free_seq(client);
2750 return Err(e.to_string());
2751 }
2752 };
2753 let result = (|| {
2754 client
2755 .buf_write(buf, &p.sealed)
2756 .map_err(|e| e.to_string())?;
2757 let op = client
2758 .state_op(linkw::StateOpKind::Restore {
2759 sequence: seq,
2760 space_id: KV_SPACE,
2761 buf_id: buf,
2762 })
2763 .map_err(|e| e.to_string())?;
2764 let done = client.wait_op_done(op).map_err(|e| e.to_string())?;
2765 if done.state != op_state::DONE || done.error != 0 {
2766 return Err(format!(
2767 "restore terminal state {} error {}",
2768 done.state, done.error
2769 ));
2770 }
2771 Ok(())
2772 })();
2773 let _ = client.buf_release(buf);
2774 if let Err(e) = result {
2775 free_seq(client);
2776 return Err(e);
2777 }
2778 match client.state_sync(linkw::StateSyncReqMsg::PublishSequence {
2779 sequence: seq,
2780 span,
2781 }) {
2782 Ok(linkw::StateSyncOkMsg::PublishSequence) => Ok(()),
2783 Ok(_) => {
2784 free_seq(client);
2785 Err("mismatched PublishSequence reply".into())
2786 }
2787 Err(e) => {
2788 free_seq(client);
2789 Err(e.to_string())
2790 }
2791 }
2792 }
2793
2794 fn relieve_pressure(&mut self, mem: &linkw::MemCountersMsg, host_short: bool) {
2795 let (high, low) = (self.pressure_high_pct as u64, self.pressure_low_pct as u64);
2796 if std::env::var_os("SUPERFLUID_DEBUG_PRESSURE").is_some() {
2797 eprintln!(
2798 "[pressure] used={} total={} evictable={}",
2799 mem.pool_blocks_used, mem.pool_blocks_total, mem.pool_bytes_evictable
2800 );
2801 }
2802 if mem.pool_blocks_total == 0 || high == 0 {
2803 return;
2804 }
2805 let Some(block_bytes) = self.host.kv_block_bytes() else {
2806 return;
2807 };
2808 let (used, total) = (mem.pool_blocks_used, mem.pool_blocks_total);
2809 if used * 100 < total * high {
2810 return;
2811 }
2812 let low_blocks = total * low / 100;
2813 let excess = used.saturating_sub(low_blocks);
2814 if excess == 0 {
2815 return;
2816 }
2817 let slack = total * block_bytes / 4;
2818 let mut evictable = mem.pool_bytes_evictable;
2819 let need = excess * block_bytes;
2820 if need > evictable && self.pins.any_held() {
2821 evictable += self
2822 .pins
2823 .yield_for(&mut self.host, &self.stats, need - evictable);
2824 }
2825 let bytes_target = (excess * block_bytes + slack).min(evictable.max(1));
2826 let all = if host_short { u64::MAX } else { !HOST_EXPORT_CLASS };
2829 let mut freed = self.evict_cache(bytes_target, all & !SHARED_PREFIX_CLASS);
2832 if freed < need {
2833 freed += self.evict_cache(need - freed, all);
2834 }
2835 if freed < need && self.pins.any_held() {
2836 let yielded = self
2837 .pins
2838 .yield_for(&mut self.host, &self.stats, need - freed);
2839 if yielded > 0 {
2840 freed += self.evict_cache((need - freed + slack).min(yielded), all);
2841 }
2842 }
2843 if freed > 0 {
2844 self.stats.pressure_evictions.fetch_add(1, Ordering::Relaxed);
2845 self.stats
2846 .pressure_bytes_evicted
2847 .fetch_add(freed, Ordering::Relaxed);
2848 }
2849 }
2850
2851 fn evict_cache(&mut self, bytes_target: u64, classes: u64) -> u64 {
2852 let res = self
2853 .host
2854 .client()
2855 .state_sync(linkw::StateSyncReqMsg::CacheEvict {
2856 bytes_target,
2857 policy: linkw::ShedPolicyMsg {
2858 victim_lanes: vec![],
2859 evictable_cache_classes: classes,
2860 protected_quota_bytes: 0,
2861 max_evict_bytes: bytes_target,
2862 },
2863 });
2864 match res {
2865 Ok(linkw::StateSyncOkMsg::CacheEvict { bytes_freed }) => bytes_freed,
2866 _ => 0,
2867 }
2868 }
2869
2870 fn retry_park_removals(&mut self) {
2871 if self.park_removals_pending.is_empty() {
2872 return;
2873 }
2874 let (Some(dir), Some(tx)) = (self.park_dir.clone(), self.park_tx.as_ref()) else {
2875 self.park_removals_pending.clear();
2876 return;
2877 };
2878 self.park_removals_pending.retain(|&session| {
2879 matches!(
2880 tx.try_send(ParkTask::Remove {
2881 dir: dir.clone(),
2882 session,
2883 }),
2884 Err(std::sync::mpsc::TrySendError::Full(_))
2885 )
2886 });
2887 }
2888
2889 fn defer_admits(&mut self, tags: &[u64]) {
2893 let mut back: Vec<GenerateJob> = Vec::new();
2894 let mut seeds: Vec<u64> = Vec::new();
2895 let mut i = 0;
2896 while i < self.lanes.len() {
2897 if tags.contains(&self.lanes[i].tag) {
2898 let lane = self.lanes.remove(i);
2899 if lane.seed_handle != 0 {
2900 seeds.push(lane.seed_handle);
2901 }
2902 back.push(lane.job);
2903 } else {
2904 i += 1;
2905 }
2906 }
2907 for seed in seeds {
2908 let _ = self.host.client().seed_release(seed);
2909 }
2910 for job in back.into_iter().rev() {
2911 let c = (job.class as usize).min(3);
2912 self.queues[c].jobs.push_front(job);
2913 }
2914 }
2915
2916 fn drop_refused(&mut self, refused: &HashMap<u64, u32>) {
2917 let mut gone: Vec<(u64, u64)> = Vec::new();
2918 let mut active = self.active.lock().expect("active set");
2919 self.lanes.retain(|lane| {
2920 let Some(&code) = refused.get(&lane.tag) else { return true };
2921 active.remove(&lane.session);
2922 let status = superfluid_abi::Status::from_raw(code as i32)
2923 .map(|s| format!("{s:?}"))
2924 .unwrap_or_else(|| format!("status {}", code as i32));
2925 tracing::warn!(session = lane.session, lane = lane.tag, status = %status, "engine refused an admit");
2926 let _ = lane.job.events.send(LaneEvent::Failed(format!("the engine refused to admit this request ({status})")));
2927 gone.push((lane.session, lane.seed_handle));
2928 false
2929 });
2930 drop(active);
2931 for (session, seed) in gone {
2932 if seed != 0 {
2933 let _ = self.host.client().seed_release(seed);
2934 }
2935 self.drop_ephemeral_artifact(session);
2936 }
2937 }
2938
2939 fn drop_ephemeral_artifact(&mut self, session: u64) {
2940 if !self.ephemeral_parked.remove(&session) {
2941 return;
2942 }
2943 let Some(dir) = self.park_dir.clone() else {
2944 return;
2945 };
2946 let Some(tx) = self.park_tx.as_ref() else {
2947 return;
2948 };
2949 if tx.try_send(ParkTask::Remove { dir, session }).is_err() {
2950 self.park_removals_pending.push(session);
2951 }
2952 }
2953
2954 fn export_lane_kv(&mut self, i: usize) -> Result<Option<(u64, u8, Vec<u8>)>, ParkRefused> {
2958 let Some(page) = self.host.page_size() else {
2959 return Ok(None);
2960 };
2961 let (tag, session) = (self.lanes[i].tag, self.lanes[i].session);
2962 let n = self.lanes[i].all_tokens.len() as u64;
2963 if n <= 1 {
2964 return Ok(None);
2965 }
2966 let mut aligned = ((n - 1) / page) * page;
2967 if aligned < page {
2968 return Ok(None);
2969 }
2970 let park_encoding = if self.park_lossy {
2971 encoding::Q8
2972 } else {
2973 encoding::LOSSLESS
2974 };
2975 let client = self.host.client();
2976 let seq = match client.state_sync(linkw::StateSyncReqMsg::LaneSequence { lane_tag: tag }) {
2977 Ok(linkw::StateSyncOkMsg::LaneSequence { sequence }) => sequence,
2978 Ok(_) => return Err("mismatched LaneSequence reply".into()),
2979 Err(e) => return Err(e.to_string().into()),
2980 };
2981 if let Ok(linkw::StateSyncOkMsg::SnapshotBoundary { boundary }) =
2982 client.state_sync(linkw::StateSyncReqMsg::SnapshotBoundary { sequence: seq, space_id: KV_SPACE, cap: n - 1 })
2983 {
2984 aligned = aligned.min(boundary);
2985 if aligned < page {
2986 return Ok(None);
2987 }
2988 }
2989 let _span = tracing::info_span!(
2990 "park",
2991 session,
2992 model = self.trace_scope.model.as_str(),
2993 store_ns = self.trace_scope.ns,
2994 covered = aligned,
2995 lossy = self.park_lossy
2996 )
2997 .entered();
2998 let (bytes, sizing_gen) = match client.state_sync(linkw::StateSyncReqMsg::ExportSize {
2999 sequence: seq,
3000 space_id: KV_SPACE,
3001 range: linkw::TokenRangeMsg {
3002 start: 0,
3003 end: aligned,
3004 },
3005 encoding: park_encoding,
3006 }) {
3007 Ok(linkw::StateSyncOkMsg::ExportSize {
3008 required_bytes,
3009 sizing_gen,
3010 }) => (required_bytes, sizing_gen),
3011 Ok(_) => return Err("mismatched ExportSize reply".into()),
3012 Err(e) => return Err(e.to_string().into()),
3013 };
3014 let buf = client.buf_create(bytes).map_err(|e| ParkRefused::from(e.to_string()))?;
3015 let result = (|| {
3016 let op_kind = if self.park_lossy {
3017 linkw::StateOpKind::Demote {
3018 sequence: seq,
3019 space_id: KV_SPACE,
3020 range: linkw::TokenRangeMsg {
3021 start: 0,
3022 end: aligned,
3023 },
3024 encoding: encoding::Q8,
3025 buf_id: buf,
3026 sizing_gen,
3027 }
3028 } else {
3029 linkw::StateOpKind::Snapshot {
3030 sequence: seq,
3031 space_id: KV_SPACE,
3032 boundary_pos: aligned,
3033 buf_id: buf,
3034 sizing_gen,
3035 }
3036 };
3037 let op = client.state_op(op_kind).map_err(|e| e.to_string())?;
3038 let done = client.wait_op_done(op).map_err(|e| e.to_string())?;
3039 if done.state != op_state::DONE || done.error != 0 {
3040 return Err(format!(
3041 "snapshot terminal state {} error {}",
3042 done.state, done.error
3043 ));
3044 }
3045 let take = if done.bytes_moved > 0 {
3046 done.bytes_moved.min(bytes)
3047 } else {
3048 bytes
3049 };
3050 client
3051 .buf_read(buf, take as usize)
3052 .map_err(|e| e.to_string())
3053 })();
3054 let _ = client.buf_release(buf);
3055 let sealed = result.map_err(ParkRefused::from)?;
3056 Ok(Some((aligned, park_encoding, sealed)))
3057 }
3058
3059 fn park_lane(&mut self, i: usize) -> Result<(), ParkRefused> {
3060 let Some(dir) = self.park_dir.clone() else {
3061 return Ok(());
3062 };
3063 if self.lanes[i].job.completion {
3064 return Ok(());
3065 }
3066 if (self.host.multi_space() || !self.lanes[i].job.media.is_empty()) && !self.host.has_unservable_space() {
3067 return self.park_lane_multi(i, dir);
3068 }
3069 if self.host.has_unservable_space() {
3070 self.stats
3071 .unservable_park_refusals
3072 .fetch_add(1, Ordering::Relaxed);
3073 return Err(ParkRefused::SpaceUnservable);
3074 }
3075 let session = self.lanes[i].session;
3076 let Some((aligned, park_encoding, sealed)) = self.export_lane_kv(i)? else {
3077 return Ok(());
3078 };
3079 let digest = park::stream_digest(&self.lanes[i].all_tokens[..aligned as usize]);
3080 let job = ParkJob {
3081 dir,
3082 session,
3083 covered: aligned,
3084 digest,
3085 encoding: park_encoding,
3086 sealed,
3087 spaces: None,
3088 };
3089 match self
3090 .park_tx
3091 .as_ref()
3092 .expect("writer exists with park_dir")
3093 .try_send(ParkTask::Write(Box::new(job)))
3094 {
3095 Ok(()) => {
3096 if self.lanes[i].job.extras.ephemeral {
3097 self.ephemeral_parked.insert(session);
3098 }
3099 if self.park_lossy {
3100 self.stats.parks_lossy.fetch_add(1, Ordering::Relaxed);
3101 } else {
3102 self.stats.parks_lossless.fetch_add(1, Ordering::Relaxed);
3103 }
3104 }
3105 Err(std::sync::mpsc::TrySendError::Full(_)) => {
3106 return Err("park writer backlogged; artifact dropped".into());
3107 }
3108 Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {
3109 return Err("park writer gone".into());
3110 }
3111 }
3112 self.lanes[i].parked_lossy = self.park_lossy;
3113 Ok(())
3114 }
3115
3116 fn park_lane_multi(&mut self, i: usize, dir: PathBuf) -> Result<(), ParkRefused> {
3117 let (tag, session) = (self.lanes[i].tag, self.lanes[i].session);
3118 let n = self.lanes[i].all_tokens.len() as u64;
3119 if n <= 1 {
3120 return Ok(());
3121 }
3122 let spaces = self.host.parkable_spaces();
3123 let lossy = self.park_lossy;
3124 let fallback_boundary = self.host.park_boundary(n - 1);
3125 let client = self.host.client();
3126 let seq = match client.state_sync(linkw::StateSyncReqMsg::LaneSequence { lane_tag: tag }) {
3127 Ok(linkw::StateSyncOkMsg::LaneSequence { sequence }) => sequence,
3128 Ok(_) => return Err("mismatched LaneSequence reply".into()),
3129 Err(e) => return Err(e.to_string().into()),
3130 };
3131 let cap = n - 1;
3132 let mut covered = cap;
3133 let mut queried = false;
3134 for &(space_id, _, _, _) in &spaces {
3135 match client.state_sync(linkw::StateSyncReqMsg::SnapshotBoundary { sequence: seq, space_id, cap }) {
3136 Ok(linkw::StateSyncOkMsg::SnapshotBoundary { boundary }) => {
3137 queried = true;
3138 covered = covered.min(boundary);
3139 }
3140 Ok(_) => return Err("mismatched SnapshotBoundary reply".into()),
3141 Err(_) => break,
3142 }
3143 }
3144 if !queried {
3145 covered = fallback_boundary;
3146 }
3147 if covered == 0 {
3148 return Ok(());
3149 }
3150 let digest = park::stream_digest(&self.lanes[i].all_tokens[..covered as usize]);
3151 let _span = tracing::info_span!(
3152 "park",
3153 session,
3154 model = self.trace_scope.model.as_str(),
3155 store_ns = self.trace_scope.ns,
3156 covered,
3157 lossy = self.park_lossy,
3158 multi = true
3159 )
3160 .entered();
3161 let mut parked: Vec<park::ParkedSpace> = Vec::with_capacity(spaces.len());
3162 for &(space_id, kind, _, _) in &spaces {
3163 let paged = kind == superfluid_abi::space_kind::PAGED_TOKEN_KV;
3164 let encoding = if paged && lossy { encoding::Q8 } else { encoding::LOSSLESS };
3165 let range = if paged {
3166 linkw::TokenRangeMsg { start: 0, end: covered }
3167 } else {
3168 linkw::TokenRangeMsg { start: covered, end: covered }
3169 };
3170 let (bytes, sizing_gen) = match client.state_sync(linkw::StateSyncReqMsg::ExportSize {
3171 sequence: seq,
3172 space_id,
3173 range,
3174 encoding,
3175 }) {
3176 Ok(linkw::StateSyncOkMsg::ExportSize { required_bytes, sizing_gen }) => {
3177 (required_bytes, sizing_gen)
3178 }
3179 Ok(_) => return Err("mismatched ExportSize reply".into()),
3180 Err(e) => return Err(format!("export size (space {space_id}): {e}").into()),
3181 };
3182 let buf = client.buf_create(bytes).map_err(|e| ParkRefused::from(e.to_string()))?;
3183 let result = (|| {
3184 let op_kind = if paged && lossy {
3185 linkw::StateOpKind::Demote {
3186 sequence: seq,
3187 space_id,
3188 range,
3189 encoding: encoding::Q8,
3190 buf_id: buf,
3191 sizing_gen,
3192 }
3193 } else {
3194 linkw::StateOpKind::Snapshot {
3195 sequence: seq,
3196 space_id,
3197 boundary_pos: covered,
3198 buf_id: buf,
3199 sizing_gen,
3200 }
3201 };
3202 let op = client.state_op(op_kind).map_err(|e| e.to_string())?;
3203 let done = client.wait_op_done(op).map_err(|e| e.to_string())?;
3204 if done.state != op_state::DONE || done.error != 0 {
3205 return Err(format!(
3206 "space {space_id} export terminal state {} error {}",
3207 done.state, done.error
3208 ));
3209 }
3210 let take = if done.bytes_moved > 0 { done.bytes_moved.min(bytes) } else { bytes };
3211 client.buf_read(buf, take as usize).map_err(|e| e.to_string())
3212 })();
3213 let _ = client.buf_release(buf);
3214 let sealed = result.map_err(ParkRefused::from)?;
3215 parked.push(park::ParkedSpace {
3216 space_id,
3217 kind,
3218 encoding,
3219 sealed,
3220 });
3221 }
3222 let job = ParkJob {
3223 dir,
3224 session,
3225 covered,
3226 digest,
3227 encoding: if lossy { encoding::Q8 } else { encoding::LOSSLESS },
3228 sealed: Vec::new(),
3229 spaces: Some(parked),
3230 };
3231 match self
3232 .park_tx
3233 .as_ref()
3234 .expect("writer exists with park_dir")
3235 .try_send(ParkTask::Write(Box::new(job)))
3236 {
3237 Ok(()) => {
3238 if self.lanes[i].job.extras.ephemeral {
3239 self.ephemeral_parked.insert(session);
3240 }
3241 if lossy {
3242 self.stats.parks_lossy.fetch_add(1, Ordering::Relaxed);
3243 } else {
3244 self.stats.parks_lossless.fetch_add(1, Ordering::Relaxed);
3245 }
3246 }
3247 Err(std::sync::mpsc::TrySendError::Full(_)) => {
3248 return Err("park writer backlogged; artifact dropped".into());
3249 }
3250 Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {
3251 return Err("park writer gone".into());
3252 }
3253 }
3254 self.lanes[i].parked_lossy = lossy;
3255 Ok(())
3256 }
3257
3258 fn resume_adopt(&mut self, job: &GenerateJob, p: &park::Parked) -> Result<(u64, u64), String> {
3259 if p.covered as usize >= job.stream.len() {
3260 return Err("artifact covers the whole stream (nothing to prefill)".into());
3261 }
3262 if park::stream_digest(&job.stream[..p.covered as usize]) != p.stream_digest {
3263 return Err("stream digest mismatch (rewritten history)".into());
3264 }
3265 let _op_span = tracing::info_span!(
3266 "restore",
3267 session = job.session,
3268 model = self.trace_scope.model.as_str(),
3269 store_ns = self.trace_scope.ns,
3270 covered = p.covered,
3271 spaces = p.spaces.len(),
3272 multi = true
3273 )
3274 .entered();
3275 let want = self.host.parkable_spaces();
3276 for &(id, _, _, _) in &want {
3277 if !p.spaces.iter().any(|s| s.space_id == id) {
3278 return Err(format!("artifact lacks space {id}"));
3279 }
3280 }
3281 let covered_span = self
3282 .stage(&job.stream[..p.covered as usize])
3283 .map_err(|e| e.to_string())?;
3284 let client = self.host.client();
3285 let seq = match client.state_sync(linkw::StateSyncReqMsg::CreateSequence) {
3286 Ok(linkw::StateSyncOkMsg::CreateSequence { sequence }) => sequence,
3287 Ok(_) => return Err("mismatched CreateSequence reply".into()),
3288 Err(e) => return Err(e.to_string()),
3289 };
3290 let free_seq = |client: &mut superfluid_agent::WorkerClient| {
3291 let _ = client.state_sync(linkw::StateSyncReqMsg::FreeSequence { sequence: seq });
3292 };
3293 let mut order: Vec<&park::ParkedSpace> = p.spaces.iter().collect();
3294 order.sort_by_key(|s| (s.kind != superfluid_abi::space_kind::PAGED_TOKEN_KV, s.space_id));
3295 for sp in order {
3296 let buf = match client.buf_create(sp.sealed.len() as u64) {
3297 Ok(b) => b,
3298 Err(e) => {
3299 free_seq(client);
3300 return Err(e.to_string());
3301 }
3302 };
3303 let result = (|| {
3304 client.buf_write(buf, &sp.sealed).map_err(|e| e.to_string())?;
3305 let op = client
3306 .state_op(linkw::StateOpKind::Restore {
3307 sequence: seq,
3308 space_id: sp.space_id,
3309 buf_id: buf,
3310 })
3311 .map_err(|e| e.to_string())?;
3312 let done = client.wait_op_done(op).map_err(|e| e.to_string())?;
3313 if done.state != op_state::DONE || done.error != 0 {
3314 return Err(format!(
3315 "restore space {} terminal state {} error {}",
3316 sp.space_id, done.state, done.error
3317 ));
3318 }
3319 Ok(())
3320 })();
3321 let _ = client.buf_release(buf);
3322 if let Err(e) = result {
3323 free_seq(client);
3324 return Err(e);
3325 }
3326 }
3327 match client.state_sync(linkw::StateSyncReqMsg::AdoptSequence {
3328 sequence: seq,
3329 span: covered_span,
3330 }) {
3331 Ok(linkw::StateSyncOkMsg::AdoptSequence { seed_handle }) => Ok((seed_handle, p.covered)),
3332 Ok(_) => {
3333 free_seq(client);
3334 Err("mismatched AdoptSequence reply".into())
3335 }
3336 Err(e) => {
3337 free_seq(client);
3338 Err(e.to_string())
3339 }
3340 }
3341 }
3342
3343 fn commit_tool_block(
3344 codec: &Arc<dyn TextCodec + Send + Sync>,
3345 store: &Arc<Mutex<SessionStore>>,
3346 tool_lease_ms: u64,
3347 session: u64,
3348 raw: String,
3349 schemas: Option<&crate::codec::ToolSchemas>,
3350 ) -> Result<Option<crate::CommittedEvent>, DaemonError> {
3351 match codec.parse_tool_call_with(&raw, schemas) {
3352 Some((name, args)) => {
3353 let mut store = store.lock().expect("store");
3354 let e = store.commit_tool_use(session, name, args)?;
3355 if tool_lease_ms > 0 {
3356 let deadline = crate::wal::now_unix_ms() + tool_lease_ms;
3357 let _ = store.commit_tool_lease(session, e.event_id, deadline);
3358 }
3359 Ok(Some(e))
3360 }
3361 None => Ok(Some(
3362 store.lock().expect("store").commit_tool_parse_failure(session, raw)?,
3363 )),
3364 }
3365 }
3366
3367 fn finish_failed(&mut self, job: &GenerateJob, msg: &str) {
3368 self.drop_ephemeral_artifact(job.session);
3369 self.active.lock().expect("active set").remove(&job.session);
3370 let _ = job.events.send(LaneEvent::Failed(msg.to_string()));
3371 }
3372
3373 fn publish_occupancy(&self) {
3374 self.stats.clear_tick_in_flight();
3375 self.stats
3376 .lanes_active
3377 .store(self.lanes.len() as u64, Ordering::Relaxed);
3378 for (c, q) in self.queues.iter().enumerate() {
3379 self.stats.queue_depth[c].store(q.jobs.len() as u64, Ordering::Relaxed);
3380 }
3381 }
3382
3383 fn tick_once(&mut self, tick_span: &tracing::Span) -> Result<(), DaemonError> {
3384 self.publish_occupancy();
3385 if self.lanes.is_empty() {
3386 return Ok(());
3387 }
3388
3389 let mut locally_done: Vec<u64> = Vec::new();
3390 for i in 0..self.lanes.len() {
3391 let session = self.lanes[i].session;
3392 let is_active = !matches!(self.lanes[i].phase, LanePhase::Finishing);
3393 if is_active && self.cancels.lock().expect("cancels").remove(&session) {
3394 let sink = OutputSink { codec: &self.codec, store: &self.store, bus: &self.bus, tool_lease_ms: self.tool_lease_ms };
3395 let lane = &mut self.lanes[i];
3396 lane.fin = finish::CANCELLED;
3397 lane.out.close(&sink, session, &lane.job.events, finish::CANCELLED, false)?;
3398 if lane.admitted {
3399 lane.phase = LanePhase::Finishing;
3400 } else {
3401 if lane.seed_handle != 0 {
3402 let _ = self.host.client().seed_release(lane.seed_handle);
3403 }
3404 locally_done.push(lane.tag);
3405 }
3406 }
3407 }
3408 let now_dl = std::time::Instant::now();
3409 for i in 0..self.lanes.len() {
3410 if matches!(self.lanes[i].phase, LanePhase::Finishing) {
3411 continue;
3412 }
3413 let Some(deadline) = self.lanes[i].job.deadline else {
3414 continue;
3415 };
3416 if now_dl <= deadline {
3417 continue;
3418 }
3419 let completion = self.lanes[i].job.completion;
3420 let session = self.lanes[i].session;
3421 if !completion {
3422 let sink = OutputSink { codec: &self.codec, store: &self.store, bus: &self.bus, tool_lease_ms: self.tool_lease_ms };
3423 let lane = &mut self.lanes[i];
3424 lane.out.close(&sink, session, &lane.job.events, finish::LENGTH, false)?;
3425 lane.fin = finish::LENGTH;
3426 } else {
3427 let _ = self.lanes[i].job.events.send(LaneEvent::Expired);
3428 }
3429 let lane = &mut self.lanes[i];
3430 if lane.admitted {
3431 lane.phase = LanePhase::Finishing;
3432 } else {
3433 if lane.seed_handle != 0 {
3434 let _ = self.host.client().seed_release(lane.seed_handle);
3435 }
3436 locally_done.push(lane.tag);
3437 }
3438 }
3439 if !locally_done.is_empty() {
3440 let mut locally_retired: Vec<u64> = Vec::new();
3441 let mut active = self.active.lock().expect("active set");
3442 self.lanes.retain(|lane| {
3443 if locally_done.contains(&lane.tag) {
3444 active.remove(&lane.session);
3445 let _ = lane.job.events.send(LaneEvent::Done {
3446 tokens_generated: lane.produced,
3447 finish: lane.fin,
3448 warm_prefix: lane.warm_prefix,
3449 spec: lane.spec,
3450 decode: lane.decode_measured(),
3451 });
3452 locally_retired.push(lane.session);
3453 false
3454 } else {
3455 true
3456 }
3457 });
3458 drop(active);
3459 for session in locally_retired {
3460 self.drop_ephemeral_artifact(session);
3461 }
3462 }
3463 self.publish_occupancy();
3464 if self.lanes.is_empty() {
3465 return Ok(());
3466 }
3467
3468 self.tick_count += 1;
3469 self.stats.ticks.fetch_add(1, Ordering::Relaxed);
3470
3471 if self.park_lossy {
3472 for lane in &mut self.lanes {
3473 if lane.preempted
3474 && !lane.parked
3475 && lane.job.params.seed != 0
3476 && !lane.job.extras.seed_drawn
3477 {
3478 lane.parked = true;
3479 }
3480 }
3481 }
3482 if self.park_lossy && !self.pins.is_empty() && self.pins.enforceable(&self.host) {
3483 for lane in &mut self.lanes {
3484 if matches!(lane.phase, LanePhase::Finishing)
3485 && !lane.parked
3486 && lane.job.media.is_empty()
3487 && !lane.job.completion
3488 && self.pins.awaits_lease(&self.host, lane.session)
3489 {
3490 lane.parked = true;
3491 }
3492 }
3493 }
3494
3495 if !self.park_removals_pending.is_empty() {
3496 if let (Some(dir), Some(tx)) = (self.park_dir.clone(), self.park_tx.as_ref()) {
3497 self.park_removals_pending.retain(|&session| {
3498 matches!(
3499 tx.try_send(ParkTask::Remove {
3500 dir: dir.clone(),
3501 session,
3502 }),
3503 Err(std::sync::mpsc::TrySendError::Full(_))
3504 )
3505 });
3506 } else {
3507 self.park_removals_pending.clear();
3508 }
3509 }
3510
3511 if self.park_dir.is_some() {
3512 for i in 0..self.lanes.len() {
3513 let keeps_prefix = self.lanes[i].job.completion && self.lanes[i].job.extras.keep_prefix;
3514 if matches!(self.lanes[i].phase, LanePhase::Finishing)
3515 && self.lanes[i].admitted
3516 && !self.lanes[i].parked
3517 && !keeps_prefix
3518 && (!self.lanes[i].job.extras.ephemeral || self.lanes[i].preempted)
3519 {
3520 self.lanes[i].parked = true;
3521 if let Err(msg) = self.park_lane(i) {
3522 eprintln!(
3523 "[superfluid] park skipped (session {}): {msg}",
3524 self.lanes[i].session
3525 );
3526 }
3527 } else if matches!(self.lanes[i].phase, LanePhase::Finishing)
3528 && self.lanes[i].admitted
3529 && !self.lanes[i].parked
3530 {
3531 self.lanes[i].parked = true;
3532 if let Err(msg) = self.persist_prefix(i) {
3533 eprintln!("[superfluid] prefix not kept (session {}): {msg}", self.lanes[i].session);
3534 }
3535 }
3536 }
3537 }
3538
3539 let rate_cap = |ceiling: u32| match self.prefill_secs_per_token_bulk {
3540 Some(spt) if spt > 0.0 => Some((MAX_PREFILL_TICK_SECS / spt).clamp(512.0, ceiling as f64) as u32),
3541 _ => None,
3542 };
3543 let solo = !self.prefill_budget_fixed
3544 && self.lanes.len() == 1
3545 && self.queued() == 0
3546 && self
3547 .lanes
3548 .iter()
3549 .all(|l| matches!(l.phase, LanePhase::Prefilling));
3550 let solo_budget = if solo { rate_cap(SOLO_PREFILL_BUDGET) } else { None };
3551 let open = self.joiner_waits_on_this_tick();
3552 let tick_prefill_budget = if self
3553 .lanes
3554 .iter()
3555 .any(|l| !l.job.media.is_empty() && matches!(l.phase, LanePhase::Prefilling))
3556 {
3557 PREFILL_BUDGET
3558 } else if open {
3559 let spt = match (self.prefill_secs_per_token_bulk, self.prefill_secs_per_token_last) {
3562 (Some(a), Some(b)) => Some(a.max(b)),
3563 (a, b) => a.or(b),
3564 };
3565 let cap = match spt.filter(|&t| t > 0.0 && self.prefill_rate_samples >= 2) {
3566 Some(spt) => {
3567 let yields = matches!(
3568 self.host.capabilities().read().expect("capability record").get("serving", "tick_yield"),
3569 crate::capabilities::Cap::Yes
3570 );
3571 open_prefill_tick_cap(spt, open_prefill_tick_target(yields, self.latency_sensitive()))
3572 }
3573 None => 512,
3574 };
3575 self.prefill_budget.min(cap)
3576 } else if let Some(cap) = solo_budget.filter(|&cap| cap > self.prefill_budget) {
3577 cap
3578 } else {
3579 match self.prefill_secs_per_token_bulk {
3580 Some(spt) if spt > 0.0 => {
3581 let lone = self.lanes.len() == 1
3582 && self.queues.iter().all(|q| q.jobs.is_empty())
3583 && self.lanes.iter().all(|l| {
3584 Self::prefill_end_of(&self.speculation, l).saturating_sub(l.prefilled)
3585 > prefill_tick_cap(spt, true) as u64
3586 });
3587 self.prefill_budget.min(prefill_tick_cap(spt, lone))
3588 }
3589 _ => self.prefill_budget,
3590 }
3591 };
3592 let mut plan = linkw::TickPlanMsg {
3593 plan_seq: {
3594 self.plan_seq += 1;
3595 self.plan_seq
3596 },
3597 flags: 0,
3598 prefill_token_budget: tick_prefill_budget,
3599 max_decode_lanes: self.max_lanes as u32,
3600 admits: vec![],
3601 commits: vec![],
3602 prefills: vec![],
3603 decodes: vec![],
3604 retires: vec![],
3605 shed_policy: linkw::ShedPolicyMsg {
3606 victim_lanes: vec![],
3607 evictable_cache_classes: u64::MAX,
3608 protected_quota_bytes: 0,
3609 max_evict_bytes: u64::MAX,
3610 },
3611 };
3612
3613 let prefill_demand_lanes = self
3614 .lanes
3615 .iter()
3616 .filter(|l| {
3617 !matches!(l.phase, LanePhase::Finishing)
3618 && l.prefilled < Self::prefill_end_of(&self.speculation, l)
3619 })
3620 .count()
3621 + self.queues.iter().map(|q| q.jobs.len()).sum::<usize>();
3622 self.plan_prefix_checkpoints();
3623 let mut budget = PrefillBudget::new(tick_prefill_budget as u64);
3624 let mut retiring: Vec<u64> = Vec::new();
3625 let starvation = self.starvation_ticks;
3626 let mut order: Vec<usize> = (0..self.lanes.len()).collect();
3627 order.sort_by_key(|&i| {
3628 let l = &self.lanes[i];
3629 let remaining = Self::prefill_end_of(&self.speculation, l).saturating_sub(l.prefilled);
3630 prefill_rank(l.job.class, remaining, l.tag, l.prefill_starved_ticks, starvation)
3631 });
3632 let mut decode_candidates: Vec<usize> = Vec::new();
3633 let speculation = Arc::clone(&self.speculation);
3634 for &li in &order {
3635 let lane = &mut self.lanes[li];
3636 if matches!(lane.phase, LanePhase::Finishing) {
3637 plan.retires.push(linkw::LaneRetireMsg {
3638 lane_tag: lane.tag,
3639 publish_to_cache: !lane.parked_lossy
3640 && lane.job.media.is_empty()
3641 && (!lane.job.completion || lane.job.extras.keep_prefix),
3642 });
3643 retiring.push(lane.tag);
3644 continue;
3645 }
3646 if !lane.admitted {
3647 let strategy_slot = Self::strategy_slot_for(&speculation, &lane.job);
3648 plan.admits.push(linkw::LaneAdmitMsg {
3649 lane_tag: lane.tag,
3650 prompt: lane.span,
3651 seed_handle: lane.seed_handle,
3652 sampling: if lane.job.params.temperature == 0.0 { 0 } else { 1 },
3653 params: linkw::SamplingParamsMsg {
3654 temperature: lane.job.params.temperature,
3655 top_p: lane.job.params.top_p,
3656 min_p: lane.job.params.min_p,
3657 top_k: lane.job.params.top_k,
3658 freq_penalty: lane.job.extras.frequency_penalty,
3659 presence_penalty: lane.job.extras.presence_penalty,
3660 repeat_penalty: lane.job.extras.repeat_penalty,
3661 flags: if lane.job.extras.ignore_eos {
3662 superfluid_abi::sampling_flag::IGNORE_EOS
3663 } else {
3664 0
3665 },
3666 },
3667 rng_counter_base: lane
3668 .job
3669 .params
3670 .seed
3671 .wrapping_add(lane.job.stream.len() as u64),
3672 grammar_handle: lane.job.grammar_handle,
3673 grammar_replay: if lane.job.grammar_handle != 0 { lane.produced_before } else { 0 },
3674 decode_replay: Self::decode_replay_for(strategy_slot, lane.produced_before),
3675 logit_bias_handle: lane.job.extras.logit_bias_handle,
3676 want_logprobs: lane.job.extras.want_logprobs,
3677 top_logprobs: lane.job.extras.top_logprobs,
3678 strategy_slot,
3679 determinism_class: 0,
3680 minimum_exactness: 0,
3681 allow_approximate: false,
3682 required_cert_id: 0,
3683 host_sampler_identity: [0; 32],
3684 });
3685 }
3686 let total = Self::prefill_end_of(&speculation, lane);
3687 if lane.prefilled < total {
3688 let mut chunk = budget.offer_to(total - lane.prefilled, lane.job.class) as u32;
3689 if lane.checkpoint_at > lane.prefilled {
3690 chunk = chunk.min((lane.checkpoint_at - lane.prefilled) as u32);
3691 }
3692 for &(off, _, n) in &lane.job.media {
3693 let (rs, re) = (off, off + n as u64);
3694 let (cs, ce) = (lane.prefilled, lane.prefilled + chunk as u64);
3695 if rs < ce && re > ce {
3696 chunk = if rs > cs { (rs - cs) as u32 } else { 0 };
3697 }
3698 }
3699 if chunk > 0 {
3700 lane.prefill_starved_ticks = 0;
3701 }
3702 if chunk > 0 {
3703 plan.prefills.push(linkw::LanePrefillMsg {
3704 lane_tag: lane.tag,
3705 token_offset: lane.prefilled as u32,
3706 token_count: chunk,
3707 });
3708 lane.prefilled += chunk as u64;
3709 lane.checkpoint_due = lane.checkpoint_at > 0 && lane.prefilled == lane.checkpoint_at;
3710 budget.take(chunk as u64, lane.prefilled >= total);
3711 if lane.prefilled >= total && lane.job.class <= crate::qos::INLINE_COMPLETION {
3712 budget.hold_below(lane.job.class);
3713 }
3714 }
3715 }
3716 if lane.prefilled >= total {
3717 lane.phase = LanePhase::Decoding;
3718 if lane.job.max_tokens > lane.produced {
3719 decode_candidates.push(li);
3720 }
3721 }
3722 }
3723 let youngest_served = plan.prefills.iter().map(|p| p.lane_tag).max();
3725 for lane in self.lanes.iter_mut() {
3726 let waiting = lane.prefilled < Self::prefill_end_of(&self.speculation, lane)
3727 && !matches!(lane.phase, LanePhase::Finishing)
3728 && !plan.prefills.iter().any(|p| p.lane_tag == lane.tag);
3729 lane.prefill_starved_ticks = starved_after(lane.prefill_starved_ticks, waiting, youngest_served, lane.tag);
3730 }
3731 let divisor = self.background_tick_divisor.load(Ordering::Relaxed).max(1);
3732 let tick_no = self.tick_count;
3733 decode_candidates.sort_by_key(|&i| {
3734 let l = &self.lanes[i];
3735 let starving = l.starved_ticks >= starvation;
3736 (!starving, l.job.class, l.tag)
3737 });
3738 let planned_prefill_tokens: f64 = plan.prefills.iter().map(|p| p.token_count as f64).sum();
3739 let finishing = self.lanes.iter().any(|l| matches!(l.phase, LanePhase::Finishing));
3740 let shape = self.tick_shape;
3741 let target_secs = tick_target(&shape, self.tick_target_secs, self.lanes.len(), self.max_lanes, finishing);
3742 self.stats.tick_target_live_ms.store((target_secs * 1000.0).round() as u64, Ordering::Relaxed);
3743 let grant_cap: u16 = match (plan.prefills.len(), self.decode_round_secs) {
3744 (p, Some(c)) if c > 0.0 && !decode_candidates.is_empty() => {
3745 if p > 0 {
3746 let prefill_secs = self.prefill_secs_per_token.map_or(0.0, |s| s * planned_prefill_tokens);
3747 let interactive =
3748 decode_candidates.iter().any(|&i| self.lanes[i].job.class <= crate::qos::INLINE_COMPLETION);
3749 let share =
3750 if interactive { DECODE_SHARE_BESIDE_PREFILL_INTERACTIVE } else { DECODE_SHARE_BESIDE_PREFILL };
3751 ((prefill_secs * share / c) as u32)
3752 .clamp(shape.beside_prefill as u32, TICK_DECODE_MAX as u32) as u16
3753 } else {
3754 let rounds = ((target_secs / c) as u32).clamp(shape.decode_floor as u32, TICK_DECODE_MAX as u32);
3755 let first_out = decode_candidates
3756 .iter()
3757 .map(|&i| self.lanes[i].job.max_tokens.saturating_sub(self.lanes[i].produced))
3758 .min()
3759 .unwrap_or(rounds);
3760 let rounds = rounds.min(first_out.max(shape.beside_prefill as u32));
3761 if self.just_retired {
3762 rounds.min(shape.beside_prefill as u32) as u16
3763 } else {
3764 rounds as u16
3765 }
3766 }
3767 }
3768 _ => shape.decode_floor,
3769 };
3770 self.stats.decode_grant_live.store(grant_cap as u64, Ordering::Relaxed);
3771 let mut decode_left = if self.tick_decode_budget_fixed {
3772 self.tick_decode_budget
3773 } else {
3774 (self.max_lanes as u32).saturating_mul(grant_cap as u32)
3775 };
3776 let retire_only = plan.prefills.is_empty() && !plan.retires.is_empty();
3777 if retire_only {
3778 decode_candidates.clear();
3779 }
3780 self.just_retired = retire_only;
3781 for &li in &decode_candidates {
3782 let lane = &mut self.lanes[li];
3783 let background = lane.job.class >= crate::qos::FOREGROUND_AGENT;
3784 let throttled = background && divisor > 1 && !tick_no.is_multiple_of(divisor as u64);
3785 let want = if throttled {
3786 0
3787 } else {
3788 (lane.job.max_tokens - lane.produced)
3789 .min(grant_cap as u32)
3790 .min(decode_left) as u16
3791 };
3792 if want > 0 {
3793 if lane.starved_ticks >= starvation {
3794 self.stats.starvation_grants.fetch_add(1, Ordering::Relaxed);
3795 }
3796 lane.starved_ticks = 0;
3797 decode_left -= want as u32;
3798 let overshoot = (lane.job.max_tokens - lane.produced - want as u32).min(16) as u16;
3799 plan.decodes.push(linkw::LaneDecodeMsg {
3800 lane_tag: lane.tag,
3801 max_new_tokens: want,
3802 overshoot,
3803 });
3804 } else {
3805 lane.starved_ticks = lane.starved_ticks.saturating_add(1);
3806 }
3807 }
3808
3809 let wants_partial = self.lanes.iter().any(|l| {
3810 !matches!(l.phase, LanePhase::Finishing)
3811 && (l.job.extras.wants_deltas
3812 || l.job.completion
3813 || self.bus.has_provisional_subscriber(l.session))
3814 });
3815 if wants_partial {
3816 plan.flags |= superfluid_abi::tick_flags::PARTIAL_EMITS;
3817 }
3818 let mut plan_prefill_tokens: u32 = plan.prefills.iter().map(|p| p.token_count).sum();
3819 let prefill_by_tag: HashMap<u64, u32> =
3820 plan.prefills.iter().map(|p| (p.lane_tag, p.token_count)).collect();
3821 let plan_decode_rounds: u32 =
3822 plan.decodes.iter().map(|d| d.max_new_tokens as u32).max().unwrap_or(0);
3823 let plan_prefill_lanes_planned = plan.prefills.len();
3824 let plan_prefill_lanes = plan.prefills.len().max(prefill_demand_lanes);
3825 let plan_decode_lanes = plan.decodes.len();
3826 if !tick_span.is_disabled() {
3827 let planned: HashSet<u64> = plan
3828 .admits
3829 .iter()
3830 .map(|a| a.lane_tag)
3831 .chain(plan.prefills.iter().map(|p| p.lane_tag))
3832 .chain(plan.decodes.iter().map(|d| d.lane_tag))
3833 .chain(plan.retires.iter().map(|r| r.lane_tag))
3834 .collect();
3835 tick_span.record("lanes", planned.len());
3836 for lane in self.lanes.iter().filter(|l| planned.contains(&l.tag)) {
3837 tick_span.follows_from(&lane.job.trace);
3838 }
3839 }
3840 let tick_started = std::time::Instant::now();
3841 self.stats.set_tick_in_flight(TickInFlight {
3842 since_ms: uptime_ms().max(1),
3843 prefill_tokens: plan_prefill_tokens as u64,
3844 prefill_lanes: plan_prefill_lanes_planned as u64,
3845 decode_lanes: plan_decode_lanes as u64,
3846 admits: plan.admits.len() as u64,
3847 retires: plan.retires.len() as u64,
3848 target_ms: (target_secs * 1000.0).round() as u64,
3849 });
3850 let mut deferred: Vec<u64> = Vec::new();
3851 let ev = {
3852 let Scheduler {
3853 host,
3854 lanes,
3855 bus,
3856 codec,
3857 pins,
3858 stats,
3859 ..
3860 } = self;
3861 let by_tag: std::collections::HashMap<u64, usize> =
3862 lanes.iter().enumerate().map(|(i, l)| (l.tag, i)).collect();
3863 let fewer = (!plan.admits.is_empty()
3867 && plan.admits.iter().all(|a| {
3868 by_tag.get(&a.lane_tag).is_some_and(|&i| lanes[i].job.media.is_empty())
3869 }))
3870 .then(|| plan.clone());
3871 let mut on_emit = |em: linkw::TickEmitMsg| {
3872 let Some(&i) = by_tag.get(&em.lane_tag) else { return };
3873 let lane = &mut lanes[i];
3874 lane.frame_fed += em.tokens.len();
3875 stats
3876 .decode_tokens
3877 .fetch_add(em.tokens.len() as u64, Ordering::Relaxed);
3878 if !lane.ttft_observed && lane.ttft_frame_ms.is_none() && !em.tokens.is_empty() {
3879 lane.ttft_frame_ms = Some(lane.admitted_at.elapsed().as_millis() as u64);
3880 }
3881 if lane.job.completion {
3882 if !lane.emit_dead
3883 && lane.job.events.send(LaneEvent::CompletionTokens(em.tokens)).is_err()
3884 {
3885 lane.emit_dead = true;
3886 }
3887 } else {
3888 lane.out.provisional(codec, bus, lane.session, &lane.job.events, &em.tokens);
3889 }
3890 };
3891 let retry = pins.any_held().then(|| plan.clone());
3892 let needs_replan = |r: &Result<linkw::TickEventsMsg, superfluid_agent::AgentError>| {
3893 matches!(r, Err(superfluid_agent::AgentError::Rejected(s)) if *s == superfluid_abi::Status::NeedsReplan as i32)
3894 };
3895 let mut res = match (host.client().tick_streaming(plan, &mut on_emit), retry) {
3896 (Err(superfluid_agent::AgentError::Rejected(status)), Some(plan))
3897 if status == superfluid_abi::Status::NeedsReplan as i32
3898 && pins.yield_all(host, stats) =>
3899 {
3900 host.client().tick_streaming(plan, &mut on_emit)
3901 }
3902 (res, _) => res,
3903 };
3904 if let Some(full) = fewer.filter(|_| needs_replan(&res)) {
3905 let mut keep = full.admits.len();
3906 while keep > 0 && needs_replan(&res) {
3907 keep /= 2;
3908 res = host.client().tick_streaming(with_first_admits(&full, keep), &mut on_emit);
3909 }
3910 if res.is_ok() {
3911 deferred = full.admits[keep..].iter().map(|a| a.lane_tag).collect();
3912 }
3913 }
3914 stats.clear_tick_in_flight();
3915 res?
3916 };
3917 if !deferred.is_empty() {
3918 plan_prefill_tokens = plan_prefill_tokens
3919 .saturating_sub(deferred.iter().filter_map(|t| prefill_by_tag.get(t)).sum());
3920 self.stats.admits_deferred.fetch_add(deferred.len() as u64, Ordering::Relaxed);
3921 self.defer_admits(&deferred);
3922 }
3923 let tick_secs = tick_started.elapsed().as_secs_f64().max(1e-3);
3924 let tick_done = std::time::Instant::now();
3925 for e in ev.shed.entries.iter().filter(|e| e.kind == superfluid_abi::shed_kind::PREFILL_CHUNK) {
3928 if let Some(lane) = self.lanes.iter_mut().find(|l| l.tag == e.lane_tag) {
3929 lane.prefilled = lane.prefilled.saturating_sub(e.tokens);
3930 if matches!(lane.phase, LanePhase::Decoding) {
3931 lane.phase = LanePhase::Prefilling;
3932 }
3933 if lane.checkpoint_due && lane.prefilled != lane.checkpoint_at {
3934 lane.checkpoint_due = false;
3935 }
3936 }
3937 }
3938 let prefill_tokens_ran: u32 = plan_prefill_tokens
3939 .saturating_sub(
3940 ev.admit_results
3941 .iter()
3942 .filter(|a| a.status != superfluid_abi::admit_status::ADMITTED)
3943 .filter_map(|a| prefill_by_tag.get(&a.lane_tag))
3944 .sum(),
3945 )
3946 .saturating_sub(
3947 ev.faults
3948 .iter()
3949 .filter(|f| f.code == superfluid_abi::fault::MEDIA)
3950 .filter_map(|f| prefill_by_tag.get(&f.lane_tag))
3951 .sum(),
3952 )
3953 .saturating_sub(
3954 ev.shed
3955 .entries
3956 .iter()
3957 .filter(|e| e.kind == superfluid_abi::shed_kind::PREFILL_CHUNK)
3958 .map(|e| e.tokens.min(u32::MAX as u64) as u32)
3959 .sum(),
3960 );
3961 let tick_rounds_exec: u32 = ev
3962 .emits
3963 .iter()
3964 .filter(|e| e.n_tokens > 0)
3965 .map(|e| e.n_tokens.saturating_sub(e.spec.accepted).max(1))
3966 .max()
3967 .unwrap_or(0);
3968 let first_from_prefill: HashSet<u64> = self
3972 .lanes
3973 .iter()
3974 .filter(|l| l.produced == 0 && prefill_by_tag.contains_key(&l.tag))
3975 .map(|l| l.tag)
3976 .collect();
3977 let tick_decode_rounds: Option<u32> = ev
3978 .emits
3979 .iter()
3980 .filter(|e| e.n_tokens > 0)
3981 .map(|e| {
3982 let rounds = e.n_tokens.saturating_sub(e.spec.accepted).max(1);
3983 rounds - u32::from(first_from_prefill.contains(&e.lane_tag))
3984 })
3985 .max();
3986 self.stats
3987 .prefill_tokens
3988 .fetch_add(prefill_tokens_ran as u64, Ordering::Relaxed);
3989 let phase_split = plan_prefill_tokens == 0 || ev.timings.prefill_ns > 0;
3990 let (prefill_secs_meas, decode_secs_meas) = if phase_split {
3991 (ev.timings.prefill_ns as f64 / 1e9, ev.timings.decode_ns as f64 / 1e9)
3992 } else {
3993 let decode_est = self.decode_round_secs.unwrap_or(0.0) * tick_decode_rounds.unwrap_or(0) as f64;
3994 ((tick_secs - decode_est).max(1e-3), decode_est)
3995 };
3996 let decode_rounds_ran = match tick_decode_rounds {
3999 Some(rounds) => rounds.min(plan_decode_rounds),
4000 None => plan_decode_rounds,
4001 };
4002 if decode_rounds_ran > 0 && phase_split && decode_secs_meas > 0.0 {
4003 let sample = decode_secs_meas / decode_rounds_ran as f64;
4004 self.decode_round_secs = Some(match self.decode_round_secs {
4005 Some(c) => 0.7 * c + 0.3 * sample,
4006 None => sample,
4007 });
4008 } else if plan_prefill_tokens == 0 && decode_rounds_ran > 0 {
4009 let sample = tick_secs / decode_rounds_ran as f64;
4010 self.decode_round_secs = Some(match self.decode_round_secs {
4011 Some(c) => 0.7 * c + 0.3 * sample,
4012 None => sample,
4013 });
4014 }
4015 if prefill_tokens_ran >= PREFILL_RATE_MIN_TOKENS && phase_split {
4016 let sample = prefill_secs_meas.max(1e-3) / prefill_tokens_ran as f64;
4017 self.prefill_secs_per_token = Some(match self.prefill_secs_per_token {
4018 Some(c) => 0.7 * c + 0.3 * sample,
4019 None => sample,
4020 });
4021 if prefill_tokens_ran >= 512 {
4022 self.prefill_secs_per_token_last = Some(sample);
4023 self.prefill_rate_samples = self.prefill_rate_samples.saturating_add(1);
4024 self.prefill_secs_per_token_bulk = Some(match self.prefill_secs_per_token_bulk {
4025 Some(c) => 0.7 * c + 0.3 * sample,
4026 None => sample,
4027 });
4028 }
4029 }
4030 if !self.prefill_budget_fixed && prefill_tokens_ran > 0 {
4031 let prefill_secs = prefill_secs_meas.max(1e-3) * plan_prefill_tokens as f64 / prefill_tokens_ran as f64;
4034 let share = plan_prefill_lanes as f64
4035 / (plan_prefill_lanes as f64 + plan_decode_lanes as f64);
4036 let factor = (self.tick_target_secs * share / prefill_secs).clamp(0.5, 2.0);
4037 self.prefill_budget =
4038 ((self.prefill_budget as f64 * factor) as u32).clamp(256, PREFILL_BUDGET);
4039 }
4040 if std::env::var_os("SUPERFLUID_DEBUG_TICK").is_some() {
4041 eprintln!(
4042 "[tick] secs={:.3} pf_secs={:.3} dec_secs={:.3} prefill_tok={} rounds={} P_eff={} P_plan={} D={} round_secs={:?} pf_secs_per_tok={:?} budget={} grant={}",
4043 tick_secs, prefill_secs_meas, decode_secs_meas, plan_prefill_tokens, plan_decode_rounds, plan_prefill_lanes, plan_prefill_lanes_planned,
4044 plan_decode_lanes, self.decode_round_secs, self.prefill_secs_per_token, self.prefill_budget, grant_cap
4045 );
4046 }
4047 self.stats
4048 .prefill_budget_live
4049 .store(self.prefill_budget as u64, Ordering::Relaxed);
4050 if !self.pins.is_empty() {
4051 let retired: Vec<(u64, Vec<u32>)> = self
4052 .lanes
4053 .iter()
4054 .filter(|l| retiring.contains(&l.tag) && !l.job.completion)
4055 .map(|l| (l.session, l.all_tokens.clone()))
4056 .collect();
4057 self.pins
4058 .after_tick(&mut self.host, &self.stats, self.tick_count, &retired);
4059 }
4060 self.publish_prefix_checkpoints();
4061 self.relieve_pressure(&ev.mem, false);
4062 self.stats
4063 .pool_blocks_used
4064 .store(ev.mem.pool_blocks_used, Ordering::Relaxed);
4065 self.stats
4066 .pool_blocks_total
4067 .store(ev.mem.pool_blocks_total, Ordering::Relaxed);
4068
4069 let refused: HashMap<u64, u32> = ev
4070 .admit_results
4071 .iter()
4072 .filter(|r| r.status == superfluid_abi::admit_status::REJECTED)
4073 .map(|r| (r.lane_tag, r.reject_code))
4074 .collect();
4075 if !refused.is_empty() {
4076 self.drop_refused(&refused);
4077 }
4078
4079 for f in &ev.faults {
4080 if let Some(lane) = self.lanes.iter_mut().find(|l| l.tag == f.lane_tag) {
4081 if lane.fault_code.is_some() {
4082 continue;
4083 }
4084 lane.fault_code = Some(f.code);
4085 lane.fin = finish::ERROR;
4086 let sink = OutputSink { codec: &self.codec, store: &self.store, bus: &self.bus, tool_lease_ms: self.tool_lease_ms };
4087 lane.out.close(&sink, lane.session, &lane.job.events, finish::ERROR, false)?;
4088 lane.phase = LanePhase::Finishing;
4089 tracing::warn!(session = lane.session, lane = lane.tag, code = f.code, "engine lane fault");
4090 }
4091 }
4092
4093 #[allow(clippy::type_complexity)]
4094 let mut logprobs_by_lane: HashMap<u64, Vec<(f32, Vec<(u32, f32)>)>> = HashMap::new();
4095 for lp in &ev.logprobs {
4096 let top: Vec<(u32, f32)> = lp
4097 .top_ids
4098 .iter()
4099 .zip(&lp.top_logprob_bits)
4100 .map(|(id, b)| (*id, f32::from_bits(*b)))
4101 .collect();
4102 logprobs_by_lane
4103 .entry(lp.lane_tag)
4104 .or_default()
4105 .push((f32::from_bits(lp.logprob_bits), top));
4106 }
4107 let mut emitted: HashMap<u64, (Vec<u32>, u32)> = HashMap::new();
4108 for emit in &ev.emits {
4109 if emit.spec.proposed > 0 {
4110 self.stats
4111 .spec_proposed
4112 .fetch_add(emit.spec.proposed as u64, Ordering::Relaxed);
4113 self.stats
4114 .spec_accepted
4115 .fetch_add(emit.spec.accepted as u64, Ordering::Relaxed);
4116 if let Some(lane) = self.lanes.iter_mut().find(|l| l.tag == emit.lane_tag) {
4117 lane.spec.0 += emit.spec.proposed as u64;
4118 lane.spec.1 += emit.spec.accepted as u64;
4119 }
4120 }
4121 if emit.n_tokens > 0 {
4122 let tokens = self.host.client().read_tokens(&emit.token_ref)?;
4123 emitted.insert(emit.lane_tag, (tokens, emit.finish));
4124 } else if emit.finish != finish::NONE {
4125 emitted.insert(emit.lane_tag, (Vec::new(), emit.finish));
4126 }
4127 }
4128 let tick_decode_tokens: u64 = emitted.values().map(|(t, _)| t.len() as u64).sum();
4129 let ttft_tail_secs = if tick_rounds_exec > 1 {
4130 decode_secs_meas * (tick_rounds_exec as f64 - 1.0) / tick_rounds_exec as f64
4131 } else {
4132 0.0
4133 };
4134 if prefill_tokens_ran > 0 || tick_decode_tokens > 0 {
4135 self.stats
4136 .tick_wall_ms_live
4137 .store((tick_secs * 1e3) as u64, Ordering::Relaxed);
4138 self.stats.prefill_rate_live.store(
4139 per_sec(prefill_tokens_ran as u64, prefill_secs_meas),
4140 Ordering::Relaxed,
4141 );
4142 let decode_secs_for_rate =
4143 if decode_secs_meas > 0.0 { decode_secs_meas } else { tick_secs };
4144 self.stats.decode_rate_live.store(
4145 per_sec(tick_decode_tokens, decode_secs_for_rate),
4146 Ordering::Relaxed,
4147 );
4148 self.stats.work_ticks.fetch_add(1, Ordering::Relaxed);
4149 }
4150 let terminators = self.codec.turn_terminators();
4151 let sink = OutputSink { codec: &self.codec, store: &self.store, bus: &self.bus, tool_lease_ms: self.tool_lease_ms };
4152 for lane in &mut self.lanes {
4153 lane.admitted = true;
4154 let Some((tokens, fin)) = emitted.remove(&lane.tag) else {
4155 continue;
4156 };
4157 if tokens.is_empty() {
4158 lane.fin = fin;
4159 if !lane.job.completion {
4160 lane.out.close(&sink, lane.session, &lane.job.events, fin, false)?;
4161 }
4162 lane.phase = LanePhase::Finishing;
4163 continue;
4164 }
4165 lane.all_tokens.extend_from_slice(&tokens);
4166 lane.produced += tokens.len() as u32;
4167 lane.fin = fin;
4168 lane.decode_ticks = Some(match lane.decode_ticks {
4169 None => (tick_done, lane.produced, tick_done, lane.produced),
4170 Some((at, by, _, _)) => (at, by, tick_done, lane.produced),
4171 });
4172 let frame_fed = std::mem::take(&mut lane.frame_fed).min(tokens.len());
4173 if !lane.job.completion {
4174 lane.out.provisional(sink.codec, sink.bus, lane.session, &lane.job.events, &tokens[frame_fed..]);
4175 }
4176 if lane.fin == finish::NONE
4177 && !lane.job.extras.ignore_eos
4178 && !terminators.is_empty()
4179 && tokens.iter().any(|t| terminators.contains(t))
4180 {
4181 lane.fin = finish::EOS;
4182 }
4183 if lane.job.extras.want_logprobs {
4184 if let Some(recs) = logprobs_by_lane.remove(&lane.tag) {
4185 if !recs.is_empty() {
4186 let items: Vec<TokenLogprob> = tokens
4187 .iter()
4188 .zip(recs)
4189 .map(|(t, (lp, top))| TokenLogprob { token: *t, logprob: lp, top })
4190 .collect();
4191 let _ = lane.job.events.send(LaneEvent::TokenLogprobs(items));
4192 }
4193 }
4194 }
4195 self.stats
4196 .decode_tokens
4197 .fetch_add((tokens.len() - frame_fed) as u64, Ordering::Relaxed);
4198 if !lane.ttft_observed && !tokens.is_empty() {
4199 lane.ttft_observed = true;
4200 let ms = lane.ttft_frame_ms.take().unwrap_or_else(|| {
4201 let elapsed = lane.admitted_at.elapsed().as_secs_f64();
4202 ((elapsed - ttft_tail_secs).max(0.0) * 1e3) as u64
4203 });
4204 self.stats.observe_ttft(ms);
4205 }
4206 let terminal_now = lane.fin != finish::NONE || lane.produced >= lane.job.max_tokens;
4207 if lane.job.completion {
4208 let leftover = tokens[frame_fed..].to_vec();
4209 if lane.emit_dead
4210 || (!leftover.is_empty()
4211 && lane.job.events.send(LaneEvent::CompletionTokens(leftover)).is_err())
4212 {
4213 lane.phase = LanePhase::Finishing;
4214 continue;
4215 }
4216 if terminal_now {
4217 lane.phase = LanePhase::Finishing;
4218 }
4219 continue;
4220 }
4221 let delivered = lane.out.commit(
4222 &sink,
4223 lane.session,
4224 &lane.job.events,
4225 &tokens,
4226 fin,
4227 terminal_now,
4228 lane.job.extras.tool_schemas.as_deref(),
4229 )?;
4230 if !delivered {
4231 lane.fin = finish::CANCELLED;
4232 lane.out.close(&sink, lane.session, &lane.job.events, finish::CANCELLED, true)?;
4233 lane.phase = LanePhase::Finishing;
4234 continue;
4235 }
4236 if terminal_now {
4237 lane.phase = LanePhase::Finishing;
4238 }
4239 }
4240
4241 if !retiring.is_empty() {
4242 let mut i = 0;
4243 let mut requeue: Vec<GenerateJob> = Vec::new();
4244 while i < self.lanes.len() {
4245 if !retiring.contains(&self.lanes[i].tag) {
4246 i += 1;
4247 continue;
4248 }
4249 let lane = self.lanes.remove(i);
4250 if lane.preempted && lane.fin == finish::NONE {
4251 let decode_before = lane.decode_measured();
4252 let mut job = lane.job;
4253 job.stream = lane.all_tokens;
4254 job.max_tokens -= lane.produced;
4255 job.carry = Some(LaneCarry {
4256 out: lane.out,
4257 produced_before: lane.produced_before + lane.produced,
4258 spec_before: lane.spec,
4259 decode_before,
4260 warm_prefix_first: Some(lane.warm_prefix_first.unwrap_or(lane.warm_prefix)),
4261 });
4262 requeue.push(job);
4263 } else {
4264 self.drop_ephemeral_artifact(lane.session);
4265 self.active.lock().expect("active set").remove(&lane.session);
4266 if let Some(code) = lane.fault_code {
4267 let _ = lane
4268 .job
4269 .events
4270 .send(LaneEvent::Failed(format!("engine lane fault (code {code})")));
4271 } else {
4272 let _ = lane.job.events.send(LaneEvent::Done {
4273 tokens_generated: lane.produced_before + lane.produced,
4274 finish: lane.fin,
4275 warm_prefix: lane.warm_prefix_first.unwrap_or(lane.warm_prefix),
4276 spec: lane.spec,
4277 decode: lane.decode_measured(),
4278 });
4279 }
4280 }
4281 }
4282 for job in requeue {
4283 let c = (job.class as usize).min(3);
4284 self.queues[c].jobs.push_front(job);
4285 }
4286 }
4287 self.publish_occupancy();
4288 Ok(())
4289 }
4290}
4291
4292fn draft_weights_identity(path: &str) -> Result<(u64, u64), DaemonError> {
4293 use std::io::Read;
4294 const SEED: u64 = 0xCBF2_9CE4_8422_2325;
4295 let mut f = std::fs::File::open(path).map_err(|_| {
4296 DaemonError::Config(
4297 "draft-model weights are not readable by the daemon; an exact draft identity \
4298 requires the daemon to read the weights file",
4299 )
4300 })?;
4301 let size = f.metadata().map(|m| m.len()).unwrap_or(0);
4302 let mut h = SEED;
4303 let mut buf = vec![0u8; 1 << 20];
4304 loop {
4305 match f.read(&mut buf) {
4306 Ok(0) => break,
4307 Ok(n) => h = superfluid_fingerprint::fnv1a64(&buf[..n], h),
4308 Err(_) => {
4309 return Err(DaemonError::Config(
4310 "failed reading draft-model weights for identity",
4311 ))
4312 }
4313 }
4314 }
4315 Ok((h, size))
4316}
4317
4318pub fn run_completion(
4319 tx: &Sender<SchedCmd>,
4320 cancels: &std::sync::Arc<crate::CancelSet>,
4321 stream: Vec<u32>,
4322 params: GenParams,
4323 max_tokens: u32,
4324 deadline_ms: u64,
4325 grammar_handle: u32,
4326) -> Result<(Vec<u32>, u32, bool), DaemonError> {
4327 run_completion_streaming(
4328 tx,
4329 cancels,
4330 stream,
4331 params,
4332 max_tokens,
4333 deadline_ms,
4334 GenExtras { grammar_handle, ..GenExtras::default() },
4335 |_| std::ops::ControlFlow::Continue(()),
4336 )
4337}
4338
4339#[allow(clippy::too_many_arguments)]
4340pub fn run_completion_streaming(
4341 tx: &Sender<SchedCmd>,
4342 cancels: &std::sync::Arc<crate::CancelSet>,
4343 stream: Vec<u32>,
4344 params: GenParams,
4345 max_tokens: u32,
4346 deadline_ms: u64,
4347 extras: GenExtras,
4348 on_tokens: impl FnMut(&[u32]) -> std::ops::ControlFlow<()>,
4349) -> Result<(Vec<u32>, u32, bool), DaemonError> {
4350 run_completion_streaming_lp(tx, cancels, stream, params, max_tokens, deadline_ms, crate::qos::INLINE_COMPLETION, extras, Vec::new(), on_tokens, |_| {})
4351}
4352
4353#[allow(clippy::too_many_arguments)]
4357pub fn run_completion_streaming_lp(
4358 tx: &Sender<SchedCmd>,
4359 cancels: &std::sync::Arc<crate::CancelSet>,
4360 stream: Vec<u32>,
4361 params: GenParams,
4362 max_tokens: u32,
4363 deadline_ms: u64,
4364 class: u8,
4365 extras: GenExtras,
4366 media: Vec<(u64, std::path::PathBuf, u32)>,
4367 mut on_tokens: impl FnMut(&[u32]) -> std::ops::ControlFlow<()>,
4368 mut on_logprobs: impl FnMut(&[TokenLogprob]),
4369) -> Result<(Vec<u32>, u32, bool), DaemonError> {
4370 let (etx, erx) = std::sync::mpsc::channel();
4371 tx.send(SchedCmd::Submit(Box::new(GenerateJob {
4372 session: 0,
4373 stream,
4374 input_end: 0,
4375 params,
4376 max_tokens,
4377 events: etx,
4378 class,
4379 batch_invariant: false,
4380 carry: None,
4381 queued_at: 0,
4382 media,
4383 completion: true,
4384 deadline: (deadline_ms > 0).then(|| std::time::Instant::now() + std::time::Duration::from_millis(deadline_ms)),
4385 grammar_handle: extras.grammar_handle,
4386 extras,
4387 trace: tracing::Span::current(),
4388 })))
4389 .map_err(|_| DaemonError::Protocol("scheduler is gone"))?;
4390 if class <= crate::qos::INLINE_COMPLETION {
4391 cancels.latency_request();
4393 }
4394 let mut tokens = Vec::new();
4395 loop {
4396 match erx.recv() {
4397 Ok(LaneEvent::Provisional { .. }) => {}
4398 Ok(LaneEvent::CompletionTokens(t)) => {
4399 if on_tokens(&t).is_break() {
4400 tokens.extend(t);
4401 return Ok((tokens, finish::CANCELLED, false));
4402 }
4403 tokens.extend(t);
4404 }
4405 Ok(LaneEvent::Done { finish, .. }) => return Ok((tokens, finish, false)),
4406 Ok(LaneEvent::Expired) => return Ok((Vec::new(), finish::NONE, true)),
4407 Ok(LaneEvent::Committed(_)) => {}
4408 Ok(LaneEvent::TokenLogprobs(items)) => on_logprobs(&items),
4409 Ok(LaneEvent::Failed(msg)) => return Err(DaemonError::Generation(msg)),
4410 Ok(LaneEvent::Error(e)) => return Err(e),
4411 Err(_) => return Err(DaemonError::Protocol("scheduler dropped the lane")),
4412 }
4413 }
4414}
4415
4416pub type RunJobResult = (Vec<CommittedEvent>, u32, u32, u64, Vec<TokenLogprob>, (u64, u64), (u64, f64));
4417
4418#[allow(clippy::too_many_arguments)]
4419pub fn run_job(
4420 tx: &Sender<SchedCmd>,
4421 job_session: u64,
4422 stream: Vec<u32>,
4423 input_end: u64,
4424 params: GenParams,
4425 max_tokens: u32,
4426 class: u8,
4427 batch_invariant: bool,
4428 media: Vec<(u64, std::path::PathBuf, u32)>,
4429 extras: GenExtras,
4430 cancels: &Arc<crate::CancelSet>,
4431 mut on_event: impl FnMut(&CommittedEvent) -> Result<(), DaemonError>,
4432 mut on_logprobs: impl FnMut(&[TokenLogprob]),
4433 mut on_delta: impl FnMut(u32, &str, u32, bool),
4434) -> Result<RunJobResult, DaemonError> {
4435 let (etx, erx) = std::sync::mpsc::channel();
4436 let grammar_handle = extras.grammar_handle;
4437 tx.send(SchedCmd::Submit(Box::new(GenerateJob {
4438 session: job_session,
4439 stream,
4440 input_end,
4441 params,
4442 max_tokens,
4443 events: etx,
4444 class,
4445 batch_invariant,
4446 carry: None,
4447 queued_at: 0,
4448 media,
4449 completion: false,
4450 deadline: None,
4451 grammar_handle,
4452 extras,
4453 trace: tracing::Span::current(),
4454 })))
4455 .map_err(|_| DaemonError::Protocol("scheduler is gone"))?;
4456 if class <= crate::qos::INLINE_COMPLETION {
4457 cancels.latency_request();
4461 }
4462 let mut events = Vec::new();
4463 let mut logprobs: Vec<TokenLogprob> = Vec::new();
4464 let mut delivery_err: Option<DaemonError> = None;
4465 loop {
4466 match erx.recv() {
4467 Ok(LaneEvent::Provisional { channel, text, produced, closes }) => {
4468 if delivery_err.is_none() {
4469 on_delta(channel, &text, produced, closes);
4470 }
4471 }
4472 Ok(LaneEvent::TokenLogprobs(items)) => {
4473 if delivery_err.is_none() {
4474 on_logprobs(&items);
4475 }
4476 logprobs.extend(items);
4477 }
4478 Ok(LaneEvent::Committed(e)) => {
4479 if delivery_err.is_none() {
4480 if let Err(err) = on_event(&e) {
4481 cancels.cancel(job_session);
4482 delivery_err = Some(err);
4483 } else {
4484 events.push(e);
4485 }
4486 }
4487 }
4488 Ok(LaneEvent::Done {
4489 tokens_generated,
4490 finish,
4491 warm_prefix,
4492 spec,
4493 decode,
4494 }) => {
4495 return match delivery_err {
4496 Some(err) => Err(err),
4497 None => Ok((events, tokens_generated, finish, warm_prefix, logprobs, spec, decode)),
4498 };
4499 }
4500 Ok(LaneEvent::Failed(msg)) => {
4501 return Err(delivery_err.unwrap_or(DaemonError::Generation(msg)));
4502 }
4503 Ok(LaneEvent::Error(e)) => {
4504 return Err(delivery_err.unwrap_or(e));
4505 }
4506 Ok(LaneEvent::CompletionTokens(_)) | Ok(LaneEvent::Expired) => {
4507 return Err(DaemonError::Protocol("completion event on a generate lane"));
4508 }
4509 Err(_) => {
4510 return Err(delivery_err.unwrap_or(DaemonError::Protocol("scheduler dropped the lane")));
4511 }
4512 }
4513 }
4514}