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