metrale_scheduler/
scripted.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: [`ScriptedDeviceIo`], a test router wrapped around an inner router.
4//! It answers the decode lane's readback from a script (per sequence, in
5//! order), fails a launch on cue, reports a configurable depth, and records
6//! every effect and LoRA command, while the inner router still performs the
7//! device bookkeeping the rows need.
8//!
9//! A row is matched to its script line by the key function the test passes
10//! (the server's trace-harness tests key by session hash).
11//!
12//! Depth: a launch whose readback is a device argmax (plain, masked or a
13//! fed step) is forwarded and left in flight with the inner router; the
14//! script is substituted when the ticket is awaited, so an inner router
15//! that pipelines keeps pipelining under the script. A host-logits launch
16//! is settled at once, because its answer is written into the plan's own
17//! buffer, which only exists for the duration of the call.
18//!
19//! Owner: scheduler.
20//! Invariants: none beyond the types.
21
22use std::cell::{Cell, RefCell};
23use std::collections::{HashMap, VecDeque};
24
25use crate::device::DeviceIo;
26use crate::effect::{Effect, EffectOutcome};
27use crate::plan::{DecodeRows, Label, Readback, StepOutcome, StepPlan, StepResult};
28use crate::types::{DeviceError, DeviceFacts, Ticket};
29
30/// 2026-09-25: A scripted launch failure.
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum Fault {
33    /// 2026-09-25: A forward error that reads as KV exhaustion
34    /// ([`crate::DeviceError::is_kv_exhausted`]).
35    KvExhausted,
36    /// 2026-09-25: Any other forward error.
37    Fatal,
38}
39
40/// 2026-09-25: A BF16 one-hot logits row: 30.0 (`0x41F0`) at `hit`, 0.0 elsewhere.
41fn one_hot_bf16(dst: &mut [u8], hit: u32) {
42    dst.fill(0);
43    let i = hit as usize * 2;
44    if i + 1 < dst.len() {
45        dst[i] = 0xF0;
46        dst[i + 1] = 0x41;
47    }
48}
49
50/// 2026-09-25: How a sequence row is matched to its script line.
51pub type KeyFn<S> = Box<dyn Fn(&S) -> u64 + Send>;
52
53/// 2026-09-25: A forwarded launch whose script substitution waits for `await_result`.
54struct InFlight {
55    ticket: Ticket,
56    keys: Vec<u64>,
57    /// 2026-09-25: The inputs the plan carried (a plain step) — the fallback when a row
58    /// has no script line and the inner answered host logits.
59    fallback: Vec<u32>,
60}
61
62pub struct ScriptedDeviceIo<D: DeviceIo> {
63    inner: D,
64    key: KeyFn<D::Seq>,
65    vocab: usize,
66    max_depth: usize,
67    script: RefCell<HashMap<u64, VecDeque<u32>>>,
68    /// 2026-09-25: Faults to raise on the next launches, consumed in order.
69    faults: RefCell<VecDeque<Fault>>,
70    /// 2026-09-25: The tokens the last settled step answered, per row, for `ReadLogits`.
71    last_rows: RefCell<Vec<u32>>,
72    launches: Cell<usize>,
73    effects: RefCell<Vec<String>>,
74    /// 2026-09-25: Host-logits steps, already settled and substituted at launch.
75    settled: RefCell<VecDeque<StepResult<D::Ptr>>>,
76    /// 2026-09-25: Device-argmax steps still with the inner router, in launch order.
77    in_flight: RefCell<VecDeque<InFlight>>,
78}
79
80impl<D: DeviceIo> ScriptedDeviceIo<D> {
81    /// 2026-09-25: `key` matches a row to its script; `vocab` sizes the host-logits rows;
82    /// `max_depth` is what the router reports (never more than the inner
83    /// router's own).
84    pub fn new(
85        inner: D,
86        key: impl Fn(&D::Seq) -> u64 + Send + 'static,
87        vocab: usize,
88        max_depth: usize,
89    ) -> Self {
90        Self {
91            inner,
92            key: Box::new(key),
93            vocab,
94            max_depth,
95            script: RefCell::new(HashMap::new()),
96            faults: RefCell::new(VecDeque::new()),
97            last_rows: RefCell::new(Vec::new()),
98            launches: Cell::new(0),
99            effects: RefCell::new(Vec::new()),
100            settled: RefCell::new(VecDeque::new()),
101            in_flight: RefCell::new(VecDeque::new()),
102        }
103    }
104
105    /// 2026-09-25: The tokens the decode lane will read back for the row keyed `key`,
106    /// in order. A row with no line left answers what the inner router
107    /// answered.
108    pub fn script(&self, key: u64, tokens: impl IntoIterator<Item = u32>) {
109        self.script
110            .borrow_mut()
111            .insert(key, tokens.into_iter().collect());
112    }
113
114    /// 2026-09-25: Raise `fault` on the next launch (queued, in order).
115    pub fn fail_next(&self, fault: Fault) {
116        self.faults.borrow_mut().push_back(fault);
117    }
118
119    /// 2026-09-25: Every effect and LoRA command applied so far, by label, in order.
120    pub fn effects(&self) -> Vec<String> {
121        self.effects.borrow().clone()
122    }
123
124    pub fn launches(&self) -> usize {
125        self.launches.get()
126    }
127
128    pub fn inner(&self) -> &D {
129        &self.inner
130    }
131
132    fn take_line(&self, key: u64) -> Option<u32> {
133        self.script
134            .borrow_mut()
135            .get_mut(&key)
136            .and_then(VecDeque::pop_front)
137    }
138
139    fn next_for(&self, key: u64, fallback: u32) -> u32 {
140        self.take_line(key).unwrap_or(fallback)
141    }
142
143    fn one_hot_rows(&self, tokens: &[u32], into: &mut Vec<u8>) {
144        into.resize(tokens.len() * self.vocab * 2, 0);
145        for (i, &tok) in tokens.iter().enumerate() {
146            one_hot_bf16(&mut into[i * self.vocab * 2..(i + 1) * self.vocab * 2], tok);
147        }
148    }
149
150    /// 2026-09-25: Substitute the script into the rows the inner router answered.
151    fn substitute(&self, keys: &[u64], fallback: &[u32], model_rows: &DecodeRows) -> Vec<u32> {
152        let inner_rows: Vec<u32> = match model_rows {
153            DecodeRows::Tokens(t) => t.clone(),
154            DecodeRows::HostLogits { .. } => fallback.to_vec(),
155        };
156        let answered: Vec<u32> = keys
157            .iter()
158            .zip(&inner_rows)
159            .map(|(&k, &f)| self.next_for(k, f))
160            .collect();
161        *self.last_rows.borrow_mut() = answered.clone();
162        answered
163    }
164}
165
166impl<D: DeviceIo> DeviceIo for ScriptedDeviceIo<D>
167where
168    D::LoraCmd: Label,
169{
170    type Seq = D::Seq;
171    type Ptr = D::Ptr;
172    type Model = D::Model;
173    type LoraCmd = D::LoraCmd;
174    type LoraAck = D::LoraAck;
175
176    fn model(&self) -> &D::Model {
177        self.inner.model()
178    }
179    fn facts(&self) -> DeviceFacts {
180        self.inner.facts()
181    }
182    fn max_depth(&self) -> usize {
183        self.max_depth.min(self.inner.max_depth())
184    }
185
186    fn launch(&self, plan: StepPlan<'_>, rows: &mut [&mut D::Seq]) -> Result<Ticket, DeviceError> {
187        self.launches.set(self.launches.get() + 1);
188        if let Some(fault) = self.faults.borrow_mut().pop_front() {
189            return Err(DeviceError::Forward(match fault {
190                Fault::KvExhausted => anyhow::anyhow!("KV cache exhausted (scripted)"),
191                Fault::Fatal => anyhow::anyhow!("scripted device fault"),
192            }));
193        }
194        let keys: Vec<u64> = rows.iter().map(|r| (self.key)(r)).collect();
195        match plan {
196            StepPlan::Decode {
197                tokens,
198                ctx,
199                readback: Readback::HostLogits { into },
200            } => {
201                // 2026-09-25: The inner router does the row bookkeeping and reads its
202                // block into `scratch`, which becomes the plan's buffer; rows with a
203                // script line are overwritten below.
204                let mut scratch = Vec::new();
205                let ticket = self.inner.launch(
206                    StepPlan::Decode {
207                        tokens: tokens.clone(),
208                        ctx,
209                        readback: Readback::HostLogits { into: &mut scratch },
210                    },
211                    rows,
212                )?;
213                let inner = self.inner.await_result(ticket)?;
214                let StepOutcome::Decode { logits, .. } = inner.outcome;
215                // 2026-09-25: A row with a script line answers one-hot at that token; a
216                // row without one keeps the block the inner router read.
217                *into = scratch;
218                let row_bytes = self.vocab * 2;
219                let mut answered = Vec::with_capacity(keys.len());
220                for (i, &k) in keys.iter().enumerate() {
221                    match self.take_line(k) {
222                        Some(tok) => {
223                            if into.len() >= (i + 1) * row_bytes {
224                                one_hot_bf16(&mut into[i * row_bytes..(i + 1) * row_bytes], tok);
225                            }
226                            answered.push(tok);
227                        }
228                        None => answered.push(tokens[i]),
229                    }
230                }
231                *self.last_rows.borrow_mut() = answered;
232                self.settled.borrow_mut().push_back(StepResult {
233                    ticket,
234                    outcome: StepOutcome::Decode {
235                        logits,
236                        rows: DecodeRows::HostLogits { elem_bytes: 2 },
237                    },
238                });
239                Ok(ticket)
240            }
241            StepPlan::Decode {
242                tokens,
243                ctx,
244                readback,
245            } => {
246                let fallback = tokens.clone();
247                let ticket = self.inner.launch(
248                    StepPlan::Decode {
249                        tokens,
250                        ctx,
251                        readback,
252                    },
253                    rows,
254                )?;
255                self.in_flight.borrow_mut().push_back(InFlight {
256                    ticket,
257                    keys,
258                    fallback,
259                });
260                Ok(ticket)
261            }
262            StepPlan::DecodeFed {
263                sources,
264                masks,
265                ctx,
266            } => {
267                let ticket = self.inner.launch(
268                    StepPlan::DecodeFed {
269                        sources,
270                        masks,
271                        ctx,
272                    },
273                    rows,
274                )?;
275                self.in_flight.borrow_mut().push_back(InFlight {
276                    ticket,
277                    keys,
278                    fallback: Vec::new(),
279                });
280                Ok(ticket)
281            }
282        }
283    }
284
285    fn await_result(&self, ticket: Ticket) -> Result<StepResult<D::Ptr>, DeviceError> {
286        let settled_at = self
287            .settled
288            .borrow()
289            .iter()
290            .position(|r| r.ticket == ticket);
291        if let Some(pos) = settled_at {
292            return Ok(self.settled.borrow_mut().remove(pos).expect("position"));
293        }
294        let pending = {
295            let mut q = self.in_flight.borrow_mut();
296            let pos = q.iter().position(|f| f.ticket == ticket).ok_or_else(|| {
297                DeviceError::Effect(anyhow::anyhow!(
298                    "await_result({ticket:?}): no scripted step in flight"
299                ))
300            })?;
301            q.remove(pos).expect("position")
302        };
303        let inner = self.inner.await_result(ticket)?;
304        let StepOutcome::Decode {
305            logits,
306            rows: model_rows,
307        } = inner.outcome;
308        let answered = self.substitute(&pending.keys, &pending.fallback, &model_rows);
309        Ok(StepResult {
310            ticket,
311            outcome: StepOutcome::Decode {
312                logits,
313                rows: DecodeRows::Tokens(answered),
314            },
315        })
316    }
317
318    fn apply(&self, effect: Effect<'_, D::Seq, D::Ptr>) -> Result<EffectOutcome, DeviceError> {
319        self.effects.borrow_mut().push(effect.label());
320        match effect {
321            Effect::ReadLogits { rows, into, .. } => {
322                let answered: Vec<u32> =
323                    self.last_rows.borrow().iter().take(rows).copied().collect();
324                self.one_hot_rows(&answered, into);
325                into.resize(rows * self.vocab * 2, 0);
326                Ok(EffectOutcome::HostLogits { elem_bytes: 2 })
327            }
328            other => self.inner.apply(other),
329        }
330    }
331
332    fn lora(&mut self, cmd: D::LoraCmd) -> Result<D::LoraAck, String> {
333        self.effects
334            .borrow_mut()
335            .push(format!("lora {}", cmd.label()));
336        self.inner.lora(cmd)
337    }
338
339    fn teardown(self: Box<Self>) -> anyhow::Result<()> {
340        Box::new(self.inner).teardown()
341    }
342}
metrale_scheduler/scripted.rs

metrale_scheduler/
scripted.rs
