metrale_scheduler/
effect.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: [`Effect`], the device-side bookkeeping the core asks a router for
4//! between steps, and [`EffectOutcome`], the router's answer. Generic over the
5//! sequence state `S` and the logits handle `P`.
6//!
7//! Owner: scheduler.
8//! Invariants: none beyond the types.
9
10use crate::plan::Label;
11
12pub enum Effect<'a, S, P> {
13    /// 2026-09-25: Release a sequence's device resources: optionally offer its KV to the
14    /// prefix cache, free it, then tell the EP worker.
15    ReleaseSeq {
16        seq: &'a mut S,
17        cache: bool,
18        /// 2026-09-25: The caller's label, for the failure log lines.
19        what: &'static str,
20    },
21    /// 2026-09-25: Migrate a live sequence onto slot `target`.
22    CompactSlot {
23        seq: &'a mut S,
24        target: usize,
25    },
26    /// 2026-09-25: Disown a sequence's slot after it was migrated to another
27    /// sequence, so a later free or drop cannot release it twice.
28    DetachSlot {
29        seq: &'a mut S,
30    },
31    /// 2026-09-25: Stream a sequence's device state into `writer` (spill).
32    SaveSequenceState {
33        seq: &'a S,
34        writer: &'a mut dyn std::io::Write,
35    },
36    /// 2026-09-25: Copy the logits block `logits` (`rows` rows) to the host.
37    ReadLogits {
38        logits: P,
39        rows: usize,
40        into: &'a mut Vec<u8>,
41    },
42    /// 2026-09-25: Block-aligned Marconi SSM checkpoint after a decode step.
43    MarconiCheckpoint {
44        seq: &'a mut S,
45    },
46    SsmSnapshotSave {
47        seq: &'a S,
48        slot: usize,
49    },
50    SsmSnapshotRestore {
51        seq: &'a S,
52        slot: usize,
53    },
54    /// 2026-09-25: `save_hidden_for_catchup(row, pos)`: copy row `row`'s hidden into the
55    /// MTP catch-up ring at label `pos`.
56    SaveHiddenCatchup {
57        row: usize,
58        pos: usize,
59    },
60    /// 2026-09-25: Make the default stream wait for the secondary stream.
61    SyncSecondary,
62    /// 2026-09-25: Every worker slot down.
63    EpShutdown,
64    /// 2026-09-25: Synchronise `stream` before teardown.
65    Quiesce {
66        stream: u64,
67    },
68    /// 2026-09-25: Make sure the KV block for the sequence's next decode position exists
69    /// before a step is launched ahead of the host; answers `Reserved` with
70    /// the blocks it added (0 when the block was already there) or
71    /// `Exhausted` when the allocation fails.
72    ReserveKv {
73        seq: &'a mut S,
74    },
75    /// 2026-09-25: The over-run of a fed step was discarded for this sequence: release
76    /// the blocks the latest `ReserveKv` added for it. Only a router that
77    /// reserves ahead supports it; the synchronous router returns an error.
78    Rollback {
79        seq: &'a mut S,
80    },
81}
82
83impl<S, P> Label for Effect<'_, S, P> {
84    fn label(&self) -> String {
85        match self {
86            Self::ReleaseSeq { cache, what, .. } => format!("ReleaseSeq{{cache={cache}, {what}}}"),
87            Self::CompactSlot { target, .. } => format!("CompactSlot{{target={target}}}"),
88            Self::DetachSlot { .. } => "DetachSlot".into(),
89            Self::SaveSequenceState { .. } => "SaveSequenceState".into(),
90            Self::ReadLogits { rows, .. } => format!("ReadLogits{{rows={rows}}}"),
91            Self::MarconiCheckpoint { .. } => "MarconiCheckpoint".into(),
92            Self::SsmSnapshotSave { slot, .. } => format!("SsmSnapshotSave{{slot={slot}}}"),
93            Self::SsmSnapshotRestore { slot, .. } => format!("SsmSnapshotRestore{{slot={slot}}}"),
94            Self::SaveHiddenCatchup { row, pos } => {
95                format!("SaveHiddenCatchup{{row={row}, pos={pos}}}")
96            }
97            Self::SyncSecondary => "SyncSecondary".into(),
98            Self::EpShutdown => "EpShutdown".into(),
99            Self::Quiesce { stream } => format!("Quiesce{{stream={stream}}}"),
100            Self::ReserveKv { .. } => "ReserveKv".into(),
101            Self::Rollback { .. } => "Rollback".into(),
102        }
103    }
104}
105
106#[derive(Debug, PartialEq, Eq)]
107pub enum EffectOutcome {
108    Unit,
109    /// 2026-09-25: `ReadLogits`: the element width the model writes.
110    HostLogits {
111        elem_bytes: usize,
112    },
113    /// 2026-09-25: `ReserveKv`: the blocks added for the next position.
114    Reserved {
115        blocks: usize,
116    },
117    /// 2026-09-25: `ReserveKv`: no block could be added; the core does not launch ahead.
118    Exhausted,
119}
metrale_scheduler/effect.rs

metrale_scheduler/
effect.rs
