From 1a7423a92b71105a1d164810de2d714b6551448b Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 5 Jul 2026 16:59:51 +0100 Subject: [PATCH 01/33] support event driven system Signed-off-by: kerthcet --- Cargo.lock | 1 + Cargo.toml | 1 + docs/fsm_architecture.md | 394 +++++++++++++++++++++++++++ examples/fsm_usage.rs | 184 +++++++++++++ src/block_manager/allocator.rs | 70 +++++ src/block_manager/manager.rs | 161 +++++++++++ src/block_manager/mod.rs | 7 + src/block_manager/types.rs | 99 +++++++ src/lib.rs | 10 + src/main.rs | 2 + src/sequence_manager/events.rs | 68 +++++ src/sequence_manager/fsm_events.rs | 344 +++++++++++++++++++++++ src/sequence_manager/fsm_manager.rs | 341 +++++++++++++++++++++++ src/sequence_manager/id_generator.rs | 70 +++++ src/sequence_manager/mod.rs | 13 + src/sequence_manager/states.rs | 144 ++++++++++ 16 files changed, 1909 insertions(+) create mode 100644 docs/fsm_architecture.md create mode 100644 examples/fsm_usage.rs create mode 100644 src/block_manager/allocator.rs create mode 100644 src/block_manager/manager.rs create mode 100644 src/block_manager/mod.rs create mode 100644 src/block_manager/types.rs create mode 100644 src/lib.rs create mode 100644 src/sequence_manager/events.rs create mode 100644 src/sequence_manager/fsm_events.rs create mode 100644 src/sequence_manager/fsm_manager.rs create mode 100644 src/sequence_manager/id_generator.rs create mode 100644 src/sequence_manager/mod.rs create mode 100644 src/sequence_manager/states.rs diff --git a/Cargo.lock b/Cargo.lock index 22e04fb..cc200dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1543,6 +1543,7 @@ dependencies = [ "serde_json", "sysinfo", "tempfile", + "thiserror 2.0.18", "tokio", "tokio-stream", "tower 0.4.13", diff --git a/Cargo.toml b/Cargo.toml index 622000e..46b9364 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,7 @@ sysinfo = "0.32" rusqlite = { version = "0.32", features = ["bundled"] } rusqlite_migration = "1.3" regex = "1.11" +thiserror = "2.0" # Web server axum = "0.7" diff --git a/docs/fsm_architecture.md b/docs/fsm_architecture.md new file mode 100644 index 0000000..5b7896c --- /dev/null +++ b/docs/fsm_architecture.md @@ -0,0 +1,394 @@ +# FSM Architecture for PUMA + +Finite State Machine pattern for sequence lifecycle management. + +## Core Concepts + +### State Machine per Sequence + +Each sequence is an independent state machine: + +```rust +pub enum SequenceState { + Waiting(WaitingState), + Prefilling(PrefillingState), + Decoding(DecodingState), + Preempted(PreemptedState), + Finished(FinishedState), +} + +// Multiple sequences, different states +sequences: HashMap +``` + +### States Own Resources + +```rust +pub struct WaitingState { + seq_id: SequenceId, + prompt_tokens: usize, + max_tokens: usize, + // No blocks - not allocated yet +} + +pub struct DecodingState { + seq_id: SequenceId, + blocks: Vec, // ← State owns blocks + num_tokens: usize, + max_tokens: usize, +} + +pub struct PreemptedState { + seq_id: SequenceId, + num_tokens: usize, + max_tokens: usize, + // No blocks - they were freed +} +``` + +**Benefit:** Type system enforces "Waiting has no blocks, Decoding must have blocks" + +### Events as State Transformations + +```rust +pub struct ScheduleEvent<'a> { + block_manager: &'a mut BlockManager, + tokens_per_block: usize, +} + +impl ScheduleEvent { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Waiting(s) => { + // Allocate blocks + let blocks = allocate_blocks(s.prompt_tokens)?; + + // Transform to new state + Ok(SequenceState::Prefilling(PrefillingState { + seq_id: s.seq_id, + blocks, // Ownership transferred + tokens_filled: 0, + tokens_total: s.prompt_tokens, + max_tokens: s.max_tokens, + })) + } + _ => Err(Error::invalid_transition("ScheduleEvent requires Waiting")), + } + } +} +``` + +**Pattern:** `Event + OldState → NewState` + +## State Diagram + +``` + AddRequest + ↓ + ┌─────────────→ Waiting + │ ↓ ScheduleEvent + │ Prefilling + │ ↓ AppendTokens + │ ↓ (tokens_filled >= tokens_total) + │ Decoding + │ ├─→ AppendTokens (continue) + │ │ ↓ + │ │ ├─ Decoding (more tokens) + │ │ └─ Finished (max_tokens reached) + │ │ + │ ├─→ ForkEvent + │ │ ├─ Parent: Decoding + │ │ └─ Child: Decoding + │ │ + │ ├─→ PreemptEvent (OOM) + │ │ ↓ + │ │ Preempted + │ │ ↓ ResumeEvent + │ └───┘ + │ + └─ CompleteEvent → Finished +``` + +## State Definitions + +```rust +/// Waiting in queue for scheduling +pub struct WaitingState { + pub seq_id: SequenceId, + pub prompt_tokens: usize, + pub max_tokens: usize, +} + +/// Running prefill phase +pub struct PrefillingState { + pub seq_id: SequenceId, + pub blocks: Vec, // Owns blocks + pub tokens_filled: usize, + pub tokens_total: usize, + pub max_tokens: usize, +} + +/// Running decode phase +pub struct DecodingState { + pub seq_id: SequenceId, + pub blocks: Vec, // Owns blocks + pub num_tokens: usize, + pub max_tokens: usize, +} + +/// Preempted due to OOM +pub struct PreemptedState { + pub seq_id: SequenceId, + pub num_tokens: usize, + pub max_tokens: usize, + // Blocks freed +} + +/// Finished generation +pub struct FinishedState { + pub seq_id: SequenceId, + pub finish_reason: FinishReason, + // All resources freed +} +``` + +## Event Definitions + +### ScheduleEvent: Waiting → Prefilling + +```rust +let event = ScheduleEvent { + block_manager: &mut block_manager, + tokens_per_block: 16, +}; + +match event.apply(state) { + Ok(new_state) => { + // Waiting → Prefilling (with blocks allocated) + } + Err(Error::OutOfMemory) => { + // Can't allocate, stay in queue + } +} +``` + +### AppendTokensEvent: Prefilling → Decoding or Decoding → Decoding + +```rust +let event = AppendTokensEvent { + num_tokens: 100, + block_manager: &mut block_manager, + tokens_per_block: 16, +}; + +// From Prefilling +Prefilling(tokens_filled: 0) + AppendTokens(100) + → Decoding(num_tokens: 100) // Transition + +// From Decoding +Decoding(num_tokens: 100) + AppendTokens(1) + → Decoding(num_tokens: 101) // Stay in state +``` + +### ForkEvent: Decoding → (Decoding, Decoding) + +```rust +let event = ForkEvent { + child_id: SequenceId(2), + block_manager: &mut block_manager, +}; + +// Copy-on-write: both sequences share blocks +match event.apply(parent_state) { + Ok((parent, child)) => { + // Both in Decoding state + // Blocks ref-counted (shared) + } +} +``` + +### PreemptEvent: Decoding → Preempted + +```rust +let event = PreemptEvent { + block_manager: &mut block_manager, +}; + +// Frees blocks +Decoding(blocks: [1,2,3]) → Preempted(blocks: []) +``` + +### ResumeEvent: Preempted → Waiting + +```rust +let event = ResumeEvent; + +// Back to waiting for rescheduling +Preempted → Waiting → (Schedule) → Decoding +``` + +## Sequence Manager + +```rust +pub struct SequenceManager { + block_manager: BlockManager, + sequences: HashMap, + tokens_per_block: usize, + event_rx: SequenceEventReceiver, +} + +impl SequenceManager { + pub async fn run(mut self) { + while let Some(event) = self.event_rx.recv().await { + match event { + SequenceEvent::AppendTokens { seq_id, num_tokens } => { + // Get current state + let state = self.sequences.remove(&seq_id).unwrap(); + + // Apply event transformation + let event = AppendTokensEvent { + num_tokens, + block_manager: &mut self.block_manager, + tokens_per_block: self.tokens_per_block, + }; + + // Transform state + match event.apply(state) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + } + Err(e) => { + error!("Transition failed: {:?}", e); + } + } + } + // ... other events + } + } + } +} +``` + +## Type Safety + +**Compile-time checks:** +```rust +// ✅ Valid +let waiting = SequenceState::Waiting(...); +let event = ScheduleEvent { ... }; +event.apply(waiting)?; // OK + +// ❌ Invalid (caught at runtime, enforced by match) +let decoding = SequenceState::Decoding(...); +let event = ScheduleEvent { ... }; +event.apply(decoding)?; // Error: invalid_transition +``` + +**Resource ownership:** +```rust +fn transition(state: DecodingState) -> FinishedState { + // state.blocks moved/consumed + for block in state.blocks { + block_manager.free(block); + } + + FinishedState { + seq_id: state.seq_id, + finish_reason: FinishReason::Stop, + } + // Old state dropped, can't access blocks anymore +} +``` + +## Benefits + +1. **Type Safety** + - States enforce resource invariants + - Invalid transitions caught explicitly + +2. **Clear Ownership** + - Resources belong to states + - Automatic cleanup on state drop + +3. **Testability** + - Events are pure functions + - Test state transitions in isolation + +4. **Explicitness** + - Every transition is a function call + - State diagram maps directly to code + +## Example Flow + +```rust +// 1. Add request → Waiting +let state = SequenceState::Waiting(WaitingState { + seq_id: SequenceId(1), + prompt_tokens: 100, + max_tokens: 150, +}); + +// 2. Schedule → Prefilling (allocate 7 blocks) +let event = ScheduleEvent { ... }; +let state = event.apply(state)?; +// state = Prefilling(blocks: [1,2,3,4,5,6,7], tokens_filled: 0) + +// 3. Append tokens → Decoding (transition) +let event = AppendTokensEvent { num_tokens: 100 }; +let state = event.apply(state)?; +// state = Decoding(blocks: [1,2,3,4,5,6,7], num_tokens: 100) + +// 4. Fork → 2 sequences +let event = ForkEvent { child_id: SequenceId(2) }; +let (parent, child) = event.apply(state)?; +// parent = Decoding(blocks: [1,2,3,4,5,6,7], ref_count: 2) +// child = Decoding(blocks: [1,2,3,4,5,6,7], ref_count: 2) + +// 5. Continue generating +let event = AppendTokensEvent { num_tokens: 10 }; +let parent = event.apply(parent)?; +// parent = Decoding(blocks: [1,2,3,4,5,6,7,8], num_tokens: 110) + +// 6. Complete → Finished (free blocks) +let event = CompleteEvent { ... }; +let state = event.apply(parent)?; +// state = Finished, blocks freed +``` + +## Comparison with Simple Design + +| Simple | FSM | +|--------|-----| +| `seq.state = Running` | `state = ScheduleEvent.apply(state)` | +| Resources in struct | Resources in state variant | +| Manual cleanup | Automatic cleanup | +| Runtime checks | Type + runtime checks | + +## Usage + +```rust +use puma::block_manager::{BlockManager, CpuAllocator}; +use puma::sequence_manager::{FSMSequenceManager, SequenceEvent, SequenceId}; + +// Setup +let allocator = Box::new(CpuAllocator::new(100_000_000)); +let block_manager = BlockManager::new(allocator, 4096); +let (fsm_manager, event_tx) = FSMSequenceManager::new(block_manager, 16); + +// Run event loop +tokio::spawn(async move { fsm_manager.run().await }); + +// Send events +event_tx.send(SequenceEvent::AddRequest { + seq_id: SequenceId(1), + prompt_tokens: 100, + max_tokens: 150, +}).unwrap(); + +event_tx.send(SequenceEvent::AppendTokens { + seq_id: SequenceId(1), + num_tokens: 100, +}).unwrap(); +``` + +See `examples/fsm_usage.rs` for complete example. diff --git a/examples/fsm_usage.rs b/examples/fsm_usage.rs new file mode 100644 index 0000000..38db0fa --- /dev/null +++ b/examples/fsm_usage.rs @@ -0,0 +1,184 @@ +use puma::block_manager::{BlockManager, CpuAllocator}; +use puma::sequence_manager::{SequenceManager, SequenceEvent, SequenceIdGenerator}; + +#[tokio::main] +async fn main() { + // Initialize tracing + tracing_subscriber::fmt() + .with_env_filter("info,puma=debug") + .init(); + + println!("=== FSM-Based Sequence Manager Demo ===\n"); + + // Setup block manager + let allocator = Box::new(CpuAllocator::new(100_000_000)); // 100MB + let block_size = 4096; // 4KB per block + let block_manager = BlockManager::new(allocator, block_size); + + // Setup sequence manager + let tokens_per_block = 16; + let (seq_manager, event_tx) = SequenceManager::new(block_manager, tokens_per_block); + + // Spawn event loop + let manager_handle = tokio::spawn(async move { + seq_manager.run().await; + }); + + // ID generator (starts from 0) + let id_gen = SequenceIdGenerator::new(); + + println!("1. Adding sequence (Waiting state)"); + let seq_id = id_gen.next(); + event_tx + .send(SequenceEvent::AddRequest { + seq_id, + prompt_tokens: 100, + max_tokens: 150, + }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + println!("\n2. Appending tokens during prefill (Prefilling → Decoding)"); + event_tx + .send(SequenceEvent::AppendTokens { + seq_id, + num_tokens: 100, // Complete prefill + }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + println!("\n3. Generating tokens (Decoding state)"); + for i in 0..10 { + event_tx + .send(SequenceEvent::AppendTokens { + seq_id, + num_tokens: 1, + }) + .unwrap(); + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + } + + println!("\n4. Forking sequence for beam search"); + let child_id = id_gen.next(); + event_tx + .send(SequenceEvent::ForkSequence { + parent_id: seq_id, + child_id, + }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + println!("\n5. Continuing both sequences"); + event_tx + .send(SequenceEvent::AppendTokens { + seq_id, + num_tokens: 5, + }) + .unwrap(); + event_tx + .send(SequenceEvent::AppendTokens { + seq_id: child_id, + num_tokens: 5, + }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + // Query stats + let (tx, rx) = tokio::sync::oneshot::channel(); + event_tx + .send(SequenceEvent::GetStats { response: tx }) + .unwrap(); + + let stats = rx.await.unwrap(); + println!("\n=== Stats ==="); + println!("Total sequences: {}", stats.num_sequences); + println!("Running: {}", stats.num_running); + println!("Waiting: {}", stats.num_waiting); + println!("Preempted: {}", stats.num_preempted); + + for (block_type, block_stats) in &stats.block_stats { + println!("\nBlock type: {}", block_type); + println!(" Total blocks: {}", block_stats.total_blocks); + println!(" Allocated: {}", block_stats.allocated_blocks); + println!(" Free: {}", block_stats.free_blocks); + println!(" Total memory: {} bytes", block_stats.total_memory); + } + + println!("\n6. Completing sequences (Finished state)"); + event_tx + .send(SequenceEvent::CompleteSequence { seq_id }) + .unwrap(); + event_tx + .send(SequenceEvent::CompleteSequence { + seq_id: child_id, + }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + // Final stats + let (tx, rx) = tokio::sync::oneshot::channel(); + event_tx + .send(SequenceEvent::GetStats { response: tx }) + .unwrap(); + + let stats = rx.await.unwrap(); + println!("\n=== Final Stats ==="); + println!("Total sequences: {}", stats.num_sequences); + println!("Running: {}", stats.num_running); + + println!("\n7. Testing preemption"); + let seq3 = id_gen.next(); + event_tx + .send(SequenceEvent::AddRequest { + seq_id: seq3, + prompt_tokens: 50, + max_tokens: 100, + }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + event_tx + .send(SequenceEvent::AppendTokens { + seq_id: seq3, + num_tokens: 50, + }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + println!("\n Preempting sequence (Decoding → Preempted)"); + event_tx + .send(SequenceEvent::PreemptSequence { seq_id: seq3 }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + println!(" Resuming sequence (Preempted → Waiting → Decoding)"); + event_tx + .send(SequenceEvent::ResumeSequence { seq_id: seq3 }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + event_tx + .send(SequenceEvent::CompleteSequence { seq_id: seq3 }) + .unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + println!("\n=== Sequence Manager Demo Complete! ==="); + println!("\nState transitions demonstrated:"); + println!(" Waiting → Prefilling → Decoding → Finished"); + println!(" Decoding → Fork → Two Decoding sequences"); + println!(" Decoding → Preempted → Waiting → Decoding"); + + // Cleanup + drop(event_tx); + manager_handle.await.unwrap(); +} diff --git a/src/block_manager/allocator.rs b/src/block_manager/allocator.rs new file mode 100644 index 0000000..5a2c9b7 --- /dev/null +++ b/src/block_manager/allocator.rs @@ -0,0 +1,70 @@ +use super::types::*; + +/// Memory allocator trait +pub trait MemoryAllocator: Send + Sync { + fn allocate(&mut self, size_bytes: usize) -> Result; + fn free(&mut self, addr: MemoryAddress) -> Result<()>; + fn get_total_memory(&self) -> usize; + fn get_available_memory(&self) -> usize; +} + +/// Simple CPU memory allocator (for testing/CPU inference) +pub struct CpuAllocator { + total_memory: usize, + used_memory: usize, +} + +impl CpuAllocator { + pub fn new(total_memory: usize) -> Self { + Self { + total_memory, + used_memory: 0, + } + } +} + +impl MemoryAllocator for CpuAllocator { + fn allocate(&mut self, size_bytes: usize) -> Result { + if self.used_memory + size_bytes > self.total_memory { + return Err(Error::OutOfMemory); + } + + // Allocate aligned memory + let layout = std::alloc::Layout::from_size_align(size_bytes, 64) + .map_err(|e| Error::AllocationError(e.to_string()))?; + + let ptr = unsafe { std::alloc::alloc(layout) }; + + if ptr.is_null() { + return Err(Error::OutOfMemory); + } + + self.used_memory += size_bytes; + + Ok(MemoryAddress { ptr, size: size_bytes }) + } + + fn free(&mut self, addr: MemoryAddress) -> Result<()> { + let layout = std::alloc::Layout::from_size_align(addr.size, 64) + .map_err(|e| Error::AllocationError(e.to_string()))?; + + unsafe { + std::alloc::dealloc(addr.ptr, layout); + } + + self.used_memory -= addr.size; + + Ok(()) + } + + fn get_total_memory(&self) -> usize { + self.total_memory + } + + fn get_available_memory(&self) -> usize { + self.total_memory - self.used_memory + } +} + +// TODO: Implement CudaAllocator later +// pub struct CudaAllocator { ... } diff --git a/src/block_manager/manager.rs b/src/block_manager/manager.rs new file mode 100644 index 0000000..20967cb --- /dev/null +++ b/src/block_manager/manager.rs @@ -0,0 +1,161 @@ +use super::allocator::*; +use super::types::*; +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; + +pub struct BlockManager { + allocator: Box, + + // Block registry + block_table: HashMap, + + // Free pools per type + free_pools: HashMap>, + + // Block configurations + block_configs: HashMap, + + // Default type + default_block_type: BlockType, + + next_block_id: AtomicU64, +} + +impl BlockManager { + pub fn new( + allocator: Box, + default_block_size: usize, + ) -> Self { + let mut manager = Self { + allocator, + block_table: HashMap::new(), + free_pools: HashMap::new(), + block_configs: HashMap::new(), + default_block_type: BlockType::StandardKV, + next_block_id: AtomicU64::new(0), + }; + + // Register default type + manager.register_block_type(BlockConfig { + block_type: BlockType::StandardKV, + size_bytes: default_block_size, + }); + + manager + } + + pub fn register_block_type(&mut self, config: BlockConfig) { + self.free_pools.insert(config.block_type.clone(), Vec::new()); + self.block_configs.insert(config.block_type.clone(), config); + } + + pub fn allocate_typed(&mut self, block_type: &BlockType) -> Result { + // Try free pool first + if let Some(block_id) = self.free_pools + .get_mut(block_type) + .and_then(|pool| pool.pop()) + { + let block = self.block_table.get_mut(&block_id).unwrap(); + block.ref_count = 1; + return Ok(block_id); + } + + // Allocate new physical memory + let config = self.block_configs.get(block_type) + .ok_or_else(|| Error::UnknownBlockType(block_type.as_str().to_string()))?; + + let mem_addr = self.allocator.allocate(config.size_bytes)?; + + let block_id = BlockId(self.next_block_id.fetch_add(1, Ordering::Relaxed)); + + self.block_table.insert(block_id, Block { + block_id, + block_type: block_type.clone(), + mem_addr, + ref_count: 1, + }); + + Ok(block_id) + } + + pub fn allocate(&mut self) -> Result { + let default_type = self.default_block_type.clone(); + self.allocate_typed(&default_type) + } + + pub fn free(&mut self, block_id: BlockId) -> Result<()> { + let block = self.block_table.get_mut(&block_id) + .ok_or(Error::InvalidBlockId(block_id))?; + + if block.ref_count == 0 { + return Err(Error::DoubleFree(block_id)); + } + + block.ref_count -= 1; + + if block.ref_count == 0 { + self.free_pools + .get_mut(&block.block_type) + .unwrap() + .push(block_id); + } + + Ok(()) + } + + pub fn add_ref(&mut self, block_id: BlockId) -> Result<()> { + let block = self.block_table.get_mut(&block_id) + .ok_or(Error::InvalidBlockId(block_id))?; + block.ref_count += 1; + Ok(()) + } + + pub fn get_memory_address(&self, block_id: BlockId) -> Result { + self.block_table.get(&block_id) + .map(|b| b.mem_addr) + .ok_or(Error::InvalidBlockId(block_id)) + } + + pub fn get_block_type(&self, block_id: BlockId) -> Result<&BlockType> { + self.block_table.get(&block_id) + .map(|b| &b.block_type) + .ok_or(Error::InvalidBlockId(block_id)) + } + + pub fn get_stats(&self) -> HashMap { + let mut stats = HashMap::new(); + + for (type_id, pool) in &self.free_pools { + let config = &self.block_configs[type_id]; + let free_count = pool.len(); + + let allocated_count = self.block_table.values() + .filter(|b| b.block_type == *type_id && b.ref_count > 0) + .count(); + + stats.insert(type_id.clone(), BlockStats { + total_blocks: free_count + allocated_count, + allocated_blocks: allocated_count, + free_blocks: free_count, + block_size: config.size_bytes, + total_memory: (free_count + allocated_count) * config.size_bytes, + }); + } + + stats + } + + pub fn can_allocate(&self, block_type: &BlockType) -> bool { + if let Some(pool) = self.free_pools.get(block_type) { + if !pool.is_empty() { + return true; + } + } + + if let Some(config) = self.block_configs.get(block_type) { + return self.allocator.get_available_memory() >= config.size_bytes; + } + + false + } +} diff --git a/src/block_manager/mod.rs b/src/block_manager/mod.rs new file mode 100644 index 0000000..2d08c80 --- /dev/null +++ b/src/block_manager/mod.rs @@ -0,0 +1,7 @@ +pub mod allocator; +pub mod manager; +pub mod types; + +pub use allocator::*; +pub use manager::BlockManager; +pub use types::*; diff --git a/src/block_manager/types.rs b/src/block_manager/types.rs new file mode 100644 index 0000000..92a1d13 --- /dev/null +++ b/src/block_manager/types.rs @@ -0,0 +1,99 @@ +/// Block type - defines the kind of memory block +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum BlockType { + /// Standard KV cache blocks (for normal attention) + StandardKV, + + /// Compressed KV cache blocks (for DeepSeek-style compression) + CompressedKV, + + /// Compressor state blocks (for DeepSeek RNN compressor) + CompressorState, +} + +impl BlockType { + pub fn as_str(&self) -> &'static str { + match self { + BlockType::StandardKV => "standard_kv", + BlockType::CompressedKV => "compressed_kv", + BlockType::CompressorState => "compressor_state", + } + } +} + +impl std::fmt::Display for BlockType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct BlockId(pub u64); + +impl BlockId { + pub fn new(id: u64) -> Self { + Self(id) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct SequenceId(pub u64); + +impl SequenceId { + pub fn new(id: u64) -> Self { + Self(id) + } +} + +#[derive(Debug, Clone, Copy)] +pub struct MemoryAddress { + pub ptr: *mut u8, + pub size: usize, +} + +unsafe impl Send for MemoryAddress {} +unsafe impl Sync for MemoryAddress {} + +pub struct BlockConfig { + pub block_type: BlockType, + pub size_bytes: usize, +} + +pub struct Block { + pub block_id: BlockId, + pub block_type: BlockType, + pub mem_addr: MemoryAddress, + pub ref_count: usize, +} + +#[derive(Debug, Clone)] +pub struct BlockStats { + pub total_blocks: usize, + pub allocated_blocks: usize, + pub free_blocks: usize, + pub block_size: usize, + pub total_memory: usize, +} + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("Unknown block type: {0}")] + UnknownBlockType(String), + + #[error("Invalid block ID: {0:?}")] + InvalidBlockId(BlockId), + + #[error("Double free detected for block: {0:?}")] + DoubleFree(BlockId), + + #[error("Unknown sequence: {0:?}")] + UnknownSequence(SequenceId), + + #[error("Out of memory")] + OutOfMemory, + + #[error("Allocation error: {0}")] + AllocationError(String), +} + +pub type Result = std::result::Result; diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..ae83ace --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,10 @@ +pub mod api; +pub mod backend; +pub mod block_manager; +pub mod cli; +pub mod downloader; +pub mod registry; +pub mod sequence_manager; +pub mod storage; +pub mod system; +pub mod utils; diff --git a/src/main.rs b/src/main.rs index 02929ea..fd485d9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,10 @@ mod api; mod backend; +mod block_manager; mod cli; mod downloader; mod registry; +mod sequence_manager; mod storage; mod system; mod utils; diff --git a/src/sequence_manager/events.rs b/src/sequence_manager/events.rs new file mode 100644 index 0000000..f40fbae --- /dev/null +++ b/src/sequence_manager/events.rs @@ -0,0 +1,68 @@ +use crate::block_manager::types::*; +use std::collections::HashMap; +use tokio::sync::mpsc; + +/// Events that drive the sequence manager +#[derive(Debug)] +pub enum SequenceEvent { + // Request lifecycle + AddRequest { + seq_id: SequenceId, + prompt_tokens: usize, + max_tokens: usize, + }, + + // Token generation + AppendTokens { + seq_id: SequenceId, + num_tokens: usize, + }, + + // Sequence forking (for beam search, parallel sampling) + ForkSequence { + parent_id: SequenceId, + child_id: SequenceId, + }, + + // Completion + CompleteSequence { + seq_id: SequenceId, + }, + + // Preemption (when OOM) + PreemptSequence { + seq_id: SequenceId, + }, + + // Resume preempted sequence + ResumeSequence { + seq_id: SequenceId, + }, + + // Query + GetBlocks { + seq_id: SequenceId, + response: tokio::sync::oneshot::Sender>>, + }, + + // Stats + GetStats { + response: tokio::sync::oneshot::Sender, + }, +} + +#[derive(Debug, Clone)] +pub struct SequenceManagerStats { + pub num_sequences: usize, + pub num_running: usize, + pub num_waiting: usize, + pub num_preempted: usize, + pub block_stats: HashMap, +} + +pub type SequenceEventSender = mpsc::UnboundedSender; +pub type SequenceEventReceiver = mpsc::UnboundedReceiver; + +pub fn create_event_channel() -> (SequenceEventSender, SequenceEventReceiver) { + mpsc::unbounded_channel() +} diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs new file mode 100644 index 0000000..0ffae9c --- /dev/null +++ b/src/sequence_manager/fsm_events.rs @@ -0,0 +1,344 @@ +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; +use super::states::*; +use tracing::{debug, warn}; + +/// FSM Events - transform states with type-safe transitions +/// +/// Pattern: Event takes old state, returns new state +/// Invalid transitions return Error::invalid_transition + +/// Schedule a waiting sequence (allocate blocks) +pub struct ScheduleEvent<'a> { + pub block_manager: &'a mut BlockManager, + pub tokens_per_block: usize, +} + +impl<'a> ScheduleEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Waiting(s) => self.from_waiting(s), + _ => Err(Error::invalid_transition("ScheduleEvent requires Waiting state")), + } + } + + fn from_waiting(self, state: WaitingState) -> Result { + let blocks_needed = + (state.prompt_tokens + self.tokens_per_block - 1) / self.tokens_per_block; + + let mut blocks = Vec::new(); + for _ in 0..blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => blocks.push(block_id), + Err(Error::OutOfMemory) => { + // OOM during scheduling - free what we allocated + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(Error::OutOfMemory); + } + Err(e) => { + // Other error - cleanup and propagate + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(e); + } + } + } + + debug!( + "Scheduled seq {:?}: allocated {} blocks for {} tokens", + state.seq_id, blocks.len(), state.prompt_tokens + ); + + Ok(SequenceState::Prefilling(PrefillingState { + seq_id: state.seq_id, + blocks, + tokens_filled: 0, + tokens_total: state.prompt_tokens, + max_tokens: state.max_tokens, + })) + } +} + +/// Append tokens to a running sequence +pub struct AppendTokensEvent<'a> { + pub num_tokens: usize, + pub block_manager: &'a mut BlockManager, + pub tokens_per_block: usize, +} + +impl<'a> AppendTokensEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Prefilling(s) => self.from_prefilling(s), + SequenceState::Decoding(s) => self.from_decoding(s), + _ => Err(Error::invalid_transition("AppendTokensEvent requires Prefilling or Decoding state")), + } + } + + fn from_prefilling(self, mut state: PrefillingState) -> Result { + state.tokens_filled += self.num_tokens; + + debug!( + "Prefilling seq {:?}: {}/{} tokens", + state.seq_id, state.tokens_filled, state.tokens_total + ); + + if state.tokens_filled >= state.tokens_total { + // Transition to decode + debug!("Seq {:?} transitioning to Decoding", state.seq_id); + Ok(SequenceState::Decoding(DecodingState { + seq_id: state.seq_id, + blocks: state.blocks, + num_tokens: state.tokens_filled, + max_tokens: state.max_tokens, + })) + } else { + Ok(SequenceState::Prefilling(state)) + } + } + + fn from_decoding(self, mut state: DecodingState) -> Result { + state.num_tokens += self.num_tokens; + + // Check if finished + if state.num_tokens >= state.max_tokens { + // Free blocks + for block_id in state.blocks { + self.block_manager.free(block_id)?; + } + + debug!( + "Seq {:?} finished: reached max_tokens ({})", + state.seq_id, state.max_tokens + ); + + return Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: FinishReason::MaxTokens, + })); + } + + // Check if need more blocks + let blocks_needed = (state.num_tokens + self.tokens_per_block - 1) / self.tokens_per_block; + + while state.blocks.len() < blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => { + state.blocks.push(block_id); + debug!( + "Seq {:?}: allocated block, total blocks: {}", + state.seq_id, + state.blocks.len() + ); + } + Err(Error::OutOfMemory) => { + warn!("OOM while appending tokens to {:?}", state.seq_id); + return Err(Error::OutOfMemory); + } + Err(e) => return Err(e), + } + } + + Ok(SequenceState::Decoding(state)) + } +} + +/// Preempt a running sequence (free blocks) +pub struct PreemptEvent<'a> { + pub block_manager: &'a mut BlockManager, +} + +impl<'a> PreemptEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Decoding(s) => self.from_decoding(s), + SequenceState::Prefilling(s) => self.from_prefilling(s), + _ => Err(Error::invalid_transition("PreemptEvent requires running state")), + } + } + + fn from_decoding(self, state: DecodingState) -> Result { + // Free all blocks + for block_id in state.blocks { + self.block_manager.free(block_id)?; + } + + debug!("Preempted seq {:?}", state.seq_id); + + Ok(SequenceState::Preempted(PreemptedState { + seq_id: state.seq_id, + num_tokens: state.num_tokens, + max_tokens: state.max_tokens, + })) + } + + fn from_prefilling(self, state: PrefillingState) -> Result { + // Free all blocks + for block_id in state.blocks { + self.block_manager.free(block_id)?; + } + + debug!("Preempted seq {:?} during prefill", state.seq_id); + + Ok(SequenceState::Preempted(PreemptedState { + seq_id: state.seq_id, + num_tokens: state.tokens_filled, + max_tokens: state.max_tokens, + })) + } +} + +/// Resume a preempted sequence +pub struct ResumeEvent; + +impl ResumeEvent { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Preempted(s) => self.from_preempted(s), + _ => Err(Error::invalid_transition("ResumeEvent requires Preempted state")), + } + } + + fn from_preempted(self, state: PreemptedState) -> Result { + debug!("Resuming seq {:?}", state.seq_id); + + // Transition back to waiting for rescheduling + Ok(SequenceState::Waiting(WaitingState { + seq_id: state.seq_id, + prompt_tokens: state.num_tokens, + max_tokens: state.max_tokens, + })) + } +} + +/// Fork a sequence (copy-on-write) +pub struct ForkEvent<'a> { + pub child_id: SequenceId, + pub block_manager: &'a mut BlockManager, +} + +impl<'a> ForkEvent<'a> { + pub fn apply( + self, + state: SequenceState, + ) -> Result<(SequenceState, SequenceState)> { + match state { + SequenceState::Decoding(s) => self.from_decoding(s), + _ => Err(Error::invalid_transition("ForkEvent requires Decoding state")), + } + } + + fn from_decoding( + self, + state: DecodingState, + ) -> Result<(SequenceState, SequenceState)> { + // Copy-on-write: increment ref counts + for &block_id in &state.blocks { + self.block_manager.add_ref(block_id)?; + } + + let child = DecodingState { + seq_id: self.child_id, + blocks: state.blocks.clone(), + num_tokens: state.num_tokens, + max_tokens: state.max_tokens, + }; + + debug!( + "Forked seq {:?} → {:?}", + state.seq_id, child.seq_id + ); + + Ok(( + SequenceState::Decoding(state), + SequenceState::Decoding(child), + )) + } +} + +/// Complete a sequence +pub struct CompleteEvent<'a> { + pub reason: FinishReason, + pub block_manager: &'a mut BlockManager, +} + +impl<'a> CompleteEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Decoding(s) => self.from_decoding(s), + SequenceState::Prefilling(s) => self.from_prefilling(s), + _ => Err(Error::invalid_transition("CompleteEvent requires running state")), + } + } + + fn from_decoding(self, state: DecodingState) -> Result { + // Free all blocks + for block_id in state.blocks { + self.block_manager.free(block_id)?; + } + + debug!( + "Completed seq {:?}: {:?}", + state.seq_id, self.reason + ); + + Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: self.reason, + })) + } + + fn from_prefilling(self, state: PrefillingState) -> Result { + // Free all blocks + for block_id in state.blocks { + self.block_manager.free(block_id)?; + } + + debug!( + "Completed seq {:?} during prefill: {:?}", + state.seq_id, self.reason + ); + + Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: self.reason, + })) + } +} + +/// Abort a sequence (error or cancellation) +pub struct AbortEvent<'a> { + pub reason: String, + pub block_manager: &'a mut BlockManager, +} + +impl<'a> AbortEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + // Can abort from any state + let seq_id = state.seq_id(); + + // Free blocks if any + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + self.block_manager.free(block_id)?; + } + } + + warn!("Aborted seq {:?}: {}", seq_id, self.reason); + + Ok(SequenceState::Aborted(AbortedState { + seq_id, + reason: self.reason, + })) + } +} + +/// Invalid transition error helper +impl Error { + pub fn invalid_transition(msg: &'static str) -> Self { + Error::AllocationError(format!("Invalid state transition: {}", msg)) + } +} diff --git a/src/sequence_manager/fsm_manager.rs b/src/sequence_manager/fsm_manager.rs new file mode 100644 index 0000000..283e879 --- /dev/null +++ b/src/sequence_manager/fsm_manager.rs @@ -0,0 +1,341 @@ +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; +use super::events::*; +use super::states::*; +use super::fsm_events::*; +use std::collections::{HashMap, VecDeque}; +use tracing::{debug, info, warn}; + +/// Sequence Manager with FSM-based state transitions +pub struct SequenceManager { + block_manager: BlockManager, + + // Each sequence is a state machine + sequences: HashMap, + + tokens_per_block: usize, + + // Event-driven + event_rx: SequenceEventReceiver, + event_tx: SequenceEventSender, + + // Scheduling queues + waiting_queue: VecDeque, +} + +impl SequenceManager { + pub fn new( + block_manager: BlockManager, + tokens_per_block: usize, + ) -> (Self, SequenceEventSender) { + let (event_tx, event_rx) = create_event_channel(); + + let manager = Self { + block_manager, + sequences: HashMap::new(), + tokens_per_block, + event_rx, + event_tx: event_tx.clone(), + waiting_queue: VecDeque::new(), + }; + + (manager, event_tx) + } + + /// Main event loop + pub async fn run(mut self) { + info!("FSM SequenceManager event loop started"); + + while let Some(event) = self.event_rx.recv().await { + match event { + SequenceEvent::AddRequest { + seq_id, + prompt_tokens, + max_tokens, + } => { + self.handle_add_request(seq_id, prompt_tokens, max_tokens); + } + + SequenceEvent::AppendTokens { seq_id, num_tokens } => { + self.handle_append_tokens(seq_id, num_tokens); + } + + SequenceEvent::ForkSequence { + parent_id, + child_id, + } => { + self.handle_fork_sequence(parent_id, child_id); + } + + SequenceEvent::CompleteSequence { seq_id } => { + self.handle_complete_sequence(seq_id); + } + + SequenceEvent::PreemptSequence { seq_id } => { + self.handle_preempt_sequence(seq_id); + } + + SequenceEvent::ResumeSequence { seq_id } => { + self.handle_resume_sequence(seq_id); + } + + SequenceEvent::GetBlocks { seq_id, response } => { + let result = self.get_blocks(seq_id); + let _ = response.send(result); + } + + SequenceEvent::GetStats { response } => { + let stats = self.get_stats(); + let _ = response.send(stats); + } + } + } + + info!("FSM SequenceManager event loop stopped"); + } + + fn handle_add_request( + &mut self, + seq_id: SequenceId, + prompt_tokens: usize, + max_tokens: usize, + ) { + debug!( + "Adding request: seq_id={:?}, prompt_tokens={}, max_tokens={}", + seq_id, prompt_tokens, max_tokens + ); + + let state = SequenceState::Waiting(WaitingState { + seq_id, + prompt_tokens, + max_tokens, + }); + + self.sequences.insert(seq_id, state); + self.waiting_queue.push_back(seq_id); + + // Try to schedule immediately + self.try_schedule(); + } + + fn handle_append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let event = AppendTokensEvent { + num_tokens, + block_manager: &mut self.block_manager, + tokens_per_block: self.tokens_per_block, + }; + + match event.apply(state) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + } + Err(Error::OutOfMemory) => { + warn!("OOM while appending tokens to {:?}", seq_id); + self.handle_oom(); + } + Err(e) => { + warn!("Failed to append tokens to {:?}: {:?}", seq_id, e); + } + } + } + + fn handle_fork_sequence(&mut self, parent_id: SequenceId, child_id: SequenceId) { + let parent_state = match self.sequences.remove(&parent_id) { + Some(s) => s, + None => { + warn!("Parent sequence {:?} not found", parent_id); + return; + } + }; + + let event = ForkEvent { + child_id, + block_manager: &mut self.block_manager, + }; + + match event.apply(parent_state) { + Ok((parent_state, child_state)) => { + self.sequences.insert(parent_id, parent_state); + self.sequences.insert(child_id, child_state); + debug!("Forked sequence {:?} → {:?}", parent_id, child_id); + } + Err(e) => { + warn!("Failed to fork {:?}: {:?}", parent_id, e); + } + } + } + + fn handle_complete_sequence(&mut self, seq_id: SequenceId) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let event = CompleteEvent { + reason: FinishReason::Stop, + block_manager: &mut self.block_manager, + }; + + match event.apply(state) { + Ok(new_state) => { + info!("Completed sequence {:?}", seq_id); + // Keep finished state for stats/debugging + self.sequences.insert(seq_id, new_state); + + // Try to schedule waiting sequences + self.try_schedule(); + } + Err(e) => { + warn!("Failed to complete {:?}: {:?}", seq_id, e); + } + } + } + + fn handle_preempt_sequence(&mut self, seq_id: SequenceId) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let event = PreemptEvent { + block_manager: &mut self.block_manager, + }; + + match event.apply(state) { + Ok(new_state) => { + info!("Preempted sequence {:?}", seq_id); + self.sequences.insert(seq_id, new_state); + } + Err(e) => { + warn!("Failed to preempt {:?}: {:?}", seq_id, e); + } + } + } + + fn handle_resume_sequence(&mut self, seq_id: SequenceId) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let event = ResumeEvent; + + match event.apply(state) { + Ok(new_state) => { + debug!("Resumed sequence {:?}", seq_id); + self.sequences.insert(seq_id, new_state); + self.waiting_queue.push_back(seq_id); + self.try_schedule(); + } + Err(e) => { + warn!("Failed to resume {:?}: {:?}", seq_id, e); + } + } + } + + fn try_schedule(&mut self) { + while let Some(&seq_id) = self.waiting_queue.front() { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + self.waiting_queue.pop_front(); + continue; + } + }; + + // Only schedule if in Waiting state + if !matches!(state, SequenceState::Waiting(_)) { + self.sequences.insert(seq_id, state); + self.waiting_queue.pop_front(); + continue; + } + + let event = ScheduleEvent { + block_manager: &mut self.block_manager, + tokens_per_block: self.tokens_per_block, + }; + + match event.apply(state) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + self.waiting_queue.pop_front(); + info!("Scheduled sequence {:?}", seq_id); + } + Err(Error::OutOfMemory) => { + // Can't schedule - leave in waiting queue + // Put state back + if let Some(old_state) = self.sequences.get(&seq_id) { + // State was already removed, this shouldn't happen + // but handle gracefully + } + break; + } + Err(e) => { + warn!("Failed to schedule {:?}: {:?}", seq_id, e); + self.waiting_queue.pop_front(); + } + } + } + } + + fn handle_oom(&mut self) { + warn!("Handling OOM - preemption policy not yet implemented"); + // TODO: Implement preemption policy + // - Find lowest priority sequence + // - Preempt it + // - Retry scheduling + } + + fn get_blocks(&self, seq_id: SequenceId) -> Result> { + self.sequences + .get(&seq_id) + .and_then(|state| state.blocks()) + .map(|blocks| blocks.to_vec()) + .ok_or(Error::UnknownSequence(seq_id)) + } + + fn get_stats(&self) -> SequenceManagerStats { + let num_running = self + .sequences + .values() + .filter(|s| s.is_running()) + .count(); + + let num_waiting = self + .sequences + .values() + .filter(|s| s.is_waiting()) + .count(); + + let num_preempted = self + .sequences + .values() + .filter(|s| matches!(s, SequenceState::Preempted(_))) + .count(); + + SequenceManagerStats { + num_sequences: self.sequences.len(), + num_running, + num_waiting, + num_preempted, + block_stats: self.block_manager.get_stats(), + } + } +} diff --git a/src/sequence_manager/id_generator.rs b/src/sequence_manager/id_generator.rs new file mode 100644 index 0000000..ee9e490 --- /dev/null +++ b/src/sequence_manager/id_generator.rs @@ -0,0 +1,70 @@ +use crate::block_manager::types::SequenceId; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Sequential ID generator for sequences +/// Thread-safe, starts from 0 +pub struct SequenceIdGenerator { + next_id: AtomicU64, +} + +impl SequenceIdGenerator { + /// Create a new ID generator starting from 0 + pub fn new() -> Self { + Self { + next_id: AtomicU64::new(0), + } + } + + /// Create an ID generator starting from a specific value + pub fn with_start(start: u64) -> Self { + Self { + next_id: AtomicU64::new(start), + } + } + + /// Generate the next sequential ID + /// Thread-safe, can be called concurrently + pub fn next(&self) -> SequenceId { + SequenceId(self.next_id.fetch_add(1, Ordering::Relaxed)) + } + + /// Peek at the next ID without consuming it + pub fn peek(&self) -> SequenceId { + SequenceId(self.next_id.load(Ordering::Relaxed)) + } +} + +impl Default for SequenceIdGenerator { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sequential_generation() { + let gen = SequenceIdGenerator::new(); + assert_eq!(gen.next().0, 0); + assert_eq!(gen.next().0, 1); + assert_eq!(gen.next().0, 2); + } + + #[test] + fn test_with_start() { + let gen = SequenceIdGenerator::with_start(100); + assert_eq!(gen.next().0, 100); + assert_eq!(gen.next().0, 101); + } + + #[test] + fn test_peek() { + let gen = SequenceIdGenerator::new(); + assert_eq!(gen.peek().0, 0); + assert_eq!(gen.peek().0, 0); // Doesn't consume + assert_eq!(gen.next().0, 0); // Now consume + assert_eq!(gen.peek().0, 1); + } +} diff --git a/src/sequence_manager/mod.rs b/src/sequence_manager/mod.rs new file mode 100644 index 0000000..83dcb8f --- /dev/null +++ b/src/sequence_manager/mod.rs @@ -0,0 +1,13 @@ +pub mod events; +pub mod states; +pub mod fsm_events; +pub mod fsm_manager; +pub mod id_generator; + +pub use events::*; +pub use states::*; +pub use fsm_manager::SequenceManager; +pub use id_generator::SequenceIdGenerator; + +// Re-export from block_manager for convenience +pub use crate::block_manager::{BlockId, SequenceId}; diff --git a/src/sequence_manager/states.rs b/src/sequence_manager/states.rs new file mode 100644 index 0000000..cb6b10d --- /dev/null +++ b/src/sequence_manager/states.rs @@ -0,0 +1,144 @@ +use crate::block_manager::types::*; + +/// FSM States - each state owns its resources +/// +/// State transitions: +/// Waiting → Scheduling → Prefilling → Decoding → Finished +/// ↓ ↓ +/// └→ Preempted ←┘ + +#[derive(Debug)] +pub enum SequenceState { + /// Waiting in queue for scheduling + Waiting(WaitingState), + + /// Being scheduled (allocating resources) + Scheduling(SchedulingState), + + /// Running prefill phase + Prefilling(PrefillingState), + + /// Running decode phase (generating tokens) + Decoding(DecodingState), + + /// Preempted due to OOM + Preempted(PreemptedState), + + /// Finished generation + Finished(FinishedState), + + /// Aborted (error or cancelled) + Aborted(AbortedState), +} + +/// Waiting state - no resources allocated yet +#[derive(Debug)] +pub struct WaitingState { + pub seq_id: SequenceId, + pub prompt_tokens: usize, + pub max_tokens: usize, +} + +/// Scheduling state - in process of allocating blocks +#[derive(Debug)] +pub struct SchedulingState { + pub seq_id: SequenceId, + pub prompt_tokens: usize, + pub max_tokens: usize, + // Partially allocated blocks (if any) + pub blocks: Vec, +} + +/// Prefilling state - processing prompt +#[derive(Debug)] +pub struct PrefillingState { + pub seq_id: SequenceId, + pub blocks: Vec, + pub tokens_filled: usize, + pub tokens_total: usize, + pub max_tokens: usize, +} + +/// Decoding state - generating tokens +#[derive(Debug)] +pub struct DecodingState { + pub seq_id: SequenceId, + pub blocks: Vec, + pub num_tokens: usize, + pub max_tokens: usize, +} + +/// Preempted state - blocks freed, waiting to resume +#[derive(Debug)] +pub struct PreemptedState { + pub seq_id: SequenceId, + pub num_tokens: usize, + pub max_tokens: usize, + // No blocks - they were freed +} + +/// Finished state - all resources freed +#[derive(Debug)] +pub struct FinishedState { + pub seq_id: SequenceId, + pub finish_reason: FinishReason, +} + +/// Aborted state - error or cancellation +#[derive(Debug)] +pub struct AbortedState { + pub seq_id: SequenceId, + pub reason: String, +} + +#[derive(Debug, Clone)] +pub enum FinishReason { + /// Reached max_tokens limit + MaxTokens, + /// Natural completion (EOS token) + Stop, + /// User cancelled + Cancelled, +} + +// Helper methods for SequenceState +impl SequenceState { + pub fn seq_id(&self) -> SequenceId { + match self { + SequenceState::Waiting(s) => s.seq_id, + SequenceState::Scheduling(s) => s.seq_id, + SequenceState::Prefilling(s) => s.seq_id, + SequenceState::Decoding(s) => s.seq_id, + SequenceState::Preempted(s) => s.seq_id, + SequenceState::Finished(s) => s.seq_id, + SequenceState::Aborted(s) => s.seq_id, + } + } + + pub fn is_running(&self) -> bool { + matches!( + self, + SequenceState::Prefilling(_) | SequenceState::Decoding(_) + ) + } + + pub fn is_waiting(&self) -> bool { + matches!( + self, + SequenceState::Waiting(_) | SequenceState::Scheduling(_) + ) + } + + pub fn is_finished(&self) -> bool { + matches!(self, SequenceState::Finished(_) | SequenceState::Aborted(_)) + } + + pub fn blocks(&self) -> Option<&[BlockId]> { + match self { + SequenceState::Prefilling(s) => Some(&s.blocks), + SequenceState::Decoding(s) => Some(&s.blocks), + SequenceState::Scheduling(s) if !s.blocks.is_empty() => Some(&s.blocks), + _ => None, + } + } +} From ba1a17b8e2ebd0770158dc05b98d1ff29b1cd749 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 5 Jul 2026 17:43:33 +0100 Subject: [PATCH 02/33] rebase Main Signed-off-by: kerthcet --- examples/fsm_usage.rs | 6 +-- src/block_manager/allocator.rs | 5 ++- src/block_manager/manager.rs | 66 +++++++++++++++++------------ src/block_manager/mod.rs | 3 -- src/sequence_manager/fsm_events.rs | 55 ++++++++++++------------ src/sequence_manager/fsm_manager.rs | 27 +++--------- src/sequence_manager/mod.rs | 8 +--- 7 files changed, 82 insertions(+), 88 deletions(-) diff --git a/examples/fsm_usage.rs b/examples/fsm_usage.rs index 38db0fa..5271291 100644 --- a/examples/fsm_usage.rs +++ b/examples/fsm_usage.rs @@ -1,5 +1,5 @@ use puma::block_manager::{BlockManager, CpuAllocator}; -use puma::sequence_manager::{SequenceManager, SequenceEvent, SequenceIdGenerator}; +use puma::sequence_manager::{SequenceEvent, SequenceIdGenerator, SequenceManager}; #[tokio::main] async fn main() { @@ -113,9 +113,7 @@ async fn main() { .send(SequenceEvent::CompleteSequence { seq_id }) .unwrap(); event_tx - .send(SequenceEvent::CompleteSequence { - seq_id: child_id, - }) + .send(SequenceEvent::CompleteSequence { seq_id: child_id }) .unwrap(); tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; diff --git a/src/block_manager/allocator.rs b/src/block_manager/allocator.rs index 5a2c9b7..b048e38 100644 --- a/src/block_manager/allocator.rs +++ b/src/block_manager/allocator.rs @@ -41,7 +41,10 @@ impl MemoryAllocator for CpuAllocator { self.used_memory += size_bytes; - Ok(MemoryAddress { ptr, size: size_bytes }) + Ok(MemoryAddress { + ptr, + size: size_bytes, + }) } fn free(&mut self, addr: MemoryAddress) -> Result<()> { diff --git a/src/block_manager/manager.rs b/src/block_manager/manager.rs index 20967cb..95321a7 100644 --- a/src/block_manager/manager.rs +++ b/src/block_manager/manager.rs @@ -22,10 +22,7 @@ pub struct BlockManager { } impl BlockManager { - pub fn new( - allocator: Box, - default_block_size: usize, - ) -> Self { + pub fn new(allocator: Box, default_block_size: usize) -> Self { let mut manager = Self { allocator, block_table: HashMap::new(), @@ -45,13 +42,14 @@ impl BlockManager { } pub fn register_block_type(&mut self, config: BlockConfig) { - self.free_pools.insert(config.block_type.clone(), Vec::new()); - self.block_configs.insert(config.block_type.clone(), config); + self.free_pools.insert(config.block_type, Vec::new()); + self.block_configs.insert(config.block_type, config); } pub fn allocate_typed(&mut self, block_type: &BlockType) -> Result { // Try free pool first - if let Some(block_id) = self.free_pools + if let Some(block_id) = self + .free_pools .get_mut(block_type) .and_then(|pool| pool.pop()) { @@ -61,30 +59,37 @@ impl BlockManager { } // Allocate new physical memory - let config = self.block_configs.get(block_type) + let config = self + .block_configs + .get(block_type) .ok_or_else(|| Error::UnknownBlockType(block_type.as_str().to_string()))?; let mem_addr = self.allocator.allocate(config.size_bytes)?; let block_id = BlockId(self.next_block_id.fetch_add(1, Ordering::Relaxed)); - self.block_table.insert(block_id, Block { + self.block_table.insert( block_id, - block_type: block_type.clone(), - mem_addr, - ref_count: 1, - }); + Block { + block_id, + block_type: *block_type, + mem_addr, + ref_count: 1, + }, + ); Ok(block_id) } pub fn allocate(&mut self) -> Result { - let default_type = self.default_block_type.clone(); + let default_type = self.default_block_type; self.allocate_typed(&default_type) } pub fn free(&mut self, block_id: BlockId) -> Result<()> { - let block = self.block_table.get_mut(&block_id) + let block = self + .block_table + .get_mut(&block_id) .ok_or(Error::InvalidBlockId(block_id))?; if block.ref_count == 0 { @@ -104,20 +109,24 @@ impl BlockManager { } pub fn add_ref(&mut self, block_id: BlockId) -> Result<()> { - let block = self.block_table.get_mut(&block_id) + let block = self + .block_table + .get_mut(&block_id) .ok_or(Error::InvalidBlockId(block_id))?; block.ref_count += 1; Ok(()) } pub fn get_memory_address(&self, block_id: BlockId) -> Result { - self.block_table.get(&block_id) + self.block_table + .get(&block_id) .map(|b| b.mem_addr) .ok_or(Error::InvalidBlockId(block_id)) } pub fn get_block_type(&self, block_id: BlockId) -> Result<&BlockType> { - self.block_table.get(&block_id) + self.block_table + .get(&block_id) .map(|b| &b.block_type) .ok_or(Error::InvalidBlockId(block_id)) } @@ -129,17 +138,22 @@ impl BlockManager { let config = &self.block_configs[type_id]; let free_count = pool.len(); - let allocated_count = self.block_table.values() + let allocated_count = self + .block_table + .values() .filter(|b| b.block_type == *type_id && b.ref_count > 0) .count(); - stats.insert(type_id.clone(), BlockStats { - total_blocks: free_count + allocated_count, - allocated_blocks: allocated_count, - free_blocks: free_count, - block_size: config.size_bytes, - total_memory: (free_count + allocated_count) * config.size_bytes, - }); + stats.insert( + *type_id, + BlockStats { + total_blocks: free_count + allocated_count, + allocated_blocks: allocated_count, + free_blocks: free_count, + block_size: config.size_bytes, + total_memory: (free_count + allocated_count) * config.size_bytes, + }, + ); } stats diff --git a/src/block_manager/mod.rs b/src/block_manager/mod.rs index 2d08c80..7425962 100644 --- a/src/block_manager/mod.rs +++ b/src/block_manager/mod.rs @@ -2,6 +2,3 @@ pub mod allocator; pub mod manager; pub mod types; -pub use allocator::*; -pub use manager::BlockManager; -pub use types::*; diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs index 0ffae9c..66d606b 100644 --- a/src/sequence_manager/fsm_events.rs +++ b/src/sequence_manager/fsm_events.rs @@ -1,6 +1,6 @@ +use super::states::*; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::*; -use super::states::*; use tracing::{debug, warn}; /// FSM Events - transform states with type-safe transitions @@ -18,13 +18,14 @@ impl<'a> ScheduleEvent<'a> { pub fn apply(self, state: SequenceState) -> Result { match state { SequenceState::Waiting(s) => self.from_waiting(s), - _ => Err(Error::invalid_transition("ScheduleEvent requires Waiting state")), + _ => Err(Error::invalid_transition( + "ScheduleEvent requires Waiting state", + )), } } fn from_waiting(self, state: WaitingState) -> Result { - let blocks_needed = - (state.prompt_tokens + self.tokens_per_block - 1) / self.tokens_per_block; + let blocks_needed = state.prompt_tokens.div_ceil(self.tokens_per_block); let mut blocks = Vec::new(); for _ in 0..blocks_needed { @@ -49,7 +50,9 @@ impl<'a> ScheduleEvent<'a> { debug!( "Scheduled seq {:?}: allocated {} blocks for {} tokens", - state.seq_id, blocks.len(), state.prompt_tokens + state.seq_id, + blocks.len(), + state.prompt_tokens ); Ok(SequenceState::Prefilling(PrefillingState { @@ -74,7 +77,9 @@ impl<'a> AppendTokensEvent<'a> { match state { SequenceState::Prefilling(s) => self.from_prefilling(s), SequenceState::Decoding(s) => self.from_decoding(s), - _ => Err(Error::invalid_transition("AppendTokensEvent requires Prefilling or Decoding state")), + _ => Err(Error::invalid_transition( + "AppendTokensEvent requires Prefilling or Decoding state", + )), } } @@ -122,7 +127,7 @@ impl<'a> AppendTokensEvent<'a> { } // Check if need more blocks - let blocks_needed = (state.num_tokens + self.tokens_per_block - 1) / self.tokens_per_block; + let blocks_needed = state.num_tokens.div_ceil(self.tokens_per_block); while state.blocks.len() < blocks_needed { match self.block_manager.allocate() { @@ -156,7 +161,9 @@ impl<'a> PreemptEvent<'a> { match state { SequenceState::Decoding(s) => self.from_decoding(s), SequenceState::Prefilling(s) => self.from_prefilling(s), - _ => Err(Error::invalid_transition("PreemptEvent requires running state")), + _ => Err(Error::invalid_transition( + "PreemptEvent requires running state", + )), } } @@ -198,7 +205,9 @@ impl ResumeEvent { pub fn apply(self, state: SequenceState) -> Result { match state { SequenceState::Preempted(s) => self.from_preempted(s), - _ => Err(Error::invalid_transition("ResumeEvent requires Preempted state")), + _ => Err(Error::invalid_transition( + "ResumeEvent requires Preempted state", + )), } } @@ -221,20 +230,16 @@ pub struct ForkEvent<'a> { } impl<'a> ForkEvent<'a> { - pub fn apply( - self, - state: SequenceState, - ) -> Result<(SequenceState, SequenceState)> { + pub fn apply(self, state: SequenceState) -> Result<(SequenceState, SequenceState)> { match state { SequenceState::Decoding(s) => self.from_decoding(s), - _ => Err(Error::invalid_transition("ForkEvent requires Decoding state")), + _ => Err(Error::invalid_transition( + "ForkEvent requires Decoding state", + )), } } - fn from_decoding( - self, - state: DecodingState, - ) -> Result<(SequenceState, SequenceState)> { + fn from_decoding(self, state: DecodingState) -> Result<(SequenceState, SequenceState)> { // Copy-on-write: increment ref counts for &block_id in &state.blocks { self.block_manager.add_ref(block_id)?; @@ -247,10 +252,7 @@ impl<'a> ForkEvent<'a> { max_tokens: state.max_tokens, }; - debug!( - "Forked seq {:?} → {:?}", - state.seq_id, child.seq_id - ); + debug!("Forked seq {:?} → {:?}", state.seq_id, child.seq_id); Ok(( SequenceState::Decoding(state), @@ -270,7 +272,9 @@ impl<'a> CompleteEvent<'a> { match state { SequenceState::Decoding(s) => self.from_decoding(s), SequenceState::Prefilling(s) => self.from_prefilling(s), - _ => Err(Error::invalid_transition("CompleteEvent requires running state")), + _ => Err(Error::invalid_transition( + "CompleteEvent requires running state", + )), } } @@ -280,10 +284,7 @@ impl<'a> CompleteEvent<'a> { self.block_manager.free(block_id)?; } - debug!( - "Completed seq {:?}: {:?}", - state.seq_id, self.reason - ); + debug!("Completed seq {:?}: {:?}", state.seq_id, self.reason); Ok(SequenceState::Finished(FinishedState { seq_id: state.seq_id, diff --git a/src/sequence_manager/fsm_manager.rs b/src/sequence_manager/fsm_manager.rs index 283e879..fabf74e 100644 --- a/src/sequence_manager/fsm_manager.rs +++ b/src/sequence_manager/fsm_manager.rs @@ -1,8 +1,8 @@ -use crate::block_manager::manager::BlockManager; -use crate::block_manager::types::*; use super::events::*; -use super::states::*; use super::fsm_events::*; +use super::states::*; +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; use std::collections::{HashMap, VecDeque}; use tracing::{debug, info, warn}; @@ -94,12 +94,7 @@ impl SequenceManager { info!("FSM SequenceManager event loop stopped"); } - fn handle_add_request( - &mut self, - seq_id: SequenceId, - prompt_tokens: usize, - max_tokens: usize, - ) { + fn handle_add_request(&mut self, seq_id: SequenceId, prompt_tokens: usize, max_tokens: usize) { debug!( "Adding request: seq_id={:?}, prompt_tokens={}, max_tokens={}", seq_id, prompt_tokens, max_tokens @@ -281,7 +276,7 @@ impl SequenceManager { Err(Error::OutOfMemory) => { // Can't schedule - leave in waiting queue // Put state back - if let Some(old_state) = self.sequences.get(&seq_id) { + if let Some(_old_state) = self.sequences.get(&seq_id) { // State was already removed, this shouldn't happen // but handle gracefully } @@ -312,17 +307,9 @@ impl SequenceManager { } fn get_stats(&self) -> SequenceManagerStats { - let num_running = self - .sequences - .values() - .filter(|s| s.is_running()) - .count(); + let num_running = self.sequences.values().filter(|s| s.is_running()).count(); - let num_waiting = self - .sequences - .values() - .filter(|s| s.is_waiting()) - .count(); + let num_waiting = self.sequences.values().filter(|s| s.is_waiting()).count(); let num_preempted = self .sequences diff --git a/src/sequence_manager/mod.rs b/src/sequence_manager/mod.rs index 83dcb8f..0e0f032 100644 --- a/src/sequence_manager/mod.rs +++ b/src/sequence_manager/mod.rs @@ -1,13 +1,7 @@ pub mod events; -pub mod states; pub mod fsm_events; pub mod fsm_manager; pub mod id_generator; - -pub use events::*; -pub use states::*; -pub use fsm_manager::SequenceManager; -pub use id_generator::SequenceIdGenerator; +pub mod states; // Re-export from block_manager for convenience -pub use crate::block_manager::{BlockId, SequenceId}; From b121fa69d3ef67b4c14a008012b3d2698c5938c8 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 5 Jul 2026 17:46:30 +0100 Subject: [PATCH 03/33] fix lint Signed-off-by: kerthcet --- src/block_manager/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/block_manager/mod.rs b/src/block_manager/mod.rs index 7425962..97a6594 100644 --- a/src/block_manager/mod.rs +++ b/src/block_manager/mod.rs @@ -1,4 +1,3 @@ pub mod allocator; pub mod manager; pub mod types; - From 9390d3521ad33dc28859b23f61d755e84d432f05 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 00:05:56 +0100 Subject: [PATCH 04/33] add backup state Signed-off-by: kerthcet --- docs/fsm_architecture.md | 56 ++++++++++++++++++++++++++--- src/block_manager/manager.rs | 14 ++++++++ src/sequence_manager/fsm_manager.rs | 21 +++++++---- src/sequence_manager/states.rs | 16 ++++----- 4 files changed, 88 insertions(+), 19 deletions(-) diff --git a/docs/fsm_architecture.md b/docs/fsm_architecture.md index 5b7896c..01a3506 100644 --- a/docs/fsm_architecture.md +++ b/docs/fsm_architecture.md @@ -11,10 +11,12 @@ Each sequence is an independent state machine: ```rust pub enum SequenceState { Waiting(WaitingState), + Scheduling(SchedulingState), Prefilling(PrefillingState), Decoding(DecodingState), Preempted(PreemptedState), Finished(FinishedState), + Aborted(AbortedState), } // Multiple sequences, different states @@ -244,6 +246,9 @@ impl SequenceManager { SequenceEvent::AppendTokens { seq_id, num_tokens } => { // Get current state let state = self.sequences.remove(&seq_id).unwrap(); + + // Clone state before transition (prevents loss on failure) + let backup = state.clone(); // Apply event transformation let event = AppendTokensEvent { @@ -259,6 +264,8 @@ impl SequenceManager { } Err(e) => { error!("Transition failed: {:?}", e); + // Restore original state to prevent sequence loss + self.sequences.insert(seq_id, backup); } } } @@ -364,29 +371,68 @@ let state = event.apply(parent)?; | Manual cleanup | Automatic cleanup | | Runtime checks | Type + runtime checks | +## Error Handling and State Recovery + +On transition failure, the state is cloned before applying the event: + +```rust +fn handle_append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { + let state = self.sequences.remove(&seq_id).unwrap(); + let backup = state.clone(); // Clone before transition + + let event = AppendTokensEvent { ... }; + + match event.apply(state) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + } + Err(Error::OutOfMemory) => { + // Restore state and trigger preemption + self.sequences.insert(seq_id, backup); + self.handle_oom(); + } + Err(e) => { + // Restore state on any error (prevents block leaks) + self.sequences.insert(seq_id, backup); + warn!("Transition failed: {:?}", e); + } + } +} +``` + +**Benefits:** +- No sequence loss on transition failure +- No block leaks (blocks remain owned by restored state) +- Can retry or handle errors without losing context +- Clone cost is negligible (~50-100 bytes per state) + ## Usage ```rust use puma::block_manager::{BlockManager, CpuAllocator}; -use puma::sequence_manager::{FSMSequenceManager, SequenceEvent, SequenceId}; +use puma::sequence_manager::{SequenceManager, SequenceEvent, SequenceIdGenerator}; // Setup let allocator = Box::new(CpuAllocator::new(100_000_000)); let block_manager = BlockManager::new(allocator, 4096); -let (fsm_manager, event_tx) = FSMSequenceManager::new(block_manager, 16); +let (seq_manager, event_tx) = SequenceManager::new(block_manager, 16); + +// ID generator for sequential IDs starting from 0 +let id_gen = SequenceIdGenerator::new(); +let seq_id = id_gen.next(); // SequenceId(0) // Run event loop -tokio::spawn(async move { fsm_manager.run().await }); +tokio::spawn(async move { seq_manager.run().await }); // Send events event_tx.send(SequenceEvent::AddRequest { - seq_id: SequenceId(1), + seq_id, prompt_tokens: 100, max_tokens: 150, }).unwrap(); event_tx.send(SequenceEvent::AppendTokens { - seq_id: SequenceId(1), + seq_id, num_tokens: 100, }).unwrap(); ``` diff --git a/src/block_manager/manager.rs b/src/block_manager/manager.rs index 95321a7..449236e 100644 --- a/src/block_manager/manager.rs +++ b/src/block_manager/manager.rs @@ -173,3 +173,17 @@ impl BlockManager { false } } + +impl Drop for BlockManager { + fn drop(&mut self) { + // Free all backing memory allocations to prevent leaks + for block in self.block_table.values() { + if let Err(e) = self.allocator.free(block.mem_addr) { + eprintln!( + "Warning: failed to free block {:?} memory: {:?}", + block.block_id, e + ); + } + } + } +} diff --git a/src/sequence_manager/fsm_manager.rs b/src/sequence_manager/fsm_manager.rs index fabf74e..95aa46d 100644 --- a/src/sequence_manager/fsm_manager.rs +++ b/src/sequence_manager/fsm_manager.rs @@ -122,6 +122,7 @@ impl SequenceManager { } }; + let backup = state.clone(); let event = AppendTokensEvent { num_tokens, block_manager: &mut self.block_manager, @@ -134,10 +135,12 @@ impl SequenceManager { } Err(Error::OutOfMemory) => { warn!("OOM while appending tokens to {:?}", seq_id); + self.sequences.insert(seq_id, backup); self.handle_oom(); } Err(e) => { warn!("Failed to append tokens to {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); } } } @@ -151,6 +154,7 @@ impl SequenceManager { } }; + let backup = parent_state.clone(); let event = ForkEvent { child_id, block_manager: &mut self.block_manager, @@ -164,6 +168,7 @@ impl SequenceManager { } Err(e) => { warn!("Failed to fork {:?}: {:?}", parent_id, e); + self.sequences.insert(parent_id, backup); } } } @@ -177,6 +182,7 @@ impl SequenceManager { } }; + let backup = state.clone(); let event = CompleteEvent { reason: FinishReason::Stop, block_manager: &mut self.block_manager, @@ -193,6 +199,7 @@ impl SequenceManager { } Err(e) => { warn!("Failed to complete {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); } } } @@ -206,6 +213,7 @@ impl SequenceManager { } }; + let backup = state.clone(); let event = PreemptEvent { block_manager: &mut self.block_manager, }; @@ -217,6 +225,7 @@ impl SequenceManager { } Err(e) => { warn!("Failed to preempt {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); } } } @@ -230,6 +239,7 @@ impl SequenceManager { } }; + let backup = state.clone(); let event = ResumeEvent; match event.apply(state) { @@ -241,6 +251,7 @@ impl SequenceManager { } Err(e) => { warn!("Failed to resume {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); } } } @@ -262,6 +273,7 @@ impl SequenceManager { continue; } + let backup = state.clone(); let event = ScheduleEvent { block_manager: &mut self.block_manager, tokens_per_block: self.tokens_per_block, @@ -274,16 +286,13 @@ impl SequenceManager { info!("Scheduled sequence {:?}", seq_id); } Err(Error::OutOfMemory) => { - // Can't schedule - leave in waiting queue - // Put state back - if let Some(_old_state) = self.sequences.get(&seq_id) { - // State was already removed, this shouldn't happen - // but handle gracefully - } + // Can't schedule - restore state and leave in waiting queue + self.sequences.insert(seq_id, backup); break; } Err(e) => { warn!("Failed to schedule {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); self.waiting_queue.pop_front(); } } diff --git a/src/sequence_manager/states.rs b/src/sequence_manager/states.rs index cb6b10d..6543d83 100644 --- a/src/sequence_manager/states.rs +++ b/src/sequence_manager/states.rs @@ -7,7 +7,7 @@ use crate::block_manager::types::*; /// ↓ ↓ /// └→ Preempted ←┘ -#[derive(Debug)] +#[derive(Debug, Clone)] pub enum SequenceState { /// Waiting in queue for scheduling Waiting(WaitingState), @@ -32,7 +32,7 @@ pub enum SequenceState { } /// Waiting state - no resources allocated yet -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct WaitingState { pub seq_id: SequenceId, pub prompt_tokens: usize, @@ -40,7 +40,7 @@ pub struct WaitingState { } /// Scheduling state - in process of allocating blocks -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct SchedulingState { pub seq_id: SequenceId, pub prompt_tokens: usize, @@ -50,7 +50,7 @@ pub struct SchedulingState { } /// Prefilling state - processing prompt -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct PrefillingState { pub seq_id: SequenceId, pub blocks: Vec, @@ -60,7 +60,7 @@ pub struct PrefillingState { } /// Decoding state - generating tokens -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct DecodingState { pub seq_id: SequenceId, pub blocks: Vec, @@ -69,7 +69,7 @@ pub struct DecodingState { } /// Preempted state - blocks freed, waiting to resume -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct PreemptedState { pub seq_id: SequenceId, pub num_tokens: usize, @@ -78,14 +78,14 @@ pub struct PreemptedState { } /// Finished state - all resources freed -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct FinishedState { pub seq_id: SequenceId, pub finish_reason: FinishReason, } /// Aborted state - error or cancellation -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct AbortedState { pub seq_id: SequenceId, pub reason: String, From d4841ca8cc9902d006e847f33c737dbee2f96a97 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 00:20:29 +0100 Subject: [PATCH 05/33] optimize log Signed-off-by: kerthcet --- src/cli/serve.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cli/serve.rs b/src/cli/serve.rs index c9f7a31..5b44ca6 100644 --- a/src/cli/serve.rs +++ b/src/cli/serve.rs @@ -54,7 +54,7 @@ pub async fn execute( info!(" GET /health"); // Start server - debug!("Starting axum server"); + debug!("Starting server"); axum::serve(listener, app).await?; info!("Server shutdown"); From eaceb5fad77498f7a6d6bfc3d46202cfe5c4a9ec Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 01:05:03 +0100 Subject: [PATCH 06/33] optimize the prompt Signed-off-by: kerthcet --- Cargo.lock | 119 ++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 2 + src/cli/chat.rs | 118 +++++++++++++++++++++++++++++++++++++++++++ src/cli/commands.rs | 30 +++++------ src/cli/mod.rs | 1 + 5 files changed, 256 insertions(+), 14 deletions(-) create mode 100644 src/cli/chat.rs diff --git a/Cargo.lock b/Cargo.lock index cc200dc..e9b5cbd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -224,6 +224,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd16c4719339c4530435d38e511904438d07cce7950afa3718a84ac36c10e89e" + [[package]] name = "chrono" version = "0.4.44" @@ -277,6 +283,15 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" +[[package]] +name = "clipboard-win" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bde03770d3df201d4fb868f2c9c59e66a3e4e2bd06692a0fe701e7103c7e84d4" +dependencies = [ + "error-code", +] + [[package]] name = "colorchoice" version = "1.0.5" @@ -517,6 +532,12 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "endian-type" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d" + [[package]] name = "equivalent" version = "1.0.2" @@ -533,6 +554,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "error-code" +version = "3.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59" + [[package]] name = "fallible-iterator" version = "0.3.0" @@ -551,6 +578,17 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +[[package]] +name = "fd-lock" +version = "4.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce92ff622d6dadf7349484f42c93271a0d49b7cc4d466a936405bacbe10aa78" +dependencies = [ + "cfg-if", + "rustix", + "windows-sys 0.59.0", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -815,6 +853,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "http" version = "1.4.0" @@ -1283,6 +1330,27 @@ dependencies = [ "tempfile", ] +[[package]] +name = "nibble_vec" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43" +dependencies = [ + "smallvec", +] + +[[package]] +name = "nix" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab2156c4fce2f8df6c499cc1c763e4394b7482525bf2a9701c9d79d215f519e4" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "ntapi" version = "0.4.3" @@ -1538,6 +1606,8 @@ dependencies = [ "reqwest", "rusqlite", "rusqlite_migration", + "rustyline", + "rustyline-derive", "serde", "serde_derive", "serde_json", @@ -1574,6 +1644,16 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "radix_trie" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd" +dependencies = [ + "endian-type", + "nibble_vec", +] + [[package]] name = "rand" version = "0.9.4" @@ -1818,6 +1898,39 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +[[package]] +name = "rustyline" +version = "14.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7803e8936da37efd9b6d4478277f4b2b9bb5cdb37a113e8d63222e58da647e63" +dependencies = [ + "bitflags", + "cfg-if", + "clipboard-win", + "fd-lock", + "home", + "libc", + "log", + "memchr", + "nix", + "radix_trie", + "unicode-segmentation", + "unicode-width 0.1.14", + "utf8parse", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustyline-derive" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5af959c8bf6af1aff6d2b463a57f71aae53d1332da58419e30ad8dc7011d951" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "ryu" version = "1.0.23" @@ -2420,6 +2533,12 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + [[package]] name = "unicode-width" version = "0.1.14" diff --git a/Cargo.toml b/Cargo.toml index 46b9364..595935b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,6 +40,8 @@ tower-http = { version = "0.5", features = ["cors", "trace"] } uuid = { version = "1.0", features = ["v4", "serde"] } futures = "0.3" tokio-stream = "0.1" +rustyline = "14.0" +rustyline-derive = "0.10" [dev-dependencies] tempfile = "3.12" diff --git a/src/cli/chat.rs b/src/cli/chat.rs new file mode 100644 index 0000000..f9f52e8 --- /dev/null +++ b/src/cli/chat.rs @@ -0,0 +1,118 @@ +use colored::Colorize; +use rustyline::completion::Completer; +use rustyline::error::ReadlineError; +use rustyline::highlight::Highlighter; +use rustyline::hint::{Hint, Hinter}; +use rustyline::validate::Validator; +use rustyline::{Context, Editor}; +use rustyline_derive::{Completer, Helper, Highlighter, Validator}; +use std::borrow::Cow; +use std::io::{self, Write}; +use tokio_stream::StreamExt; + +use crate::backend::InferenceEngine; + +#[derive(Clone)] +struct PlaceholderHint { + display: String, +} + +impl Hint for PlaceholderHint { + fn display(&self) -> &str { + &self.display + } + + fn completion(&self) -> Option<&str> { + None + } +} + +/// Hint helper that shows placeholder when input is empty +#[derive(Helper, Completer, Highlighter, Validator)] +struct PlaceholderHinter { + placeholder: String, +} + +impl Hinter for PlaceholderHinter { + type Hint = PlaceholderHint; + + fn hint(&self, line: &str, _pos: usize, _ctx: &Context<'_>) -> Option { + if line.is_empty() { + // Grey/dimmed color: \x1b[2m ... \x1b[0m + Some(PlaceholderHint { + display: format!("\x1b[2m{}\x1b[0m", self.placeholder), + }) + } else { + None + } + } +} + +/// Interactive chat loop for puma run +pub async fn interactive_chat( + engine: &E, + model: &str, +) -> Result<(), io::Error> { + let mut conversation_history = Vec::new(); + + // Setup editor with placeholder hinter + let helper = PlaceholderHinter { + placeholder: "Send a message (Ctrl-C or 'exit' to quit)".to_string(), + }; + let mut rl = Editor::new().map_err(|e| io::Error::new(io::ErrorKind::Other, e))?; + rl.set_helper(Some(helper)); + + loop { + let readline = rl.readline("> "); + + let input = match readline { + Ok(line) => line.trim().to_string(), + Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => { + println!("\nGoodbye!"); + break; + } + Err(err) => { + return Err(io::Error::new(io::ErrorKind::Other, err)); + } + }; + + // Exit commands + if input.is_empty() { + continue; + } + if input == "exit" { + println!("Goodbye!"); + break; + } + + // Add user message to history + conversation_history.push(format!("User: {}", input)); + + // Build prompt from conversation history + let prompt = conversation_history.join("\n") + "\nAssistant:"; + + // Generate response with streaming + match engine.generate_stream(model, &prompt, 512, 0.7).await { + Ok(mut stream) => { + let mut full_response = String::new(); + + // Display tokens as they arrive + while let Some(token) = stream.next().await { + print!("{}", token); + io::stdout().flush()?; + full_response.push_str(&token); + } + + println!("\n"); // Double newline after response + + // Add assistant response to history + conversation_history.push(format!("Assistant: {}", full_response.trim())); + } + Err(e) => { + eprintln!("Error: {}\n", e); + } + } + } + + Ok(()) +} diff --git a/src/cli/commands.rs b/src/cli/commands.rs index 8542ca7..a129308 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -2,7 +2,8 @@ use clap::{Parser, Subcommand}; use colored::Colorize; use prettytable::{format, row, Table}; -use crate::cli::{inspect, ls, rm}; +use crate::backend::mock::MockEngine; +use crate::cli::{chat, inspect, ls, rm}; use crate::downloader::{self, Provider}; use crate::registry::model_registry::ModelRegistry; use crate::system::system_info::SystemInfo; @@ -188,16 +189,9 @@ pub async fn run(cli: Cli) { Commands::RUN(args) => { let registry = ModelRegistry::new(None); - // Check if model exists + // Ensure model exists locally, download if needed match registry.get_model(&args.model) { - Ok(Some(_)) => { - // Model exists, proceed - println!("Running model: {}", args.model); - // TODO: Implement actual model execution - println!("Model execution not yet implemented"); - } Ok(None) => { - // Model not found, download it first println!( "Model {} not found locally. Downloading...", args.model.cyan().bold() @@ -207,16 +201,24 @@ pub async fn run(cli: Cli) { eprintln!("❌ Error: {}", e); std::process::exit(1); } - - // Now run the model - println!("Running model: {}", args.model.cyan().bold()); - // TODO: Implement actual model execution - println!("Model execution not yet implemented"); } Err(e) => { eprintln!("❌ Error checking model: {}", e); std::process::exit(1); } + Ok(Some(_)) => {} + } + + // Load inference engine with the model + // TODO: Replace MockEngine with real engine that loads model files + // Real engine will use: registry.get_model(&args.model)?.metadata.cache.path + let engine = MockEngine::new(); + + // Start interactive chat (no instruction message, just start prompting) + + if let Err(e) = chat::interactive_chat(&engine, &args.model).await { + eprintln!("❌ Chat error: {}", e); + std::process::exit(1); } } diff --git a/src/cli/mod.rs b/src/cli/mod.rs index fbfc92f..2e09052 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -1,3 +1,4 @@ +pub mod chat; pub mod commands; pub mod inspect; pub mod ls; From b8a98cca5ada9e66380cc7945102beac3e357b30 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 01:08:07 +0100 Subject: [PATCH 07/33] optimize the prompt Signed-off-by: kerthcet --- src/cli/chat.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/cli/chat.rs b/src/cli/chat.rs index f9f52e8..3017ac6 100644 --- a/src/cli/chat.rs +++ b/src/cli/chat.rs @@ -68,7 +68,6 @@ pub async fn interactive_chat( let input = match readline { Ok(line) => line.trim().to_string(), Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => { - println!("\nGoodbye!"); break; } Err(err) => { @@ -81,7 +80,6 @@ pub async fn interactive_chat( continue; } if input == "exit" { - println!("Goodbye!"); break; } @@ -91,6 +89,9 @@ pub async fn interactive_chat( // Build prompt from conversation history let prompt = conversation_history.join("\n") + "\nAssistant:"; + // Empty line before response + println!(); + // Generate response with streaming match engine.generate_stream(model, &prompt, 512, 0.7).await { Ok(mut stream) => { From 8b0ccdf0b9ef7c23189a852f8a9e54600c304953 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 20:39:55 +0100 Subject: [PATCH 08/33] add tests Signed-off-by: kerthcet --- Cargo.toml | 2 +- README.md | 2 +- src/block_manager/manager_tests.rs | 132 ++++++++++++ src/block_manager/mod.rs | 5 + src/cli/chat.rs | 9 +- src/sequence_manager/fsm_events.rs | 12 +- src/sequence_manager/fsm_tests.rs | 326 +++++++++++++++++++++++++++++ src/sequence_manager/mod.rs | 5 +- src/sequence_manager/states.rs | 2 +- 9 files changed, 480 insertions(+), 15 deletions(-) create mode 100644 src/block_manager/manager_tests.rs create mode 100644 src/sequence_manager/fsm_tests.rs diff --git a/Cargo.toml b/Cargo.toml index 595935b..162101c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ name = "puma" version = "0.0.3" edition = "2021" -description = "A lightweight, high-performance inference engine for local AI." +description = "A lightweight, high-performance model engine for local AI." license = "Apache-2.0" repository = "https://github.com/InftyAI/PUMA" homepage = "https://github.com/InftyAI/PUMA" diff --git a/README.md b/README.md index 846ad16..b2558c2 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ PUMA Logo -**A lightweight, high-performance inference engine for local AI** +**A lightweight, high-performance model engine for local AI** [![Stability: Active](https://img.shields.io/badge/stability-active-brightgreen.svg)](https://github.com/InftyAI/PUMA) [![Latest Release](https://img.shields.io/github/v/release/InftyAI/PUMA)](https://github.com/InftyAI/PUMA/releases) diff --git a/src/block_manager/manager_tests.rs b/src/block_manager/manager_tests.rs new file mode 100644 index 0000000..57bc219 --- /dev/null +++ b/src/block_manager/manager_tests.rs @@ -0,0 +1,132 @@ +#[cfg(test)] +mod tests { + use super::super::allocator::CpuAllocator; + use super::super::manager::BlockManager; + use super::super::types::*; + + #[test] + fn test_allocate_and_free() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + // Allocate a block + let block_id = manager.allocate().unwrap(); + assert_eq!(block_id, BlockId(0)); + + // Free it + manager.free(block_id).unwrap(); + + // Should be back in free pool + let stats = manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.free_blocks, 1); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_ref_counting() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + let block_id = manager.allocate().unwrap(); + + // Add ref + manager.add_ref(block_id).unwrap(); + + // Free once - should still be allocated (ref_count = 1) + manager.free(block_id).unwrap(); + let stats = manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 1); + assert_eq!(block_stats.free_blocks, 0); + + // Free again - now should be in free pool (ref_count = 0) + manager.free(block_id).unwrap(); + let stats = manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.free_blocks, 1); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_double_free_error() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + let block_id = manager.allocate().unwrap(); + manager.free(block_id).unwrap(); + + // Should fail on double free + let result = manager.free(block_id); + assert!(matches!(result, Err(Error::DoubleFree(_)))); + } + + #[test] + fn test_oom() { + // Small allocator - only enough for 2 blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let mut manager = BlockManager::new(allocator, 1024); + + // First block succeeds + let block1 = manager.allocate().unwrap(); + assert_eq!(block1, BlockId(0)); + + // Second block succeeds + let block2 = manager.allocate().unwrap(); + assert_eq!(block2, BlockId(1)); + + // Third block should fail (OOM) + let result = manager.allocate(); + assert!(matches!(result, Err(Error::OutOfMemory))); + } + + #[test] + fn test_free_pool_reuse() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + // Allocate and free + let block1 = manager.allocate().unwrap(); + manager.free(block1).unwrap(); + + // Next allocation should reuse from free pool (same ID) + let block2 = manager.allocate().unwrap(); + assert_eq!(block1, block2); + } + + #[test] + fn test_get_memory_address() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + let block_id = manager.allocate().unwrap(); + let addr = manager.get_memory_address(block_id); + assert!(addr.is_ok()); + } + + #[test] + fn test_invalid_block_id() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let manager = BlockManager::new(allocator, 1024); + + let invalid_id = BlockId(999); + let result = manager.get_memory_address(invalid_id); + assert!(matches!(result, Err(Error::InvalidBlockId(_)))); + } + + #[test] + fn test_can_allocate() { + let allocator = Box::new(CpuAllocator::new(2048)); + let mut manager = BlockManager::new(allocator, 1024); + + // Should be able to allocate + assert!(manager.can_allocate(&BlockType::StandardKV)); + + // Allocate both blocks + manager.allocate().unwrap(); + manager.allocate().unwrap(); + + // Should not be able to allocate more + assert!(!manager.can_allocate(&BlockType::StandardKV)); + } +} diff --git a/src/block_manager/mod.rs b/src/block_manager/mod.rs index 97a6594..7279f14 100644 --- a/src/block_manager/mod.rs +++ b/src/block_manager/mod.rs @@ -1,3 +1,8 @@ pub mod allocator; pub mod manager; pub mod types; + +// Re-export for convenience + +#[cfg(test)] +mod manager_tests; diff --git a/src/cli/chat.rs b/src/cli/chat.rs index 3017ac6..76f2ca7 100644 --- a/src/cli/chat.rs +++ b/src/cli/chat.rs @@ -1,12 +1,7 @@ -use colored::Colorize; -use rustyline::completion::Completer; use rustyline::error::ReadlineError; -use rustyline::highlight::Highlighter; use rustyline::hint::{Hint, Hinter}; -use rustyline::validate::Validator; use rustyline::{Context, Editor}; use rustyline_derive::{Completer, Helper, Highlighter, Validator}; -use std::borrow::Cow; use std::io::{self, Write}; use tokio_stream::StreamExt; @@ -59,7 +54,7 @@ pub async fn interactive_chat( let helper = PlaceholderHinter { placeholder: "Send a message (Ctrl-C or 'exit' to quit)".to_string(), }; - let mut rl = Editor::new().map_err(|e| io::Error::new(io::ErrorKind::Other, e))?; + let mut rl = Editor::new().map_err(io::Error::other)?; rl.set_helper(Some(helper)); loop { @@ -71,7 +66,7 @@ pub async fn interactive_chat( break; } Err(err) => { - return Err(io::Error::new(io::ErrorKind::Other, err)); + return Err(io::Error::other(err)); } }; diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs index 66d606b..a0a543d 100644 --- a/src/sequence_manager/fsm_events.rs +++ b/src/sequence_manager/fsm_events.rs @@ -128,6 +128,7 @@ impl<'a> AppendTokensEvent<'a> { // Check if need more blocks let blocks_needed = state.num_tokens.div_ceil(self.tokens_per_block); + let initial_block_count = state.blocks.len(); while state.blocks.len() < blocks_needed { match self.block_manager.allocate() { @@ -139,11 +140,14 @@ impl<'a> AppendTokensEvent<'a> { state.blocks.len() ); } - Err(Error::OutOfMemory) => { - warn!("OOM while appending tokens to {:?}", state.seq_id); - return Err(Error::OutOfMemory); + Err(e) => { + // Free any blocks we allocated in this call + warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); + for block_id in state.blocks.drain(initial_block_count..) { + let _ = self.block_manager.free(block_id); + } + return Err(e); } - Err(e) => return Err(e), } } diff --git a/src/sequence_manager/fsm_tests.rs b/src/sequence_manager/fsm_tests.rs new file mode 100644 index 0000000..51dad15 --- /dev/null +++ b/src/sequence_manager/fsm_tests.rs @@ -0,0 +1,326 @@ +#[cfg(test)] +mod tests { + use crate::block_manager::allocator::CpuAllocator; + use crate::block_manager::manager::BlockManager; + use crate::block_manager::types::*; + use crate::sequence_manager::fsm_events::*; + use crate::sequence_manager::states::*; + + #[test] + fn test_schedule_event_waiting_to_prefilling() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let waiting = SequenceState::Waiting(WaitingState { + seq_id: SequenceId(1), + prompt_tokens: 100, + max_tokens: 150, + }); + + let event = ScheduleEvent { + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(waiting); + assert!(result.is_ok()); + + let new_state = result.unwrap(); + assert!(matches!(new_state, SequenceState::Prefilling(_))); + + if let SequenceState::Prefilling(state) = new_state { + assert_eq!(state.seq_id, SequenceId(1)); + assert_eq!(state.tokens_filled, 0); + assert_eq!(state.tokens_total, 100); + // Should allocate ceil(100/16) = 7 blocks + assert_eq!(state.blocks.len(), 7); + } + } + + #[test] + fn test_schedule_event_oom() { + // Small allocator - only 2 blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let waiting = SequenceState::Waiting(WaitingState { + seq_id: SequenceId(1), + prompt_tokens: 100, // Needs 7 blocks + max_tokens: 150, + }); + + let event = ScheduleEvent { + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(waiting); + assert!(matches!(result, Err(Error::OutOfMemory))); + + // Should have cleaned up - no blocks leaked + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_append_tokens_prefilling_to_decoding() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + // Allocate blocks + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let prefilling = SequenceState::Prefilling(PrefillingState { + seq_id: SequenceId(1), + blocks, + tokens_filled: 0, + tokens_total: 100, + max_tokens: 150, + }); + + let event = AppendTokensEvent { + num_tokens: 100, + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(prefilling).unwrap(); + + // Should transition to Decoding + assert!(matches!(result, SequenceState::Decoding(_))); + + if let SequenceState::Decoding(state) = result { + assert_eq!(state.num_tokens, 100); + assert_eq!(state.blocks.len(), 2); + } + } + + #[test] + fn test_append_tokens_decoding_needs_more_blocks() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![block_manager.allocate().unwrap()]; // 1 block + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks, + num_tokens: 16, // 1 block worth + max_tokens: 150, + }); + + let event = AppendTokensEvent { + num_tokens: 20, // Now need 3 blocks total + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(decoding).unwrap(); + + if let SequenceState::Decoding(state) = result { + assert_eq!(state.num_tokens, 36); + // Should allocate 2 more blocks (total 3) + assert_eq!(state.blocks.len(), 3); + } + } + + #[test] + fn test_append_tokens_decoding_oom_cleanup() { + // Small allocator - only enough for initial blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + let initial_block_count = blocks.len(); + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks, + num_tokens: 32, + max_tokens: 150, + }); + + // Try to append tokens that would need more blocks (will fail - OOM) + let event = AppendTokensEvent { + num_tokens: 50, // Would need 6 blocks total + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(decoding); + assert!(matches!(result, Err(Error::OutOfMemory))); + + // Original 2 blocks should still be allocated (not freed by error path) + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, initial_block_count); + } + + #[test] + fn test_append_tokens_reaches_max_tokens() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![block_manager.allocate().unwrap()]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks, + num_tokens: 140, + max_tokens: 150, + }); + + let event = AppendTokensEvent { + num_tokens: 10, + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(decoding).unwrap(); + + // Should transition to Finished + assert!(matches!(result, SequenceState::Finished(_))); + + if let SequenceState::Finished(state) = result { + assert_eq!(state.seq_id, SequenceId(1)); + assert_eq!(state.finish_reason, FinishReason::MaxTokens); + } + + // Blocks should be freed + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_preempt_event() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks, + num_tokens: 32, + max_tokens: 150, + }); + + let event = PreemptEvent { + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + + // Should transition to Preempted + assert!(matches!(result, SequenceState::Preempted(_))); + + if let SequenceState::Preempted(state) = result { + assert_eq!(state.num_tokens, 32); + } + + // Blocks should be freed + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_resume_event() { + let preempted = SequenceState::Preempted(PreemptedState { + seq_id: SequenceId(1), + num_tokens: 32, + max_tokens: 150, + }); + + let event = ResumeEvent; + let result = event.apply(preempted).unwrap(); + + // Should transition back to Waiting + assert!(matches!(result, SequenceState::Waiting(_))); + + if let SequenceState::Waiting(state) = result { + assert_eq!(state.prompt_tokens, 32); + assert_eq!(state.max_tokens, 150); + } + } + + #[test] + fn test_fork_event() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks, + num_tokens: 32, + max_tokens: 150, + }); + + let event = ForkEvent { + child_id: SequenceId(2), + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + let (parent, child) = result; + + // Both should be in Decoding state + assert!(matches!(parent, SequenceState::Decoding(_))); + assert!(matches!(child, SequenceState::Decoding(_))); + + if let (SequenceState::Decoding(p), SequenceState::Decoding(c)) = (parent, child) { + assert_eq!(p.seq_id, SequenceId(1)); + assert_eq!(c.seq_id, SequenceId(2)); + assert_eq!(p.blocks, c.blocks); // Shared blocks + } + } + + #[test] + fn test_complete_event() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![block_manager.allocate().unwrap()]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks, + num_tokens: 100, + max_tokens: 150, + }); + + let event = CompleteEvent { + reason: FinishReason::Stop, + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + + assert!(matches!(result, SequenceState::Finished(_))); + + if let SequenceState::Finished(state) = result { + assert_eq!(state.finish_reason, FinishReason::Stop); + } + + // Blocks should be freed + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } +} diff --git a/src/sequence_manager/mod.rs b/src/sequence_manager/mod.rs index 0e0f032..63e628c 100644 --- a/src/sequence_manager/mod.rs +++ b/src/sequence_manager/mod.rs @@ -4,4 +4,7 @@ pub mod fsm_manager; pub mod id_generator; pub mod states; -// Re-export from block_manager for convenience +#[cfg(test)] +mod fsm_tests; + +// Re-export public API for convenience diff --git a/src/sequence_manager/states.rs b/src/sequence_manager/states.rs index 6543d83..160cadf 100644 --- a/src/sequence_manager/states.rs +++ b/src/sequence_manager/states.rs @@ -91,7 +91,7 @@ pub struct AbortedState { pub reason: String, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub enum FinishReason { /// Reached max_tokens limit MaxTokens, From 812406d5cf9b9e270da029d8b8cd4b815802e81d Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 20:53:49 +0100 Subject: [PATCH 09/33] fix test Signed-off-by: kerthcet --- examples/fsm_usage.rs | 182 ------------------------------------------ src/downloader/mod.rs | 5 +- 2 files changed, 4 insertions(+), 183 deletions(-) delete mode 100644 examples/fsm_usage.rs diff --git a/examples/fsm_usage.rs b/examples/fsm_usage.rs deleted file mode 100644 index 5271291..0000000 --- a/examples/fsm_usage.rs +++ /dev/null @@ -1,182 +0,0 @@ -use puma::block_manager::{BlockManager, CpuAllocator}; -use puma::sequence_manager::{SequenceEvent, SequenceIdGenerator, SequenceManager}; - -#[tokio::main] -async fn main() { - // Initialize tracing - tracing_subscriber::fmt() - .with_env_filter("info,puma=debug") - .init(); - - println!("=== FSM-Based Sequence Manager Demo ===\n"); - - // Setup block manager - let allocator = Box::new(CpuAllocator::new(100_000_000)); // 100MB - let block_size = 4096; // 4KB per block - let block_manager = BlockManager::new(allocator, block_size); - - // Setup sequence manager - let tokens_per_block = 16; - let (seq_manager, event_tx) = SequenceManager::new(block_manager, tokens_per_block); - - // Spawn event loop - let manager_handle = tokio::spawn(async move { - seq_manager.run().await; - }); - - // ID generator (starts from 0) - let id_gen = SequenceIdGenerator::new(); - - println!("1. Adding sequence (Waiting state)"); - let seq_id = id_gen.next(); - event_tx - .send(SequenceEvent::AddRequest { - seq_id, - prompt_tokens: 100, - max_tokens: 150, - }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - println!("\n2. Appending tokens during prefill (Prefilling → Decoding)"); - event_tx - .send(SequenceEvent::AppendTokens { - seq_id, - num_tokens: 100, // Complete prefill - }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - println!("\n3. Generating tokens (Decoding state)"); - for i in 0..10 { - event_tx - .send(SequenceEvent::AppendTokens { - seq_id, - num_tokens: 1, - }) - .unwrap(); - tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; - } - - println!("\n4. Forking sequence for beam search"); - let child_id = id_gen.next(); - event_tx - .send(SequenceEvent::ForkSequence { - parent_id: seq_id, - child_id, - }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - println!("\n5. Continuing both sequences"); - event_tx - .send(SequenceEvent::AppendTokens { - seq_id, - num_tokens: 5, - }) - .unwrap(); - event_tx - .send(SequenceEvent::AppendTokens { - seq_id: child_id, - num_tokens: 5, - }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - // Query stats - let (tx, rx) = tokio::sync::oneshot::channel(); - event_tx - .send(SequenceEvent::GetStats { response: tx }) - .unwrap(); - - let stats = rx.await.unwrap(); - println!("\n=== Stats ==="); - println!("Total sequences: {}", stats.num_sequences); - println!("Running: {}", stats.num_running); - println!("Waiting: {}", stats.num_waiting); - println!("Preempted: {}", stats.num_preempted); - - for (block_type, block_stats) in &stats.block_stats { - println!("\nBlock type: {}", block_type); - println!(" Total blocks: {}", block_stats.total_blocks); - println!(" Allocated: {}", block_stats.allocated_blocks); - println!(" Free: {}", block_stats.free_blocks); - println!(" Total memory: {} bytes", block_stats.total_memory); - } - - println!("\n6. Completing sequences (Finished state)"); - event_tx - .send(SequenceEvent::CompleteSequence { seq_id }) - .unwrap(); - event_tx - .send(SequenceEvent::CompleteSequence { seq_id: child_id }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - // Final stats - let (tx, rx) = tokio::sync::oneshot::channel(); - event_tx - .send(SequenceEvent::GetStats { response: tx }) - .unwrap(); - - let stats = rx.await.unwrap(); - println!("\n=== Final Stats ==="); - println!("Total sequences: {}", stats.num_sequences); - println!("Running: {}", stats.num_running); - - println!("\n7. Testing preemption"); - let seq3 = id_gen.next(); - event_tx - .send(SequenceEvent::AddRequest { - seq_id: seq3, - prompt_tokens: 50, - max_tokens: 100, - }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - event_tx - .send(SequenceEvent::AppendTokens { - seq_id: seq3, - num_tokens: 50, - }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - println!("\n Preempting sequence (Decoding → Preempted)"); - event_tx - .send(SequenceEvent::PreemptSequence { seq_id: seq3 }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - println!(" Resuming sequence (Preempted → Waiting → Decoding)"); - event_tx - .send(SequenceEvent::ResumeSequence { seq_id: seq3 }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - event_tx - .send(SequenceEvent::CompleteSequence { seq_id: seq3 }) - .unwrap(); - - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - println!("\n=== Sequence Manager Demo Complete! ==="); - println!("\nState transitions demonstrated:"); - println!(" Waiting → Prefilling → Decoding → Finished"); - println!(" Decoding → Fork → Two Decoding sequences"); - println!(" Decoding → Preempted → Waiting → Decoding"); - - // Cleanup - drop(event_tx); - manager_handle.await.unwrap(); -} diff --git a/src/downloader/mod.rs b/src/downloader/mod.rs index 0067561..5ecc922 100644 --- a/src/downloader/mod.rs +++ b/src/downloader/mod.rs @@ -27,7 +27,10 @@ impl fmt::Display for DownloadError { } pub trait Downloader { - async fn download_model(&self, name: &str) -> Result<(), DownloadError>; + fn download_model( + &self, + name: &str, + ) -> impl std::future::Future> + Send; } /// Provider for downloading models From be8355c6858061bedc32002e389e08ded2d22e21 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 21:02:21 +0100 Subject: [PATCH 10/33] add freeError Signed-off-by: kerthcet --- src/block_manager/allocator.rs | 7 ++++++- src/block_manager/types.rs | 3 +++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/block_manager/allocator.rs b/src/block_manager/allocator.rs index b048e38..0a199eb 100644 --- a/src/block_manager/allocator.rs +++ b/src/block_manager/allocator.rs @@ -49,12 +49,17 @@ impl MemoryAllocator for CpuAllocator { fn free(&mut self, addr: MemoryAddress) -> Result<()> { let layout = std::alloc::Layout::from_size_align(addr.size, 64) - .map_err(|e| Error::AllocationError(e.to_string()))?; + .map_err(|e| Error::FreeError(e.to_string()))?; unsafe { std::alloc::dealloc(addr.ptr, layout); } + if addr.size > self.used_memory { + return Err(Error::FreeError( + "free() called with size larger than used_memory".to_string(), + )); + } self.used_memory -= addr.size; Ok(()) diff --git a/src/block_manager/types.rs b/src/block_manager/types.rs index 92a1d13..56e8427 100644 --- a/src/block_manager/types.rs +++ b/src/block_manager/types.rs @@ -94,6 +94,9 @@ pub enum Error { #[error("Allocation error: {0}")] AllocationError(String), + + #[error("Free error: {0}")] + FreeError(String), } pub type Result = std::result::Result; From 01ecb5dc216d5b548607947c2e299c04aafbac86 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 21:15:17 +0100 Subject: [PATCH 11/33] address comments Signed-off-by: kerthcet --- src/block_manager/allocator.rs | 11 ++++++----- src/cli/chat.rs | 3 ++- src/sequence_manager/fsm_manager.rs | 12 ++++++------ 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src/block_manager/allocator.rs b/src/block_manager/allocator.rs index 0a199eb..9927e85 100644 --- a/src/block_manager/allocator.rs +++ b/src/block_manager/allocator.rs @@ -48,6 +48,12 @@ impl MemoryAllocator for CpuAllocator { } fn free(&mut self, addr: MemoryAddress) -> Result<()> { + if addr.size > self.used_memory { + return Err(Error::FreeError( + "free() called with size larger than used_memory".to_string(), + )); + } + let layout = std::alloc::Layout::from_size_align(addr.size, 64) .map_err(|e| Error::FreeError(e.to_string()))?; @@ -55,11 +61,6 @@ impl MemoryAllocator for CpuAllocator { std::alloc::dealloc(addr.ptr, layout); } - if addr.size > self.used_memory { - return Err(Error::FreeError( - "free() called with size larger than used_memory".to_string(), - )); - } self.used_memory -= addr.size; Ok(()) diff --git a/src/cli/chat.rs b/src/cli/chat.rs index 76f2ca7..6d7c5a9 100644 --- a/src/cli/chat.rs +++ b/src/cli/chat.rs @@ -54,7 +54,8 @@ pub async fn interactive_chat( let helper = PlaceholderHinter { placeholder: "Send a message (Ctrl-C or 'exit' to quit)".to_string(), }; - let mut rl = Editor::new().map_err(io::Error::other)?; + let mut rl = Editor::::new() + .map_err(io::Error::other)?; rl.set_helper(Some(helper)); loop { diff --git a/src/sequence_manager/fsm_manager.rs b/src/sequence_manager/fsm_manager.rs index 95aa46d..a452e36 100644 --- a/src/sequence_manager/fsm_manager.rs +++ b/src/sequence_manager/fsm_manager.rs @@ -28,8 +28,9 @@ impl SequenceManager { block_manager: BlockManager, tokens_per_block: usize, ) -> (Self, SequenceEventSender) { - let (event_tx, event_rx) = create_event_channel(); + assert!(tokens_per_block > 0, "tokens_per_block must be > 0"); + let (event_tx, event_rx) = create_event_channel(); let manager = Self { block_manager, sequences: HashMap::new(), @@ -308,11 +309,10 @@ impl SequenceManager { } fn get_blocks(&self, seq_id: SequenceId) -> Result> { - self.sequences - .get(&seq_id) - .and_then(|state| state.blocks()) - .map(|blocks| blocks.to_vec()) - .ok_or(Error::UnknownSequence(seq_id)) + match self.sequences.get(&seq_id) { + None => Err(Error::UnknownSequence(seq_id)), + Some(state) => Ok(state.blocks().map(|b| b.to_vec()).unwrap_or_default()), + } } fn get_stats(&self) -> SequenceManagerStats { From d4b9bb386984cc7ed9de02a5c0305f9ff0c0c56d Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 21:15:53 +0100 Subject: [PATCH 12/33] fix doc Signed-off-by: kerthcet --- docs/fsm_architecture.md | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/docs/fsm_architecture.md b/docs/fsm_architecture.md index 01a3506..2bc3a91 100644 --- a/docs/fsm_architecture.md +++ b/docs/fsm_architecture.md @@ -246,7 +246,7 @@ impl SequenceManager { SequenceEvent::AppendTokens { seq_id, num_tokens } => { // Get current state let state = self.sequences.remove(&seq_id).unwrap(); - + // Clone state before transition (prevents loss on failure) let backup = state.clone(); @@ -379,9 +379,9 @@ On transition failure, the state is cloned before applying the event: fn handle_append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { let state = self.sequences.remove(&seq_id).unwrap(); let backup = state.clone(); // Clone before transition - + let event = AppendTokensEvent { ... }; - + match event.apply(state) { Ok(new_state) => { self.sequences.insert(seq_id, new_state); @@ -436,5 +436,3 @@ event_tx.send(SequenceEvent::AppendTokens { num_tokens: 100, }).unwrap(); ``` - -See `examples/fsm_usage.rs` for complete example. From ae8df8dca7cfdd12a77fdf62447d9ba36d138f61 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 6 Jul 2026 23:39:59 +0100 Subject: [PATCH 13/33] address comments Signed-off-by: kerthcet --- src/sequence_manager/fsm_events.rs | 45 ++++++++++++++++++++++++----- src/sequence_manager/fsm_manager.rs | 2 -- 2 files changed, 38 insertions(+), 9 deletions(-) diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs index a0a543d..b09a497 100644 --- a/src/sequence_manager/fsm_events.rs +++ b/src/sequence_manager/fsm_events.rs @@ -1,3 +1,4 @@ + use super::states::*; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::*; @@ -112,7 +113,12 @@ impl<'a> AppendTokensEvent<'a> { if state.num_tokens >= state.max_tokens { // Free blocks for block_id in state.blocks { - self.block_manager.free(block_id)?; + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } } debug!( @@ -174,7 +180,12 @@ impl<'a> PreemptEvent<'a> { fn from_decoding(self, state: DecodingState) -> Result { // Free all blocks for block_id in state.blocks { - self.block_manager.free(block_id)?; + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } } debug!("Preempted seq {:?}", state.seq_id); @@ -187,9 +198,14 @@ impl<'a> PreemptEvent<'a> { } fn from_prefilling(self, state: PrefillingState) -> Result { - // Free all blocks + // Free all blocks (best effort) for block_id in state.blocks { - self.block_manager.free(block_id)?; + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } } debug!("Preempted seq {:?} during prefill", state.seq_id); @@ -285,7 +301,12 @@ impl<'a> CompleteEvent<'a> { fn from_decoding(self, state: DecodingState) -> Result { // Free all blocks for block_id in state.blocks { - self.block_manager.free(block_id)?; + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } } debug!("Completed seq {:?}: {:?}", state.seq_id, self.reason); @@ -299,7 +320,12 @@ impl<'a> CompleteEvent<'a> { fn from_prefilling(self, state: PrefillingState) -> Result { // Free all blocks for block_id in state.blocks { - self.block_manager.free(block_id)?; + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } } debug!( @@ -328,7 +354,12 @@ impl<'a> AbortEvent<'a> { // Free blocks if any if let Some(blocks) = state.blocks() { for &block_id in blocks { - self.block_manager.free(block_id)?; + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } } } diff --git a/src/sequence_manager/fsm_manager.rs b/src/sequence_manager/fsm_manager.rs index a452e36..c8d1099 100644 --- a/src/sequence_manager/fsm_manager.rs +++ b/src/sequence_manager/fsm_manager.rs @@ -17,7 +17,6 @@ pub struct SequenceManager { // Event-driven event_rx: SequenceEventReceiver, - event_tx: SequenceEventSender, // Scheduling queues waiting_queue: VecDeque, @@ -36,7 +35,6 @@ impl SequenceManager { sequences: HashMap::new(), tokens_per_block, event_rx, - event_tx: event_tx.clone(), waiting_queue: VecDeque::new(), }; From 7ff6c990b9b084881f005941e963668e4b6ef046 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 7 Jul 2026 08:23:47 +0100 Subject: [PATCH 14/33] allow deadcode now Signed-off-by: kerthcet --- src/lib.rs | 2 ++ src/main.rs | 2 ++ src/sequence_manager/fsm_events.rs | 36 +++++++++++++++--------------- 3 files changed, 22 insertions(+), 18 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index ae83ace..43546f1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ +#![allow(dead_code)] + pub mod api; pub mod backend; pub mod block_manager; diff --git a/src/main.rs b/src/main.rs index fd485d9..aa4feeb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,5 @@ +#![allow(dead_code)] + mod api; mod backend; mod block_manager; diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs index b09a497..d9dbe96 100644 --- a/src/sequence_manager/fsm_events.rs +++ b/src/sequence_manager/fsm_events.rs @@ -18,14 +18,14 @@ pub struct ScheduleEvent<'a> { impl<'a> ScheduleEvent<'a> { pub fn apply(self, state: SequenceState) -> Result { match state { - SequenceState::Waiting(s) => self.from_waiting(s), + SequenceState::Waiting(s) => self.apply_to_waiting(s), _ => Err(Error::invalid_transition( "ScheduleEvent requires Waiting state", )), } } - fn from_waiting(self, state: WaitingState) -> Result { + fn apply_to_waiting(self, state: WaitingState) -> Result { let blocks_needed = state.prompt_tokens.div_ceil(self.tokens_per_block); let mut blocks = Vec::new(); @@ -76,15 +76,15 @@ pub struct AppendTokensEvent<'a> { impl<'a> AppendTokensEvent<'a> { pub fn apply(self, state: SequenceState) -> Result { match state { - SequenceState::Prefilling(s) => self.from_prefilling(s), - SequenceState::Decoding(s) => self.from_decoding(s), + SequenceState::Prefilling(s) => self.apply_to_prefilling(s), + SequenceState::Decoding(s) => self.apply_to_decoding(s), _ => Err(Error::invalid_transition( "AppendTokensEvent requires Prefilling or Decoding state", )), } } - fn from_prefilling(self, mut state: PrefillingState) -> Result { + fn apply_to_prefilling(self, mut state: PrefillingState) -> Result { state.tokens_filled += self.num_tokens; debug!( @@ -106,7 +106,7 @@ impl<'a> AppendTokensEvent<'a> { } } - fn from_decoding(self, mut state: DecodingState) -> Result { + fn apply_to_decoding(self, mut state: DecodingState) -> Result { state.num_tokens += self.num_tokens; // Check if finished @@ -169,15 +169,15 @@ pub struct PreemptEvent<'a> { impl<'a> PreemptEvent<'a> { pub fn apply(self, state: SequenceState) -> Result { match state { - SequenceState::Decoding(s) => self.from_decoding(s), - SequenceState::Prefilling(s) => self.from_prefilling(s), + SequenceState::Decoding(s) => self.apply_to_decoding(s), + SequenceState::Prefilling(s) => self.apply_to_prefilling(s), _ => Err(Error::invalid_transition( "PreemptEvent requires running state", )), } } - fn from_decoding(self, state: DecodingState) -> Result { + fn apply_to_decoding(self, state: DecodingState) -> Result { // Free all blocks for block_id in state.blocks { if let Err(e) = self.block_manager.free(block_id) { @@ -197,7 +197,7 @@ impl<'a> PreemptEvent<'a> { })) } - fn from_prefilling(self, state: PrefillingState) -> Result { + fn apply_to_prefilling(self, state: PrefillingState) -> Result { // Free all blocks (best effort) for block_id in state.blocks { if let Err(e) = self.block_manager.free(block_id) { @@ -224,14 +224,14 @@ pub struct ResumeEvent; impl ResumeEvent { pub fn apply(self, state: SequenceState) -> Result { match state { - SequenceState::Preempted(s) => self.from_preempted(s), + SequenceState::Preempted(s) => self.apply_to_preempted(s), _ => Err(Error::invalid_transition( "ResumeEvent requires Preempted state", )), } } - fn from_preempted(self, state: PreemptedState) -> Result { + fn apply_to_preempted(self, state: PreemptedState) -> Result { debug!("Resuming seq {:?}", state.seq_id); // Transition back to waiting for rescheduling @@ -252,14 +252,14 @@ pub struct ForkEvent<'a> { impl<'a> ForkEvent<'a> { pub fn apply(self, state: SequenceState) -> Result<(SequenceState, SequenceState)> { match state { - SequenceState::Decoding(s) => self.from_decoding(s), + SequenceState::Decoding(s) => self.apply_to_decoding(s), _ => Err(Error::invalid_transition( "ForkEvent requires Decoding state", )), } } - fn from_decoding(self, state: DecodingState) -> Result<(SequenceState, SequenceState)> { + fn apply_to_decoding(self, state: DecodingState) -> Result<(SequenceState, SequenceState)> { // Copy-on-write: increment ref counts for &block_id in &state.blocks { self.block_manager.add_ref(block_id)?; @@ -290,15 +290,15 @@ pub struct CompleteEvent<'a> { impl<'a> CompleteEvent<'a> { pub fn apply(self, state: SequenceState) -> Result { match state { - SequenceState::Decoding(s) => self.from_decoding(s), - SequenceState::Prefilling(s) => self.from_prefilling(s), + SequenceState::Decoding(s) => self.apply_to_decoding(s), + SequenceState::Prefilling(s) => self.apply_to_prefilling(s), _ => Err(Error::invalid_transition( "CompleteEvent requires running state", )), } } - fn from_decoding(self, state: DecodingState) -> Result { + fn apply_to_decoding(self, state: DecodingState) -> Result { // Free all blocks for block_id in state.blocks { if let Err(e) = self.block_manager.free(block_id) { @@ -317,7 +317,7 @@ impl<'a> CompleteEvent<'a> { })) } - fn from_prefilling(self, state: PrefillingState) -> Result { + fn apply_to_prefilling(self, state: PrefillingState) -> Result { // Free all blocks for block_id in state.blocks { if let Err(e) = self.block_manager.free(block_id) { From 5b611b3fbfd94021cbc2d831cc26847f7a8973ef Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 7 Jul 2026 08:24:03 +0100 Subject: [PATCH 15/33] fix lint Signed-off-by: kerthcet --- src/sequence_manager/fsm_events.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs index d9dbe96..b474cf5 100644 --- a/src/sequence_manager/fsm_events.rs +++ b/src/sequence_manager/fsm_events.rs @@ -1,4 +1,3 @@ - use super::states::*; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::*; From ba367f6cfccc668bc470034a20c99d3301bc25b6 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 7 Jul 2026 08:30:53 +0100 Subject: [PATCH 16/33] fix lint Signed-off-by: kerthcet --- src/block_manager/allocator.rs | 6 +++--- src/block_manager/types.rs | 8 ++++---- src/sequence_manager/fsm_events.rs | 3 +-- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/src/block_manager/allocator.rs b/src/block_manager/allocator.rs index 9927e85..0db8aab 100644 --- a/src/block_manager/allocator.rs +++ b/src/block_manager/allocator.rs @@ -31,7 +31,7 @@ impl MemoryAllocator for CpuAllocator { // Allocate aligned memory let layout = std::alloc::Layout::from_size_align(size_bytes, 64) - .map_err(|e| Error::AllocationError(e.to_string()))?; + .map_err(|e| Error::AllocationFailed(e.to_string()))?; let ptr = unsafe { std::alloc::alloc(layout) }; @@ -49,13 +49,13 @@ impl MemoryAllocator for CpuAllocator { fn free(&mut self, addr: MemoryAddress) -> Result<()> { if addr.size > self.used_memory { - return Err(Error::FreeError( + return Err(Error::FreeFailed( "free() called with size larger than used_memory".to_string(), )); } let layout = std::alloc::Layout::from_size_align(addr.size, 64) - .map_err(|e| Error::FreeError(e.to_string()))?; + .map_err(|e| Error::FreeFailed(e.to_string()))?; unsafe { std::alloc::dealloc(addr.ptr, layout); diff --git a/src/block_manager/types.rs b/src/block_manager/types.rs index 56e8427..0aa33f7 100644 --- a/src/block_manager/types.rs +++ b/src/block_manager/types.rs @@ -92,11 +92,11 @@ pub enum Error { #[error("Out of memory")] OutOfMemory, - #[error("Allocation error: {0}")] - AllocationError(String), + #[error("Allocation failed: {0}")] + AllocationFailed(String), - #[error("Free error: {0}")] - FreeError(String), + #[error("Free failed: {0}")] + FreeFailed(String), } pub type Result = std::result::Result; diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs index b474cf5..433eb60 100644 --- a/src/sequence_manager/fsm_events.rs +++ b/src/sequence_manager/fsm_events.rs @@ -7,7 +7,6 @@ use tracing::{debug, warn}; /// /// Pattern: Event takes old state, returns new state /// Invalid transitions return Error::invalid_transition - /// Schedule a waiting sequence (allocate blocks) pub struct ScheduleEvent<'a> { pub block_manager: &'a mut BlockManager, @@ -374,6 +373,6 @@ impl<'a> AbortEvent<'a> { /// Invalid transition error helper impl Error { pub fn invalid_transition(msg: &'static str) -> Self { - Error::AllocationError(format!("Invalid state transition: {}", msg)) + Error::AllocationFailed(format!("Invalid state transition: {}", msg)) } } From 56d03b98accb5bd533e1eefd36d72d0c7c77fcca Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 7 Jul 2026 08:54:06 +0100 Subject: [PATCH 17/33] usr Arc to solve performance issue Signed-off-by: kerthcet --- src/sequence_manager/events.rs | 5 +++ src/sequence_manager/fsm_events.rs | 52 +++++++++++++++------------- src/sequence_manager/fsm_manager.rs | 9 +++++ src/sequence_manager/fsm_tests.rs | 53 +++++++++++++++++++++++++---- src/sequence_manager/states.rs | 13 +++---- 5 files changed, 96 insertions(+), 36 deletions(-) diff --git a/src/sequence_manager/events.rs b/src/sequence_manager/events.rs index f40fbae..7398b8e 100644 --- a/src/sequence_manager/events.rs +++ b/src/sequence_manager/events.rs @@ -63,6 +63,11 @@ pub struct SequenceManagerStats { pub type SequenceEventSender = mpsc::UnboundedSender; pub type SequenceEventReceiver = mpsc::UnboundedReceiver; +/// Create event channel for sequence manager +/// +/// TODO: Use bounded channel to prevent OOM when producers outpace the event loop. +/// Capacity should be calculated based on available memory and average event size. +/// Consider: capacity = (available_memory * 0.1) / sizeof(SequenceEvent) pub fn create_event_channel() -> (SequenceEventSender, SequenceEventReceiver) { mpsc::unbounded_channel() } diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs index 433eb60..39d83d9 100644 --- a/src/sequence_manager/fsm_events.rs +++ b/src/sequence_manager/fsm_events.rs @@ -1,6 +1,7 @@ use super::states::*; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::*; +use std::sync::Arc; use tracing::{debug, warn}; /// FSM Events - transform states with type-safe transitions @@ -56,7 +57,7 @@ impl<'a> ScheduleEvent<'a> { Ok(SequenceState::Prefilling(PrefillingState { seq_id: state.seq_id, - blocks, + blocks: Arc::new(blocks), tokens_filled: 0, tokens_total: state.prompt_tokens, max_tokens: state.max_tokens, @@ -110,7 +111,7 @@ impl<'a> AppendTokensEvent<'a> { // Check if finished if state.num_tokens >= state.max_tokens { // Free blocks - for block_id in state.blocks { + for &block_id in state.blocks.iter() { if let Err(e) = self.block_manager.free(block_id) { warn!( "Failed to free block {:?} for {:?}: {:?}", @@ -134,23 +135,28 @@ impl<'a> AppendTokensEvent<'a> { let blocks_needed = state.num_tokens.div_ceil(self.tokens_per_block); let initial_block_count = state.blocks.len(); - while state.blocks.len() < blocks_needed { - match self.block_manager.allocate() { - Ok(block_id) => { - state.blocks.push(block_id); - debug!( - "Seq {:?}: allocated block, total blocks: {}", - state.seq_id, - state.blocks.len() - ); - } - Err(e) => { - // Free any blocks we allocated in this call - warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); - for block_id in state.blocks.drain(initial_block_count..) { - let _ = self.block_manager.free(block_id); + if state.blocks.len() < blocks_needed { + // Need to allocate more blocks - use Arc::make_mut for copy-on-write + let blocks = Arc::make_mut(&mut state.blocks); + + while blocks.len() < blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => { + blocks.push(block_id); + debug!( + "Seq {:?}: allocated block, total blocks: {}", + state.seq_id, + blocks.len() + ); + } + Err(e) => { + // Free any blocks we allocated in this call + warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); + for block_id in blocks.drain(initial_block_count..) { + let _ = self.block_manager.free(block_id); + } + return Err(e); } - return Err(e); } } } @@ -177,7 +183,7 @@ impl<'a> PreemptEvent<'a> { fn apply_to_decoding(self, state: DecodingState) -> Result { // Free all blocks - for block_id in state.blocks { + for &block_id in state.blocks.iter() { if let Err(e) = self.block_manager.free(block_id) { warn!( "Failed to free block {:?} for {:?}: {:?}", @@ -197,7 +203,7 @@ impl<'a> PreemptEvent<'a> { fn apply_to_prefilling(self, state: PrefillingState) -> Result { // Free all blocks (best effort) - for block_id in state.blocks { + for &block_id in state.blocks.iter() { if let Err(e) = self.block_manager.free(block_id) { warn!( "Failed to free block {:?} for {:?}: {:?}", @@ -259,7 +265,7 @@ impl<'a> ForkEvent<'a> { fn apply_to_decoding(self, state: DecodingState) -> Result<(SequenceState, SequenceState)> { // Copy-on-write: increment ref counts - for &block_id in &state.blocks { + for &block_id in state.blocks.iter() { self.block_manager.add_ref(block_id)?; } @@ -298,7 +304,7 @@ impl<'a> CompleteEvent<'a> { fn apply_to_decoding(self, state: DecodingState) -> Result { // Free all blocks - for block_id in state.blocks { + for &block_id in state.blocks.iter() { if let Err(e) = self.block_manager.free(block_id) { warn!( "Failed to free block {:?} for {:?}: {:?}", @@ -317,7 +323,7 @@ impl<'a> CompleteEvent<'a> { fn apply_to_prefilling(self, state: PrefillingState) -> Result { // Free all blocks - for block_id in state.blocks { + for &block_id in state.blocks.iter() { if let Err(e) = self.block_manager.free(block_id) { warn!( "Failed to free block {:?} for {:?}: {:?}", diff --git a/src/sequence_manager/fsm_manager.rs b/src/sequence_manager/fsm_manager.rs index c8d1099..ae74297 100644 --- a/src/sequence_manager/fsm_manager.rs +++ b/src/sequence_manager/fsm_manager.rs @@ -145,6 +145,15 @@ impl SequenceManager { } fn handle_fork_sequence(&mut self, parent_id: SequenceId, child_id: SequenceId) { + // Check if child_id already exists + if self.sequences.contains_key(&child_id) { + warn!( + "Cannot fork {:?} → {:?}: child ID already exists", + parent_id, child_id + ); + return; + } + let parent_state = match self.sequences.remove(&parent_id) { Some(s) => s, None => { diff --git a/src/sequence_manager/fsm_tests.rs b/src/sequence_manager/fsm_tests.rs index 51dad15..5ec2bb0 100644 --- a/src/sequence_manager/fsm_tests.rs +++ b/src/sequence_manager/fsm_tests.rs @@ -5,6 +5,7 @@ mod tests { use crate::block_manager::types::*; use crate::sequence_manager::fsm_events::*; use crate::sequence_manager::states::*; + use std::sync::Arc; #[test] fn test_schedule_event_waiting_to_prefilling() { @@ -76,7 +77,7 @@ mod tests { let prefilling = SequenceState::Prefilling(PrefillingState { seq_id: SequenceId(1), - blocks, + blocks: Arc::new(blocks), tokens_filled: 0, tokens_total: 100, max_tokens: 150, @@ -108,7 +109,7 @@ mod tests { let decoding = SequenceState::Decoding(DecodingState { seq_id: SequenceId(1), - blocks, + blocks: Arc::new(blocks), num_tokens: 16, // 1 block worth max_tokens: 150, }); @@ -142,7 +143,7 @@ mod tests { let decoding = SequenceState::Decoding(DecodingState { seq_id: SequenceId(1), - blocks, + blocks: Arc::new(blocks), num_tokens: 32, max_tokens: 150, }); @@ -172,7 +173,7 @@ mod tests { let decoding = SequenceState::Decoding(DecodingState { seq_id: SequenceId(1), - blocks, + blocks: Arc::new(blocks), num_tokens: 140, max_tokens: 150, }); @@ -211,7 +212,7 @@ mod tests { let decoding = SequenceState::Decoding(DecodingState { seq_id: SequenceId(1), - blocks, + blocks: Arc::new(blocks), num_tokens: 32, max_tokens: 150, }); @@ -267,7 +268,7 @@ mod tests { let decoding = SequenceState::Decoding(DecodingState { seq_id: SequenceId(1), - blocks, + blocks: Arc::new(blocks), num_tokens: 32, max_tokens: 150, }); @@ -300,7 +301,7 @@ mod tests { let decoding = SequenceState::Decoding(DecodingState { seq_id: SequenceId(1), - blocks, + blocks: Arc::new(blocks), num_tokens: 100, max_tokens: 150, }); @@ -323,4 +324,42 @@ mod tests { let block_stats = stats.get(&BlockType::StandardKV).unwrap(); assert_eq!(block_stats.allocated_blocks, 0); } + + #[test] + fn test_fork_duplicate_child_id() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 32, + max_tokens: 150, + }); + + // Fork with child_id = SequenceId(2) + let event = ForkEvent { + child_id: SequenceId(2), + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + let (parent, child) = result; + + // Verify fork succeeded + if let SequenceState::Decoding(c) = &child { + assert_eq!(c.seq_id, SequenceId(2)); + } + + // Now try to fork again with the SAME child_id (should fail in manager) + // This test verifies the ForkEvent itself works, but the manager + // should check for duplicate child_id before calling this + drop(parent); + drop(child); + } } diff --git a/src/sequence_manager/states.rs b/src/sequence_manager/states.rs index 160cadf..2128f1a 100644 --- a/src/sequence_manager/states.rs +++ b/src/sequence_manager/states.rs @@ -1,4 +1,5 @@ use crate::block_manager::types::*; +use std::sync::Arc; /// FSM States - each state owns its resources /// @@ -46,14 +47,14 @@ pub struct SchedulingState { pub prompt_tokens: usize, pub max_tokens: usize, // Partially allocated blocks (if any) - pub blocks: Vec, + pub blocks: Arc>, } /// Prefilling state - processing prompt #[derive(Debug, Clone)] pub struct PrefillingState { pub seq_id: SequenceId, - pub blocks: Vec, + pub blocks: Arc>, // Arc for cheap cloning (O(1)) pub tokens_filled: usize, pub tokens_total: usize, pub max_tokens: usize, @@ -63,7 +64,7 @@ pub struct PrefillingState { #[derive(Debug, Clone)] pub struct DecodingState { pub seq_id: SequenceId, - pub blocks: Vec, + pub blocks: Arc>, // Arc for cheap cloning (O(1)) pub num_tokens: usize, pub max_tokens: usize, } @@ -135,9 +136,9 @@ impl SequenceState { pub fn blocks(&self) -> Option<&[BlockId]> { match self { - SequenceState::Prefilling(s) => Some(&s.blocks), - SequenceState::Decoding(s) => Some(&s.blocks), - SequenceState::Scheduling(s) if !s.blocks.is_empty() => Some(&s.blocks), + SequenceState::Prefilling(s) => Some(s.blocks.as_ref()), + SequenceState::Decoding(s) => Some(s.blocks.as_ref()), + SequenceState::Scheduling(s) if !s.blocks.is_empty() => Some(s.blocks.as_ref()), _ => None, } } From ea5670bed63264b6537ae50f26982cc61fffae95 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 7 Jul 2026 09:11:54 +0100 Subject: [PATCH 18/33] fix issue Signed-off-by: kerthcet --- src/block_manager/types.rs | 3 +++ src/sequence_manager/fsm_events.rs | 19 ++++++------------- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/src/block_manager/types.rs b/src/block_manager/types.rs index 0aa33f7..3076802 100644 --- a/src/block_manager/types.rs +++ b/src/block_manager/types.rs @@ -97,6 +97,9 @@ pub enum Error { #[error("Free failed: {0}")] FreeFailed(String), + + #[error("Invalid state transition: {0}")] + InvalidTransition(&'static str), } pub type Result = std::result::Result; diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs index 39d83d9..c41f2cb 100644 --- a/src/sequence_manager/fsm_events.rs +++ b/src/sequence_manager/fsm_events.rs @@ -18,7 +18,7 @@ impl<'a> ScheduleEvent<'a> { pub fn apply(self, state: SequenceState) -> Result { match state { SequenceState::Waiting(s) => self.apply_to_waiting(s), - _ => Err(Error::invalid_transition( + _ => Err(Error::InvalidTransition( "ScheduleEvent requires Waiting state", )), } @@ -77,7 +77,7 @@ impl<'a> AppendTokensEvent<'a> { match state { SequenceState::Prefilling(s) => self.apply_to_prefilling(s), SequenceState::Decoding(s) => self.apply_to_decoding(s), - _ => Err(Error::invalid_transition( + _ => Err(Error::InvalidTransition( "AppendTokensEvent requires Prefilling or Decoding state", )), } @@ -175,7 +175,7 @@ impl<'a> PreemptEvent<'a> { match state { SequenceState::Decoding(s) => self.apply_to_decoding(s), SequenceState::Prefilling(s) => self.apply_to_prefilling(s), - _ => Err(Error::invalid_transition( + _ => Err(Error::InvalidTransition( "PreemptEvent requires running state", )), } @@ -229,7 +229,7 @@ impl ResumeEvent { pub fn apply(self, state: SequenceState) -> Result { match state { SequenceState::Preempted(s) => self.apply_to_preempted(s), - _ => Err(Error::invalid_transition( + _ => Err(Error::InvalidTransition( "ResumeEvent requires Preempted state", )), } @@ -257,7 +257,7 @@ impl<'a> ForkEvent<'a> { pub fn apply(self, state: SequenceState) -> Result<(SequenceState, SequenceState)> { match state { SequenceState::Decoding(s) => self.apply_to_decoding(s), - _ => Err(Error::invalid_transition( + _ => Err(Error::InvalidTransition( "ForkEvent requires Decoding state", )), } @@ -296,7 +296,7 @@ impl<'a> CompleteEvent<'a> { match state { SequenceState::Decoding(s) => self.apply_to_decoding(s), SequenceState::Prefilling(s) => self.apply_to_prefilling(s), - _ => Err(Error::invalid_transition( + _ => Err(Error::InvalidTransition( "CompleteEvent requires running state", )), } @@ -375,10 +375,3 @@ impl<'a> AbortEvent<'a> { })) } } - -/// Invalid transition error helper -impl Error { - pub fn invalid_transition(msg: &'static str) -> Self { - Error::AllocationFailed(format!("Invalid state transition: {}", msg)) - } -} From 6dfe52d64fdb841cddecca262174738f7083d692 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 7 Jul 2026 10:25:43 +0100 Subject: [PATCH 19/33] fix tests Signed-off-by: kerthcet --- src/sequence_manager/fsm_events.rs | 363 +++++++++++++++++++++++++++ src/sequence_manager/fsm_manager.rs | 74 ++++++ src/sequence_manager/fsm_tests.rs | 365 ---------------------------- src/sequence_manager/mod.rs | 3 - 4 files changed, 437 insertions(+), 368 deletions(-) delete mode 100644 src/sequence_manager/fsm_tests.rs diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs index c41f2cb..0c4904d 100644 --- a/src/sequence_manager/fsm_events.rs +++ b/src/sequence_manager/fsm_events.rs @@ -375,3 +375,366 @@ impl<'a> AbortEvent<'a> { })) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + use crate::block_manager::manager::BlockManager; + + #[test] + fn test_schedule_event_waiting_to_prefilling() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let waiting = SequenceState::Waiting(WaitingState { + seq_id: SequenceId(1), + prompt_tokens: 100, + max_tokens: 150, + }); + + let event = ScheduleEvent { + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(waiting); + assert!(result.is_ok()); + + let new_state = result.unwrap(); + assert!(matches!(new_state, SequenceState::Prefilling(_))); + + if let SequenceState::Prefilling(state) = new_state { + assert_eq!(state.seq_id, SequenceId(1)); + assert_eq!(state.tokens_filled, 0); + assert_eq!(state.tokens_total, 100); + // Should allocate ceil(100/16) = 7 blocks + assert_eq!(state.blocks.len(), 7); + } + } + + #[test] + fn test_schedule_event_oom() { + // Small allocator - only 2 blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let waiting = SequenceState::Waiting(WaitingState { + seq_id: SequenceId(1), + prompt_tokens: 100, // Needs 7 blocks + max_tokens: 150, + }); + + let event = ScheduleEvent { + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(waiting); + assert!(matches!(result, Err(Error::OutOfMemory))); + + // Should have cleaned up - no blocks leaked + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_append_tokens_prefilling_to_decoding() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + // Allocate blocks + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let prefilling = SequenceState::Prefilling(PrefillingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: 100, + max_tokens: 150, + }); + + let event = AppendTokensEvent { + num_tokens: 100, + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(prefilling).unwrap(); + + // Should transition to Decoding + assert!(matches!(result, SequenceState::Decoding(_))); + + if let SequenceState::Decoding(state) = result { + assert_eq!(state.num_tokens, 100); + assert_eq!(state.blocks.len(), 2); + } + } + + #[test] + fn test_append_tokens_decoding_needs_more_blocks() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![block_manager.allocate().unwrap()]; // 1 block + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 16, // 1 block worth + max_tokens: 150, + }); + + let event = AppendTokensEvent { + num_tokens: 20, // Now need 3 blocks total + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(decoding).unwrap(); + + if let SequenceState::Decoding(state) = result { + assert_eq!(state.num_tokens, 36); + // Should allocate 2 more blocks (total 3) + assert_eq!(state.blocks.len(), 3); + } + } + + #[test] + fn test_append_tokens_decoding_oom_cleanup() { + // Small allocator - only enough for initial blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + let initial_block_count = blocks.len(); + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 32, + max_tokens: 150, + }); + + // Try to append tokens that would need more blocks (will fail - OOM) + let event = AppendTokensEvent { + num_tokens: 50, // Would need 6 blocks total + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(decoding); + assert!(matches!(result, Err(Error::OutOfMemory))); + + // Original 2 blocks should still be allocated (not freed by error path) + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, initial_block_count); + } + + #[test] + fn test_append_tokens_reaches_max_tokens() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![block_manager.allocate().unwrap()]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 140, + max_tokens: 150, + }); + + let event = AppendTokensEvent { + num_tokens: 10, + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(decoding).unwrap(); + + // Should transition to Finished + assert!(matches!(result, SequenceState::Finished(_))); + + if let SequenceState::Finished(state) = result { + assert_eq!(state.seq_id, SequenceId(1)); + assert_eq!(state.finish_reason, FinishReason::MaxTokens); + } + + // Blocks should be freed + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_preempt_event() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 32, + max_tokens: 150, + }); + + let event = PreemptEvent { + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + + // Should transition to Preempted + assert!(matches!(result, SequenceState::Preempted(_))); + + if let SequenceState::Preempted(state) = result { + assert_eq!(state.num_tokens, 32); + } + + // Blocks should be freed + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_resume_event() { + let preempted = SequenceState::Preempted(PreemptedState { + seq_id: SequenceId(1), + num_tokens: 32, + max_tokens: 150, + }); + + let event = ResumeEvent; + let result = event.apply(preempted).unwrap(); + + // Should transition back to Waiting + assert!(matches!(result, SequenceState::Waiting(_))); + + if let SequenceState::Waiting(state) = result { + assert_eq!(state.prompt_tokens, 32); + assert_eq!(state.max_tokens, 150); + } + } + + #[test] + fn test_fork_event() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 32, + max_tokens: 150, + }); + + let event = ForkEvent { + child_id: SequenceId(2), + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + let (parent, child) = result; + + // Both should be in Decoding state + assert!(matches!(parent, SequenceState::Decoding(_))); + assert!(matches!(child, SequenceState::Decoding(_))); + + if let (SequenceState::Decoding(p), SequenceState::Decoding(c)) = (parent, child) { + assert_eq!(p.seq_id, SequenceId(1)); + assert_eq!(c.seq_id, SequenceId(2)); + assert_eq!(p.blocks, c.blocks); // Shared blocks + } + } + + #[test] + fn test_complete_event() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![block_manager.allocate().unwrap()]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 100, + max_tokens: 150, + }); + + let event = CompleteEvent { + reason: FinishReason::Stop, + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + + assert!(matches!(result, SequenceState::Finished(_))); + + if let SequenceState::Finished(state) = result { + assert_eq!(state.finish_reason, FinishReason::Stop); + } + + // Blocks should be freed + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_fork_duplicate_child_id() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 32, + max_tokens: 150, + }); + + // Fork with child_id = SequenceId(2) + let event = ForkEvent { + child_id: SequenceId(2), + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + let (parent, child) = result; + + // Verify fork succeeded + if let SequenceState::Decoding(c) = &child { + assert_eq!(c.seq_id, SequenceId(2)); + } + + // Now try to fork again with the SAME child_id (should fail in manager) + // This test verifies the ForkEvent itself works, but the manager + // should check for duplicate child_id before calling this + drop(parent); + drop(child); + } +} diff --git a/src/sequence_manager/fsm_manager.rs b/src/sequence_manager/fsm_manager.rs index ae74297..762f93b 100644 --- a/src/sequence_manager/fsm_manager.rs +++ b/src/sequence_manager/fsm_manager.rs @@ -99,6 +99,15 @@ impl SequenceManager { seq_id, prompt_tokens, max_tokens ); + // Check for duplicate seq_id + if self.sequences.contains_key(&seq_id) { + warn!( + "Duplicate AddRequest for {:?} - sequence already exists, ignoring", + seq_id + ); + return; + } + let state = SequenceState::Waiting(WaitingState { seq_id, prompt_tokens, @@ -342,3 +351,68 @@ impl SequenceManager { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + + fn create_test_manager() -> (SequenceManager, SequenceEventSender) { + let allocator = Box::new(CpuAllocator::new(1024 * 1024)); // 1MB + let block_manager = BlockManager::new(allocator, 256); // 256 byte blocks + SequenceManager::new(block_manager, 16) // 16 tokens per block + } + + #[test] + fn test_duplicate_add_request_ignored() { + let (mut manager, _tx) = create_test_manager(); + + let seq_id = SequenceId(0); + + // Add first request + manager.handle_add_request(seq_id, 32, 128); + + // Verify it was added + assert!(manager.sequences.contains_key(&seq_id)); + let first_state = manager.sequences.get(&seq_id).cloned(); + + // Try to add duplicate - should be ignored + manager.handle_add_request(seq_id, 64, 256); // Different params + + // Verify the original state is unchanged + let current_state = manager.sequences.get(&seq_id); + assert_eq!(format!("{:?}", first_state), format!("{:?}", current_state)); + + // Verify original params (may have transitioned to Prefilling via try_schedule) + match manager.sequences.get(&seq_id) { + Some(SequenceState::Waiting(s)) => { + assert_eq!(s.prompt_tokens, 32); // Original value + assert_eq!(s.max_tokens, 128); // Original value + } + Some(SequenceState::Prefilling(s)) => { + // Scheduled automatically, check original params + assert_eq!(s.tokens_total, 32); // Original prompt_tokens + assert_eq!(s.max_tokens, 128); // Original value + } + _ => panic!( + "Expected Waiting or Prefilling state, got: {:?}", + manager.sequences.get(&seq_id) + ), + } + } + + #[test] + fn test_add_different_sequences() { + let (mut manager, _tx) = create_test_manager(); + + // Add multiple different sequences - all should succeed + manager.handle_add_request(SequenceId(0), 32, 128); + manager.handle_add_request(SequenceId(1), 64, 256); + manager.handle_add_request(SequenceId(2), 16, 64); + + assert_eq!(manager.sequences.len(), 3); + assert!(manager.sequences.contains_key(&SequenceId(0))); + assert!(manager.sequences.contains_key(&SequenceId(1))); + assert!(manager.sequences.contains_key(&SequenceId(2))); + } +} diff --git a/src/sequence_manager/fsm_tests.rs b/src/sequence_manager/fsm_tests.rs deleted file mode 100644 index 5ec2bb0..0000000 --- a/src/sequence_manager/fsm_tests.rs +++ /dev/null @@ -1,365 +0,0 @@ -#[cfg(test)] -mod tests { - use crate::block_manager::allocator::CpuAllocator; - use crate::block_manager::manager::BlockManager; - use crate::block_manager::types::*; - use crate::sequence_manager::fsm_events::*; - use crate::sequence_manager::states::*; - use std::sync::Arc; - - #[test] - fn test_schedule_event_waiting_to_prefilling() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let waiting = SequenceState::Waiting(WaitingState { - seq_id: SequenceId(1), - prompt_tokens: 100, - max_tokens: 150, - }); - - let event = ScheduleEvent { - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(waiting); - assert!(result.is_ok()); - - let new_state = result.unwrap(); - assert!(matches!(new_state, SequenceState::Prefilling(_))); - - if let SequenceState::Prefilling(state) = new_state { - assert_eq!(state.seq_id, SequenceId(1)); - assert_eq!(state.tokens_filled, 0); - assert_eq!(state.tokens_total, 100); - // Should allocate ceil(100/16) = 7 blocks - assert_eq!(state.blocks.len(), 7); - } - } - - #[test] - fn test_schedule_event_oom() { - // Small allocator - only 2 blocks - let allocator = Box::new(CpuAllocator::new(2048)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let waiting = SequenceState::Waiting(WaitingState { - seq_id: SequenceId(1), - prompt_tokens: 100, // Needs 7 blocks - max_tokens: 150, - }); - - let event = ScheduleEvent { - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(waiting); - assert!(matches!(result, Err(Error::OutOfMemory))); - - // Should have cleaned up - no blocks leaked - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_append_tokens_prefilling_to_decoding() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - // Allocate blocks - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - - let prefilling = SequenceState::Prefilling(PrefillingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - tokens_filled: 0, - tokens_total: 100, - max_tokens: 150, - }); - - let event = AppendTokensEvent { - num_tokens: 100, - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(prefilling).unwrap(); - - // Should transition to Decoding - assert!(matches!(result, SequenceState::Decoding(_))); - - if let SequenceState::Decoding(state) = result { - assert_eq!(state.num_tokens, 100); - assert_eq!(state.blocks.len(), 2); - } - } - - #[test] - fn test_append_tokens_decoding_needs_more_blocks() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![block_manager.allocate().unwrap()]; // 1 block - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 16, // 1 block worth - max_tokens: 150, - }); - - let event = AppendTokensEvent { - num_tokens: 20, // Now need 3 blocks total - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(decoding).unwrap(); - - if let SequenceState::Decoding(state) = result { - assert_eq!(state.num_tokens, 36); - // Should allocate 2 more blocks (total 3) - assert_eq!(state.blocks.len(), 3); - } - } - - #[test] - fn test_append_tokens_decoding_oom_cleanup() { - // Small allocator - only enough for initial blocks - let allocator = Box::new(CpuAllocator::new(2048)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - let initial_block_count = blocks.len(); - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 32, - max_tokens: 150, - }); - - // Try to append tokens that would need more blocks (will fail - OOM) - let event = AppendTokensEvent { - num_tokens: 50, // Would need 6 blocks total - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(decoding); - assert!(matches!(result, Err(Error::OutOfMemory))); - - // Original 2 blocks should still be allocated (not freed by error path) - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, initial_block_count); - } - - #[test] - fn test_append_tokens_reaches_max_tokens() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![block_manager.allocate().unwrap()]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 140, - max_tokens: 150, - }); - - let event = AppendTokensEvent { - num_tokens: 10, - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(decoding).unwrap(); - - // Should transition to Finished - assert!(matches!(result, SequenceState::Finished(_))); - - if let SequenceState::Finished(state) = result { - assert_eq!(state.seq_id, SequenceId(1)); - assert_eq!(state.finish_reason, FinishReason::MaxTokens); - } - - // Blocks should be freed - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_preempt_event() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 32, - max_tokens: 150, - }); - - let event = PreemptEvent { - block_manager: &mut block_manager, - }; - - let result = event.apply(decoding).unwrap(); - - // Should transition to Preempted - assert!(matches!(result, SequenceState::Preempted(_))); - - if let SequenceState::Preempted(state) = result { - assert_eq!(state.num_tokens, 32); - } - - // Blocks should be freed - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_resume_event() { - let preempted = SequenceState::Preempted(PreemptedState { - seq_id: SequenceId(1), - num_tokens: 32, - max_tokens: 150, - }); - - let event = ResumeEvent; - let result = event.apply(preempted).unwrap(); - - // Should transition back to Waiting - assert!(matches!(result, SequenceState::Waiting(_))); - - if let SequenceState::Waiting(state) = result { - assert_eq!(state.prompt_tokens, 32); - assert_eq!(state.max_tokens, 150); - } - } - - #[test] - fn test_fork_event() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 32, - max_tokens: 150, - }); - - let event = ForkEvent { - child_id: SequenceId(2), - block_manager: &mut block_manager, - }; - - let result = event.apply(decoding).unwrap(); - let (parent, child) = result; - - // Both should be in Decoding state - assert!(matches!(parent, SequenceState::Decoding(_))); - assert!(matches!(child, SequenceState::Decoding(_))); - - if let (SequenceState::Decoding(p), SequenceState::Decoding(c)) = (parent, child) { - assert_eq!(p.seq_id, SequenceId(1)); - assert_eq!(c.seq_id, SequenceId(2)); - assert_eq!(p.blocks, c.blocks); // Shared blocks - } - } - - #[test] - fn test_complete_event() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![block_manager.allocate().unwrap()]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 100, - max_tokens: 150, - }); - - let event = CompleteEvent { - reason: FinishReason::Stop, - block_manager: &mut block_manager, - }; - - let result = event.apply(decoding).unwrap(); - - assert!(matches!(result, SequenceState::Finished(_))); - - if let SequenceState::Finished(state) = result { - assert_eq!(state.finish_reason, FinishReason::Stop); - } - - // Blocks should be freed - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_fork_duplicate_child_id() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 32, - max_tokens: 150, - }); - - // Fork with child_id = SequenceId(2) - let event = ForkEvent { - child_id: SequenceId(2), - block_manager: &mut block_manager, - }; - - let result = event.apply(decoding).unwrap(); - let (parent, child) = result; - - // Verify fork succeeded - if let SequenceState::Decoding(c) = &child { - assert_eq!(c.seq_id, SequenceId(2)); - } - - // Now try to fork again with the SAME child_id (should fail in manager) - // This test verifies the ForkEvent itself works, but the manager - // should check for duplicate child_id before calling this - drop(parent); - drop(child); - } -} diff --git a/src/sequence_manager/mod.rs b/src/sequence_manager/mod.rs index 63e628c..138bdca 100644 --- a/src/sequence_manager/mod.rs +++ b/src/sequence_manager/mod.rs @@ -4,7 +4,4 @@ pub mod fsm_manager; pub mod id_generator; pub mod states; -#[cfg(test)] -mod fsm_tests; - // Re-export public API for convenience From f93278b9009504687577c74ee37212dc8c850a68 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sat, 11 Jul 2026 10:37:11 +0100 Subject: [PATCH 20/33] add comments Signed-off-by: kerthcet --- docs/fsm_architecture.md | 472 ++++++----------------------- src/block_manager/manager.rs | 181 ++++++++++- src/block_manager/manager_tests.rs | 132 -------- src/block_manager/mod.rs | 3 - 4 files changed, 268 insertions(+), 520 deletions(-) delete mode 100644 src/block_manager/manager_tests.rs diff --git a/docs/fsm_architecture.md b/docs/fsm_architecture.md index 2bc3a91..92bd011 100644 --- a/docs/fsm_architecture.md +++ b/docs/fsm_architecture.md @@ -6,81 +6,44 @@ Finite State Machine pattern for sequence lifecycle management. ### State Machine per Sequence -Each sequence is an independent state machine: - -```rust -pub enum SequenceState { - Waiting(WaitingState), - Scheduling(SchedulingState), - Prefilling(PrefillingState), - Decoding(DecodingState), - Preempted(PreemptedState), - Finished(FinishedState), - Aborted(AbortedState), -} - -// Multiple sequences, different states -sequences: HashMap -``` +Each sequence is an independent state machine. Multiple sequences can be in different states simultaneously. + +**Implementation:** `src/sequence_manager/states.rs` + +**States:** +- `Waiting` - In queue, no resources allocated +- `Scheduling` - Being scheduled (transitional) +- `Prefilling` - Processing prompt +- `Decoding` - Generating tokens +- `Preempted` - Temporarily suspended (blocks freed) +- `Finished` - Completed successfully +- `Aborted` - Error or cancelled ### States Own Resources -```rust -pub struct WaitingState { - seq_id: SequenceId, - prompt_tokens: usize, - max_tokens: usize, - // No blocks - not allocated yet -} - -pub struct DecodingState { - seq_id: SequenceId, - blocks: Vec, // ← State owns blocks - num_tokens: usize, - max_tokens: usize, -} - -pub struct PreemptedState { - seq_id: SequenceId, - num_tokens: usize, - max_tokens: usize, - // No blocks - they were freed -} -``` +Each state owns the resources it needs: +- `WaitingState` has no blocks (not allocated yet) +- `DecodingState` owns blocks via `Arc>` +- `PreemptedState` has no blocks (freed) -**Benefit:** Type system enforces "Waiting has no blocks, Decoding must have blocks" +**Benefit:** Type system enforces resource invariants - can't access blocks that don't exist. ### Events as State Transformations -```rust -pub struct ScheduleEvent<'a> { - block_manager: &'a mut BlockManager, - tokens_per_block: usize, -} - -impl ScheduleEvent { - pub fn apply(self, state: SequenceState) -> Result { - match state { - SequenceState::Waiting(s) => { - // Allocate blocks - let blocks = allocate_blocks(s.prompt_tokens)?; - - // Transform to new state - Ok(SequenceState::Prefilling(PrefillingState { - seq_id: s.seq_id, - blocks, // Ownership transferred - tokens_filled: 0, - tokens_total: s.prompt_tokens, - max_tokens: s.max_tokens, - })) - } - _ => Err(Error::invalid_transition("ScheduleEvent requires Waiting")), - } - } -} -``` +Events transform states: `Event + OldState → NewState` + +**Implementation:** `src/sequence_manager/fsm_events.rs` + +**Events:** +- `ScheduleEvent` - Waiting → Prefilling (allocate blocks) +- `AppendTokensEvent` - Prefilling → Decoding or Decoding → Decoding/Finished +- `ForkEvent` - Decoding → (Decoding, Decoding) for beam search +- `PreemptEvent` - Decoding → Preempted (free blocks on OOM) +- `ResumeEvent` - Preempted → Waiting (re-queue) +- `CompleteEvent` - * → Finished (free blocks) +- `AbortEvent` - * → Aborted (cleanup) -**Pattern:** `Event + OldState → NewState` +Invalid transitions return `Error::InvalidTransition`. ## State Diagram @@ -91,348 +54,101 @@ impl ScheduleEvent { │ ↓ ScheduleEvent │ Prefilling │ ↓ AppendTokens - │ ↓ (tokens_filled >= tokens_total) │ Decoding │ ├─→ AppendTokens (continue) - │ │ ↓ │ │ ├─ Decoding (more tokens) - │ │ └─ Finished (max_tokens reached) + │ │ └─ Finished (max_tokens) │ │ - │ ├─→ ForkEvent + │ ├─→ ForkEvent (beam search) │ │ ├─ Parent: Decoding │ │ └─ Child: Decoding │ │ │ ├─→ PreemptEvent (OOM) - │ │ ↓ - │ │ Preempted + │ │ ↓ Preempted │ │ ↓ ResumeEvent │ └───┘ │ └─ CompleteEvent → Finished ``` -## State Definitions - -```rust -/// Waiting in queue for scheduling -pub struct WaitingState { - pub seq_id: SequenceId, - pub prompt_tokens: usize, - pub max_tokens: usize, -} - -/// Running prefill phase -pub struct PrefillingState { - pub seq_id: SequenceId, - pub blocks: Vec, // Owns blocks - pub tokens_filled: usize, - pub tokens_total: usize, - pub max_tokens: usize, -} - -/// Running decode phase -pub struct DecodingState { - pub seq_id: SequenceId, - pub blocks: Vec, // Owns blocks - pub num_tokens: usize, - pub max_tokens: usize, -} - -/// Preempted due to OOM -pub struct PreemptedState { - pub seq_id: SequenceId, - pub num_tokens: usize, - pub max_tokens: usize, - // Blocks freed -} - -/// Finished generation -pub struct FinishedState { - pub seq_id: SequenceId, - pub finish_reason: FinishReason, - // All resources freed -} -``` - -## Event Definitions - -### ScheduleEvent: Waiting → Prefilling - -```rust -let event = ScheduleEvent { - block_manager: &mut block_manager, - tokens_per_block: 16, -}; - -match event.apply(state) { - Ok(new_state) => { - // Waiting → Prefilling (with blocks allocated) - } - Err(Error::OutOfMemory) => { - // Can't allocate, stay in queue - } -} -``` - -### AppendTokensEvent: Prefilling → Decoding or Decoding → Decoding +## Sequence Manager -```rust -let event = AppendTokensEvent { - num_tokens: 100, - block_manager: &mut block_manager, - tokens_per_block: 16, -}; +Event-driven architecture with async event loop. -// From Prefilling -Prefilling(tokens_filled: 0) + AppendTokens(100) - → Decoding(num_tokens: 100) // Transition +**Implementation:** `src/sequence_manager/fsm_manager.rs` -// From Decoding -Decoding(num_tokens: 100) + AppendTokens(1) - → Decoding(num_tokens: 101) // Stay in state -``` +**Key features:** +- Receives events via async channel +- Maintains `HashMap` +- Clones state before transitions (prevents loss on failure) +- Restores backup on error (no block leaks) +- Duplicate `seq_id` check prevents overwrites +- Duplicate `child_id` check prevents fork overwrites -### ForkEvent: Decoding → (Decoding, Decoding) - -```rust -let event = ForkEvent { - child_id: SequenceId(2), - block_manager: &mut block_manager, -}; - -// Copy-on-write: both sequences share blocks -match event.apply(parent_state) { - Ok((parent, child)) => { - // Both in Decoding state - // Blocks ref-counted (shared) - } -} -``` +**Tests:** Inline in `fsm_manager.rs` and `fsm_events.rs` -### PreemptEvent: Decoding → Preempted +## Arc Optimization -```rust -let event = PreemptEvent { - block_manager: &mut block_manager, -}; +States use `Arc>` for efficient cloning: -// Frees blocks -Decoding(blocks: [1,2,3]) → Preempted(blocks: []) -``` +- **State cloning** (for backup/fork): **O(1)** - just increment ref count +- **Modifying blocks**: Copy-on-write via `Arc::make_mut()` - only copies if shared +- **Fork sequences**: Share blocks via Arc, diverge only when modified -### ResumeEvent: Preempted → Waiting +**Benefit:** Hot path performance - state cloning costs ~50 bytes instead of O(n) vector copy. -```rust -let event = ResumeEvent; +## Error Handling -// Back to waiting for rescheduling -Preempted → Waiting → (Schedule) → Decoding -``` +All state transitions are wrapped with backup/restore: -## Sequence Manager +1. Remove state from HashMap +2. Clone state as backup +3. Apply event transformation +4. On success: insert new state +5. On failure: restore backup, log warning -```rust -pub struct SequenceManager { - block_manager: BlockManager, - sequences: HashMap, - tokens_per_block: usize, - event_rx: SequenceEventReceiver, -} - -impl SequenceManager { - pub async fn run(mut self) { - while let Some(event) = self.event_rx.recv().await { - match event { - SequenceEvent::AppendTokens { seq_id, num_tokens } => { - // Get current state - let state = self.sequences.remove(&seq_id).unwrap(); - - // Clone state before transition (prevents loss on failure) - let backup = state.clone(); - - // Apply event transformation - let event = AppendTokensEvent { - num_tokens, - block_manager: &mut self.block_manager, - tokens_per_block: self.tokens_per_block, - }; - - // Transform state - match event.apply(state) { - Ok(new_state) => { - self.sequences.insert(seq_id, new_state); - } - Err(e) => { - error!("Transition failed: {:?}", e); - // Restore original state to prevent sequence loss - self.sequences.insert(seq_id, backup); - } - } - } - // ... other events - } - } - } -} -``` +**Prevents:** +- Sequence loss on transition failure +- Block leaks (blocks remain owned by restored state) +- Duplicate sequences (check before insert) ## Type Safety -**Compile-time checks:** -```rust -// ✅ Valid -let waiting = SequenceState::Waiting(...); -let event = ScheduleEvent { ... }; -event.apply(waiting)?; // OK - -// ❌ Invalid (caught at runtime, enforced by match) -let decoding = SequenceState::Decoding(...); -let event = ScheduleEvent { ... }; -event.apply(decoding)?; // Error: invalid_transition -``` +**Runtime checks:** +- Invalid transitions caught by pattern matching +- Return `Error::InvalidTransition` with clear message **Resource ownership:** -```rust -fn transition(state: DecodingState) -> FinishedState { - // state.blocks moved/consumed - for block in state.blocks { - block_manager.free(block); - } - - FinishedState { - seq_id: state.seq_id, - finish_reason: FinishReason::Stop, - } - // Old state dropped, can't access blocks anymore -} -``` +- States own their blocks via Arc +- Type system prevents accessing freed blocks +- Drop trait ensures cleanup ## Benefits -1. **Type Safety** - - States enforce resource invariants - - Invalid transitions caught explicitly - -2. **Clear Ownership** - - Resources belong to states - - Automatic cleanup on state drop - -3. **Testability** - - Events are pure functions - - Test state transitions in isolation - -4. **Explicitness** - - Every transition is a function call - - State diagram maps directly to code - -## Example Flow - -```rust -// 1. Add request → Waiting -let state = SequenceState::Waiting(WaitingState { - seq_id: SequenceId(1), - prompt_tokens: 100, - max_tokens: 150, -}); - -// 2. Schedule → Prefilling (allocate 7 blocks) -let event = ScheduleEvent { ... }; -let state = event.apply(state)?; -// state = Prefilling(blocks: [1,2,3,4,5,6,7], tokens_filled: 0) - -// 3. Append tokens → Decoding (transition) -let event = AppendTokensEvent { num_tokens: 100 }; -let state = event.apply(state)?; -// state = Decoding(blocks: [1,2,3,4,5,6,7], num_tokens: 100) - -// 4. Fork → 2 sequences -let event = ForkEvent { child_id: SequenceId(2) }; -let (parent, child) = event.apply(state)?; -// parent = Decoding(blocks: [1,2,3,4,5,6,7], ref_count: 2) -// child = Decoding(blocks: [1,2,3,4,5,6,7], ref_count: 2) - -// 5. Continue generating -let event = AppendTokensEvent { num_tokens: 10 }; -let parent = event.apply(parent)?; -// parent = Decoding(blocks: [1,2,3,4,5,6,7,8], num_tokens: 110) - -// 6. Complete → Finished (free blocks) -let event = CompleteEvent { ... }; -let state = event.apply(parent)?; -// state = Finished, blocks freed -``` - -## Comparison with Simple Design - -| Simple | FSM | -|--------|-----| -| `seq.state = Running` | `state = ScheduleEvent.apply(state)` | -| Resources in struct | Resources in state variant | -| Manual cleanup | Automatic cleanup | -| Runtime checks | Type + runtime checks | - -## Error Handling and State Recovery - -On transition failure, the state is cloned before applying the event: - -```rust -fn handle_append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { - let state = self.sequences.remove(&seq_id).unwrap(); - let backup = state.clone(); // Clone before transition - - let event = AppendTokensEvent { ... }; - - match event.apply(state) { - Ok(new_state) => { - self.sequences.insert(seq_id, new_state); - } - Err(Error::OutOfMemory) => { - // Restore state and trigger preemption - self.sequences.insert(seq_id, backup); - self.handle_oom(); - } - Err(e) => { - // Restore state on any error (prevents block leaks) - self.sequences.insert(seq_id, backup); - warn!("Transition failed: {:?}", e); - } - } -} -``` - -**Benefits:** -- No sequence loss on transition failure -- No block leaks (blocks remain owned by restored state) -- Can retry or handle errors without losing context -- Clone cost is negligible (~50-100 bytes per state) - -## Usage - -```rust -use puma::block_manager::{BlockManager, CpuAllocator}; -use puma::sequence_manager::{SequenceManager, SequenceEvent, SequenceIdGenerator}; - -// Setup -let allocator = Box::new(CpuAllocator::new(100_000_000)); -let block_manager = BlockManager::new(allocator, 4096); -let (seq_manager, event_tx) = SequenceManager::new(block_manager, 16); - -// ID generator for sequential IDs starting from 0 -let id_gen = SequenceIdGenerator::new(); -let seq_id = id_gen.next(); // SequenceId(0) - -// Run event loop -tokio::spawn(async move { seq_manager.run().await }); - -// Send events -event_tx.send(SequenceEvent::AddRequest { - seq_id, - prompt_tokens: 100, - max_tokens: 150, -}).unwrap(); - -event_tx.send(SequenceEvent::AppendTokens { - seq_id, - num_tokens: 100, -}).unwrap(); -``` +1. **Type Safety** - States enforce resource invariants, invalid transitions caught explicitly +2. **Clear Ownership** - Resources belong to states, automatic cleanup on drop +3. **Testability** - Events are pure functions, test state transitions in isolation +4. **Explicitness** - Every transition is explicit, state diagram maps to code +5. **Performance** - Arc enables O(1) cloning, copy-on-write for efficiency +6. **Reliability** - Backup/restore pattern prevents data loss on errors + +## Key Files + +- `src/sequence_manager/states.rs` - State definitions +- `src/sequence_manager/fsm_events.rs` - Event implementations + tests +- `src/sequence_manager/fsm_manager.rs` - SequenceManager + tests +- `src/sequence_manager/events.rs` - Event channel types +- `src/sequence_manager/id_generator.rs` - Sequential ID generator + +## Usage Example + +See `src/sequence_manager/fsm_manager.rs` tests for complete examples. + +Basic flow: +1. Create BlockManager and SequenceManager +2. Get event sender channel +3. Spawn event loop in background +4. Send `AddRequest` event +5. Send `AppendTokens` events as tokens generate +6. Sequence automatically transitions through states +7. Finished/Aborted when done diff --git a/src/block_manager/manager.rs b/src/block_manager/manager.rs index 449236e..b8176e8 100644 --- a/src/block_manager/manager.rs +++ b/src/block_manager/manager.rs @@ -3,21 +3,31 @@ use super::types::*; use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; +/// Manages memory blocks with reference counting and pooling. +/// +/// Features: +/// - Reference counting: blocks can be shared, freed only when ref_count = 0 +/// - Free pools: reuse freed blocks instead of allocating new memory +/// - Type-specific blocks: different block types with different sizes +/// - Thread-safe ID generation: atomic counter for BlockId pub struct BlockManager { + /// Underlying memory allocator (CPU or GPU) allocator: Box, - // Block registry + /// All blocks (allocated + free) - tracks metadata and physical memory block_table: HashMap, - // Free pools per type + /// Free block pools per type - blocks ready for reuse (ref_count = 0) + /// Keyed by BlockType, values are BlockIds ready to be allocated free_pools: HashMap>, - // Block configurations + /// Configuration for each block type (size, etc.) block_configs: HashMap, - // Default type + /// Default block type for allocate() calls default_block_type: BlockType, + /// Atomic counter for generating unique BlockIds next_block_id: AtomicU64, } @@ -41,13 +51,20 @@ impl BlockManager { manager } + /// Register a new block type with specific configuration. pub fn register_block_type(&mut self, config: BlockConfig) { self.free_pools.insert(config.block_type, Vec::new()); self.block_configs.insert(config.block_type, config); } + /// Allocate a block of specific type. + /// + /// Flow: + /// 1. Check free pool first (reuse if available) + /// 2. If pool empty, allocate new physical memory + /// 3. Return BlockId with ref_count = 1 pub fn allocate_typed(&mut self, block_type: &BlockType) -> Result { - // Try free pool first + // Try free pool first - reuse freed blocks if let Some(block_id) = self .free_pools .get_mut(block_type) @@ -58,7 +75,7 @@ impl BlockManager { return Ok(block_id); } - // Allocate new physical memory + // No free blocks - allocate new physical memory let config = self .block_configs .get(block_type) @@ -81,11 +98,16 @@ impl BlockManager { Ok(block_id) } + /// Allocate a block using the default type. pub fn allocate(&mut self) -> Result { let default_type = self.default_block_type; self.allocate_typed(&default_type) } + /// Decrement reference count. When ref_count reaches 0, block goes to free pool. + /// + /// Note: Physical memory is NOT freed - block is just marked as reusable. + /// This enables fast reallocation without syscalls. pub fn free(&mut self, block_id: BlockId) -> Result<()> { let block = self .block_table @@ -98,6 +120,7 @@ impl BlockManager { block.ref_count -= 1; + // When ref_count hits 0, move to free pool for reuse if block.ref_count == 0 { self.free_pools .get_mut(&block.block_type) @@ -108,6 +131,9 @@ impl BlockManager { Ok(()) } + /// Increment reference count (for sharing blocks between sequences). + /// + /// Used for fork/beam search where multiple sequences share the same blocks. pub fn add_ref(&mut self, block_id: BlockId) -> Result<()> { let block = self .block_table @@ -117,6 +143,7 @@ impl BlockManager { Ok(()) } + /// Get the physical memory address for a block. pub fn get_memory_address(&self, block_id: BlockId) -> Result { self.block_table .get(&block_id) @@ -124,6 +151,7 @@ impl BlockManager { .ok_or(Error::InvalidBlockId(block_id)) } + /// Get the block type for a block. pub fn get_block_type(&self, block_id: BlockId) -> Result<&BlockType> { self.block_table .get(&block_id) @@ -131,6 +159,7 @@ impl BlockManager { .ok_or(Error::InvalidBlockId(block_id)) } + /// Get statistics for all block types (allocated, free, total memory, etc.). pub fn get_stats(&self) -> HashMap { let mut stats = HashMap::new(); @@ -159,13 +188,16 @@ impl BlockManager { stats } + /// Check if a block can be allocated (either from free pool or new memory). pub fn can_allocate(&self, block_type: &BlockType) -> bool { + // Check free pool first if let Some(pool) = self.free_pools.get(block_type) { if !pool.is_empty() { return true; } } + // Check if allocator has enough memory if let Some(config) = self.block_configs.get(block_type) { return self.allocator.get_available_memory() >= config.size_bytes; } @@ -175,8 +207,11 @@ impl BlockManager { } impl Drop for BlockManager { + /// Free all physical memory when BlockManager is dropped. + /// + /// This ensures no memory leaks even if blocks weren't explicitly freed. fn drop(&mut self) { - // Free all backing memory allocations to prevent leaks + // Free all backing memory allocations for block in self.block_table.values() { if let Err(e) = self.allocator.free(block.mem_addr) { eprintln!( @@ -187,3 +222,135 @@ impl Drop for BlockManager { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + + #[test] + fn test_allocate_and_free() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + // Allocate a block + let block_id = manager.allocate().unwrap(); + assert_eq!(block_id, BlockId(0)); + + // Free it + manager.free(block_id).unwrap(); + + // Should be back in free pool + let stats = manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.free_blocks, 1); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_ref_counting() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + let block_id = manager.allocate().unwrap(); + + // Add ref + manager.add_ref(block_id).unwrap(); + + // Free once - should still be allocated (ref_count = 1) + manager.free(block_id).unwrap(); + let stats = manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 1); + assert_eq!(block_stats.free_blocks, 0); + + // Free again - now should be in free pool (ref_count = 0) + manager.free(block_id).unwrap(); + let stats = manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.free_blocks, 1); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_double_free_error() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + let block_id = manager.allocate().unwrap(); + manager.free(block_id).unwrap(); + + // Should fail on double free + let result = manager.free(block_id); + assert!(matches!(result, Err(Error::DoubleFree(_)))); + } + + #[test] + fn test_oom() { + // Small allocator - only enough for 2 blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let mut manager = BlockManager::new(allocator, 1024); + + // First block succeeds + let block1 = manager.allocate().unwrap(); + assert_eq!(block1, BlockId(0)); + + // Second block succeeds + let block2 = manager.allocate().unwrap(); + assert_eq!(block2, BlockId(1)); + + // Third block should fail (OOM) + let result = manager.allocate(); + assert!(matches!(result, Err(Error::OutOfMemory))); + } + + #[test] + fn test_free_pool_reuse() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + // Allocate and free + let block1 = manager.allocate().unwrap(); + manager.free(block1).unwrap(); + + // Next allocation should reuse from free pool (same ID) + let block2 = manager.allocate().unwrap(); + assert_eq!(block1, block2); + } + + #[test] + fn test_get_memory_address() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let mut manager = BlockManager::new(allocator, 1024); + + let block_id = manager.allocate().unwrap(); + let addr = manager.get_memory_address(block_id); + assert!(addr.is_ok()); + } + + #[test] + fn test_invalid_block_id() { + let allocator = Box::new(CpuAllocator::new(10_000)); + let manager = BlockManager::new(allocator, 1024); + + let invalid_id = BlockId(999); + let result = manager.get_memory_address(invalid_id); + assert!(matches!(result, Err(Error::InvalidBlockId(_)))); + } + + #[test] + fn test_can_allocate() { + let allocator = Box::new(CpuAllocator::new(2048)); + let mut manager = BlockManager::new(allocator, 1024); + + // Should be able to allocate + assert!(manager.can_allocate(&BlockType::StandardKV)); + + // Allocate both blocks + manager.allocate().unwrap(); + manager.allocate().unwrap(); + + // Should not be able to allocate more + assert!(!manager.can_allocate(&BlockType::StandardKV)); + } +} diff --git a/src/block_manager/manager_tests.rs b/src/block_manager/manager_tests.rs deleted file mode 100644 index 57bc219..0000000 --- a/src/block_manager/manager_tests.rs +++ /dev/null @@ -1,132 +0,0 @@ -#[cfg(test)] -mod tests { - use super::super::allocator::CpuAllocator; - use super::super::manager::BlockManager; - use super::super::types::*; - - #[test] - fn test_allocate_and_free() { - let allocator = Box::new(CpuAllocator::new(10_000)); - let mut manager = BlockManager::new(allocator, 1024); - - // Allocate a block - let block_id = manager.allocate().unwrap(); - assert_eq!(block_id, BlockId(0)); - - // Free it - manager.free(block_id).unwrap(); - - // Should be back in free pool - let stats = manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.free_blocks, 1); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_ref_counting() { - let allocator = Box::new(CpuAllocator::new(10_000)); - let mut manager = BlockManager::new(allocator, 1024); - - let block_id = manager.allocate().unwrap(); - - // Add ref - manager.add_ref(block_id).unwrap(); - - // Free once - should still be allocated (ref_count = 1) - manager.free(block_id).unwrap(); - let stats = manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, 1); - assert_eq!(block_stats.free_blocks, 0); - - // Free again - now should be in free pool (ref_count = 0) - manager.free(block_id).unwrap(); - let stats = manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.free_blocks, 1); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_double_free_error() { - let allocator = Box::new(CpuAllocator::new(10_000)); - let mut manager = BlockManager::new(allocator, 1024); - - let block_id = manager.allocate().unwrap(); - manager.free(block_id).unwrap(); - - // Should fail on double free - let result = manager.free(block_id); - assert!(matches!(result, Err(Error::DoubleFree(_)))); - } - - #[test] - fn test_oom() { - // Small allocator - only enough for 2 blocks - let allocator = Box::new(CpuAllocator::new(2048)); - let mut manager = BlockManager::new(allocator, 1024); - - // First block succeeds - let block1 = manager.allocate().unwrap(); - assert_eq!(block1, BlockId(0)); - - // Second block succeeds - let block2 = manager.allocate().unwrap(); - assert_eq!(block2, BlockId(1)); - - // Third block should fail (OOM) - let result = manager.allocate(); - assert!(matches!(result, Err(Error::OutOfMemory))); - } - - #[test] - fn test_free_pool_reuse() { - let allocator = Box::new(CpuAllocator::new(10_000)); - let mut manager = BlockManager::new(allocator, 1024); - - // Allocate and free - let block1 = manager.allocate().unwrap(); - manager.free(block1).unwrap(); - - // Next allocation should reuse from free pool (same ID) - let block2 = manager.allocate().unwrap(); - assert_eq!(block1, block2); - } - - #[test] - fn test_get_memory_address() { - let allocator = Box::new(CpuAllocator::new(10_000)); - let mut manager = BlockManager::new(allocator, 1024); - - let block_id = manager.allocate().unwrap(); - let addr = manager.get_memory_address(block_id); - assert!(addr.is_ok()); - } - - #[test] - fn test_invalid_block_id() { - let allocator = Box::new(CpuAllocator::new(10_000)); - let manager = BlockManager::new(allocator, 1024); - - let invalid_id = BlockId(999); - let result = manager.get_memory_address(invalid_id); - assert!(matches!(result, Err(Error::InvalidBlockId(_)))); - } - - #[test] - fn test_can_allocate() { - let allocator = Box::new(CpuAllocator::new(2048)); - let mut manager = BlockManager::new(allocator, 1024); - - // Should be able to allocate - assert!(manager.can_allocate(&BlockType::StandardKV)); - - // Allocate both blocks - manager.allocate().unwrap(); - manager.allocate().unwrap(); - - // Should not be able to allocate more - assert!(!manager.can_allocate(&BlockType::StandardKV)); - } -} diff --git a/src/block_manager/mod.rs b/src/block_manager/mod.rs index 7279f14..a4468f5 100644 --- a/src/block_manager/mod.rs +++ b/src/block_manager/mod.rs @@ -3,6 +3,3 @@ pub mod manager; pub mod types; // Re-export for convenience - -#[cfg(test)] -mod manager_tests; From 6dec6350aac7b080b8038c1f92b5fc73ce908669 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sat, 11 Jul 2026 19:57:27 +0100 Subject: [PATCH 21/33] scheduler Signed-off-by: kerthcet --- src/lib.rs | 2 +- src/main.rs | 2 +- src/scheduler/core.rs | 806 ++++++++++++++++++ src/scheduler/events.rs | 58 ++ .../fsm_events.rs | 0 .../id_generator.rs | 0 src/{sequence_manager => scheduler}/mod.rs | 4 +- src/{sequence_manager => scheduler}/states.rs | 0 src/sequence_manager/events.rs | 73 -- src/sequence_manager/fsm_manager.rs | 418 --------- 10 files changed, 867 insertions(+), 496 deletions(-) create mode 100644 src/scheduler/core.rs create mode 100644 src/scheduler/events.rs rename src/{sequence_manager => scheduler}/fsm_events.rs (100%) rename src/{sequence_manager => scheduler}/id_generator.rs (100%) rename src/{sequence_manager => scheduler}/mod.rs (54%) rename src/{sequence_manager => scheduler}/states.rs (100%) delete mode 100644 src/sequence_manager/events.rs delete mode 100644 src/sequence_manager/fsm_manager.rs diff --git a/src/lib.rs b/src/lib.rs index 43546f1..0f9ad10 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,7 +6,7 @@ pub mod block_manager; pub mod cli; pub mod downloader; pub mod registry; -pub mod sequence_manager; +pub mod scheduler; pub mod storage; pub mod system; pub mod utils; diff --git a/src/main.rs b/src/main.rs index aa4feeb..45a5cd7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,7 +6,7 @@ mod block_manager; mod cli; mod downloader; mod registry; -mod sequence_manager; +mod scheduler; mod storage; mod system; mod utils; diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs new file mode 100644 index 0000000..e250e78 --- /dev/null +++ b/src/scheduler/core.rs @@ -0,0 +1,806 @@ +use super::events::*; +use super::fsm_events::*; +use super::states::*; +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; +use std::collections::{HashMap, VecDeque}; +use std::sync::Arc; +use tokio::sync::mpsc; +use tracing::{debug, info, warn}; + +/// Scheduler - central controller +/// +/// Architecture: +/// - Scheduler DRIVES inference: builds batches → calls backend → updates state +/// - External events (via channel): AddRequest, CancelRequest, GetBlocks, GetStats +/// - Internal operations: token generation, state transitions, OOM handling +/// +/// Responsibilities: +/// - Owns BlockManager (memory allocation) +/// - Owns sequence state storage (FSM states) +/// - Manages batches (waiting/prefill/decode) +/// - Scheduling policy (what to schedule when) +/// - Uses FSM events for type-safe state transitions +pub struct Scheduler { + /// Block manager for memory allocation + block_manager: BlockManager, + + /// All sequences and their states + sequences: HashMap, + + /// Waiting sequences (not yet scheduled) + waiting_queue: VecDeque, + + /// Currently running sequences in prefill phase + prefill_batch: Vec, + + /// Currently running sequences in decode phase + decode_batch: Vec, + + /// Event receiver + event_rx: mpsc::UnboundedReceiver, + + /// Configuration + max_batch_size: usize, + tokens_per_block: usize, +} + +impl Scheduler { + pub fn new( + block_manager: BlockManager, + max_batch_size: usize, + tokens_per_block: usize, + ) -> (Self, mpsc::UnboundedSender) { + let (event_tx, event_rx) = mpsc::unbounded_channel(); + + let scheduler = Self { + block_manager, + sequences: HashMap::new(), + waiting_queue: VecDeque::new(), + prefill_batch: Vec::new(), + decode_batch: Vec::new(), + event_rx, + max_batch_size, + tokens_per_block, + }; + + (scheduler, event_tx) + } + + /// Main scheduler loop - drives inference and handles external events + /// + /// Loop structure: + /// 1. Build batch from ready sequences + /// 2. Run inference (TODO: integrate backend) + /// 3. Update state based on inference results (internal) + /// 4. Handle external events (non-blocking) + /// 5. Try scheduling new sequences + pub async fn run(mut self) { + info!("Scheduler started"); + + loop { + // 1. Build batch for inference + // TODO: Call backend.forward(batch) when backend is integrated + + // 2. Handle external events (non-blocking) + // Use try_recv to not block - we want to keep generating tokens + match self.event_rx.try_recv() { + Ok(event) => { + match event { + SchedulerEvent::AddRequest { + seq_id, + prompt_tokens, + max_tokens, + } => { + self.add_request(seq_id, prompt_tokens, max_tokens); + } + + SchedulerEvent::CancelRequest { seq_id } => { + self.cancel_request(seq_id); + } + + SchedulerEvent::GetBlocks { seq_id, response } => { + let result = self.get_blocks(seq_id); + let _ = response.send(result); + } + + SchedulerEvent::GetStats { response } => { + let stats = self.get_stats(); + let _ = response.send(stats); + } + } + + // After event, try scheduling + self.try_schedule(); + } + Err(mpsc::error::TryRecvError::Empty) => { + // No events - continue to next iteration + } + Err(mpsc::error::TryRecvError::Disconnected) => { + info!("Event channel closed, shutting down scheduler"); + break; + } + } + + // Small yield to prevent busy loop + // TODO: Remove once backend integration drives the timing + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + } + + info!("Scheduler stopped"); + } + + // ===== External Event Handlers ===== + + fn add_request(&mut self, seq_id: SequenceId, prompt_tokens: usize, max_tokens: usize) { + debug!( + "Adding request: seq_id={:?}, prompt_tokens={}, max_tokens={}", + seq_id, prompt_tokens, max_tokens + ); + + // Check duplicate + if self.sequences.contains_key(&seq_id) { + warn!("Duplicate seq_id {:?}", seq_id); + return; + } + + // Create waiting state + let state = SequenceState::Waiting(WaitingState { + seq_id, + prompt_tokens, + max_tokens, + }); + + self.sequences.insert(seq_id, state); + self.waiting_queue.push_back(seq_id); + } + + fn cancel_request(&mut self, seq_id: SequenceId) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Cannot cancel - sequence {:?} not found", seq_id); + return; + } + }; + + let event = AbortEvent { + reason: "User cancelled".to_string(), + block_manager: &mut self.block_manager, + }; + + match event.apply(state) { + Ok(new_state) => { + info!("Cancelled sequence {:?}", seq_id); + self.sequences.insert(seq_id, new_state); + + // Remove from queues/batches + self.waiting_queue.retain(|&id| id != seq_id); + self.prefill_batch.retain(|&id| id != seq_id); + self.decode_batch.retain(|&id| id != seq_id); + } + Err(e) => { + warn!("Failed to cancel {:?}: {:?}", seq_id, e); + } + } + } + + // ===== Internal Methods (called by scheduler loop during inference) ===== + // These are NOT exposed as events - scheduler calls them directly when: + // - Generating tokens (append_tokens) + // - Detecting stop tokens (complete_sequence) + // - Handling OOM (preempt_sequence, resume_sequence) + // - Beam search (fork_sequence) + + fn append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let backup = state.clone(); + let event = AppendTokensEvent { + num_tokens, + block_manager: &mut self.block_manager, + tokens_per_block: self.tokens_per_block, + }; + + match event.apply(state) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + } + Err(Error::OutOfMemory) => { + warn!("OOM while appending tokens to {:?}", seq_id); + self.sequences.insert(seq_id, backup); + self.handle_oom(); + } + Err(e) => { + warn!("Failed to append tokens to {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); + } + } + } + + fn complete_sequence(&mut self, seq_id: SequenceId, reason: FinishReason) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let backup = state.clone(); + let event = CompleteEvent { + reason, + block_manager: &mut self.block_manager, + }; + + match event.apply(state) { + Ok(new_state) => { + info!("Completed sequence {:?}", seq_id); + self.sequences.insert(seq_id, new_state); + + // Remove from batches + self.prefill_batch.retain(|&id| id != seq_id); + self.decode_batch.retain(|&id| id != seq_id); + } + Err(e) => { + warn!("Failed to complete {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); + } + } + } + + fn fork_sequence(&mut self, parent_id: SequenceId, child_id: SequenceId) { + // Check duplicate child_id + if self.sequences.contains_key(&child_id) { + warn!("Child ID {:?} already exists", child_id); + return; + } + + let parent_state = match self.sequences.remove(&parent_id) { + Some(s) => s, + None => { + warn!("Parent sequence {:?} not found", parent_id); + return; + } + }; + + let backup = parent_state.clone(); + let event = ForkEvent { + child_id, + block_manager: &mut self.block_manager, + }; + + match event.apply(parent_state) { + Ok((parent_state, child_state)) => { + self.sequences.insert(parent_id, parent_state); + self.sequences.insert(child_id, child_state); + debug!("Forked sequence {:?} → {:?}", parent_id, child_id); + } + Err(e) => { + warn!("Failed to fork {:?}: {:?}", parent_id, e); + self.sequences.insert(parent_id, backup); + } + } + } + + fn preempt_sequence(&mut self, seq_id: SequenceId) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let backup = state.clone(); + let event = PreemptEvent { + block_manager: &mut self.block_manager, + }; + + match event.apply(state) { + Ok(new_state) => { + info!("Preempted sequence {:?}", seq_id); + self.sequences.insert(seq_id, new_state); + + // Remove from batches + self.prefill_batch.retain(|&id| id != seq_id); + self.decode_batch.retain(|&id| id != seq_id); + } + Err(e) => { + warn!("Failed to preempt {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); + } + } + } + + fn resume_sequence(&mut self, seq_id: SequenceId) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let backup = state.clone(); + let event = ResumeEvent; + + match event.apply(state) { + Ok(new_state) => { + debug!("Resumed sequence {:?}", seq_id); + self.sequences.insert(seq_id, new_state); + self.waiting_queue.push_back(seq_id); + } + Err(e) => { + warn!("Failed to resume {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); + } + } + } + + // ===== Scheduling Logic ===== + + fn try_schedule(&mut self) { + self.schedule_prefill(); + self.schedule_decode(); + } + + fn schedule_prefill(&mut self) { + while let Some(seq_id) = self.waiting_queue.pop_front() { + // Check batch limit + if self.prefill_batch.len() >= self.max_batch_size { + self.waiting_queue.push_front(seq_id); + debug!("Prefill batch full"); + break; + } + + // Get sequence state + let state = match self.sequences.get(&seq_id) { + Some(SequenceState::Waiting(s)) => s, + _ => { + warn!("Sequence {:?} not in Waiting state", seq_id); + continue; + } + }; + + // Calculate blocks needed + let blocks_needed = state.prompt_tokens.div_ceil(self.tokens_per_block); + + // Check memory + if !self.block_manager.can_allocate(&BlockType::StandardKV) { + self.waiting_queue.push_front(seq_id); + debug!("OOM - cannot schedule {:?}", seq_id); + break; + } + + // Allocate blocks + let mut blocks = Vec::new(); + for _ in 0..blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => blocks.push(block_id), + Err(Error::OutOfMemory) => { + // Cleanup and stop + for b in blocks { + let _ = self.block_manager.free(b); + } + self.waiting_queue.push_front(seq_id); + warn!("OOM during allocation for {:?}", seq_id); + return; + } + Err(e) => { + for b in blocks { + let _ = self.block_manager.free(b); + } + warn!("Allocation error: {:?}", e); + return; + } + } + } + + // Transition: Waiting → Prefilling + let state = self.sequences.remove(&seq_id).unwrap(); + if let SequenceState::Waiting(w) = state { + let new_state = SequenceState::Prefilling(PrefillingState { + seq_id, + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: w.prompt_tokens, + max_tokens: w.max_tokens, + }); + + self.sequences.insert(seq_id, new_state); + self.prefill_batch.push(seq_id); + + info!("Scheduled {:?} for prefill", seq_id); + } + } + } + + fn schedule_decode(&mut self) { + // TODO: Continuous batching + // - Move completed prefill → decode batch + // - Remove finished sequences + // - Respect max_batch_size + + debug!("Decode batch size: {}", self.decode_batch.len()); + } + + fn handle_oom(&mut self) { + warn!("Handling OOM"); + // TODO: Preemption policy + // Find lowest priority sequence + // Preempt it + // Retry scheduling + } + + // ===== Query Methods ===== + + fn get_blocks(&self, seq_id: SequenceId) -> Result> { + match self.sequences.get(&seq_id) { + Some(state) => Ok(state.blocks().map(|b| b.to_vec()).unwrap_or_default()), + None => Err(Error::UnknownSequence(seq_id)), + } + } + + fn get_stats(&self) -> SchedulerStats { + let num_running = self.sequences.values().filter(|s| s.is_running()).count(); + let num_waiting = self.waiting_queue.len(); + let num_preempted = self + .sequences + .values() + .filter(|s| matches!(s, SequenceState::Preempted(_))) + .count(); + + SchedulerStats { + num_sequences: self.sequences.len(), + num_running, + num_waiting, + num_preempted, + block_stats: self.block_manager.get_stats(), + } + } +} + +/// Create scheduler event channel +pub fn create_scheduler_channel() -> ( + mpsc::UnboundedSender, + mpsc::UnboundedReceiver, +) { + mpsc::unbounded_channel() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + + fn create_test_scheduler() -> (Scheduler, mpsc::UnboundedSender) { + let allocator = Box::new(CpuAllocator::new(100_000)); + let block_manager = BlockManager::new(allocator, 1024); + Scheduler::new(block_manager, 10, 16) + } + + #[test] + fn test_add_request() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + + assert_eq!(scheduler.sequences.len(), 1); + assert_eq!(scheduler.waiting_queue.len(), 1); + assert!(matches!( + scheduler.sequences.get(&SequenceId(1)), + Some(SequenceState::Waiting(_)) + )); + } + + #[test] + fn test_add_duplicate_request() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), 100, 200); // Duplicate + + // Should not add duplicate + assert_eq!(scheduler.sequences.len(), 1); + assert_eq!(scheduler.waiting_queue.len(), 1); + } + + #[test] + fn test_cancel_request_waiting() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.cancel_request(SequenceId(1)); + + // Should transition to Aborted + assert!(matches!( + scheduler.sequences.get(&SequenceId(1)), + Some(SequenceState::Aborted(_)) + )); + assert_eq!(scheduler.waiting_queue.len(), 0); + } + + #[test] + fn test_cancel_request_not_found() { + let (mut scheduler, _tx) = create_test_scheduler(); + + // Should not panic + scheduler.cancel_request(SequenceId(999)); + assert_eq!(scheduler.sequences.len(), 0); + } + + #[test] + fn test_schedule_prefill_success() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + + // Should move to prefilling + assert!(matches!( + scheduler.sequences.get(&SequenceId(1)), + Some(SequenceState::Prefilling(_)) + )); + assert_eq!(scheduler.prefill_batch.len(), 1); + assert_eq!(scheduler.waiting_queue.len(), 0); + } + + #[test] + fn test_schedule_prefill_multiple() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(2), 50, 100); + scheduler.schedule_prefill(); + + // Both should be scheduled + assert_eq!(scheduler.prefill_batch.len(), 2); + assert_eq!(scheduler.waiting_queue.len(), 0); + } + + #[test] + fn test_schedule_prefill_batch_limit() { + let (mut scheduler, _tx) = create_test_scheduler(); + + // Add more than max_batch_size (10) + for i in 0..15 { + scheduler.add_request(SequenceId(i), 100, 200); + } + scheduler.schedule_prefill(); + + // Only 10 should be scheduled + assert_eq!(scheduler.prefill_batch.len(), 10); + assert_eq!(scheduler.waiting_queue.len(), 5); + } + + #[test] + fn test_schedule_prefill_oom() { + // Small allocator - only 2 blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let block_manager = BlockManager::new(allocator, 1024); + let (mut scheduler, _tx) = Scheduler::new(block_manager, 10, 16); + + // First request needs 7 blocks (100 tokens / 16 tokens_per_block) + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + + // Should fail - not enough memory + assert_eq!(scheduler.prefill_batch.len(), 0); + assert_eq!(scheduler.waiting_queue.len(), 1); + assert!(matches!( + scheduler.sequences.get(&SequenceId(1)), + Some(SequenceState::Waiting(_)) + )); + } + + #[test] + fn test_append_tokens_prefilling() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + + // Append all tokens + scheduler.append_tokens(SequenceId(1), 100); + + // Should transition to Decoding + assert!(matches!( + scheduler.sequences.get(&SequenceId(1)), + Some(SequenceState::Decoding(_)) + )); + } + + #[test] + fn test_append_tokens_decoding() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + // Append more tokens + scheduler.append_tokens(SequenceId(1), 10); + + if let Some(SequenceState::Decoding(state)) = scheduler.sequences.get(&SequenceId(1)) { + assert_eq!(state.num_tokens, 110); + } else { + panic!("Expected Decoding state"); + } + } + + #[test] + fn test_append_tokens_reaches_max() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + // Reach max_tokens + scheduler.append_tokens(SequenceId(1), 100); // 100 + 100 = 200 = max_tokens + + // Should transition to Finished + assert!(matches!( + scheduler.sequences.get(&SequenceId(1)), + Some(SequenceState::Finished(_)) + )); + } + + #[test] + fn test_complete_sequence() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + scheduler.complete_sequence(SequenceId(1), FinishReason::Stop); + + // Should be Finished + assert!(matches!( + scheduler.sequences.get(&SequenceId(1)), + Some(SequenceState::Finished(_)) + )); + assert_eq!(scheduler.prefill_batch.len(), 0); + } + + #[test] + fn test_fork_sequence() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + scheduler.fork_sequence(SequenceId(1), SequenceId(2)); + + // Both should exist + assert!(scheduler.sequences.contains_key(&SequenceId(1))); + assert!(scheduler.sequences.contains_key(&SequenceId(2))); + assert!(matches!( + scheduler.sequences.get(&SequenceId(2)), + Some(SequenceState::Decoding(_)) + )); + } + + #[test] + fn test_fork_duplicate_child_id() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); + + scheduler.add_request(SequenceId(2), 50, 100); // Child ID exists + scheduler.fork_sequence(SequenceId(1), SequenceId(2)); // Should fail + + // Original seq 2 should be unchanged (Waiting) + assert!(matches!( + scheduler.sequences.get(&SequenceId(2)), + Some(SequenceState::Waiting(_)) + )); + } + + #[test] + fn test_preempt_sequence() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); // → Decoding + + scheduler.preempt_sequence(SequenceId(1)); + + // Should be Preempted + assert!(matches!( + scheduler.sequences.get(&SequenceId(1)), + Some(SequenceState::Preempted(_)) + )); + assert_eq!(scheduler.decode_batch.len(), 0); + } + + #[test] + fn test_resume_sequence() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + scheduler.append_tokens(SequenceId(1), 100); + scheduler.preempt_sequence(SequenceId(1)); + + scheduler.resume_sequence(SequenceId(1)); + + // Should be back to Waiting + assert!(matches!( + scheduler.sequences.get(&SequenceId(1)), + Some(SequenceState::Waiting(_)) + )); + assert_eq!(scheduler.waiting_queue.len(), 1); + } + + #[test] + fn test_get_blocks() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + + let blocks = scheduler.get_blocks(SequenceId(1)).unwrap(); + assert!(!blocks.is_empty()); + } + + #[test] + fn test_get_blocks_not_found() { + let (scheduler, _tx) = create_test_scheduler(); + + let result = scheduler.get_blocks(SequenceId(999)); + assert!(matches!(result, Err(Error::UnknownSequence(_)))); + } + + #[test] + fn test_get_stats() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(2), 50, 100); + scheduler.schedule_prefill(); + + let stats = scheduler.get_stats(); + assert_eq!(stats.num_sequences, 2); + assert_eq!(stats.num_running, 2); // Both in prefill + assert_eq!(stats.num_waiting, 0); + } + + #[test] + fn test_try_schedule() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.try_schedule(); + + // Should schedule automatically + assert_eq!(scheduler.prefill_batch.len(), 1); + assert_eq!(scheduler.waiting_queue.len(), 0); + } + + #[test] + fn test_cancel_removes_from_batches() { + let (mut scheduler, _tx) = create_test_scheduler(); + + scheduler.add_request(SequenceId(1), 100, 200); + scheduler.schedule_prefill(); + assert_eq!(scheduler.prefill_batch.len(), 1); + + scheduler.cancel_request(SequenceId(1)); + + // Should remove from batch + assert_eq!(scheduler.prefill_batch.len(), 0); + } +} diff --git a/src/scheduler/events.rs b/src/scheduler/events.rs new file mode 100644 index 0000000..3541e45 --- /dev/null +++ b/src/scheduler/events.rs @@ -0,0 +1,58 @@ +use crate::block_manager::types::*; +use std::collections::HashMap; +use tokio::sync::mpsc; + +/// External events that clients/backend send to the scheduler +/// +/// Only includes events that come from outside the scheduler: +/// - AddRequest: Client starts a new request +/// - CancelRequest: Client cancels an ongoing request +/// - GetBlocks: Backend queries block IDs for inference +/// - GetStats: Monitoring/observability queries +/// +/// Internal operations (token generation, state transitions, OOM handling) +/// are NOT events - the scheduler handles them directly in its run loop. +#[derive(Debug)] +pub enum SchedulerEvent { + /// Client starts a new inference request + AddRequest { + seq_id: SequenceId, + prompt_tokens: usize, + max_tokens: usize, + }, + + /// Client cancels a request (user stop, disconnect, timeout) + CancelRequest { seq_id: SequenceId }, + + /// Backend queries block IDs for a sequence (for KV cache addressing) + GetBlocks { + seq_id: SequenceId, + response: tokio::sync::oneshot::Sender>>, + }, + + /// Query scheduler statistics + GetStats { + response: tokio::sync::oneshot::Sender, + }, +} + +#[derive(Debug, Clone)] +pub struct SchedulerStats { + pub num_sequences: usize, + pub num_running: usize, + pub num_waiting: usize, + pub num_preempted: usize, + pub block_stats: HashMap, +} + +pub type SchedulerEventSender = mpsc::UnboundedSender; +pub type SchedulerEventReceiver = mpsc::UnboundedReceiver; + +/// Create event channel for scheduler +/// +/// TODO: Use bounded channel to prevent OOM when producers outpace the event loop. +/// Capacity should be calculated based on available memory and average event size. +/// Consider: capacity = (available_memory * 0.1) / sizeof(SchedulerEvent) +pub fn create_event_channel() -> (SchedulerEventSender, SchedulerEventReceiver) { + mpsc::unbounded_channel() +} diff --git a/src/sequence_manager/fsm_events.rs b/src/scheduler/fsm_events.rs similarity index 100% rename from src/sequence_manager/fsm_events.rs rename to src/scheduler/fsm_events.rs diff --git a/src/sequence_manager/id_generator.rs b/src/scheduler/id_generator.rs similarity index 100% rename from src/sequence_manager/id_generator.rs rename to src/scheduler/id_generator.rs diff --git a/src/sequence_manager/mod.rs b/src/scheduler/mod.rs similarity index 54% rename from src/sequence_manager/mod.rs rename to src/scheduler/mod.rs index 138bdca..dd9bd23 100644 --- a/src/sequence_manager/mod.rs +++ b/src/scheduler/mod.rs @@ -1,7 +1,5 @@ +pub mod core; pub mod events; pub mod fsm_events; -pub mod fsm_manager; pub mod id_generator; pub mod states; - -// Re-export public API for convenience diff --git a/src/sequence_manager/states.rs b/src/scheduler/states.rs similarity index 100% rename from src/sequence_manager/states.rs rename to src/scheduler/states.rs diff --git a/src/sequence_manager/events.rs b/src/sequence_manager/events.rs deleted file mode 100644 index 7398b8e..0000000 --- a/src/sequence_manager/events.rs +++ /dev/null @@ -1,73 +0,0 @@ -use crate::block_manager::types::*; -use std::collections::HashMap; -use tokio::sync::mpsc; - -/// Events that drive the sequence manager -#[derive(Debug)] -pub enum SequenceEvent { - // Request lifecycle - AddRequest { - seq_id: SequenceId, - prompt_tokens: usize, - max_tokens: usize, - }, - - // Token generation - AppendTokens { - seq_id: SequenceId, - num_tokens: usize, - }, - - // Sequence forking (for beam search, parallel sampling) - ForkSequence { - parent_id: SequenceId, - child_id: SequenceId, - }, - - // Completion - CompleteSequence { - seq_id: SequenceId, - }, - - // Preemption (when OOM) - PreemptSequence { - seq_id: SequenceId, - }, - - // Resume preempted sequence - ResumeSequence { - seq_id: SequenceId, - }, - - // Query - GetBlocks { - seq_id: SequenceId, - response: tokio::sync::oneshot::Sender>>, - }, - - // Stats - GetStats { - response: tokio::sync::oneshot::Sender, - }, -} - -#[derive(Debug, Clone)] -pub struct SequenceManagerStats { - pub num_sequences: usize, - pub num_running: usize, - pub num_waiting: usize, - pub num_preempted: usize, - pub block_stats: HashMap, -} - -pub type SequenceEventSender = mpsc::UnboundedSender; -pub type SequenceEventReceiver = mpsc::UnboundedReceiver; - -/// Create event channel for sequence manager -/// -/// TODO: Use bounded channel to prevent OOM when producers outpace the event loop. -/// Capacity should be calculated based on available memory and average event size. -/// Consider: capacity = (available_memory * 0.1) / sizeof(SequenceEvent) -pub fn create_event_channel() -> (SequenceEventSender, SequenceEventReceiver) { - mpsc::unbounded_channel() -} diff --git a/src/sequence_manager/fsm_manager.rs b/src/sequence_manager/fsm_manager.rs deleted file mode 100644 index 762f93b..0000000 --- a/src/sequence_manager/fsm_manager.rs +++ /dev/null @@ -1,418 +0,0 @@ -use super::events::*; -use super::fsm_events::*; -use super::states::*; -use crate::block_manager::manager::BlockManager; -use crate::block_manager::types::*; -use std::collections::{HashMap, VecDeque}; -use tracing::{debug, info, warn}; - -/// Sequence Manager with FSM-based state transitions -pub struct SequenceManager { - block_manager: BlockManager, - - // Each sequence is a state machine - sequences: HashMap, - - tokens_per_block: usize, - - // Event-driven - event_rx: SequenceEventReceiver, - - // Scheduling queues - waiting_queue: VecDeque, -} - -impl SequenceManager { - pub fn new( - block_manager: BlockManager, - tokens_per_block: usize, - ) -> (Self, SequenceEventSender) { - assert!(tokens_per_block > 0, "tokens_per_block must be > 0"); - - let (event_tx, event_rx) = create_event_channel(); - let manager = Self { - block_manager, - sequences: HashMap::new(), - tokens_per_block, - event_rx, - waiting_queue: VecDeque::new(), - }; - - (manager, event_tx) - } - - /// Main event loop - pub async fn run(mut self) { - info!("FSM SequenceManager event loop started"); - - while let Some(event) = self.event_rx.recv().await { - match event { - SequenceEvent::AddRequest { - seq_id, - prompt_tokens, - max_tokens, - } => { - self.handle_add_request(seq_id, prompt_tokens, max_tokens); - } - - SequenceEvent::AppendTokens { seq_id, num_tokens } => { - self.handle_append_tokens(seq_id, num_tokens); - } - - SequenceEvent::ForkSequence { - parent_id, - child_id, - } => { - self.handle_fork_sequence(parent_id, child_id); - } - - SequenceEvent::CompleteSequence { seq_id } => { - self.handle_complete_sequence(seq_id); - } - - SequenceEvent::PreemptSequence { seq_id } => { - self.handle_preempt_sequence(seq_id); - } - - SequenceEvent::ResumeSequence { seq_id } => { - self.handle_resume_sequence(seq_id); - } - - SequenceEvent::GetBlocks { seq_id, response } => { - let result = self.get_blocks(seq_id); - let _ = response.send(result); - } - - SequenceEvent::GetStats { response } => { - let stats = self.get_stats(); - let _ = response.send(stats); - } - } - } - - info!("FSM SequenceManager event loop stopped"); - } - - fn handle_add_request(&mut self, seq_id: SequenceId, prompt_tokens: usize, max_tokens: usize) { - debug!( - "Adding request: seq_id={:?}, prompt_tokens={}, max_tokens={}", - seq_id, prompt_tokens, max_tokens - ); - - // Check for duplicate seq_id - if self.sequences.contains_key(&seq_id) { - warn!( - "Duplicate AddRequest for {:?} - sequence already exists, ignoring", - seq_id - ); - return; - } - - let state = SequenceState::Waiting(WaitingState { - seq_id, - prompt_tokens, - max_tokens, - }); - - self.sequences.insert(seq_id, state); - self.waiting_queue.push_back(seq_id); - - // Try to schedule immediately - self.try_schedule(); - } - - fn handle_append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Sequence {:?} not found", seq_id); - return; - } - }; - - let backup = state.clone(); - let event = AppendTokensEvent { - num_tokens, - block_manager: &mut self.block_manager, - tokens_per_block: self.tokens_per_block, - }; - - match event.apply(state) { - Ok(new_state) => { - self.sequences.insert(seq_id, new_state); - } - Err(Error::OutOfMemory) => { - warn!("OOM while appending tokens to {:?}", seq_id); - self.sequences.insert(seq_id, backup); - self.handle_oom(); - } - Err(e) => { - warn!("Failed to append tokens to {:?}: {:?}", seq_id, e); - self.sequences.insert(seq_id, backup); - } - } - } - - fn handle_fork_sequence(&mut self, parent_id: SequenceId, child_id: SequenceId) { - // Check if child_id already exists - if self.sequences.contains_key(&child_id) { - warn!( - "Cannot fork {:?} → {:?}: child ID already exists", - parent_id, child_id - ); - return; - } - - let parent_state = match self.sequences.remove(&parent_id) { - Some(s) => s, - None => { - warn!("Parent sequence {:?} not found", parent_id); - return; - } - }; - - let backup = parent_state.clone(); - let event = ForkEvent { - child_id, - block_manager: &mut self.block_manager, - }; - - match event.apply(parent_state) { - Ok((parent_state, child_state)) => { - self.sequences.insert(parent_id, parent_state); - self.sequences.insert(child_id, child_state); - debug!("Forked sequence {:?} → {:?}", parent_id, child_id); - } - Err(e) => { - warn!("Failed to fork {:?}: {:?}", parent_id, e); - self.sequences.insert(parent_id, backup); - } - } - } - - fn handle_complete_sequence(&mut self, seq_id: SequenceId) { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Sequence {:?} not found", seq_id); - return; - } - }; - - let backup = state.clone(); - let event = CompleteEvent { - reason: FinishReason::Stop, - block_manager: &mut self.block_manager, - }; - - match event.apply(state) { - Ok(new_state) => { - info!("Completed sequence {:?}", seq_id); - // Keep finished state for stats/debugging - self.sequences.insert(seq_id, new_state); - - // Try to schedule waiting sequences - self.try_schedule(); - } - Err(e) => { - warn!("Failed to complete {:?}: {:?}", seq_id, e); - self.sequences.insert(seq_id, backup); - } - } - } - - fn handle_preempt_sequence(&mut self, seq_id: SequenceId) { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Sequence {:?} not found", seq_id); - return; - } - }; - - let backup = state.clone(); - let event = PreemptEvent { - block_manager: &mut self.block_manager, - }; - - match event.apply(state) { - Ok(new_state) => { - info!("Preempted sequence {:?}", seq_id); - self.sequences.insert(seq_id, new_state); - } - Err(e) => { - warn!("Failed to preempt {:?}: {:?}", seq_id, e); - self.sequences.insert(seq_id, backup); - } - } - } - - fn handle_resume_sequence(&mut self, seq_id: SequenceId) { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Sequence {:?} not found", seq_id); - return; - } - }; - - let backup = state.clone(); - let event = ResumeEvent; - - match event.apply(state) { - Ok(new_state) => { - debug!("Resumed sequence {:?}", seq_id); - self.sequences.insert(seq_id, new_state); - self.waiting_queue.push_back(seq_id); - self.try_schedule(); - } - Err(e) => { - warn!("Failed to resume {:?}: {:?}", seq_id, e); - self.sequences.insert(seq_id, backup); - } - } - } - - fn try_schedule(&mut self) { - while let Some(&seq_id) = self.waiting_queue.front() { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - self.waiting_queue.pop_front(); - continue; - } - }; - - // Only schedule if in Waiting state - if !matches!(state, SequenceState::Waiting(_)) { - self.sequences.insert(seq_id, state); - self.waiting_queue.pop_front(); - continue; - } - - let backup = state.clone(); - let event = ScheduleEvent { - block_manager: &mut self.block_manager, - tokens_per_block: self.tokens_per_block, - }; - - match event.apply(state) { - Ok(new_state) => { - self.sequences.insert(seq_id, new_state); - self.waiting_queue.pop_front(); - info!("Scheduled sequence {:?}", seq_id); - } - Err(Error::OutOfMemory) => { - // Can't schedule - restore state and leave in waiting queue - self.sequences.insert(seq_id, backup); - break; - } - Err(e) => { - warn!("Failed to schedule {:?}: {:?}", seq_id, e); - self.sequences.insert(seq_id, backup); - self.waiting_queue.pop_front(); - } - } - } - } - - fn handle_oom(&mut self) { - warn!("Handling OOM - preemption policy not yet implemented"); - // TODO: Implement preemption policy - // - Find lowest priority sequence - // - Preempt it - // - Retry scheduling - } - - fn get_blocks(&self, seq_id: SequenceId) -> Result> { - match self.sequences.get(&seq_id) { - None => Err(Error::UnknownSequence(seq_id)), - Some(state) => Ok(state.blocks().map(|b| b.to_vec()).unwrap_or_default()), - } - } - - fn get_stats(&self) -> SequenceManagerStats { - let num_running = self.sequences.values().filter(|s| s.is_running()).count(); - - let num_waiting = self.sequences.values().filter(|s| s.is_waiting()).count(); - - let num_preempted = self - .sequences - .values() - .filter(|s| matches!(s, SequenceState::Preempted(_))) - .count(); - - SequenceManagerStats { - num_sequences: self.sequences.len(), - num_running, - num_waiting, - num_preempted, - block_stats: self.block_manager.get_stats(), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::block_manager::allocator::CpuAllocator; - - fn create_test_manager() -> (SequenceManager, SequenceEventSender) { - let allocator = Box::new(CpuAllocator::new(1024 * 1024)); // 1MB - let block_manager = BlockManager::new(allocator, 256); // 256 byte blocks - SequenceManager::new(block_manager, 16) // 16 tokens per block - } - - #[test] - fn test_duplicate_add_request_ignored() { - let (mut manager, _tx) = create_test_manager(); - - let seq_id = SequenceId(0); - - // Add first request - manager.handle_add_request(seq_id, 32, 128); - - // Verify it was added - assert!(manager.sequences.contains_key(&seq_id)); - let first_state = manager.sequences.get(&seq_id).cloned(); - - // Try to add duplicate - should be ignored - manager.handle_add_request(seq_id, 64, 256); // Different params - - // Verify the original state is unchanged - let current_state = manager.sequences.get(&seq_id); - assert_eq!(format!("{:?}", first_state), format!("{:?}", current_state)); - - // Verify original params (may have transitioned to Prefilling via try_schedule) - match manager.sequences.get(&seq_id) { - Some(SequenceState::Waiting(s)) => { - assert_eq!(s.prompt_tokens, 32); // Original value - assert_eq!(s.max_tokens, 128); // Original value - } - Some(SequenceState::Prefilling(s)) => { - // Scheduled automatically, check original params - assert_eq!(s.tokens_total, 32); // Original prompt_tokens - assert_eq!(s.max_tokens, 128); // Original value - } - _ => panic!( - "Expected Waiting or Prefilling state, got: {:?}", - manager.sequences.get(&seq_id) - ), - } - } - - #[test] - fn test_add_different_sequences() { - let (mut manager, _tx) = create_test_manager(); - - // Add multiple different sequences - all should succeed - manager.handle_add_request(SequenceId(0), 32, 128); - manager.handle_add_request(SequenceId(1), 64, 256); - manager.handle_add_request(SequenceId(2), 16, 64); - - assert_eq!(manager.sequences.len(), 3); - assert!(manager.sequences.contains_key(&SequenceId(0))); - assert!(manager.sequences.contains_key(&SequenceId(1))); - assert!(manager.sequences.contains_key(&SequenceId(2))); - } -} From 73298dfbe5a67bf45d073a8ca811c84bba9bbf5b Mon Sep 17 00:00:00 2001 From: kerthcet Date: Thu, 16 Jul 2026 00:07:07 +0100 Subject: [PATCH 22/33] rearch Signed-off-by: kerthcet --- Cargo.lock | 342 +++++++++++++- Cargo.toml | 1 + docs/fsm_architecture.md | 137 ++++-- src/backend/llm_engine.rs | 212 +++++++++ src/backend/mod.rs | 1 + src/block_manager/types.rs | 3 + src/cli/commands.rs | 26 +- src/fsm/events.rs | 34 ++ src/fsm/mod.rs | 45 ++ src/{scheduler => fsm}/states.rs | 23 +- src/fsm/transitions.rs | 373 ++++++++++++++++ src/lib.rs | 1 + src/main.rs | 1 + src/scheduler/core.rs | 691 +++++++++++++++++++---------- src/scheduler/events.rs | 31 +- src/scheduler/fsm_events.rs | 740 ------------------------------- src/scheduler/mod.rs | 4 +- 17 files changed, 1635 insertions(+), 1030 deletions(-) create mode 100644 src/backend/llm_engine.rs create mode 100644 src/fsm/events.rs create mode 100644 src/fsm/mod.rs rename src/{scheduler => fsm}/states.rs (78%) create mode 100644 src/fsm/transitions.rs delete mode 100644 src/scheduler/fsm_events.rs diff --git a/Cargo.lock b/Cargo.lock index e9b5cbd..e062eab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -172,6 +172,12 @@ dependencies = [ "tracing", ] +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + [[package]] name = "base64" version = "0.22.1" @@ -308,6 +314,19 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "console" +version = "0.15.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width 0.2.2", + "windows-sys 0.59.0", +] + [[package]] name = "console" version = "0.16.3" @@ -430,6 +449,41 @@ dependencies = [ "memchr", ] +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn", +] + [[package]] name = "der" version = "0.8.0" @@ -449,6 +503,37 @@ dependencies = [ "powerfmt", ] +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn", +] + [[package]] name = "dirs" version = "6.0.0" @@ -560,6 +645,15 @@ version = "3.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59" +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + [[package]] name = "fallible-iterator" version = "0.3.0" @@ -838,12 +932,12 @@ dependencies = [ "dirs", "futures", "http", - "indicatif", + "indicatif 0.18.4", "libc", "log", "native-tls", "num_cpus", - "rand", + "rand 0.9.4", "reqwest", "serde", "serde_json", @@ -966,7 +1060,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -1097,6 +1191,12 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "idna" version = "1.1.0" @@ -1130,13 +1230,26 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console 0.15.11", + "number_prefix", + "portable-atomic", + "unicode-width 0.2.2", + "web-time", +] + [[package]] name = "indicatif" version = "0.18.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" dependencies = [ - "console", + "console 0.16.3", "portable-atomic", "unicode-width 0.2.2", "unit-prefix", @@ -1176,6 +1289,24 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" +[[package]] +name = "itertools" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1265,6 +1396,22 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + [[package]] name = "matchers" version = "0.2.0" @@ -1292,6 +1439,12 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -1313,6 +1466,28 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "native-tls" version = "0.2.18" @@ -1351,6 +1526,16 @@ dependencies = [ "libc", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "ntapi" version = "0.4.3" @@ -1394,6 +1579,12 @@ dependencies = [ "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "once_cell" version = "1.21.4" @@ -1406,6 +1597,28 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "openssl" version = "0.10.78" @@ -1479,6 +1692,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "pem-rfc7468" version = "1.0.0" @@ -1600,7 +1819,7 @@ dependencies = [ "dirs", "futures", "hf-hub", - "indicatif", + "indicatif 0.18.4", "prettytable-rs", "regex", "reqwest", @@ -1614,6 +1833,7 @@ dependencies = [ "sysinfo", "tempfile", "thiserror 2.0.18", + "tokenizers", "tokio", "tokio-stream", "tower 0.4.13", @@ -1654,14 +1874,35 @@ dependencies = [ "nibble_vec", ] +[[package]] +name = "rand" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -1671,7 +1912,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", ] [[package]] @@ -1693,6 +1943,17 @@ dependencies = [ "rayon-core", ] +[[package]] +name = "rayon-cond" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "059f538b55efd2309c9794130bc149c6a553db90e9d99c2030785c82f0bd7df9" +dependencies = [ + "either", + "itertools 0.11.0", + "rayon", +] + [[package]] name = "rayon-core" version = "1.13.0" @@ -1769,7 +2030,7 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -2111,6 +2372,18 @@ dependencies = [ "winapi", ] +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom", + "serde", + "unicode-segmentation", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -2309,6 +2582,38 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tokenizers" +version = "0.20.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b08cc37428a476fc9e20ac850132a513a2e1ce32b6a31addf2b74fa7033b905" +dependencies = [ + "aho-corasick", + "derive_builder", + "esaxx-rs", + "getrandom 0.2.17", + "indicatif 0.17.11", + "itertools 0.12.1", + "lazy_static", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.8.7", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 1.0.69", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + [[package]] name = "tokio" version = "1.52.1" @@ -2533,6 +2838,15 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + [[package]] name = "unicode-segmentation" version = "1.13.3" @@ -2557,6 +2871,12 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + [[package]] name = "unit-prefix" version = "0.5.2" @@ -2575,7 +2895,7 @@ version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" dependencies = [ - "base64", + "base64 0.22.1", "cookie_store", "der", "flate2", @@ -2599,7 +2919,7 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" dependencies = [ - "base64", + "base64 0.22.1", "http", "httparse", "log", diff --git a/Cargo.toml b/Cargo.toml index 162101c..cdcccd0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,6 +42,7 @@ futures = "0.3" tokio-stream = "0.1" rustyline = "14.0" rustyline-derive = "0.10" +tokenizers = "0.20" [dev-dependencies] tempfile = "3.12" diff --git a/docs/fsm_architecture.md b/docs/fsm_architecture.md index 92bd011..1d413b2 100644 --- a/docs/fsm_architecture.md +++ b/docs/fsm_architecture.md @@ -1,6 +1,6 @@ # FSM Architecture for PUMA -Finite State Machine pattern for sequence lifecycle management. +Finite State Machine pattern for sequence lifecycle management, integrated into the Scheduler. ## Core Concepts @@ -8,7 +8,7 @@ Finite State Machine pattern for sequence lifecycle management. Each sequence is an independent state machine. Multiple sequences can be in different states simultaneously. -**Implementation:** `src/sequence_manager/states.rs` +**Implementation:** `src/fsm/states.rs` **States:** - `Waiting` - In queue, no resources allocated @@ -32,19 +32,48 @@ Each state owns the resources it needs: Events transform states: `Event + OldState → NewState` -**Implementation:** `src/sequence_manager/fsm_events.rs` +**Implementation:** `src/fsm/events.rs` **Events:** -- `ScheduleEvent` - Waiting → Prefilling (allocate blocks) -- `AppendTokensEvent` - Prefilling → Decoding or Decoding → Decoding/Finished -- `ForkEvent` - Decoding → (Decoding, Decoding) for beam search -- `PreemptEvent` - Decoding → Preempted (free blocks on OOM) -- `ResumeEvent` - Preempted → Waiting (re-queue) -- `CompleteEvent` - * → Finished (free blocks) -- `AbortEvent` - * → Aborted (cleanup) +- `Schedule` - Waiting → Prefilling (allocate blocks) +- `AppendTokens` - Prefilling → Decoding or Decoding → Decoding/Finished +- `Fork` - Decoding → (Decoding, Decoding) for beam search +- `Preempt` - Decoding → Preempted (free blocks on OOM) +- `Resume` - Preempted → Waiting (re-queue) +- `Complete` - * → Finished (free blocks) +- `Abort` - * → Aborted (cleanup) Invalid transitions return `Error::InvalidTransition`. +### FSM Integration with Scheduler + +The FSM logic is **integrated into the Scheduler** to solve the ownership boundary problem: + +**Public API:** `Scheduler::transition(state, event)` - Event-based abstraction +- Single entry point for all state transitions +- Easy to add logging, metrics, debugging +- Event replay capability for testing + +**Internal Implementation:** Private `transition_*()` methods +- Type-safe helpers that enforce correct state types +- Direct access to `self.block_manager` - no parameter passing +- Called by `transition()` after event dispatching + +```rust +impl Scheduler { + pub fn transition(&mut self, state: SequenceState, event: Event) -> Result { + match (state, event) { + (SequenceState::Waiting(s), Event::Schedule{..}) => + self.transition_schedule(s), + // ... dispatches to internal methods + } + } + + fn transition_schedule(&mut self, state: WaitingState) -> Result { + // Direct access to self.block_manager, self.tokens_per_block + } +} + ## State Diagram ``` @@ -71,21 +100,21 @@ Invalid transitions return `Error::InvalidTransition`. └─ CompleteEvent → Finished ``` -## Sequence Manager +## Scheduler Integration -Event-driven architecture with async event loop. +The Scheduler owns the FSM and manages all state transitions. -**Implementation:** `src/sequence_manager/fsm_manager.rs` +**Implementation:** `src/scheduler/core.rs` **Key features:** -- Receives events via async channel +- Synchronous API (LLMEngine handles async coordination) - Maintains `HashMap` +- Owns `BlockManager` directly - no parameter passing needed - Clones state before transitions (prevents loss on failure) - Restores backup on error (no block leaks) -- Duplicate `seq_id` check prevents overwrites -- Duplicate `child_id` check prevents fork overwrites +- Event-based `apply()` method for maintainability -**Tests:** Inline in `fsm_manager.rs` and `fsm_events.rs` +**Tests:** `cargo test --lib` (120 tests passing) ## Arc Optimization @@ -132,23 +161,67 @@ All state transitions are wrapped with backup/restore: 5. **Performance** - Arc enables O(1) cloning, copy-on-write for efficiency 6. **Reliability** - Backup/restore pattern prevents data loss on errors +## Two-Level Event Architecture + +PUMA uses a two-level event system (inspired by TokenSpeed): + +**Level 1: FSM Events (Internal)** - `src/fsm/events.rs` +- Created by Scheduler internally +- Include scheduler context (tokens_per_block, etc.) +- Used with `Scheduler::apply(state, event)` +- Examples: `Event::Schedule`, `Event::AppendTokens` + +**Level 2: Scheduler Events (External)** - `src/scheduler/events.rs` +- Created by external components (clients, GPU workers) +- Simple data, no resource pointers +- LLMEngine translates these to FSM events +- Examples: `SchedulerEvent::AddRequest`, `SchedulerEvent::CancelRequest` + +**Why two levels?** External components cannot create FSM events because: +- They don't own `BlockManager` +- They don't know scheduler policies (tokens_per_block) +- FSM events require scheduler context + +See `ARCHITECTURE.md` for detailed explanation and flow diagrams. + ## Key Files -- `src/sequence_manager/states.rs` - State definitions -- `src/sequence_manager/fsm_events.rs` - Event implementations + tests -- `src/sequence_manager/fsm_manager.rs` - SequenceManager + tests -- `src/sequence_manager/events.rs` - Event channel types -- `src/sequence_manager/id_generator.rs` - Sequential ID generator +- `src/fsm/states.rs` - State definitions and helper methods +- `src/fsm/events.rs` - FSM event enum (internal transitions) +- `src/scheduler/core.rs` - Scheduler with integrated FSM (`apply()` method) +- `src/scheduler/events.rs` - External scheduler events +- `src/backend/llm_engine.rs` - Event coordinator +- `ARCHITECTURE.md` - Comprehensive architecture documentation ## Usage Example -See `src/sequence_manager/fsm_manager.rs` tests for complete examples. - -Basic flow: -1. Create BlockManager and SequenceManager -2. Get event sender channel -3. Spawn event loop in background -4. Send `AddRequest` event -5. Send `AppendTokens` events as tokens generate -6. Sequence automatically transitions through states -7. Finished/Aborted when done +```rust +// External component sends high-level event +scheduler_tx.send(SchedulerEvent::AddRequest { + seq_id: 1, + prompt_tokens: 10, + max_tokens: 100, +})?; + +// Scheduler translates to FSM events internally +impl Scheduler { + pub fn schedule(&mut self) -> bool { + // Pop from waiting queue + let state = self.sequences.remove(&seq_id).unwrap(); + + // Create FSM event with scheduler context + let event = Event::Schedule { + tokens_per_block: self.tokens_per_block, + }; + + // Perform transition (direct access to self.block_manager) + match self.transition(state, event) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + self.prefill_batch.push(seq_id); + } + Err(e) => { /* handle error */ } + } + } +} +``` diff --git a/src/backend/llm_engine.rs b/src/backend/llm_engine.rs new file mode 100644 index 0000000..e5fb190 --- /dev/null +++ b/src/backend/llm_engine.rs @@ -0,0 +1,212 @@ +use std::io; +use std::sync::atomic::{AtomicU64, Ordering}; +use tokenizers::Tokenizer; +use tokio::sync::mpsc; + +use super::engine::{GenerateResponse, InferenceEngine}; +use crate::block_manager::allocator::CpuAllocator; +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::SequenceId; +use crate::scheduler::core::Scheduler; +use crate::scheduler::events::SchedulerEvent; + +/// LLM Engine - integrates scheduler with inference backend +/// +/// Architecture +/// - LLMEngine drives the main loop +/// - Owns tokenizer (model-specific) +/// - Scheduler is synchronous (just methods, no async) +/// - Direct ownership (no Arc/Mutex overhead) +/// - Loop: handle events → schedule() → forward() → process_outputs() +pub struct LLMEngine { + backend: B, + scheduler: Scheduler, + tokenizer: Tokenizer, + event_rx: mpsc::UnboundedReceiver, + event_tx: mpsc::UnboundedSender, + seq_id_counter: AtomicU64, + model: String, +} + +impl LLMEngine { + pub fn new(backend: B, tokenizer: Tokenizer, model: String) -> Self { + // Create block manager (100MB memory pool, 512 bytes per block) + let allocator = Box::new(CpuAllocator::new(1024 * 1024 * 100)); + let block_manager = BlockManager::new(allocator, 512); + + // Create scheduler (max 32 batch size, 16 tokens per block) + let scheduler = Scheduler::new(block_manager, 32, 16); + + // Event channel + let (event_tx, event_rx) = mpsc::unbounded_channel(); + + Self { + backend, + scheduler, + tokenizer, + event_rx, + event_tx, + seq_id_counter: AtomicU64::new(0), + model, + } + } + + fn next_seq_id(&self) -> SequenceId { + SequenceId(self.seq_id_counter.fetch_add(1, Ordering::Relaxed)) + } + + /// Main loop - drives scheduling and inference + /// + /// Pattern + /// 1. Handle events (add/cancel requests) + /// 2. Schedule (decide what to run) + /// 3. Forward (run inference if work exists) + /// 4. Process outputs (update scheduler state) + pub async fn run(mut self) { + tracing::info!("LLMEngine started"); + + loop { + // 1. Handle events (non-blocking, drain all) + while let Ok(event) = self.event_rx.try_recv() { + match event { + SchedulerEvent::AddRequest { + seq_id, + token_ids, + max_tokens, + } => { + self.scheduler.add_request(seq_id, token_ids, max_tokens); + } + + SchedulerEvent::CancelRequest { seq_id } => { + self.scheduler.cancel_request(seq_id); + } + + SchedulerEvent::GetBlocks { seq_id, response } => { + let result = self.scheduler.get_blocks(seq_id); + let _ = response.send(result); + } + + SchedulerEvent::GetStats { response } => { + let stats = self.scheduler.get_stats(); + let _ = response.send(stats); + } + } + } + + // 2. Schedule (always, not event-driven) + let has_work = self.scheduler.schedule(); + + // 3. Forward (if work exists) + if has_work { + // TODO: Build batch and call backend.forward() + // For now, just yield + } + + // 4. Small yield to prevent busy loop + tokio::task::yield_now().await; + } + } +} + +// Implement InferenceEngine directly for LLMEngine +impl InferenceEngine for LLMEngine { + async fn generate( + &self, + _model: &str, + prompt: &str, + max_tokens: usize, + _temperature: f32, + ) -> Result { + let seq_id = self.next_seq_id(); + + // Tokenize prompt using HuggingFace tokenizer + let encoding = self + .tokenizer + .encode(prompt, false) + .map_err(|e| io::Error::other(format!("Tokenization failed: {}", e)))?; + let token_ids = encoding.get_ids().to_vec(); + let num_tokens = token_ids.len(); + + // Send request to scheduler + self.event_tx + .send(SchedulerEvent::AddRequest { + seq_id, + token_ids, + max_tokens, + }) + .map_err(|e| io::Error::other(format!("Engine send failed: {}", e)))?; + + // TODO: Wait for actual result from scheduler/GPU + // For now, return dummy response + Ok(GenerateResponse { + text: format!("Generated response for: {}", prompt), + prompt_tokens: num_tokens, + completion_tokens: 0, + }) + } + + async fn generate_stream( + &self, + _model: &str, + prompt: &str, + max_tokens: usize, + _temperature: f32, + ) -> Result + Send>>, io::Error> { + let seq_id = self.next_seq_id(); + + // Tokenize prompt using HuggingFace tokenizer + let encoding = self + .tokenizer + .encode(prompt, false) + .map_err(|e| io::Error::other(format!("Tokenization failed: {}", e)))?; + let token_ids = encoding.get_ids().to_vec(); + + // Send request to scheduler + self.event_tx + .send(SchedulerEvent::AddRequest { + seq_id, + token_ids, + max_tokens, + }) + .map_err(|e| io::Error::other(format!("Engine send failed: {}", e)))?; + + // TODO: Stream results from scheduler/GPU + // For now, return dummy stream + let stream = tokio_stream::iter(vec![ + format!("Token 1 for: {}\n", prompt), + "Token 2\n".to_string(), + "Token 3\n".to_string(), + ]); + Ok(Box::pin(stream)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::mock::MockEngine; + + fn create_test_tokenizer() -> Tokenizer { + // Create a simple BPE tokenizer for testing + use tokenizers::models::bpe::BPE; + use tokenizers::Tokenizer as TokenizerBuilder; + + let bpe = BPE::default(); + TokenizerBuilder::new(bpe) + } + + #[tokio::test] + async fn test_llm_engine() { + let backend = MockEngine::new(); + let tokenizer = create_test_tokenizer(); + let engine = LLMEngine::new(backend, tokenizer, "test-model".to_string()); + + // TODO: Spawn engine loop when real backend integration is done + // tokio::spawn(async move { + // engine_copy.run().await; + // }); + + let result = engine.generate("test-model", "Hello world", 100, 0.7).await; + assert!(result.is_ok()); + } +} diff --git a/src/backend/mod.rs b/src/backend/mod.rs index 1631447..d9220b1 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -1,4 +1,5 @@ pub mod engine; +pub mod llm_engine; pub mod mock; pub use engine::*; diff --git a/src/block_manager/types.rs b/src/block_manager/types.rs index 3076802..69a5c76 100644 --- a/src/block_manager/types.rs +++ b/src/block_manager/types.rs @@ -45,6 +45,9 @@ impl SequenceId { } } +/// Token ID type - used for tokenized text +pub type TokenId = u32; + #[derive(Debug, Clone, Copy)] pub struct MemoryAddress { pub ptr: *mut u8, diff --git a/src/cli/commands.rs b/src/cli/commands.rs index a129308..4ed1b4b 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -1,7 +1,9 @@ + use clap::{Parser, Subcommand}; use colored::Colorize; use prettytable::{format, row, Table}; +use crate::backend::llm_engine::LLMEngine; use crate::backend::mock::MockEngine; use crate::cli::{chat, inspect, ls, rm}; use crate::downloader::{self, Provider}; @@ -209,13 +211,27 @@ pub async fn run(cli: Cli) { Ok(Some(_)) => {} } - // Load inference engine with the model - // TODO: Replace MockEngine with real engine that loads model files - // Real engine will use: registry.get_model(&args.model)?.metadata.cache.path - let engine = MockEngine::new(); + // Load tokenizer + // TODO: Load from model directory when model is downloaded + // For now, create a simple test tokenizer + use tokenizers::{models::bpe::BPE, Tokenizer}; + let tokenizer = Tokenizer::new(BPE::default()); + + // Load inference backend + // TODO: Replace MockEngine with real backend that loads model files + // Real backend will use: registry.get_model(&args.model)?.metadata.cache.path + let backend = MockEngine::new(); + + // Create LLMEngine (integrates scheduler + memory management + tokenizer) + let engine = LLMEngine::new(backend, tokenizer, args.model.clone()); - // Start interactive chat (no instruction message, just start prompting) + // TODO: Spawn engine loop when real backend integration is done + // For now, engine.generate() sends events but backend is MockEngine + // tokio::spawn(async move { + // engine_copy.run().await; + // }); + // Start interactive chat if let Err(e) = chat::interactive_chat(&engine, &args.model).await { eprintln!("❌ Chat error: {}", e); std::process::exit(1); diff --git a/src/fsm/events.rs b/src/fsm/events.rs new file mode 100644 index 0000000..47883b4 --- /dev/null +++ b/src/fsm/events.rs @@ -0,0 +1,34 @@ +use super::states::FinishReason; +use crate::block_manager::types::SequenceId; + +/// FSM Events - pure data that triggers state transitions +/// +/// These events represent state machine transitions and are created +/// internally by the scheduler. External components should use +/// SchedulerEvent instead. +#[derive(Debug, Clone)] +pub enum Event { + /// Allocate blocks and start prefilling + Schedule { tokens_per_block: usize }, + + /// Append tokens to a running sequence + AppendTokens { + num_tokens: usize, + tokens_per_block: usize, + }, + + /// Mark sequence as complete + Complete { reason: FinishReason }, + + /// Preempt a running sequence (free blocks) + Preempt, + + /// Resume a preempted sequence + Resume, + + /// Fork a sequence (copy-on-write) + Fork { child_id: SequenceId }, + + /// Abort a sequence with error + Abort { reason: String }, +} diff --git a/src/fsm/mod.rs b/src/fsm/mod.rs new file mode 100644 index 0000000..e945239 --- /dev/null +++ b/src/fsm/mod.rs @@ -0,0 +1,45 @@ +//! FSM - Finite State Machine for sequence lifecycle +//! +//! This module defines states and events for the sequence state machine. +//! The actual transition logic is implemented in `Scheduler::transition()` and +//! private `Scheduler::transition_*()` methods. +//! +//! # Architecture +//! +//! - `states`: State definitions (Waiting, Prefilling, Decoding, etc.) +//! - `events`: Event types that trigger transitions +//! - `transitions`: (DEPRECATED - logic moved to Scheduler) +//! +//! # Usage +//! +//! States and events are used by the Scheduler: +//! +//! ```rust,ignore +//! use crate::fsm::{Event, SequenceState}; +//! use crate::scheduler::Scheduler; +//! +//! let mut scheduler = Scheduler::new(...); +//! let state = SequenceState::Waiting(WaitingState { ... }); +//! let event = Event::Schedule { tokens_per_block: 16 }; +//! +//! // Scheduler owns BlockManager and performs transitions +//! let new_state = scheduler.transition(state, event)?; +//! ``` +//! +//! # Design Rationale +//! +//! FSM logic is **inside Scheduler** rather than a separate `fsm::apply()` because: +//! - Scheduler owns `BlockManager` - no need to pass as parameter +//! - Direct access to scheduler context (tokens_per_block, etc.) +//! - Event-based `transition()` method provides single entry point for logging/metrics +//! - Type-safe internal `transition_*()` methods enforce correct state types +//! +//! See `docs/fsm_architecture.md` for detailed documentation. + +pub mod events; +pub mod states; +// pub mod transitions; // DEPRECATED: logic moved to Scheduler + +// Re-export commonly used types +pub use events::Event; +pub use states::*; diff --git a/src/scheduler/states.rs b/src/fsm/states.rs similarity index 78% rename from src/scheduler/states.rs rename to src/fsm/states.rs index 2128f1a..50426a5 100644 --- a/src/scheduler/states.rs +++ b/src/fsm/states.rs @@ -36,7 +36,7 @@ pub enum SequenceState { #[derive(Debug, Clone)] pub struct WaitingState { pub seq_id: SequenceId, - pub prompt_tokens: usize, + pub token_ids: Arc>, // Arc for O(1) cloning, GPU gets reference pub max_tokens: usize, } @@ -54,7 +54,8 @@ pub struct SchedulingState { #[derive(Debug, Clone)] pub struct PrefillingState { pub seq_id: SequenceId, - pub blocks: Arc>, // Arc for cheap cloning (O(1)) + pub token_ids: Arc>, // Arc for O(1) cloning, needed if preempted + pub blocks: Arc>, // Arc for cheap cloning (O(1)) pub tokens_filled: usize, pub tokens_total: usize, pub max_tokens: usize, @@ -64,7 +65,8 @@ pub struct PrefillingState { #[derive(Debug, Clone)] pub struct DecodingState { pub seq_id: SequenceId, - pub blocks: Arc>, // Arc for cheap cloning (O(1)) + pub token_ids: Arc>, // Arc for O(1) cloning, needed if preempted + pub blocks: Arc>, // Arc for cheap cloning (O(1)) pub num_tokens: usize, pub max_tokens: usize, } @@ -73,7 +75,8 @@ pub struct DecodingState { #[derive(Debug, Clone)] pub struct PreemptedState { pub seq_id: SequenceId, - pub num_tokens: usize, + pub token_ids: Arc>, // Arc for O(1) cloning, needed to resume + pub num_tokens: usize, // How many tokens were generated so far pub max_tokens: usize, // No blocks - they were freed } @@ -142,4 +145,16 @@ impl SequenceState { _ => None, } } + + pub fn variant_name(&self) -> &'static str { + match self { + SequenceState::Waiting(_) => "Waiting", + SequenceState::Scheduling(_) => "Scheduling", + SequenceState::Prefilling(_) => "Prefilling", + SequenceState::Decoding(_) => "Decoding", + SequenceState::Preempted(_) => "Preempted", + SequenceState::Finished(_) => "Finished", + SequenceState::Aborted(_) => "Aborted", + } + } } diff --git a/src/fsm/transitions.rs b/src/fsm/transitions.rs new file mode 100644 index 0000000..1631658 --- /dev/null +++ b/src/fsm/transitions.rs @@ -0,0 +1,373 @@ +use super::events::Event; +use super::states::*; +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; +use std::sync::Arc; +use tracing::{debug, warn}; + +/// Apply an event to a state, producing a new state +/// +/// Pure function that handles all state transitions. +/// Pattern: (current_state, event) -> new_state +pub fn apply( + state: SequenceState, + event: Event, + block_manager: &mut BlockManager, +) -> Result { + match (state, event) { + // Waiting → Prefilling (Schedule) + (SequenceState::Waiting(s), Event::Schedule { tokens_per_block }) => { + apply_schedule(s, block_manager, tokens_per_block) + } + + // Prefilling → Decoding/Prefilling (AppendTokens) + ( + SequenceState::Prefilling(s), + Event::AppendTokens { + num_tokens, + tokens_per_block, + }, + ) => apply_append_tokens_prefilling(s, num_tokens, block_manager, tokens_per_block), + + // Decoding → Decoding/Finished (AppendTokens) + ( + SequenceState::Decoding(s), + Event::AppendTokens { + num_tokens, + tokens_per_block, + }, + ) => apply_append_tokens_decoding(s, num_tokens, block_manager, tokens_per_block), + + // Prefilling/Decoding → Finished (Complete) + (SequenceState::Prefilling(s), Event::Complete { reason }) => { + apply_complete_prefilling(s, reason, block_manager) + } + (SequenceState::Decoding(s), Event::Complete { reason }) => { + apply_complete_decoding(s, reason, block_manager) + } + + // Prefilling/Decoding → Preempted (Preempt) + (SequenceState::Prefilling(s), Event::Preempt) => { + apply_preempt_prefilling(s, block_manager) + } + (SequenceState::Decoding(s), Event::Preempt) => apply_preempt_decoding(s, block_manager), + + // Preempted → Waiting (Resume) + (SequenceState::Preempted(s), Event::Resume) => apply_resume(s), + + // Decoding → (Decoding, Decoding) (Fork) + (SequenceState::Decoding(s), Event::Fork { child_id }) => { + apply_fork(s, child_id, block_manager) + } + + // Any → Aborted (Abort) + (state, Event::Abort { reason }) => apply_abort(state, reason, block_manager), + + // Invalid transitions + (_state, _event) => Err(Error::InvalidTransition("Invalid state transition")), + } +} + +fn state_name(state: &SequenceState) -> &'static str { + match state { + SequenceState::Waiting(_) => "Waiting", + SequenceState::Scheduling(_) => "Scheduling", + SequenceState::Prefilling(_) => "Prefilling", + SequenceState::Decoding(_) => "Decoding", + SequenceState::Preempted(_) => "Preempted", + SequenceState::Finished(_) => "Finished", + SequenceState::Aborted(_) => "Aborted", + } +} + +// Helper functions for each transition + +fn apply_schedule( + state: WaitingState, + block_manager: &mut BlockManager, + tokens_per_block: usize, +) -> Result { + let blocks_needed = state.prompt_tokens.div_ceil(tokens_per_block); + + let mut blocks = Vec::new(); + for _ in 0..blocks_needed { + match block_manager.allocate() { + Ok(block_id) => blocks.push(block_id), + Err(Error::OutOfMemory) => { + // Cleanup on OOM + for block in blocks { + let _ = block_manager.free(block); + } + return Err(Error::OutOfMemory); + } + Err(e) => { + // Other error - cleanup and propagate + for block in blocks { + let _ = block_manager.free(block); + } + return Err(e); + } + } + } + + debug!( + "Scheduled seq {:?}: allocated {} blocks for {} tokens", + state.seq_id, + blocks.len(), + state.prompt_tokens + ); + + Ok(SequenceState::Prefilling(PrefillingState { + seq_id: state.seq_id, + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: state.prompt_tokens, + max_tokens: state.max_tokens, + })) +} + +fn apply_append_tokens_prefilling( + mut state: PrefillingState, + num_tokens: usize, + _block_manager: &mut BlockManager, + _tokens_per_block: usize, +) -> Result { + state.tokens_filled += num_tokens; + + debug!( + "Prefilling seq {:?}: {}/{} tokens", + state.seq_id, state.tokens_filled, state.tokens_total + ); + + if state.tokens_filled >= state.tokens_total { + debug!("Seq {:?} transitioning to Decoding", state.seq_id); + Ok(SequenceState::Decoding(DecodingState { + seq_id: state.seq_id, + blocks: state.blocks, + num_tokens: state.tokens_filled, + max_tokens: state.max_tokens, + })) + } else { + Ok(SequenceState::Prefilling(state)) + } +} + +fn apply_append_tokens_decoding( + mut state: DecodingState, + num_tokens: usize, + block_manager: &mut BlockManager, + tokens_per_block: usize, +) -> Result { + state.num_tokens += num_tokens; + + // Check if finished + if state.num_tokens >= state.max_tokens { + // Free blocks + for &block_id in state.blocks.iter() { + if let Err(e) = block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!( + "Seq {:?} finished: reached max_tokens ({})", + state.seq_id, state.max_tokens + ); + + return Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: FinishReason::MaxTokens, + })); + } + + // Check if need more blocks + let blocks_needed = state.num_tokens.div_ceil(tokens_per_block); + let initial_block_count = state.blocks.len(); + + if state.blocks.len() < blocks_needed { + // Need to allocate more blocks - use Arc::make_mut for copy-on-write + let blocks = Arc::make_mut(&mut state.blocks); + + while blocks.len() < blocks_needed { + match block_manager.allocate() { + Ok(block_id) => { + blocks.push(block_id); + debug!( + "Seq {:?}: allocated block, total blocks: {}", + state.seq_id, + blocks.len() + ); + } + Err(e) => { + // Free any blocks we allocated in this call + warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); + for block_id in blocks.drain(initial_block_count..) { + let _ = block_manager.free(block_id); + } + return Err(e); + } + } + } + } + + Ok(SequenceState::Decoding(state)) +} + +fn apply_complete_prefilling( + state: PrefillingState, + reason: FinishReason, + block_manager: &mut BlockManager, +) -> Result { + // Free all blocks + for &block_id in state.blocks.iter() { + if let Err(e) = block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!( + "Completed seq {:?} during prefill: {:?}", + state.seq_id, reason + ); + + Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: reason, + })) +} + +fn apply_complete_decoding( + state: DecodingState, + reason: FinishReason, + block_manager: &mut BlockManager, +) -> Result { + // Free all blocks + for &block_id in state.blocks.iter() { + if let Err(e) = block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!("Completed seq {:?}: {:?}", state.seq_id, reason); + + Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: reason, + })) +} + +fn apply_preempt_prefilling( + state: PrefillingState, + block_manager: &mut BlockManager, +) -> Result { + // Free all blocks + for &block_id in state.blocks.iter() { + if let Err(e) = block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!("Preempted seq {:?} during prefill", state.seq_id); + + Ok(SequenceState::Preempted(PreemptedState { + seq_id: state.seq_id, + num_tokens: state.tokens_filled, + max_tokens: state.max_tokens, + })) +} + +fn apply_preempt_decoding( + state: DecodingState, + block_manager: &mut BlockManager, +) -> Result { + // Free all blocks + for &block_id in state.blocks.iter() { + if let Err(e) = block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!("Preempted seq {:?}", state.seq_id); + + Ok(SequenceState::Preempted(PreemptedState { + seq_id: state.seq_id, + num_tokens: state.num_tokens, + max_tokens: state.max_tokens, + })) +} + +fn apply_resume(state: PreemptedState) -> Result { + debug!("Resuming seq {:?}", state.seq_id); + + // Transition back to waiting for rescheduling + Ok(SequenceState::Waiting(WaitingState { + seq_id: state.seq_id, + prompt_tokens: state.num_tokens, + max_tokens: state.max_tokens, + })) +} + +fn apply_fork( + state: DecodingState, + child_id: SequenceId, + block_manager: &mut BlockManager, +) -> Result { + // Copy-on-write: increment ref counts + for &block_id in state.blocks.iter() { + block_manager.add_ref(block_id)?; + } + + let child = DecodingState { + seq_id: child_id, + blocks: state.blocks.clone(), + num_tokens: state.num_tokens, + max_tokens: state.max_tokens, + }; + + debug!("Forked seq {:?} → {:?}", state.seq_id, child.seq_id); + + // Fork returns parent, child needs to be handled separately by caller + // For now just return error - need to handle multi-state result + Err(Error::InvalidTransition( + "Fork not yet supported in unified apply", + )) +} + +fn apply_abort( + state: SequenceState, + reason: String, + block_manager: &mut BlockManager, +) -> Result { + let seq_id = state.seq_id(); + + // Free blocks if any + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + warn!("Aborted seq {:?}: {}", seq_id, reason); + + Ok(SequenceState::Aborted(AbortedState { seq_id, reason })) +} diff --git a/src/lib.rs b/src/lib.rs index 0f9ad10..5b5fab3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,6 +5,7 @@ pub mod backend; pub mod block_manager; pub mod cli; pub mod downloader; +pub mod fsm; pub mod registry; pub mod scheduler; pub mod storage; diff --git a/src/main.rs b/src/main.rs index 45a5cd7..4ef25f1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,7 @@ mod backend; mod block_manager; mod cli; mod downloader; +mod fsm; mod registry; mod scheduler; mod storage; diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs index e250e78..552dbf7 100644 --- a/src/scheduler/core.rs +++ b/src/scheduler/core.rs @@ -1,26 +1,27 @@ -use super::events::*; -use super::fsm_events::*; -use super::states::*; +use super::events::SchedulerStats; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::*; +use crate::fsm::{ + AbortedState, DecodingState, Event, FinishReason, FinishedState, PreemptedState, + PrefillingState, SequenceState, WaitingState, +}; use std::collections::{HashMap, VecDeque}; use std::sync::Arc; -use tokio::sync::mpsc; use tracing::{debug, info, warn}; -/// Scheduler - central controller +/// Scheduler - memory and batch management (synchronous API) /// -/// Architecture: -/// - Scheduler DRIVES inference: builds batches → calls backend → updates state -/// - External events (via channel): AddRequest, CancelRequest, GetBlocks, GetStats -/// - Internal operations: token generation, state transitions, OOM handling +/// Architecture +/// - Pure synchronous methods (no async, no loop) +/// - LLMEngine calls: add_request() → schedule() → process_outputs() +/// - Owns BlockManager and performs FSM transitions directly /// /// Responsibilities: /// - Owns BlockManager (memory allocation) /// - Owns sequence state storage (FSM states) /// - Manages batches (waiting/prefill/decode) /// - Scheduling policy (what to schedule when) -/// - Uses FSM events for type-safe state transitions +/// - FSM transitions (internal methods with direct BlockManager access) pub struct Scheduler { /// Block manager for memory allocation block_manager: BlockManager, @@ -37,9 +38,6 @@ pub struct Scheduler { /// Currently running sequences in decode phase decode_batch: Vec, - /// Event receiver - event_rx: mpsc::UnboundedReceiver, - /// Configuration max_batch_size: usize, tokens_per_block: usize, @@ -50,92 +48,26 @@ impl Scheduler { block_manager: BlockManager, max_batch_size: usize, tokens_per_block: usize, - ) -> (Self, mpsc::UnboundedSender) { - let (event_tx, event_rx) = mpsc::unbounded_channel(); - - let scheduler = Self { + ) -> Self { + Self { block_manager, sequences: HashMap::new(), waiting_queue: VecDeque::new(), prefill_batch: Vec::new(), decode_batch: Vec::new(), - event_rx, max_batch_size, tokens_per_block, - }; - - (scheduler, event_tx) - } - - /// Main scheduler loop - drives inference and handles external events - /// - /// Loop structure: - /// 1. Build batch from ready sequences - /// 2. Run inference (TODO: integrate backend) - /// 3. Update state based on inference results (internal) - /// 4. Handle external events (non-blocking) - /// 5. Try scheduling new sequences - pub async fn run(mut self) { - info!("Scheduler started"); - - loop { - // 1. Build batch for inference - // TODO: Call backend.forward(batch) when backend is integrated - - // 2. Handle external events (non-blocking) - // Use try_recv to not block - we want to keep generating tokens - match self.event_rx.try_recv() { - Ok(event) => { - match event { - SchedulerEvent::AddRequest { - seq_id, - prompt_tokens, - max_tokens, - } => { - self.add_request(seq_id, prompt_tokens, max_tokens); - } - - SchedulerEvent::CancelRequest { seq_id } => { - self.cancel_request(seq_id); - } - - SchedulerEvent::GetBlocks { seq_id, response } => { - let result = self.get_blocks(seq_id); - let _ = response.send(result); - } - - SchedulerEvent::GetStats { response } => { - let stats = self.get_stats(); - let _ = response.send(stats); - } - } - - // After event, try scheduling - self.try_schedule(); - } - Err(mpsc::error::TryRecvError::Empty) => { - // No events - continue to next iteration - } - Err(mpsc::error::TryRecvError::Disconnected) => { - info!("Event channel closed, shutting down scheduler"); - break; - } - } - - // Small yield to prevent busy loop - // TODO: Remove once backend integration drives the timing - tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; } - - info!("Scheduler stopped"); } - // ===== External Event Handlers ===== + // ===== Public API (called by LLMEngine) ===== - fn add_request(&mut self, seq_id: SequenceId, prompt_tokens: usize, max_tokens: usize) { + pub fn add_request(&mut self, seq_id: SequenceId, token_ids: Vec, max_tokens: usize) { debug!( - "Adding request: seq_id={:?}, prompt_tokens={}, max_tokens={}", - seq_id, prompt_tokens, max_tokens + "Adding request: seq_id={:?}, num_tokens={}, max_tokens={}", + seq_id, + token_ids.len(), + max_tokens ); // Check duplicate @@ -144,10 +76,10 @@ impl Scheduler { return; } - // Create waiting state + // Create waiting state - tokens already provided by LLMEngine let state = SequenceState::Waiting(WaitingState { seq_id, - prompt_tokens, + token_ids: Arc::new(token_ids), max_tokens, }); @@ -155,7 +87,7 @@ impl Scheduler { self.waiting_queue.push_back(seq_id); } - fn cancel_request(&mut self, seq_id: SequenceId) { + pub fn cancel_request(&mut self, seq_id: SequenceId) { let state = match self.sequences.remove(&seq_id) { Some(s) => s, None => { @@ -164,12 +96,11 @@ impl Scheduler { } }; - let event = AbortEvent { + let event = Event::Abort { reason: "User cancelled".to_string(), - block_manager: &mut self.block_manager, }; - match event.apply(state) { + match self.transition(state, event) { Ok(new_state) => { info!("Cancelled sequence {:?}", seq_id); self.sequences.insert(seq_id, new_state); @@ -185,14 +116,22 @@ impl Scheduler { } } - // ===== Internal Methods (called by scheduler loop during inference) ===== - // These are NOT exposed as events - scheduler calls them directly when: - // - Generating tokens (append_tokens) - // - Detecting stop tokens (complete_sequence) - // - Handling OOM (preempt_sequence, resume_sequence) - // - Beam search (fork_sequence) + /// Schedule sequences - returns whether work was scheduled + /// + /// Called by LLMEngine every iteration to: + /// 1. Move waiting → prefilling + /// 2. Allocate blocks + /// 3. Build batches + pub fn schedule(&mut self) -> bool { + let before = self.prefill_batch.len(); + self.schedule_prefill(); + self.schedule_decode(); + self.prefill_batch.len() > before || !self.decode_batch.is_empty() + } + + // ===== Internal State Management (called by LLMEngine after inference) ===== - fn append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { + pub fn append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { let state = match self.sequences.remove(&seq_id) { Some(s) => s, None => { @@ -202,13 +141,12 @@ impl Scheduler { }; let backup = state.clone(); - let event = AppendTokensEvent { + let event = Event::AppendTokens { num_tokens, - block_manager: &mut self.block_manager, tokens_per_block: self.tokens_per_block, }; - match event.apply(state) { + match self.transition(state, event) { Ok(new_state) => { self.sequences.insert(seq_id, new_state); } @@ -224,7 +162,7 @@ impl Scheduler { } } - fn complete_sequence(&mut self, seq_id: SequenceId, reason: FinishReason) { + pub fn complete_sequence(&mut self, seq_id: SequenceId, reason: FinishReason) { let state = match self.sequences.remove(&seq_id) { Some(s) => s, None => { @@ -234,12 +172,9 @@ impl Scheduler { }; let backup = state.clone(); - let event = CompleteEvent { - reason, - block_manager: &mut self.block_manager, - }; + let event = Event::Complete { reason }; - match event.apply(state) { + match self.transition(state, event) { Ok(new_state) => { info!("Completed sequence {:?}", seq_id); self.sequences.insert(seq_id, new_state); @@ -255,7 +190,7 @@ impl Scheduler { } } - fn fork_sequence(&mut self, parent_id: SequenceId, child_id: SequenceId) { + pub fn fork_sequence(&mut self, parent_id: SequenceId, child_id: SequenceId) { // Check duplicate child_id if self.sequences.contains_key(&child_id) { warn!("Child ID {:?} already exists", child_id); @@ -270,22 +205,35 @@ impl Scheduler { } }; + // TODO: Fork needs special handling - returns (parent, child) + // For now, use manual implementation let backup = parent_state.clone(); - let event = ForkEvent { - child_id, - block_manager: &mut self.block_manager, - }; - match event.apply(parent_state) { - Ok((parent_state, child_state)) => { - self.sequences.insert(parent_id, parent_state); - self.sequences.insert(child_id, child_state); - debug!("Forked sequence {:?} → {:?}", parent_id, child_id); - } - Err(e) => { - warn!("Failed to fork {:?}: {:?}", parent_id, e); - self.sequences.insert(parent_id, backup); + if let SequenceState::Decoding(s) = parent_state { + // Copy-on-write: increment ref counts + for &block_id in s.blocks.iter() { + if let Err(e) = self.block_manager.add_ref(block_id) { + warn!("Failed to fork {:?}: {:?}", parent_id, e); + self.sequences.insert(parent_id, backup); + return; + } } + + let child = DecodingState { + seq_id: child_id, + token_ids: Arc::clone(&s.token_ids), + blocks: Arc::clone(&s.blocks), + num_tokens: s.num_tokens, + max_tokens: s.max_tokens, + }; + + self.sequences.insert(parent_id, SequenceState::Decoding(s)); + self.sequences + .insert(child_id, SequenceState::Decoding(child)); + debug!("Forked sequence {:?} → {:?}", parent_id, child_id); + } else { + warn!("Cannot fork {:?} - not in Decoding state", parent_id); + self.sequences.insert(parent_id, backup); } } @@ -299,11 +247,9 @@ impl Scheduler { }; let backup = state.clone(); - let event = PreemptEvent { - block_manager: &mut self.block_manager, - }; + let event = Event::Preempt; - match event.apply(state) { + match self.transition(state, event) { Ok(new_state) => { info!("Preempted sequence {:?}", seq_id); self.sequences.insert(seq_id, new_state); @@ -329,9 +275,9 @@ impl Scheduler { }; let backup = state.clone(); - let event = ResumeEvent; + let event = Event::Resume; - match event.apply(state) { + match self.transition(state, event) { Ok(new_state) => { debug!("Resumed sequence {:?}", seq_id); self.sequences.insert(seq_id, new_state); @@ -360,64 +306,43 @@ impl Scheduler { break; } - // Get sequence state - let state = match self.sequences.get(&seq_id) { - Some(SequenceState::Waiting(s)) => s, - _ => { - warn!("Sequence {:?} not in Waiting state", seq_id); + // Get and remove state + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); continue; } }; - // Calculate blocks needed - let blocks_needed = state.prompt_tokens.div_ceil(self.tokens_per_block); - - // Check memory - if !self.block_manager.can_allocate(&BlockType::StandardKV) { - self.waiting_queue.push_front(seq_id); - debug!("OOM - cannot schedule {:?}", seq_id); - break; - } + let event = Event::Schedule { + tokens_per_block: self.tokens_per_block, + }; - // Allocate blocks - let mut blocks = Vec::new(); - for _ in 0..blocks_needed { - match self.block_manager.allocate() { - Ok(block_id) => blocks.push(block_id), - Err(Error::OutOfMemory) => { - // Cleanup and stop - for b in blocks { - let _ = self.block_manager.free(b); - } - self.waiting_queue.push_front(seq_id); - warn!("OOM during allocation for {:?}", seq_id); - return; - } - Err(e) => { - for b in blocks { - let _ = self.block_manager.free(b); - } - warn!("Allocation error: {:?}", e); - return; - } + match self.transition(state.clone(), event) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + self.prefill_batch.push(seq_id); + info!("Scheduled {:?} for prefill", seq_id); + } + Err(Error::OutOfMemory) => { + // Transition already cleaned up allocated blocks + // Restore state, put back in queue and stop + self.sequences.insert(seq_id, state); + self.waiting_queue.push_front(seq_id); + debug!("OOM - cannot schedule {:?}", seq_id); + break; + } + Err(Error::InvalidTransition(_)) => { + // Not in Waiting state - restore and skip + self.sequences.insert(seq_id, state); + warn!("Cannot schedule {:?} - not in Waiting state", seq_id); + } + Err(e) => { + // Other error - restore state + self.sequences.insert(seq_id, state); + warn!("Failed to schedule {:?}: {:?}", seq_id, e); } - } - - // Transition: Waiting → Prefilling - let state = self.sequences.remove(&seq_id).unwrap(); - if let SequenceState::Waiting(w) = state { - let new_state = SequenceState::Prefilling(PrefillingState { - seq_id, - blocks: Arc::new(blocks), - tokens_filled: 0, - tokens_total: w.prompt_tokens, - max_tokens: w.max_tokens, - }); - - self.sequences.insert(seq_id, new_state); - self.prefill_batch.push(seq_id); - - info!("Scheduled {:?} for prefill", seq_id); } } } @@ -441,14 +366,14 @@ impl Scheduler { // ===== Query Methods ===== - fn get_blocks(&self, seq_id: SequenceId) -> Result> { + pub fn get_blocks(&self, seq_id: SequenceId) -> Result> { match self.sequences.get(&seq_id) { Some(state) => Ok(state.blocks().map(|b| b.to_vec()).unwrap_or_default()), None => Err(Error::UnknownSequence(seq_id)), } } - fn get_stats(&self) -> SchedulerStats { + pub fn get_stats(&self) -> SchedulerStats { let num_running = self.sequences.values().filter(|s| s.is_running()).count(); let num_waiting = self.waiting_queue.len(); let num_preempted = self @@ -465,22 +390,318 @@ impl Scheduler { block_stats: self.block_manager.get_stats(), } } -} -/// Create scheduler event channel -pub fn create_scheduler_channel() -> ( - mpsc::UnboundedSender, - mpsc::UnboundedReceiver, -) { - mpsc::unbounded_channel() + // ===== FSM Transitions ===== + + /// Transition sequence state using FSM event - single entry point for all transitions + /// + /// This method provides an event-based abstraction over the internal transition methods. + /// + /// # Design + /// + /// **Public API**: Event-based for maintainability + /// - Single entry point makes it easy to add cross-cutting concerns (logging, metrics) + /// - Event enum can be serialized for debugging/replay + /// - Clean interface for callers + /// + /// **Internal Implementation**: Type-safe methods + /// - Private `transition_*()` methods enforce correct state types at compile time + /// - Called by this method after event dispatching + /// + /// # Arguments + /// + /// * `state` - Current sequence state (will be consumed) + /// * `event` - FSM event to apply + /// + /// # Returns + /// + /// New state on success, or `Error::InvalidTransition` for invalid state/event combinations + /// + /// # Example + /// + /// ```rust,ignore + /// let event = Event::Schedule { tokens_per_block: 16 }; + /// let new_state = self.transition(state, event)?; + /// self.sequences.insert(seq_id, new_state); + /// ``` + pub fn transition(&mut self, state: SequenceState, event: Event) -> Result { + use crate::fsm::Event; + + match (state, event) { + // Waiting → Prefilling + (SequenceState::Waiting(s), Event::Schedule { .. }) => self.transition_schedule(s), + + // Prefilling → Prefilling/Decoding + (SequenceState::Prefilling(s), Event::AppendTokens { num_tokens, .. }) => { + self.transition_append_tokens_prefilling(s, num_tokens) + } + + // Decoding → Decoding/Finished + (SequenceState::Decoding(s), Event::AppendTokens { num_tokens, .. }) => { + self.transition_append_tokens_decoding(s, num_tokens) + } + + // Any running state → Finished + (state, Event::Complete { reason }) => self.transition_complete(state, reason), + + // Running → Preempted + ( + state @ (SequenceState::Prefilling(_) | SequenceState::Decoding(_)), + Event::Preempt, + ) => self.transition_preempt(state), + + // Preempted → Waiting + (SequenceState::Preempted(s), Event::Resume) => self.transition_resume(s), + + // Any → Aborted + (state, Event::Abort { reason }) => self.transition_abort(state, reason), + + // Invalid transitions + _ => Err(Error::InvalidTransition("Invalid state transition")), + } + } + + // ===== Internal Transition Methods ===== + // + // These private methods implement the actual transition logic. + // Called by apply() after event dispatching. + + /// Transition: Waiting → Prefilling (allocate blocks for tokenized prompt) + fn transition_schedule(&mut self, state: WaitingState) -> Result { + let prompt_tokens = state.token_ids.len(); + let blocks_needed = prompt_tokens.div_ceil(self.tokens_per_block); + + let mut blocks = Vec::new(); + for _ in 0..blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => blocks.push(block_id), + Err(Error::OutOfMemory) => { + // Cleanup on OOM + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(Error::OutOfMemory); + } + Err(e) => { + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(e); + } + } + } + + debug!( + "Scheduled seq {:?}: allocated {} blocks for {} tokens", + state.seq_id, + blocks.len(), + prompt_tokens + ); + + Ok(SequenceState::Prefilling(PrefillingState { + seq_id: state.seq_id, + token_ids: state.token_ids, + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: prompt_tokens, + max_tokens: state.max_tokens, + })) + } + + /// Transition: Prefilling → Decoding or Prefilling (append tokens) + fn transition_append_tokens_prefilling( + &mut self, + mut state: PrefillingState, + num_tokens: usize, + ) -> Result { + state.tokens_filled += num_tokens; + + debug!( + "Prefilling seq {:?}: {}/{} tokens", + state.seq_id, state.tokens_filled, state.tokens_total + ); + + if state.tokens_filled >= state.tokens_total { + debug!("Seq {:?} transitioning to Decoding", state.seq_id); + Ok(SequenceState::Decoding(DecodingState { + seq_id: state.seq_id, + token_ids: state.token_ids, + blocks: state.blocks, + num_tokens: state.tokens_filled, + max_tokens: state.max_tokens, + })) + } else { + Ok(SequenceState::Prefilling(state)) + } + } + + /// Transition: Decoding → Decoding or Finished (append tokens, maybe allocate more blocks) + fn transition_append_tokens_decoding( + &mut self, + mut state: DecodingState, + num_tokens: usize, + ) -> Result { + state.num_tokens += num_tokens; + + // Check if finished + if state.num_tokens >= state.max_tokens { + // Free blocks + for &block_id in state.blocks.iter() { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!( + "Seq {:?} finished: reached max_tokens ({})", + state.seq_id, state.max_tokens + ); + + return Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: FinishReason::MaxTokens, + })); + } + + // Check if need more blocks + let blocks_needed = state.num_tokens.div_ceil(self.tokens_per_block); + let initial_block_count = state.blocks.len(); + + if state.blocks.len() < blocks_needed { + let blocks = Arc::make_mut(&mut state.blocks); + + while blocks.len() < blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => { + blocks.push(block_id); + debug!( + "Seq {:?}: allocated block, total blocks: {}", + state.seq_id, + blocks.len() + ); + } + Err(e) => { + warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); + for block_id in blocks.drain(initial_block_count..) { + let _ = self.block_manager.free(block_id); + } + return Err(e); + } + } + } + } + + Ok(SequenceState::Decoding(state)) + } + + /// Transition: Prefilling/Decoding → Finished + fn transition_complete( + &mut self, + state: SequenceState, + reason: FinishReason, + ) -> Result { + let seq_id = state.seq_id(); + + // Free blocks if any + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + debug!("Completed seq {:?}: {:?}", seq_id, reason); + + Ok(SequenceState::Finished(FinishedState { + seq_id, + finish_reason: reason, + })) + } + + /// Transition: Prefilling/Decoding → Preempted (free blocks) + fn transition_preempt(&mut self, state: SequenceState) -> Result { + let seq_id = state.seq_id(); + let (token_ids, num_tokens, max_tokens) = match &state { + SequenceState::Prefilling(s) => { + (Arc::clone(&s.token_ids), s.tokens_filled, s.max_tokens) + } + SequenceState::Decoding(s) => (Arc::clone(&s.token_ids), s.num_tokens, s.max_tokens), + _ => { + return Err(Error::InvalidTransition( + "Can only preempt Prefilling/Decoding", + )) + } + }; + + // Free blocks + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + debug!("Preempted seq {:?}", seq_id); + + Ok(SequenceState::Preempted(PreemptedState { + seq_id, + token_ids, + num_tokens, + max_tokens, + })) + } + + /// Transition: Preempted → Waiting + fn transition_resume(&mut self, state: PreemptedState) -> Result { + debug!("Resuming seq {:?}", state.seq_id); + + Ok(SequenceState::Waiting(WaitingState { + seq_id: state.seq_id, + token_ids: state.token_ids, + max_tokens: state.max_tokens, + })) + } + + /// Transition: Any → Aborted (free blocks, cleanup) + fn transition_abort(&mut self, state: SequenceState, reason: String) -> Result { + let seq_id = state.seq_id(); + + // Free blocks if any + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + warn!("Aborted seq {:?}: {}", seq_id, reason); + + Ok(SequenceState::Aborted(AbortedState { seq_id, reason })) + } } +/// Create scheduler event channel #[cfg(test)] mod tests { use super::*; use crate::block_manager::allocator::CpuAllocator; - fn create_test_scheduler() -> (Scheduler, mpsc::UnboundedSender) { + fn create_test_scheduler() -> Scheduler { let allocator = Box::new(CpuAllocator::new(100_000)); let block_manager = BlockManager::new(allocator, 1024); Scheduler::new(block_manager, 10, 16) @@ -488,9 +709,9 @@ mod tests { #[test] fn test_add_request() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); assert_eq!(scheduler.sequences.len(), 1); assert_eq!(scheduler.waiting_queue.len(), 1); @@ -502,10 +723,10 @@ mod tests { #[test] fn test_add_duplicate_request() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); - scheduler.add_request(SequenceId(1), 100, 200); // Duplicate + scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); // Duplicate // Should not add duplicate assert_eq!(scheduler.sequences.len(), 1); @@ -514,9 +735,9 @@ mod tests { #[test] fn test_cancel_request_waiting() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.cancel_request(SequenceId(1)); // Should transition to Aborted @@ -529,7 +750,7 @@ mod tests { #[test] fn test_cancel_request_not_found() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); // Should not panic scheduler.cancel_request(SequenceId(999)); @@ -538,9 +759,9 @@ mod tests { #[test] fn test_schedule_prefill_success() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); // Should move to prefilling @@ -554,10 +775,10 @@ mod tests { #[test] fn test_schedule_prefill_multiple() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); - scheduler.add_request(SequenceId(2), 50, 100); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(2), vec![0; 50], 100); scheduler.schedule_prefill(); // Both should be scheduled @@ -567,11 +788,11 @@ mod tests { #[test] fn test_schedule_prefill_batch_limit() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); // Add more than max_batch_size (10) for i in 0..15 { - scheduler.add_request(SequenceId(i), 100, 200); + scheduler.add_request(SequenceId(i), vec![0; 100], 200); } scheduler.schedule_prefill(); @@ -585,10 +806,10 @@ mod tests { // Small allocator - only 2 blocks let allocator = Box::new(CpuAllocator::new(2048)); let block_manager = BlockManager::new(allocator, 1024); - let (mut scheduler, _tx) = Scheduler::new(block_manager, 10, 16); + let mut scheduler = Scheduler::new(block_manager, 10, 16); // First request needs 7 blocks (100 tokens / 16 tokens_per_block) - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); // Should fail - not enough memory @@ -602,9 +823,9 @@ mod tests { #[test] fn test_append_tokens_prefilling() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); // Append all tokens @@ -619,9 +840,9 @@ mod tests { #[test] fn test_append_tokens_decoding() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding @@ -637,9 +858,9 @@ mod tests { #[test] fn test_append_tokens_reaches_max() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding @@ -655,9 +876,9 @@ mod tests { #[test] fn test_complete_sequence() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding @@ -673,9 +894,9 @@ mod tests { #[test] fn test_fork_sequence() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding @@ -692,13 +913,13 @@ mod tests { #[test] fn test_fork_duplicate_child_id() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); - scheduler.add_request(SequenceId(2), 50, 100); // Child ID exists + scheduler.add_request(SequenceId(2), vec![0; 50], 100); // Child ID exists scheduler.fork_sequence(SequenceId(1), SequenceId(2)); // Should fail // Original seq 2 should be unchanged (Waiting) @@ -710,9 +931,9 @@ mod tests { #[test] fn test_preempt_sequence() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding @@ -728,9 +949,9 @@ mod tests { #[test] fn test_resume_sequence() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); scheduler.preempt_sequence(SequenceId(1)); @@ -747,9 +968,9 @@ mod tests { #[test] fn test_get_blocks() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); let blocks = scheduler.get_blocks(SequenceId(1)).unwrap(); @@ -758,7 +979,7 @@ mod tests { #[test] fn test_get_blocks_not_found() { - let (scheduler, _tx) = create_test_scheduler(); + let scheduler = create_test_scheduler(); let result = scheduler.get_blocks(SequenceId(999)); assert!(matches!(result, Err(Error::UnknownSequence(_)))); @@ -766,10 +987,10 @@ mod tests { #[test] fn test_get_stats() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); - scheduler.add_request(SequenceId(2), 50, 100); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(2), vec![0; 50], 100); scheduler.schedule_prefill(); let stats = scheduler.get_stats(); @@ -780,9 +1001,9 @@ mod tests { #[test] fn test_try_schedule() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.try_schedule(); // Should schedule automatically @@ -792,9 +1013,9 @@ mod tests { #[test] fn test_cancel_removes_from_batches() { - let (mut scheduler, _tx) = create_test_scheduler(); + let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), 100, 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200); scheduler.schedule_prefill(); assert_eq!(scheduler.prefill_batch.len(), 1); diff --git a/src/scheduler/events.rs b/src/scheduler/events.rs index 3541e45..f2c9382 100644 --- a/src/scheduler/events.rs +++ b/src/scheduler/events.rs @@ -1,3 +1,31 @@ +//! Scheduler Events - External API +//! +//! # Two-Level Event Architecture +//! +//! PUMA follows TokenSpeed's pattern with two distinct event levels: +//! +//! ## Level 1: FSM Events (Internal - `crate::fsm::Event`) +//! - Pure state machine transitions +//! - Created by scheduler internally +//! - Includes scheduler resources (BlockManager, allocators) +//! - Examples: Schedule, AppendTokens, Complete +//! +//! ## Level 2: Scheduler Events (External - `SchedulerEvent`) +//! - API boundary for external components +//! - Simple data, no resource pointers +//! - Can be created by clients, GPU, cache systems +//! - Examples: AddRequest, CancelRequest, TokensGenerated +//! +//! ## Why Two Levels? +//! +//! External components (GPU, clients) **cannot create** FSM events because: +//! - FSM events need `&mut BlockManager` (owned by scheduler) +//! - FSM events include policy decisions (scheduler computes) +//! - FSM events track internal state (scheduler maintains) +//! +//! The scheduler translates external events → FSM events, adding +//! the necessary context and resources. + use crate::block_manager::types::*; use std::collections::HashMap; use tokio::sync::mpsc; @@ -17,8 +45,9 @@ pub enum SchedulerEvent { /// Client starts a new inference request AddRequest { seq_id: SequenceId, - prompt_tokens: usize, + token_ids: Vec, // Tokenized by LLMEngine max_tokens: usize, + response_tx: tokio::sync::oneshot::Sender>, // Send result back }, /// Client cancels a request (user stop, disconnect, timeout) diff --git a/src/scheduler/fsm_events.rs b/src/scheduler/fsm_events.rs deleted file mode 100644 index 0c4904d..0000000 --- a/src/scheduler/fsm_events.rs +++ /dev/null @@ -1,740 +0,0 @@ -use super::states::*; -use crate::block_manager::manager::BlockManager; -use crate::block_manager::types::*; -use std::sync::Arc; -use tracing::{debug, warn}; - -/// FSM Events - transform states with type-safe transitions -/// -/// Pattern: Event takes old state, returns new state -/// Invalid transitions return Error::invalid_transition -/// Schedule a waiting sequence (allocate blocks) -pub struct ScheduleEvent<'a> { - pub block_manager: &'a mut BlockManager, - pub tokens_per_block: usize, -} - -impl<'a> ScheduleEvent<'a> { - pub fn apply(self, state: SequenceState) -> Result { - match state { - SequenceState::Waiting(s) => self.apply_to_waiting(s), - _ => Err(Error::InvalidTransition( - "ScheduleEvent requires Waiting state", - )), - } - } - - fn apply_to_waiting(self, state: WaitingState) -> Result { - let blocks_needed = state.prompt_tokens.div_ceil(self.tokens_per_block); - - let mut blocks = Vec::new(); - for _ in 0..blocks_needed { - match self.block_manager.allocate() { - Ok(block_id) => blocks.push(block_id), - Err(Error::OutOfMemory) => { - // OOM during scheduling - free what we allocated - for block in blocks { - let _ = self.block_manager.free(block); - } - return Err(Error::OutOfMemory); - } - Err(e) => { - // Other error - cleanup and propagate - for block in blocks { - let _ = self.block_manager.free(block); - } - return Err(e); - } - } - } - - debug!( - "Scheduled seq {:?}: allocated {} blocks for {} tokens", - state.seq_id, - blocks.len(), - state.prompt_tokens - ); - - Ok(SequenceState::Prefilling(PrefillingState { - seq_id: state.seq_id, - blocks: Arc::new(blocks), - tokens_filled: 0, - tokens_total: state.prompt_tokens, - max_tokens: state.max_tokens, - })) - } -} - -/// Append tokens to a running sequence -pub struct AppendTokensEvent<'a> { - pub num_tokens: usize, - pub block_manager: &'a mut BlockManager, - pub tokens_per_block: usize, -} - -impl<'a> AppendTokensEvent<'a> { - pub fn apply(self, state: SequenceState) -> Result { - match state { - SequenceState::Prefilling(s) => self.apply_to_prefilling(s), - SequenceState::Decoding(s) => self.apply_to_decoding(s), - _ => Err(Error::InvalidTransition( - "AppendTokensEvent requires Prefilling or Decoding state", - )), - } - } - - fn apply_to_prefilling(self, mut state: PrefillingState) -> Result { - state.tokens_filled += self.num_tokens; - - debug!( - "Prefilling seq {:?}: {}/{} tokens", - state.seq_id, state.tokens_filled, state.tokens_total - ); - - if state.tokens_filled >= state.tokens_total { - // Transition to decode - debug!("Seq {:?} transitioning to Decoding", state.seq_id); - Ok(SequenceState::Decoding(DecodingState { - seq_id: state.seq_id, - blocks: state.blocks, - num_tokens: state.tokens_filled, - max_tokens: state.max_tokens, - })) - } else { - Ok(SequenceState::Prefilling(state)) - } - } - - fn apply_to_decoding(self, mut state: DecodingState) -> Result { - state.num_tokens += self.num_tokens; - - // Check if finished - if state.num_tokens >= state.max_tokens { - // Free blocks - for &block_id in state.blocks.iter() { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!( - "Seq {:?} finished: reached max_tokens ({})", - state.seq_id, state.max_tokens - ); - - return Ok(SequenceState::Finished(FinishedState { - seq_id: state.seq_id, - finish_reason: FinishReason::MaxTokens, - })); - } - - // Check if need more blocks - let blocks_needed = state.num_tokens.div_ceil(self.tokens_per_block); - let initial_block_count = state.blocks.len(); - - if state.blocks.len() < blocks_needed { - // Need to allocate more blocks - use Arc::make_mut for copy-on-write - let blocks = Arc::make_mut(&mut state.blocks); - - while blocks.len() < blocks_needed { - match self.block_manager.allocate() { - Ok(block_id) => { - blocks.push(block_id); - debug!( - "Seq {:?}: allocated block, total blocks: {}", - state.seq_id, - blocks.len() - ); - } - Err(e) => { - // Free any blocks we allocated in this call - warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); - for block_id in blocks.drain(initial_block_count..) { - let _ = self.block_manager.free(block_id); - } - return Err(e); - } - } - } - } - - Ok(SequenceState::Decoding(state)) - } -} - -/// Preempt a running sequence (free blocks) -pub struct PreemptEvent<'a> { - pub block_manager: &'a mut BlockManager, -} - -impl<'a> PreemptEvent<'a> { - pub fn apply(self, state: SequenceState) -> Result { - match state { - SequenceState::Decoding(s) => self.apply_to_decoding(s), - SequenceState::Prefilling(s) => self.apply_to_prefilling(s), - _ => Err(Error::InvalidTransition( - "PreemptEvent requires running state", - )), - } - } - - fn apply_to_decoding(self, state: DecodingState) -> Result { - // Free all blocks - for &block_id in state.blocks.iter() { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!("Preempted seq {:?}", state.seq_id); - - Ok(SequenceState::Preempted(PreemptedState { - seq_id: state.seq_id, - num_tokens: state.num_tokens, - max_tokens: state.max_tokens, - })) - } - - fn apply_to_prefilling(self, state: PrefillingState) -> Result { - // Free all blocks (best effort) - for &block_id in state.blocks.iter() { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!("Preempted seq {:?} during prefill", state.seq_id); - - Ok(SequenceState::Preempted(PreemptedState { - seq_id: state.seq_id, - num_tokens: state.tokens_filled, - max_tokens: state.max_tokens, - })) - } -} - -/// Resume a preempted sequence -pub struct ResumeEvent; - -impl ResumeEvent { - pub fn apply(self, state: SequenceState) -> Result { - match state { - SequenceState::Preempted(s) => self.apply_to_preempted(s), - _ => Err(Error::InvalidTransition( - "ResumeEvent requires Preempted state", - )), - } - } - - fn apply_to_preempted(self, state: PreemptedState) -> Result { - debug!("Resuming seq {:?}", state.seq_id); - - // Transition back to waiting for rescheduling - Ok(SequenceState::Waiting(WaitingState { - seq_id: state.seq_id, - prompt_tokens: state.num_tokens, - max_tokens: state.max_tokens, - })) - } -} - -/// Fork a sequence (copy-on-write) -pub struct ForkEvent<'a> { - pub child_id: SequenceId, - pub block_manager: &'a mut BlockManager, -} - -impl<'a> ForkEvent<'a> { - pub fn apply(self, state: SequenceState) -> Result<(SequenceState, SequenceState)> { - match state { - SequenceState::Decoding(s) => self.apply_to_decoding(s), - _ => Err(Error::InvalidTransition( - "ForkEvent requires Decoding state", - )), - } - } - - fn apply_to_decoding(self, state: DecodingState) -> Result<(SequenceState, SequenceState)> { - // Copy-on-write: increment ref counts - for &block_id in state.blocks.iter() { - self.block_manager.add_ref(block_id)?; - } - - let child = DecodingState { - seq_id: self.child_id, - blocks: state.blocks.clone(), - num_tokens: state.num_tokens, - max_tokens: state.max_tokens, - }; - - debug!("Forked seq {:?} → {:?}", state.seq_id, child.seq_id); - - Ok(( - SequenceState::Decoding(state), - SequenceState::Decoding(child), - )) - } -} - -/// Complete a sequence -pub struct CompleteEvent<'a> { - pub reason: FinishReason, - pub block_manager: &'a mut BlockManager, -} - -impl<'a> CompleteEvent<'a> { - pub fn apply(self, state: SequenceState) -> Result { - match state { - SequenceState::Decoding(s) => self.apply_to_decoding(s), - SequenceState::Prefilling(s) => self.apply_to_prefilling(s), - _ => Err(Error::InvalidTransition( - "CompleteEvent requires running state", - )), - } - } - - fn apply_to_decoding(self, state: DecodingState) -> Result { - // Free all blocks - for &block_id in state.blocks.iter() { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!("Completed seq {:?}: {:?}", state.seq_id, self.reason); - - Ok(SequenceState::Finished(FinishedState { - seq_id: state.seq_id, - finish_reason: self.reason, - })) - } - - fn apply_to_prefilling(self, state: PrefillingState) -> Result { - // Free all blocks - for &block_id in state.blocks.iter() { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!( - "Completed seq {:?} during prefill: {:?}", - state.seq_id, self.reason - ); - - Ok(SequenceState::Finished(FinishedState { - seq_id: state.seq_id, - finish_reason: self.reason, - })) - } -} - -/// Abort a sequence (error or cancellation) -pub struct AbortEvent<'a> { - pub reason: String, - pub block_manager: &'a mut BlockManager, -} - -impl<'a> AbortEvent<'a> { - pub fn apply(self, state: SequenceState) -> Result { - // Can abort from any state - let seq_id = state.seq_id(); - - // Free blocks if any - if let Some(blocks) = state.blocks() { - for &block_id in blocks { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, seq_id, e - ); - } - } - } - - warn!("Aborted seq {:?}: {}", seq_id, self.reason); - - Ok(SequenceState::Aborted(AbortedState { - seq_id, - reason: self.reason, - })) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::block_manager::allocator::CpuAllocator; - use crate::block_manager::manager::BlockManager; - - #[test] - fn test_schedule_event_waiting_to_prefilling() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let waiting = SequenceState::Waiting(WaitingState { - seq_id: SequenceId(1), - prompt_tokens: 100, - max_tokens: 150, - }); - - let event = ScheduleEvent { - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(waiting); - assert!(result.is_ok()); - - let new_state = result.unwrap(); - assert!(matches!(new_state, SequenceState::Prefilling(_))); - - if let SequenceState::Prefilling(state) = new_state { - assert_eq!(state.seq_id, SequenceId(1)); - assert_eq!(state.tokens_filled, 0); - assert_eq!(state.tokens_total, 100); - // Should allocate ceil(100/16) = 7 blocks - assert_eq!(state.blocks.len(), 7); - } - } - - #[test] - fn test_schedule_event_oom() { - // Small allocator - only 2 blocks - let allocator = Box::new(CpuAllocator::new(2048)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let waiting = SequenceState::Waiting(WaitingState { - seq_id: SequenceId(1), - prompt_tokens: 100, // Needs 7 blocks - max_tokens: 150, - }); - - let event = ScheduleEvent { - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(waiting); - assert!(matches!(result, Err(Error::OutOfMemory))); - - // Should have cleaned up - no blocks leaked - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_append_tokens_prefilling_to_decoding() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - // Allocate blocks - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - - let prefilling = SequenceState::Prefilling(PrefillingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - tokens_filled: 0, - tokens_total: 100, - max_tokens: 150, - }); - - let event = AppendTokensEvent { - num_tokens: 100, - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(prefilling).unwrap(); - - // Should transition to Decoding - assert!(matches!(result, SequenceState::Decoding(_))); - - if let SequenceState::Decoding(state) = result { - assert_eq!(state.num_tokens, 100); - assert_eq!(state.blocks.len(), 2); - } - } - - #[test] - fn test_append_tokens_decoding_needs_more_blocks() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![block_manager.allocate().unwrap()]; // 1 block - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 16, // 1 block worth - max_tokens: 150, - }); - - let event = AppendTokensEvent { - num_tokens: 20, // Now need 3 blocks total - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(decoding).unwrap(); - - if let SequenceState::Decoding(state) = result { - assert_eq!(state.num_tokens, 36); - // Should allocate 2 more blocks (total 3) - assert_eq!(state.blocks.len(), 3); - } - } - - #[test] - fn test_append_tokens_decoding_oom_cleanup() { - // Small allocator - only enough for initial blocks - let allocator = Box::new(CpuAllocator::new(2048)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - let initial_block_count = blocks.len(); - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 32, - max_tokens: 150, - }); - - // Try to append tokens that would need more blocks (will fail - OOM) - let event = AppendTokensEvent { - num_tokens: 50, // Would need 6 blocks total - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(decoding); - assert!(matches!(result, Err(Error::OutOfMemory))); - - // Original 2 blocks should still be allocated (not freed by error path) - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, initial_block_count); - } - - #[test] - fn test_append_tokens_reaches_max_tokens() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![block_manager.allocate().unwrap()]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 140, - max_tokens: 150, - }); - - let event = AppendTokensEvent { - num_tokens: 10, - block_manager: &mut block_manager, - tokens_per_block: 16, - }; - - let result = event.apply(decoding).unwrap(); - - // Should transition to Finished - assert!(matches!(result, SequenceState::Finished(_))); - - if let SequenceState::Finished(state) = result { - assert_eq!(state.seq_id, SequenceId(1)); - assert_eq!(state.finish_reason, FinishReason::MaxTokens); - } - - // Blocks should be freed - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_preempt_event() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 32, - max_tokens: 150, - }); - - let event = PreemptEvent { - block_manager: &mut block_manager, - }; - - let result = event.apply(decoding).unwrap(); - - // Should transition to Preempted - assert!(matches!(result, SequenceState::Preempted(_))); - - if let SequenceState::Preempted(state) = result { - assert_eq!(state.num_tokens, 32); - } - - // Blocks should be freed - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_resume_event() { - let preempted = SequenceState::Preempted(PreemptedState { - seq_id: SequenceId(1), - num_tokens: 32, - max_tokens: 150, - }); - - let event = ResumeEvent; - let result = event.apply(preempted).unwrap(); - - // Should transition back to Waiting - assert!(matches!(result, SequenceState::Waiting(_))); - - if let SequenceState::Waiting(state) = result { - assert_eq!(state.prompt_tokens, 32); - assert_eq!(state.max_tokens, 150); - } - } - - #[test] - fn test_fork_event() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 32, - max_tokens: 150, - }); - - let event = ForkEvent { - child_id: SequenceId(2), - block_manager: &mut block_manager, - }; - - let result = event.apply(decoding).unwrap(); - let (parent, child) = result; - - // Both should be in Decoding state - assert!(matches!(parent, SequenceState::Decoding(_))); - assert!(matches!(child, SequenceState::Decoding(_))); - - if let (SequenceState::Decoding(p), SequenceState::Decoding(c)) = (parent, child) { - assert_eq!(p.seq_id, SequenceId(1)); - assert_eq!(c.seq_id, SequenceId(2)); - assert_eq!(p.blocks, c.blocks); // Shared blocks - } - } - - #[test] - fn test_complete_event() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![block_manager.allocate().unwrap()]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 100, - max_tokens: 150, - }); - - let event = CompleteEvent { - reason: FinishReason::Stop, - block_manager: &mut block_manager, - }; - - let result = event.apply(decoding).unwrap(); - - assert!(matches!(result, SequenceState::Finished(_))); - - if let SequenceState::Finished(state) = result { - assert_eq!(state.finish_reason, FinishReason::Stop); - } - - // Blocks should be freed - let stats = block_manager.get_stats(); - let block_stats = stats.get(&BlockType::StandardKV).unwrap(); - assert_eq!(block_stats.allocated_blocks, 0); - } - - #[test] - fn test_fork_duplicate_child_id() { - let allocator = Box::new(CpuAllocator::new(100_000)); - let mut block_manager = BlockManager::new(allocator, 1024); - - let blocks = vec![ - block_manager.allocate().unwrap(), - block_manager.allocate().unwrap(), - ]; - - let decoding = SequenceState::Decoding(DecodingState { - seq_id: SequenceId(1), - blocks: Arc::new(blocks), - num_tokens: 32, - max_tokens: 150, - }); - - // Fork with child_id = SequenceId(2) - let event = ForkEvent { - child_id: SequenceId(2), - block_manager: &mut block_manager, - }; - - let result = event.apply(decoding).unwrap(); - let (parent, child) = result; - - // Verify fork succeeded - if let SequenceState::Decoding(c) = &child { - assert_eq!(c.seq_id, SequenceId(2)); - } - - // Now try to fork again with the SAME child_id (should fail in manager) - // This test verifies the ForkEvent itself works, but the manager - // should check for duplicate child_id before calling this - drop(parent); - drop(child); - } -} diff --git a/src/scheduler/mod.rs b/src/scheduler/mod.rs index dd9bd23..f463eba 100644 --- a/src/scheduler/mod.rs +++ b/src/scheduler/mod.rs @@ -1,5 +1,5 @@ pub mod core; pub mod events; -pub mod fsm_events; pub mod id_generator; -pub mod states; + +// Re-export for convenience From 57251f518d7685d9e96d0535172416da8283520a Mon Sep 17 00:00:00 2001 From: kerthcet Date: Thu, 16 Jul 2026 01:07:51 +0100 Subject: [PATCH 23/33] support cli Signed-off-by: kerthcet --- Cargo.lock | 23 +++++++ Cargo.toml | 1 + Makefile | 1 + src/backend/llm_engine.rs | 76 +++++++++++++++-------- src/cli/commands.rs | 26 ++++++-- src/scheduler/core.rs | 126 ++++++++++++++++++++++++++++++-------- src/scheduler/events.rs | 11 +++- 7 files changed, 205 insertions(+), 59 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e062eab..6c71338 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -94,6 +94,28 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -1812,6 +1834,7 @@ dependencies = [ name = "puma" version = "0.0.3" dependencies = [ + "async-stream", "axum", "chrono", "clap", diff --git a/Cargo.toml b/Cargo.toml index cdcccd0..40cd30b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ tokio-stream = "0.1" rustyline = "14.0" rustyline-derive = "0.10" tokenizers = "0.20" +async-stream = "0.3" [dev-dependencies] tempfile = "3.12" diff --git a/Makefile b/Makefile index 0e72459..d48a626 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,5 @@ build: + rm -rf ./puma cargo build && cp target/debug/puma ./puma test: diff --git a/src/backend/llm_engine.rs b/src/backend/llm_engine.rs index e5fb190..1990e0c 100644 --- a/src/backend/llm_engine.rs +++ b/src/backend/llm_engine.rs @@ -8,7 +8,7 @@ use crate::block_manager::allocator::CpuAllocator; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::SequenceId; use crate::scheduler::core::Scheduler; -use crate::scheduler::events::SchedulerEvent; +use crate::scheduler::events::{ResponseSender, SchedulerEvent}; /// LLM Engine - integrates scheduler with inference backend /// @@ -73,8 +73,10 @@ impl LLMEngine { seq_id, token_ids, max_tokens, + response_tx, } => { - self.scheduler.add_request(seq_id, token_ids, max_tokens); + self.scheduler + .add_request(seq_id, token_ids, max_tokens, response_tx); } SchedulerEvent::CancelRequest { seq_id } => { @@ -120,28 +122,39 @@ impl InferenceEngine for LLMEngine { let seq_id = self.next_seq_id(); // Tokenize prompt using HuggingFace tokenizer - let encoding = self - .tokenizer - .encode(prompt, false) - .map_err(|e| io::Error::other(format!("Tokenization failed: {}", e)))?; - let token_ids = encoding.get_ids().to_vec(); + let token_ids = match self.tokenizer.encode(prompt, false) { + Ok(encoding) => encoding.get_ids().to_vec(), + Err(e) => { + // Fallback: if tokenizer has no vocab, use dummy tokens + tracing::warn!("Tokenization failed ({}), using dummy tokens", e); + vec![0u32; prompt.len().min(100)] // Dummy: 1 token per char, max 100 + } + }; let num_tokens = token_ids.len(); - // Send request to scheduler + // Create response channel (single response) + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + + // Send request to scheduler with response channel self.event_tx .send(SchedulerEvent::AddRequest { seq_id, token_ids, max_tokens, + response_tx: ResponseSender::Single(response_tx), }) .map_err(|e| io::Error::other(format!("Engine send failed: {}", e)))?; - // TODO: Wait for actual result from scheduler/GPU - // For now, return dummy response + // Wait for result from scheduler + let text = response_rx + .await + .map_err(|e| io::Error::other(format!("Response channel closed: {}", e)))? + .map_err(|e| io::Error::other(format!("Scheduler error: {:?}", e)))?; + Ok(GenerateResponse { - text: format!("Generated response for: {}", prompt), + text, prompt_tokens: num_tokens, - completion_tokens: 0, + completion_tokens: 0, // TODO: Track actual completion tokens }) } @@ -155,28 +168,41 @@ impl InferenceEngine for LLMEngine { let seq_id = self.next_seq_id(); // Tokenize prompt using HuggingFace tokenizer - let encoding = self - .tokenizer - .encode(prompt, false) - .map_err(|e| io::Error::other(format!("Tokenization failed: {}", e)))?; - let token_ids = encoding.get_ids().to_vec(); + let token_ids = match self.tokenizer.encode(prompt, false) { + Ok(encoding) => encoding.get_ids().to_vec(), + Err(e) => { + // Fallback: if tokenizer has no vocab, use dummy tokens + tracing::warn!("Tokenization failed ({}), using dummy tokens", e); + vec![0u32; prompt.len().min(100)] // Dummy: 1 token per char, max 100 + } + }; + + // Create streaming response channel + let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel(); - // Send request to scheduler + // Send request to scheduler with streaming channel self.event_tx .send(SchedulerEvent::AddRequest { seq_id, token_ids, max_tokens, + response_tx: ResponseSender::Stream(response_tx), }) .map_err(|e| io::Error::other(format!("Engine send failed: {}", e)))?; - // TODO: Stream results from scheduler/GPU - // For now, return dummy stream - let stream = tokio_stream::iter(vec![ - format!("Token 1 for: {}\n", prompt), - "Token 2\n".to_string(), - "Token 3\n".to_string(), - ]); + // Create stream that receives tokens from scheduler + let stream = async_stream::stream! { + while let Some(result) = response_rx.recv().await { + match result { + Ok(token) => yield token, + Err(e) => { + tracing::error!("Stream error: {:?}", e); + break; + } + } + } + }; + Ok(Box::pin(stream)) } } diff --git a/src/cli/commands.rs b/src/cli/commands.rs index 4ed1b4b..c4d003c 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -1,8 +1,9 @@ - use clap::{Parser, Subcommand}; use colored::Colorize; use prettytable::{format, row, Table}; +use tokenizers::Tokenizer; + use crate::backend::llm_engine::LLMEngine; use crate::backend::mock::MockEngine; use crate::cli::{chat, inspect, ls, rm}; @@ -211,11 +212,24 @@ pub async fn run(cli: Cli) { Ok(Some(_)) => {} } - // Load tokenizer - // TODO: Load from model directory when model is downloaded - // For now, create a simple test tokenizer - use tokenizers::{models::bpe::BPE, Tokenizer}; - let tokenizer = Tokenizer::new(BPE::default()); + // Load tokenizer from model directory + let model_info = registry.get_model(&args.model).unwrap().unwrap(); + let tokenizer_path = format!( + "{}/snapshots/{}/tokenizer.json", + model_info.metadata.cache.path, model_info.metadata.cache.revision + ); + + let tokenizer = match Tokenizer::from_file(&tokenizer_path) { + Ok(tok) => tok, + Err(e) => { + eprintln!( + "Warning: Could not load tokenizer from {}: {}. Using BPE tokenizer instead.", + tokenizer_path, e + ); + use tokenizers::models::bpe::BPE; + Tokenizer::new(BPE::default()) + } + }; // Load inference backend // TODO: Replace MockEngine with real backend that loads model files diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs index 552dbf7..7b713f9 100644 --- a/src/scheduler/core.rs +++ b/src/scheduler/core.rs @@ -1,4 +1,4 @@ -use super::events::SchedulerStats; +use super::events::{ResponseSender, SchedulerStats}; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::*; use crate::fsm::{ @@ -29,6 +29,9 @@ pub struct Scheduler { /// All sequences and their states sequences: HashMap, + /// Response channels to send results back to clients (supports streaming) + response_channels: HashMap, + /// Waiting sequences (not yet scheduled) waiting_queue: VecDeque, @@ -52,6 +55,7 @@ impl Scheduler { Self { block_manager, sequences: HashMap::new(), + response_channels: HashMap::new(), waiting_queue: VecDeque::new(), prefill_batch: Vec::new(), decode_batch: Vec::new(), @@ -62,7 +66,13 @@ impl Scheduler { // ===== Public API (called by LLMEngine) ===== - pub fn add_request(&mut self, seq_id: SequenceId, token_ids: Vec, max_tokens: usize) { + pub fn add_request( + &mut self, + seq_id: SequenceId, + token_ids: Vec, + max_tokens: usize, + response_tx: ResponseSender, + ) { debug!( "Adding request: seq_id={:?}, num_tokens={}, max_tokens={}", seq_id, @@ -73,9 +83,13 @@ impl Scheduler { // Check duplicate if self.sequences.contains_key(&seq_id) { warn!("Duplicate seq_id {:?}", seq_id); + self.send_error(response_tx, Error::InvalidTransition("Duplicate sequence ID")); return; } + // Store response channel + self.response_channels.insert(seq_id, response_tx); + // Create waiting state - tokens already provided by LLMEngine let state = SequenceState::Waiting(WaitingState { seq_id, @@ -162,7 +176,7 @@ impl Scheduler { } } - pub fn complete_sequence(&mut self, seq_id: SequenceId, reason: FinishReason) { + pub fn complete_sequence(&mut self, seq_id: SequenceId, reason: FinishReason, result: String) { let state = match self.sequences.remove(&seq_id) { Some(s) => s, None => { @@ -182,10 +196,18 @@ impl Scheduler { // Remove from batches self.prefill_batch.retain(|&id| id != seq_id); self.decode_batch.retain(|&id| id != seq_id); + + // Send result back to client + self.send_result(seq_id, result); } Err(e) => { warn!("Failed to complete {:?}: {:?}", seq_id, e); self.sequences.insert(seq_id, backup); + + // Send error back to client + if let Some(response_tx) = self.response_channels.remove(&seq_id) { + self.send_error(response_tx, e); + } } } } @@ -693,6 +715,50 @@ impl Scheduler { Ok(SequenceState::Aborted(AbortedState { seq_id, reason })) } + + // ===== Response Channel Helpers ===== + + /// Send a token to client (for streaming) + pub fn send_token(&mut self, seq_id: SequenceId, token: String) { + if let Some(response_tx) = self.response_channels.get(&seq_id) { + match response_tx { + ResponseSender::Stream(tx) => { + let _ = tx.send(Ok(token)); + } + ResponseSender::Single(_) => { + // Single response - can't stream individual tokens + // Will send complete result at the end + } + } + } + } + + /// Send error to client + fn send_error(&self, response_tx: ResponseSender, error: Error) { + match response_tx { + ResponseSender::Single(tx) => { + let _ = tx.send(Err(error)); + } + ResponseSender::Stream(tx) => { + let _ = tx.send(Err(error)); + } + } + } + + /// Send final result to client + fn send_result(&mut self, seq_id: SequenceId, result: String) { + if let Some(response_tx) = self.response_channels.remove(&seq_id) { + match response_tx { + ResponseSender::Single(tx) => { + let _ = tx.send(Ok(result)); + } + ResponseSender::Stream(_tx) => { + // Already streamed tokens, just close the channel + // (drop _tx automatically closes it) + } + } + } + } } /// Create scheduler event channel @@ -700,6 +766,7 @@ impl Scheduler { mod tests { use super::*; use crate::block_manager::allocator::CpuAllocator; + use crate::scheduler::events::ResponseSender; fn create_test_scheduler() -> Scheduler { let allocator = Box::new(CpuAllocator::new(100_000)); @@ -707,11 +774,16 @@ mod tests { Scheduler::new(block_manager, 10, 16) } + fn create_dummy_response() -> ResponseSender { + let (tx, _rx) = tokio::sync::oneshot::channel(); + ResponseSender::Single(tx) + } + #[test] fn test_add_request() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); assert_eq!(scheduler.sequences.len(), 1); assert_eq!(scheduler.waiting_queue.len(), 1); @@ -725,8 +797,8 @@ mod tests { fn test_add_duplicate_request() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); // Duplicate + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); // Duplicate // Should not add duplicate assert_eq!(scheduler.sequences.len(), 1); @@ -737,7 +809,7 @@ mod tests { fn test_cancel_request_waiting() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.cancel_request(SequenceId(1)); // Should transition to Aborted @@ -761,7 +833,7 @@ mod tests { fn test_schedule_prefill_success() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); // Should move to prefilling @@ -777,8 +849,8 @@ mod tests { fn test_schedule_prefill_multiple() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); - scheduler.add_request(SequenceId(2), vec![0; 50], 100); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.add_request(SequenceId(2), vec![0; 50], 100, create_dummy_response()); scheduler.schedule_prefill(); // Both should be scheduled @@ -792,7 +864,7 @@ mod tests { // Add more than max_batch_size (10) for i in 0..15 { - scheduler.add_request(SequenceId(i), vec![0; 100], 200); + scheduler.add_request(SequenceId(i), vec![0; 100], 200, create_dummy_response()); } scheduler.schedule_prefill(); @@ -809,7 +881,7 @@ mod tests { let mut scheduler = Scheduler::new(block_manager, 10, 16); // First request needs 7 blocks (100 tokens / 16 tokens_per_block) - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); // Should fail - not enough memory @@ -825,7 +897,7 @@ mod tests { fn test_append_tokens_prefilling() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); // Append all tokens @@ -842,7 +914,7 @@ mod tests { fn test_append_tokens_decoding() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding @@ -860,7 +932,7 @@ mod tests { fn test_append_tokens_reaches_max() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding @@ -878,11 +950,11 @@ mod tests { fn test_complete_sequence() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding - scheduler.complete_sequence(SequenceId(1), FinishReason::Stop); + scheduler.complete_sequence(SequenceId(1), FinishReason::Stop, "test result".to_string()); // Should be Finished assert!(matches!( @@ -896,7 +968,7 @@ mod tests { fn test_fork_sequence() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding @@ -915,11 +987,11 @@ mod tests { fn test_fork_duplicate_child_id() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); - scheduler.add_request(SequenceId(2), vec![0; 50], 100); // Child ID exists + scheduler.add_request(SequenceId(2), vec![0; 50], 100, create_dummy_response()); // Child ID exists scheduler.fork_sequence(SequenceId(1), SequenceId(2)); // Should fail // Original seq 2 should be unchanged (Waiting) @@ -933,7 +1005,7 @@ mod tests { fn test_preempt_sequence() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); // → Decoding @@ -951,7 +1023,7 @@ mod tests { fn test_resume_sequence() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); scheduler.append_tokens(SequenceId(1), 100); scheduler.preempt_sequence(SequenceId(1)); @@ -970,7 +1042,7 @@ mod tests { fn test_get_blocks() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); let blocks = scheduler.get_blocks(SequenceId(1)).unwrap(); @@ -989,8 +1061,8 @@ mod tests { fn test_get_stats() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); - scheduler.add_request(SequenceId(2), vec![0; 50], 100); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); + scheduler.add_request(SequenceId(2), vec![0; 50], 100, create_dummy_response()); scheduler.schedule_prefill(); let stats = scheduler.get_stats(); @@ -1003,7 +1075,7 @@ mod tests { fn test_try_schedule() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.try_schedule(); // Should schedule automatically @@ -1015,7 +1087,7 @@ mod tests { fn test_cancel_removes_from_batches() { let mut scheduler = create_test_scheduler(); - scheduler.add_request(SequenceId(1), vec![0; 100], 200); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); scheduler.schedule_prefill(); assert_eq!(scheduler.prefill_batch.len(), 1); diff --git a/src/scheduler/events.rs b/src/scheduler/events.rs index f2c9382..1589fc9 100644 --- a/src/scheduler/events.rs +++ b/src/scheduler/events.rs @@ -38,6 +38,15 @@ use tokio::sync::mpsc; /// - GetBlocks: Backend queries block IDs for inference /// - GetStats: Monitoring/observability queries /// +/// Response sender - supports both single response and streaming +#[derive(Debug)] +pub enum ResponseSender { + /// Single response (non-streaming) + Single(tokio::sync::oneshot::Sender>), + /// Streaming response (multiple tokens) + Stream(tokio::sync::mpsc::UnboundedSender>), +} + /// Internal operations (token generation, state transitions, OOM handling) /// are NOT events - the scheduler handles them directly in its run loop. #[derive(Debug)] @@ -47,7 +56,7 @@ pub enum SchedulerEvent { seq_id: SequenceId, token_ids: Vec, // Tokenized by LLMEngine max_tokens: usize, - response_tx: tokio::sync::oneshot::Sender>, // Send result back + response_tx: ResponseSender, // Send result(s) back }, /// Client cancels a request (user stop, disconnect, timeout) From 6900ecded3e2832598afde486d70f3b2103dab05 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Thu, 16 Jul 2026 15:27:58 +0100 Subject: [PATCH 24/33] WIP: Backend rename and tokenization refactor (partial) --- src/api/chat.rs | 2 +- src/api/completions.rs | 2 +- src/api/models.rs | 2 +- src/api/routes.rs | 2 +- src/backend/engine.rs | 32 +++++----- src/backend/llm_engine.rs | 119 ++++++++++++++++++++------------------ src/backend/mock.rs | 64 +++++++++----------- src/backend/mod.rs | 3 +- src/cli/chat.rs | 2 +- 9 files changed, 112 insertions(+), 116 deletions(-) diff --git a/src/api/chat.rs b/src/api/chat.rs index 0dfcf4f..2d8bb67 100644 --- a/src/api/chat.rs +++ b/src/api/chat.rs @@ -16,7 +16,7 @@ use crate::api::types::{ ChatChoice, ChatChoiceDelta, ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, ChatMessage, ChatMessageDelta, ErrorResponse, Usage, }; -use crate::backend::InferenceEngine; +use crate::backend::{Backend, LLMEngine}; /// Main handler for chat completions pub async fn chat_completions( diff --git a/src/api/completions.rs b/src/api/completions.rs index 497736b..cf084be 100644 --- a/src/api/completions.rs +++ b/src/api/completions.rs @@ -5,7 +5,7 @@ use crate::api::routes::AppState; use crate::api::types::{ CompletionChoice, CompletionRequest, CompletionResponse, ErrorResponse, Usage, }; -use crate::backend::InferenceEngine; +use crate::backend::{Backend, LLMEngine}; /// Handler for legacy text completions pub async fn completions( diff --git a/src/api/models.rs b/src/api/models.rs index c8304e1..8c06705 100644 --- a/src/api/models.rs +++ b/src/api/models.rs @@ -6,7 +6,7 @@ use axum::{ use crate::api::routes::AppState; use crate::api::types::{ErrorResponse, Model, ModelList}; -use crate::backend::InferenceEngine; +use crate::backend::{Backend, LLMEngine}; /// List all available models pub async fn list_models( diff --git a/src/api/routes.rs b/src/api/routes.rs index dae6c0e..b2d01b7 100644 --- a/src/api/routes.rs +++ b/src/api/routes.rs @@ -10,7 +10,7 @@ use tower_http::{ LatencyUnit, }; -use crate::backend::InferenceEngine; +use crate::backend::{Backend, LLMEngine}; use crate::registry::model_registry::ModelRegistry; use super::{chat, completions, models}; diff --git a/src/backend/engine.rs b/src/backend/engine.rs index 2d464a5..93924f6 100644 --- a/src/backend/engine.rs +++ b/src/backend/engine.rs @@ -1,34 +1,30 @@ use std::io; use std::pin::Pin; use tokio_stream::Stream; +use crate::block_manager::types::TokenId; -/// Inference engine trait -pub trait InferenceEngine: Send + Sync { - /// Generate text completion +/// Backend trait - low-level inference that works with token IDs +/// +/// LLMEngine handles tokenization (text → tokens) +/// Backend handles inference (tokens → tokens) +pub trait Backend: Send + Sync { + /// Generate tokens from input token_ids + /// Returns generated token IDs fn generate( &self, - model: &str, - prompt: &str, + token_ids: Vec, max_tokens: usize, temperature: f32, - ) -> impl std::future::Future> + Send; + ) -> impl std::future::Future, io::Error>> + Send; - /// Generate text with streaming + /// Generate tokens with streaming + /// Returns stream of token IDs as they're generated fn generate_stream( &self, - model: &str, - prompt: &str, + token_ids: Vec, max_tokens: usize, temperature: f32, ) -> impl std::future::Future< - Output = Result + Send>>, io::Error>, + Output = Result + Send>>, io::Error>, > + Send; } - -/// Generation response -#[derive(Debug, Clone)] -pub struct GenerateResponse { - pub text: String, - pub prompt_tokens: usize, - pub completion_tokens: usize, -} diff --git a/src/backend/llm_engine.rs b/src/backend/llm_engine.rs index 1990e0c..f3298c3 100644 --- a/src/backend/llm_engine.rs +++ b/src/backend/llm_engine.rs @@ -3,7 +3,15 @@ use std::sync::atomic::{AtomicU64, Ordering}; use tokenizers::Tokenizer; use tokio::sync::mpsc; -use super::engine::{GenerateResponse, InferenceEngine}; +use super::engine::Backend; + +/// User-facing response for generate() +#[derive(Debug, Clone)] +pub struct GenerateResponse { + pub text: String, + pub prompt_tokens: usize, + pub completion_tokens: usize, +} use crate::block_manager::allocator::CpuAllocator; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::SequenceId; @@ -18,7 +26,7 @@ use crate::scheduler::events::{ResponseSender, SchedulerEvent}; /// - Scheduler is synchronous (just methods, no async) /// - Direct ownership (no Arc/Mutex overhead) /// - Loop: handle events → schedule() → forward() → process_outputs() -pub struct LLMEngine { +pub struct LLMEngine { backend: B, scheduler: Scheduler, tokenizer: Tokenizer, @@ -28,7 +36,7 @@ pub struct LLMEngine { model: String, } -impl LLMEngine { +impl LLMEngine { pub fn new(backend: B, tokenizer: Tokenizer, model: String) -> Self { // Create block manager (100MB memory pool, 512 bytes per block) let allocator = Box::new(CpuAllocator::new(1024 * 1024 * 100)); @@ -55,6 +63,39 @@ impl LLMEngine { SequenceId(self.seq_id_counter.fetch_add(1, Ordering::Relaxed)) } + /// Tokenize prompt + fn tokenize(&self, prompt: &str) -> Result, io::Error> { + self.tokenizer + .encode(prompt, false) + .map(|encoding| encoding.get_ids().to_vec()) + .map_err(|e| io::Error::other(format!("Tokenization failed: {}", e))) + } + + /// Decode token IDs to text + fn decode_tokens(&self, token_ids: &[u32]) -> Result { + self.tokenizer + .decode(token_ids, false) + .map_err(|e| io::Error::other(format!("Detokenization failed: {}", e))) + } + + /// Send request helper + fn send_request( + &self, + token_ids: Vec, + max_tokens: usize, + response_tx: ResponseSender, + ) -> Result<(), io::Error> { + let seq_id = self.next_seq_id(); + self.event_tx + .send(SchedulerEvent::AddRequest { + seq_id, + token_ids, + max_tokens, + response_tx, + }) + .map_err(|e| io::Error::other(format!("Engine send failed: {}", e))) + } + /// Main loop - drives scheduling and inference /// /// Pattern @@ -110,40 +151,23 @@ impl LLMEngine { } } -// Implement InferenceEngine directly for LLMEngine -impl InferenceEngine for LLMEngine { - async fn generate( +// User-facing API methods +impl LLMEngine { + /// Generate text completion (single response) + pub async fn generate( &self, - _model: &str, prompt: &str, max_tokens: usize, - _temperature: f32, ) -> Result { - let seq_id = self.next_seq_id(); - - // Tokenize prompt using HuggingFace tokenizer - let token_ids = match self.tokenizer.encode(prompt, false) { - Ok(encoding) => encoding.get_ids().to_vec(), - Err(e) => { - // Fallback: if tokenizer has no vocab, use dummy tokens - tracing::warn!("Tokenization failed ({}), using dummy tokens", e); - vec![0u32; prompt.len().min(100)] // Dummy: 1 token per char, max 100 - } - }; - let num_tokens = token_ids.len(); + // Tokenize + let token_ids = self.tokenize(prompt)?; + let prompt_tokens = token_ids.len(); - // Create response channel (single response) + // Create response channel let (response_tx, response_rx) = tokio::sync::oneshot::channel(); - // Send request to scheduler with response channel - self.event_tx - .send(SchedulerEvent::AddRequest { - seq_id, - token_ids, - max_tokens, - response_tx: ResponseSender::Single(response_tx), - }) - .map_err(|e| io::Error::other(format!("Engine send failed: {}", e)))?; + // Send request + self.send_request(token_ids, max_tokens, ResponseSender::Single(response_tx))?; // Wait for result from scheduler let text = response_rx @@ -153,42 +177,25 @@ impl InferenceEngine for LLMEngine { Ok(GenerateResponse { text, - prompt_tokens: num_tokens, - completion_tokens: 0, // TODO: Track actual completion tokens + prompt_tokens, + completion_tokens: 0, }) } - async fn generate_stream( + /// Generate with streaming + pub async fn generate_stream( &self, - _model: &str, prompt: &str, max_tokens: usize, - _temperature: f32, ) -> Result + Send>>, io::Error> { - let seq_id = self.next_seq_id(); + // Tokenize + let token_ids = self.tokenize(prompt)?; - // Tokenize prompt using HuggingFace tokenizer - let token_ids = match self.tokenizer.encode(prompt, false) { - Ok(encoding) => encoding.get_ids().to_vec(), - Err(e) => { - // Fallback: if tokenizer has no vocab, use dummy tokens - tracing::warn!("Tokenization failed ({}), using dummy tokens", e); - vec![0u32; prompt.len().min(100)] // Dummy: 1 token per char, max 100 - } - }; - - // Create streaming response channel + // Create channel let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel(); - // Send request to scheduler with streaming channel - self.event_tx - .send(SchedulerEvent::AddRequest { - seq_id, - token_ids, - max_tokens, - response_tx: ResponseSender::Stream(response_tx), - }) - .map_err(|e| io::Error::other(format!("Engine send failed: {}", e)))?; + // Send request + self.send_request(token_ids, max_tokens, ResponseSender::Stream(response_tx))?; // Create stream that receives tokens from scheduler let stream = async_stream::stream! { diff --git a/src/backend/mock.rs b/src/backend/mock.rs index 92f056e..faa3c63 100644 --- a/src/backend/mock.rs +++ b/src/backend/mock.rs @@ -3,7 +3,8 @@ use std::io; use std::pin::Pin; use tokio_stream::Stream; -use super::engine::{GenerateResponse, InferenceEngine}; +use super::engine::Backend; +use crate::block_manager::types::TokenId; /// Mock engine for testing (replace with MLX later) #[derive(Clone)] @@ -15,52 +16,43 @@ impl MockEngine { } } -impl InferenceEngine for MockEngine { +impl Backend for MockEngine { async fn generate( &self, - model: &str, - prompt: &str, + token_ids: Vec, max_tokens: usize, _temperature: f32, - ) -> Result { - // Mock response for testing - let response_text = format!( - "This is a mock response from model '{}' for prompt: '{}' (max_tokens: {})", - model, - prompt.chars().take(50).collect::(), - max_tokens - ); - - Ok(GenerateResponse { - text: response_text, - prompt_tokens: prompt.split_whitespace().count(), - completion_tokens: 20, - }) + ) -> Result, io::Error> { + // Mock: echo input + add generated tokens + // This makes different inputs produce different outputs + let mut result = token_ids; + for i in 0..max_tokens.min(10) { + result.push(1000 + i as u32); // Append tokens 1000, 1001, 1002... + } + Ok(result) } async fn generate_stream( &self, - model: &str, - _prompt: &str, + token_ids: Vec, max_tokens: usize, _temperature: f32, - ) -> Result + Send>>, io::Error> { - // Mock streaming response - let tokens = vec![ - "This ".to_string(), - "is ".to_string(), - "a ".to_string(), - "mock ".to_string(), - "streaming ".to_string(), - "response ".to_string(), - format!("from model '{}' ", model), - format!("(max_tokens: {}).", max_tokens), - ]; + ) -> Result + Send>>, io::Error> { + // Mock: echo input tokens first, then generate new ones + // This makes different inputs produce different outputs + + // First yield all input tokens (the prompt) + let mut all_tokens = token_ids; + + // Then add generated tokens + for i in 0..max_tokens.min(10) { + all_tokens.push(1000 + i as u32); // Generate tokens 1000, 1001, 1002... + } - // Simulate delay between tokens - let stream = stream::iter(tokens).then(|token| async move { - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - token + // Simulate delay between tokens (like real GPU) + let stream = stream::iter(all_tokens).then(|token_id| async move { + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + token_id }); Ok(Box::pin(stream)) diff --git a/src/backend/mod.rs b/src/backend/mod.rs index d9220b1..68fc792 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -2,4 +2,5 @@ pub mod engine; pub mod llm_engine; pub mod mock; -pub use engine::*; +pub use engine::Backend; +pub use llm_engine::LLMEngine; diff --git a/src/cli/chat.rs b/src/cli/chat.rs index 6d7c5a9..d9bb260 100644 --- a/src/cli/chat.rs +++ b/src/cli/chat.rs @@ -5,7 +5,7 @@ use rustyline_derive::{Completer, Helper, Highlighter, Validator}; use std::io::{self, Write}; use tokio_stream::StreamExt; -use crate::backend::InferenceEngine; +use crate::backend::{Backend, LLMEngine}; #[derive(Clone)] struct PlaceholderHint { From 4a78c7d4545da6de7371fdb6efd36189d0ab1536 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Thu, 16 Jul 2026 15:30:34 +0100 Subject: [PATCH 25/33] WIP: API signature updates (broken state) --- src/api/chat.rs | 12 ++++++------ src/api/completions.rs | 5 +++-- src/api/models.rs | 9 +++++---- src/api/routes.rs | 16 ++++++++-------- 4 files changed, 22 insertions(+), 20 deletions(-) diff --git a/src/api/chat.rs b/src/api/chat.rs index 2d8bb67..061b11f 100644 --- a/src/api/chat.rs +++ b/src/api/chat.rs @@ -19,8 +19,8 @@ use crate::api::types::{ use crate::backend::{Backend, LLMEngine}; /// Main handler for chat completions -pub async fn chat_completions( - State(state): State>, +pub async fn chat_completions( + State(state): State>>>, Json(req): Json, ) -> Response { let engine = state.engine; @@ -83,8 +83,8 @@ pub async fn chat_completions( } /// Non-streaming chat completion -async fn chat_completions_non_stream( - engine: Arc, +async fn chat_completions_non_stream( + engine: Arc>>, req: ChatCompletionRequest, ) -> Result> { let id = format!("chatcmpl-{}", Uuid::new_v4()); @@ -125,8 +125,8 @@ async fn chat_completions_non_stream( } /// Streaming chat completion -async fn chat_completions_stream( - engine: Arc, +async fn chat_completions_stream( + engine: Arc>>, req: ChatCompletionRequest, ) -> Sse>> { let id = format!("chatcmpl-{}", Uuid::new_v4()); diff --git a/src/api/completions.rs b/src/api/completions.rs index cf084be..119a73c 100644 --- a/src/api/completions.rs +++ b/src/api/completions.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use axum::{extract::State, response::IntoResponse, Json}; use uuid::Uuid; @@ -8,8 +9,8 @@ use crate::api::types::{ use crate::backend::{Backend, LLMEngine}; /// Handler for legacy text completions -pub async fn completions( - State(state): State>, +pub async fn completions( + State(state): State>>>, Json(req): Json, ) -> impl IntoResponse { let engine = state.engine; diff --git a/src/api/models.rs b/src/api/models.rs index 8c06705..53486c4 100644 --- a/src/api/models.rs +++ b/src/api/models.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use axum::{ extract::{Path, State}, response::IntoResponse, @@ -9,8 +10,8 @@ use crate::api::types::{ErrorResponse, Model, ModelList}; use crate::backend::{Backend, LLMEngine}; /// List all available models -pub async fn list_models( - State(state): State>, +pub async fn list_models( + State(state): State>>>, ) -> impl IntoResponse { let registry = state.registry; match registry.load_models(None) { @@ -43,8 +44,8 @@ pub async fn list_models( } /// Get a specific model by ID -pub async fn get_model( - State(state): State>, +pub async fn get_model( + State(state): State>>>, Path(model_id): Path, ) -> impl IntoResponse { let registry = state.registry; diff --git a/src/api/routes.rs b/src/api/routes.rs index b2d01b7..55faa7e 100644 --- a/src/api/routes.rs +++ b/src/api/routes.rs @@ -17,26 +17,26 @@ use super::{chat, completions, models}; /// Shared application state #[derive(Clone)] -pub struct AppState { - pub engine: Arc, +pub struct AppState { + pub engine: Arc>>, pub registry: Arc, } /// Create the API router with all endpoints -pub fn create_router( - engine: Arc, +pub fn create_router( + engine: Arc>>, registry: Arc, ) -> Router { let state = AppState { engine, registry }; Router::new() // Chat completions (most important) - .route("/v1/chat/completions", post(chat::chat_completions::)) + .route("/v1/chat/completions", post(chat::chat_completions::>>)) // Legacy completions - .route("/v1/completions", post(completions::completions::)) + .route("/v1/completions", post(completions::completions::>>)) // Models - .route("/v1/models", get(models::list_models::)) - .route("/v1/models/:model", get(models::get_model::)) + .route("/v1/models", get(models::list_models::>>)) + .route("/v1/models/:model", get(models::get_model::>>)) // Health check .route("/health", get(health_check)) // Pass state From 3bc604cc1e693182d65292479cdc8994b8910d18 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sat, 18 Jul 2026 13:23:47 +0100 Subject: [PATCH 26/33] enable the simple pipeline Signed-off-by: kerthcet --- src/api/chat.rs | 33 ++-- src/api/completions.rs | 6 +- src/api/models.rs | 10 +- src/api/routes.rs | 19 +- src/api/tests.rs | 18 +- src/backend/llm_engine.rs | 392 ++++++++++++++++++++++++++------------ src/backend/mock.rs | 144 +++++++++++--- src/backend/mod.rs | 3 +- src/cli/chat.rs | 23 +-- src/cli/commands.rs | 30 +-- src/cli/serve.rs | 18 +- src/main.rs | 20 +- src/scheduler/core.rs | 39 +++- src/scheduler/events.rs | 7 + src/utils/mod.rs | 1 + src/utils/prompt.rs | 65 +++++++ 16 files changed, 596 insertions(+), 232 deletions(-) create mode 100644 src/utils/prompt.rs diff --git a/src/api/chat.rs b/src/api/chat.rs index 061b11f..a15e99e 100644 --- a/src/api/chat.rs +++ b/src/api/chat.rs @@ -7,7 +7,6 @@ use axum::{ Json, }; use futures::stream::StreamExt; -use std::sync::Arc; use tokio_stream::wrappers::ReceiverStream; use uuid::Uuid; @@ -16,11 +15,11 @@ use crate::api::types::{ ChatChoice, ChatChoiceDelta, ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, ChatMessage, ChatMessageDelta, ErrorResponse, Usage, }; -use crate::backend::{Backend, LLMEngine}; +use crate::backend::EngineHandle; /// Main handler for chat completions -pub async fn chat_completions( - State(state): State>>>, +pub async fn chat_completions( + State(state): State, Json(req): Json, ) -> Response { let engine = state.engine; @@ -83,8 +82,8 @@ pub async fn chat_completions( } /// Non-streaming chat completion -async fn chat_completions_non_stream( - engine: Arc>>, +async fn chat_completions_non_stream( + engine: EngineHandle, req: ChatCompletionRequest, ) -> Result> { let id = format!("chatcmpl-{}", Uuid::new_v4()); @@ -125,8 +124,8 @@ async fn chat_completions_non_stream( } /// Streaming chat completion -async fn chat_completions_stream( - engine: Arc>>, +async fn chat_completions_stream( + engine: EngineHandle, req: ChatCompletionRequest, ) -> Sse>> { let id = format!("chatcmpl-{}", Uuid::new_v4()); @@ -236,19 +235,9 @@ async fn chat_completions_stream( Sse::new(ReceiverStream::new(rx)).keep_alive(KeepAlive::default()) } -/// Format chat messages into a prompt +/// Format chat messages into a prompt via the shared prompt formatter. fn format_chat_messages(messages: &[ChatMessage]) -> String { - messages - .iter() - .map(|m| { - if m.role == "system" { - format!("System: {}", m.content) - } else if m.role == "user" { - format!("User: {}", m.content) - } else { - format!("Assistant: {}", m.content) - } - }) - .collect::>() - .join("\n") + crate::utils::prompt::format_conversation( + messages.iter().map(|m| (m.role.as_str(), m.content.as_str())), + ) } diff --git a/src/api/completions.rs b/src/api/completions.rs index 119a73c..c9bedf8 100644 --- a/src/api/completions.rs +++ b/src/api/completions.rs @@ -1,4 +1,3 @@ -use std::sync::Arc; use axum::{extract::State, response::IntoResponse, Json}; use uuid::Uuid; @@ -6,11 +5,10 @@ use crate::api::routes::AppState; use crate::api::types::{ CompletionChoice, CompletionRequest, CompletionResponse, ErrorResponse, Usage, }; -use crate::backend::{Backend, LLMEngine}; /// Handler for legacy text completions -pub async fn completions( - State(state): State>>>, +pub async fn completions( + State(state): State, Json(req): Json, ) -> impl IntoResponse { let engine = state.engine; diff --git a/src/api/models.rs b/src/api/models.rs index 53486c4..00b3be6 100644 --- a/src/api/models.rs +++ b/src/api/models.rs @@ -1,4 +1,3 @@ -use std::sync::Arc; use axum::{ extract::{Path, State}, response::IntoResponse, @@ -7,12 +6,9 @@ use axum::{ use crate::api::routes::AppState; use crate::api::types::{ErrorResponse, Model, ModelList}; -use crate::backend::{Backend, LLMEngine}; /// List all available models -pub async fn list_models( - State(state): State>>>, -) -> impl IntoResponse { +pub async fn list_models(State(state): State) -> impl IntoResponse { let registry = state.registry; match registry.load_models(None) { Ok(models) => { @@ -44,8 +40,8 @@ pub async fn list_models( } /// Get a specific model by ID -pub async fn get_model( - State(state): State>>>, +pub async fn get_model( + State(state): State, Path(model_id): Path, ) -> impl IntoResponse { let registry = state.registry; diff --git a/src/api/routes.rs b/src/api/routes.rs index 55faa7e..15875be 100644 --- a/src/api/routes.rs +++ b/src/api/routes.rs @@ -10,33 +10,30 @@ use tower_http::{ LatencyUnit, }; -use crate::backend::{Backend, LLMEngine}; +use crate::backend::EngineHandle; use crate::registry::model_registry::ModelRegistry; use super::{chat, completions, models}; /// Shared application state #[derive(Clone)] -pub struct AppState { - pub engine: Arc>>, +pub struct AppState { + pub engine: EngineHandle, pub registry: Arc, } /// Create the API router with all endpoints -pub fn create_router( - engine: Arc>>, - registry: Arc, -) -> Router { +pub fn create_router(engine: EngineHandle, registry: Arc) -> Router { let state = AppState { engine, registry }; Router::new() // Chat completions (most important) - .route("/v1/chat/completions", post(chat::chat_completions::>>)) + .route("/v1/chat/completions", post(chat::chat_completions)) // Legacy completions - .route("/v1/completions", post(completions::completions::>>)) + .route("/v1/completions", post(completions::completions)) // Models - .route("/v1/models", get(models::list_models::>>)) - .route("/v1/models/:model", get(models::get_model::>>)) + .route("/v1/models", get(models::list_models)) + .route("/v1/models/:model", get(models::get_model)) // Health check .route("/health", get(health_check)) // Pass state diff --git a/src/api/tests.rs b/src/api/tests.rs index fba05ca..fd8939e 100644 --- a/src/api/tests.rs +++ b/src/api/tests.rs @@ -13,13 +13,21 @@ use tempfile::TempDir; use tower::util::ServiceExt; // for `oneshot` and `ready` use super::routes::create_router; +use crate::backend::engine; use crate::backend::mock::MockEngine; use crate::registry::model_registry::{CacheInfo, ModelInfo, ModelMetadata, ModelRegistry}; /// Helper to create test app with a pre-registered test model /// Returns the router and the temp directory (which must be kept alive) fn create_test_app() -> (axum::Router, TempDir) { - let engine = Arc::new(MockEngine::new()); + // Build the engine and spawn its runner; the handle drives the router + let (handle, runner) = engine( + MockEngine::new(), + create_test_tokenizer(), + "test-model".to_string(), + ); + tokio::spawn(runner.serve()); + let temp_dir = TempDir::new().unwrap(); let registry = Arc::new(ModelRegistry::new(Some(temp_dir.path().to_path_buf()))); @@ -49,7 +57,13 @@ fn create_test_app() -> (axum::Router, TempDir) { .register_model(test_model) .expect("failed to register test model"); - (create_router(engine, registry), temp_dir) + (create_router(handle, registry), temp_dir) +} + +/// Simple BPE tokenizer for tests +fn create_test_tokenizer() -> tokenizers::Tokenizer { + use tokenizers::models::bpe::BPE; + tokenizers::Tokenizer::new(BPE::default()) } /// Helper to make a JSON request diff --git a/src/backend/llm_engine.rs b/src/backend/llm_engine.rs index f3298c3..db73eef 100644 --- a/src/backend/llm_engine.rs +++ b/src/backend/llm_engine.rs @@ -1,9 +1,17 @@ use std::io; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; use tokenizers::Tokenizer; use tokio::sync::mpsc; +use tokio_stream::StreamExt; use super::engine::Backend; +use crate::block_manager::allocator::CpuAllocator; +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::{SequenceId, TokenId}; +use crate::fsm::FinishReason; +use crate::scheduler::core::Scheduler; +use crate::scheduler::events::{ResponseSender, SchedulerEvent}; /// User-facing response for generate() #[derive(Debug, Clone)] @@ -12,53 +20,27 @@ pub struct GenerateResponse { pub prompt_tokens: usize, pub completion_tokens: usize, } -use crate::block_manager::allocator::CpuAllocator; -use crate::block_manager::manager::BlockManager; -use crate::block_manager::types::SequenceId; -use crate::scheduler::core::Scheduler; -use crate::scheduler::events::{ResponseSender, SchedulerEvent}; -/// LLM Engine - integrates scheduler with inference backend +/// Send-side handle to the engine (cheap to clone). /// -/// Architecture -/// - LLMEngine drives the main loop -/// - Owns tokenizer (model-specific) -/// - Scheduler is synchronous (just methods, no async) -/// - Direct ownership (no Arc/Mutex overhead) -/// - Loop: handle events → schedule() → forward() → process_outputs() -pub struct LLMEngine { - backend: B, - scheduler: Scheduler, - tokenizer: Tokenizer, - event_rx: mpsc::UnboundedReceiver, +/// Holds only what's needed to translate a user request into a +/// `SchedulerEvent` and push it onto the event channel: +/// - tokenizer (text → tokens) +/// - event sender (fire the event) +/// - sequence id counter (assign ids) +/// +/// The actual work happens in [`EngineRunner::serve`], which owns the scheduler +/// and backend and consumes these events. This handle is `Clone` so the API +/// layer and CLI can share it freely across tasks. +#[derive(Clone)] +pub struct EngineHandle { + tokenizer: Arc, event_tx: mpsc::UnboundedSender, - seq_id_counter: AtomicU64, + seq_id_counter: Arc, model: String, } -impl LLMEngine { - pub fn new(backend: B, tokenizer: Tokenizer, model: String) -> Self { - // Create block manager (100MB memory pool, 512 bytes per block) - let allocator = Box::new(CpuAllocator::new(1024 * 1024 * 100)); - let block_manager = BlockManager::new(allocator, 512); - - // Create scheduler (max 32 batch size, 16 tokens per block) - let scheduler = Scheduler::new(block_manager, 32, 16); - - // Event channel - let (event_tx, event_rx) = mpsc::unbounded_channel(); - - Self { - backend, - scheduler, - tokenizer, - event_rx, - event_tx, - seq_id_counter: AtomicU64::new(0), - model, - } - } - +impl EngineHandle { fn next_seq_id(&self) -> SequenceId { SequenceId(self.seq_id_counter.fetch_add(1, Ordering::Relaxed)) } @@ -71,14 +53,7 @@ impl LLMEngine { .map_err(|e| io::Error::other(format!("Tokenization failed: {}", e))) } - /// Decode token IDs to text - fn decode_tokens(&self, token_ids: &[u32]) -> Result { - self.tokenizer - .decode(token_ids, false) - .map_err(|e| io::Error::other(format!("Detokenization failed: {}", e))) - } - - /// Send request helper + /// Send an AddRequest event to the engine loop fn send_request( &self, token_ids: Vec, @@ -96,6 +71,193 @@ impl LLMEngine { .map_err(|e| io::Error::other(format!("Engine send failed: {}", e))) } + /// Generate text completion (single response) + pub async fn generate( + &self, + _model: &str, + prompt: &str, + max_tokens: usize, + _temperature: f32, + ) -> Result { + let token_ids = self.tokenize(prompt)?; + let prompt_tokens = token_ids.len(); + + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + self.send_request(token_ids, max_tokens, ResponseSender::Single(response_tx))?; + + let text = response_rx + .await + .map_err(|e| io::Error::other(format!("Response channel closed: {}", e)))? + .map_err(|e| io::Error::other(format!("Scheduler error: {:?}", e)))?; + + Ok(GenerateResponse { + text, + prompt_tokens, + completion_tokens: 0, + }) + } + + /// Generate with streaming + pub async fn generate_stream( + &self, + _model: &str, + prompt: &str, + max_tokens: usize, + _temperature: f32, + ) -> Result + Send>>, io::Error> { + let token_ids = self.tokenize(prompt)?; + + let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel(); + self.send_request(token_ids, max_tokens, ResponseSender::Stream(response_tx))?; + + let stream = async_stream::stream! { + while let Some(result) = response_rx.recv().await { + match result { + Ok(token) => yield token, + Err(e) => { + tracing::error!("Stream error: {:?}", e); + break; + } + } + } + }; + + Ok(Box::pin(stream)) + } + + /// Model name this handle serves + pub fn model(&self) -> &str { + &self.model + } +} + +/// The engine itself: owns the scheduler and backend, drives the event loop. +/// +/// Drains the event channel and runs inference. Spawn its +/// [`serve`](Self::serve) on a task; hold the paired [`EngineHandle`] everywhere +/// else to submit work. +pub struct EngineRunner { + backend: B, + scheduler: Scheduler, + tokenizer: Arc, + event_rx: mpsc::UnboundedReceiver, +} + +impl EngineRunner { + /// Decode token IDs to text + fn decode_tokens(&self, token_ids: &[u32]) -> Result { + self.tokenizer + .decode(token_ids, false) + .map_err(|e| io::Error::other(format!("Detokenization failed: {}", e))) + } + + /// Run one-shot inference for a non-streaming request. + /// + /// Calls [`Backend::generate`], decodes the completion, advances the FSM, + /// and delivers the full text to the client's single-response channel. + async fn forward_single( + &mut self, + seq_id: SequenceId, + token_ids: Arc>, + max_tokens: usize, + ) { + let prompt_len = token_ids.len(); + let completion_tokens = match self + .backend + .generate((*token_ids).clone(), max_tokens, 0.0) + .await + { + Ok(tokens) => tokens, + Err(e) => { + tracing::error!("Backend inference failed for {:?}: {}", seq_id, e); + self.scheduler.cancel_request(seq_id); + return; + } + }; + + // The backend returns only the newly generated completion tokens. + let num_completion = completion_tokens.len(); + let text = self.decode_tokens(&completion_tokens).unwrap_or_default(); + + // Advance FSM state (prefill → decode → finished) so blocks are freed, + // then deliver the result to the client channel. + self.scheduler.append_tokens(seq_id, prompt_len); + self.scheduler + .complete_sequence(seq_id, FinishReason::Stop, text); + tracing::debug!( + "Completed {:?}: {} completion tokens", + seq_id, + num_completion + ); + } + + /// Run streaming inference for a streaming request. + /// + /// Calls [`Backend::generate_stream`] and forwards each decoded text chunk to + /// the client via [`send_chunk`](Scheduler::send_chunk) as it arrives, then + /// advances the FSM and closes the stream once generation finishes. + async fn forward_stream( + &mut self, + seq_id: SequenceId, + token_ids: Arc>, + max_tokens: usize, + ) { + let prompt_len = token_ids.len(); + let mut stream = match self + .backend + .generate_stream((*token_ids).clone(), max_tokens, 0.0) + .await + { + Ok(stream) => stream, + Err(e) => { + tracing::error!("Backend stream failed for {:?}: {}", seq_id, e); + self.scheduler.cancel_request(seq_id); + return; + } + }; + + // The backend streams only the newly generated completion tokens. + // + // Decode incrementally rather than one token at a time: subword/byte-BPE + // tokenizers split multi-byte characters (emoji, CJK) across several + // tokens, so a lone token often decodes to `�` or "". We keep the running + // list of completion tokens and decode the whole prefix each step, which + // is always well-formed. When the decode ends in the Unicode replacement + // character the trailing multi-byte char is still incomplete, so we hold + // the output back until the next token completes it; otherwise we emit the + // newly-decoded suffix. This is the streaming detokenization scheme vLLM + // and TGI use. + let mut completion_ids: Vec = Vec::new(); + let mut sent_len = 0usize; // bytes of the decoded completion already sent + while let Some(token_id) = stream.next().await { + completion_ids.push(token_id); + + let decoded = self.decode_tokens(&completion_ids).unwrap_or_default(); + // A trailing replacement char means an incomplete multi-byte + // character; wait for the next token before emitting more. + if decoded.ends_with('\u{FFFD}') { + continue; + } + if decoded.len() > sent_len { + let delta = decoded[sent_len..].to_string(); + sent_len = decoded.len(); + self.scheduler.send_chunk(seq_id, delta); + } + } + let num_completion = completion_ids.len(); + + // Advance FSM state (prefill → decode → finished) so blocks are freed, + // then close the stream channel. + self.scheduler.append_tokens(seq_id, prompt_len); + self.scheduler + .complete_sequence(seq_id, FinishReason::Stop, String::new()); + tracing::debug!( + "Streamed {:?}: {} completion tokens", + seq_id, + num_completion + ); + } + /// Main loop - drives scheduling and inference /// /// Pattern @@ -103,8 +265,8 @@ impl LLMEngine { /// 2. Schedule (decide what to run) /// 3. Forward (run inference if work exists) /// 4. Process outputs (update scheduler state) - pub async fn run(mut self) { - tracing::info!("LLMEngine started"); + pub async fn serve(mut self) { + tracing::info!("EngineRunner started"); loop { // 1. Handle events (non-blocking, drain all) @@ -139,10 +301,24 @@ impl LLMEngine { // 2. Schedule (always, not event-driven) let has_work = self.scheduler.schedule(); - // 3. Forward (if work exists) + // 3. Forward: run inference on the scheduled batch and report + // results back to the scheduler, which fulfills client channels. + // Streaming requests use the backend's streaming path (token by + // token); single requests use the one-shot path. + // + // TODO: this processes one sequence at a time, awaiting each backend + // call in turn. For real batch inference, partition the batch by mode + // (see the `PrefillBatch` design) and hand the non-streaming bucket to + // a single `Backend::generate_batch` call instead of looping. if has_work { - // TODO: Build batch and call backend.forward() - // For now, just yield + for work in self.scheduler.take_prefill_batch() { + let (seq_id, token_ids, max_tokens, streaming) = work; + if streaming { + self.forward_stream(seq_id, token_ids, max_tokens).await; + } else { + self.forward_single(seq_id, token_ids, max_tokens).await; + } + } } // 4. Small yield to prevent busy loop @@ -151,67 +327,43 @@ impl LLMEngine { } } -// User-facing API methods -impl LLMEngine { - /// Generate text completion (single response) - pub async fn generate( - &self, - prompt: &str, - max_tokens: usize, - ) -> Result { - // Tokenize - let token_ids = self.tokenize(prompt)?; - let prompt_tokens = token_ids.len(); - - // Create response channel - let (response_tx, response_rx) = tokio::sync::oneshot::channel(); - - // Send request - self.send_request(token_ids, max_tokens, ResponseSender::Single(response_tx))?; - - // Wait for result from scheduler - let text = response_rx - .await - .map_err(|e| io::Error::other(format!("Response channel closed: {}", e)))? - .map_err(|e| io::Error::other(format!("Scheduler error: {:?}", e)))?; - - Ok(GenerateResponse { - text, - prompt_tokens, - completion_tokens: 0, - }) - } - - /// Generate with streaming - pub async fn generate_stream( - &self, - prompt: &str, - max_tokens: usize, - ) -> Result + Send>>, io::Error> { - // Tokenize - let token_ids = self.tokenize(prompt)?; - - // Create channel - let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel(); - - // Send request - self.send_request(token_ids, max_tokens, ResponseSender::Stream(response_tx))?; - - // Create stream that receives tokens from scheduler - let stream = async_stream::stream! { - while let Some(result) = response_rx.recv().await { - match result { - Ok(token) => yield token, - Err(e) => { - tracing::error!("Stream error: {:?}", e); - break; - } - } - } - }; - - Ok(Box::pin(stream)) - } +/// Construct a paired [`EngineHandle`] and [`EngineRunner`]. +/// +/// Spawn `runner.serve()` on a task and share the returned handle with the API / +/// CLI. All request submission goes through events, so the handle never +/// touches the scheduler directly. +pub fn engine( + backend: B, + tokenizer: Tokenizer, + model: String, +) -> (EngineHandle, EngineRunner) { + // Create block manager (100MB memory pool, 512 bytes per block) + let allocator = Box::new(CpuAllocator::new(1024 * 1024 * 100)); + let block_manager = BlockManager::new(allocator, 512); + + // Create scheduler (max 32 batch size, 16 tokens per block) + let scheduler = Scheduler::new(block_manager, 32, 16); + + // Event channel + let (event_tx, event_rx) = mpsc::unbounded_channel(); + + let tokenizer = Arc::new(tokenizer); + + let handle = EngineHandle { + tokenizer: tokenizer.clone(), + event_tx, + seq_id_counter: Arc::new(AtomicU64::new(1)), + model, + }; + + let runner = EngineRunner { + backend, + scheduler, + tokenizer, + event_rx, + }; + + (handle, runner) } #[cfg(test)] @@ -220,7 +372,6 @@ mod tests { use crate::backend::mock::MockEngine; fn create_test_tokenizer() -> Tokenizer { - // Create a simple BPE tokenizer for testing use tokenizers::models::bpe::BPE; use tokenizers::Tokenizer as TokenizerBuilder; @@ -232,14 +383,11 @@ mod tests { async fn test_llm_engine() { let backend = MockEngine::new(); let tokenizer = create_test_tokenizer(); - let engine = LLMEngine::new(backend, tokenizer, "test-model".to_string()); + let (handle, runner) = engine(backend, tokenizer, "test-model".to_string()); - // TODO: Spawn engine loop when real backend integration is done - // tokio::spawn(async move { - // engine_copy.run().await; - // }); + tokio::spawn(runner.serve()); - let result = engine.generate("test-model", "Hello world", 100, 0.7).await; + let result = handle.generate("test-model", "Hello world", 100, 0.7).await; assert!(result.is_ok()); } } diff --git a/src/backend/mock.rs b/src/backend/mock.rs index faa3c63..16ecfad 100644 --- a/src/backend/mock.rs +++ b/src/backend/mock.rs @@ -6,14 +6,71 @@ use tokio_stream::Stream; use super::engine::Backend; use crate::block_manager::types::TokenId; -/// Mock engine for testing (replace with MLX later) +/// Default vocab size the mock samples completion tokens from. +/// +/// Small enough that ids land in the low end of any real tokenizer's vocab +/// (so they decode to real text), large enough for varied output. +const DEFAULT_VOCAB_SIZE: u32 = 1000; + +/// Mock inference engine that behaves like a real autoregressive model. +/// +/// Unlike a naive echo, this consumes the prompt as *context* and emits only +/// **new** completion tokens — never the prompt back. Tokens are produced by a +/// deterministic pseudo-random walk seeded from the context, so: +/// - the same prompt always yields the same completion (reproducible tests), +/// - different prompts yield different completions (input-dependent), +/// - ids stay within the vocab range, so they decode to real text. +/// +/// Generation stops at `max_tokens` (the mock has no EOS concept yet). #[derive(Clone)] -pub struct MockEngine; +pub struct MockEngine { + vocab_size: u32, +} impl MockEngine { pub fn new() -> Self { - Self + Self { + vocab_size: DEFAULT_VOCAB_SIZE, + } + } + + /// Construct with an explicit vocab size (useful for tests pinning output). + pub fn with_vocab_size(vocab_size: u32) -> Self { + Self { + vocab_size: vocab_size.max(1), + } + } + + /// Deterministically generate `max_tokens` completion token ids from the + /// prompt context, mimicking an autoregressive decode. + /// + /// Each step hashes the running context (prompt + tokens produced so far) + /// into the vocab range and appends the result — exactly the shape of real + /// inference (next token conditioned on all prior tokens), just with a hash + /// standing in for a learned distribution. + fn generate_tokens(&self, prompt: &[TokenId], max_tokens: usize) -> Vec { + let mut context: Vec = prompt.to_vec(); + let mut completion = Vec::with_capacity(max_tokens); + for _ in 0..max_tokens { + let next = (hash_tokens(&context) % self.vocab_size as u64) as TokenId; + completion.push(next); + context.push(next); + } + completion + } +} + +/// FNV-1a hash over a token sequence — order-sensitive and fast, so each +/// distinct context maps to a distinct next token deterministically. +fn hash_tokens(tokens: &[TokenId]) -> u64 { + let mut hash: u64 = 0xcbf29ce484222325; + for &tok in tokens { + for byte in tok.to_le_bytes() { + hash ^= byte as u64; + hash = hash.wrapping_mul(0x100000001b3); + } } + hash } impl Backend for MockEngine { @@ -23,13 +80,8 @@ impl Backend for MockEngine { max_tokens: usize, _temperature: f32, ) -> Result, io::Error> { - // Mock: echo input + add generated tokens - // This makes different inputs produce different outputs - let mut result = token_ids; - for i in 0..max_tokens.min(10) { - result.push(1000 + i as u32); // Append tokens 1000, 1001, 1002... - } - Ok(result) + // Return ONLY the new completion tokens, as a real model does. + Ok(self.generate_tokens(&token_ids, max_tokens)) } async fn generate_stream( @@ -38,23 +90,13 @@ impl Backend for MockEngine { max_tokens: usize, _temperature: f32, ) -> Result + Send>>, io::Error> { - // Mock: echo input tokens first, then generate new ones - // This makes different inputs produce different outputs - - // First yield all input tokens (the prompt) - let mut all_tokens = token_ids; - - // Then add generated tokens - for i in 0..max_tokens.min(10) { - all_tokens.push(1000 + i as u32); // Generate tokens 1000, 1001, 1002... - } - - // Simulate delay between tokens (like real GPU) - let stream = stream::iter(all_tokens).then(|token_id| async move { - tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + // Same completion, but yielded one token at a time with a small delay to + // mimic a real GPU's inter-token latency. + let completion = self.generate_tokens(&token_ids, max_tokens); + let stream = stream::iter(completion).then(|token_id| async move { + tokio::time::sleep(tokio::time::Duration::from_millis(20)).await; token_id }); - Ok(Box::pin(stream)) } } @@ -64,3 +106,55 @@ impl Default for MockEngine { Self::new() } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn returns_only_completion_tokens() { + let engine = MockEngine::new(); + let prompt = vec![5, 9, 2]; + let out = engine.generate(prompt.clone(), 4, 0.0).await.unwrap(); + // Completion-only: exactly max_tokens, and it does not start with the + // prompt (a real model returns a continuation, not an echo). + assert_eq!(out.len(), 4); + assert_ne!(&out[..prompt.len().min(out.len())], &prompt[..]); + } + + #[tokio::test] + async fn is_deterministic() { + let engine = MockEngine::new(); + let a = engine.generate(vec![1, 2, 3], 8, 0.0).await.unwrap(); + let b = engine.generate(vec![1, 2, 3], 8, 0.0).await.unwrap(); + assert_eq!(a, b, "same prompt must yield same completion"); + } + + #[tokio::test] + async fn is_input_dependent() { + let engine = MockEngine::new(); + let a = engine.generate(vec![1, 2, 3], 8, 0.0).await.unwrap(); + let b = engine.generate(vec![3, 2, 1], 8, 0.0).await.unwrap(); + assert_ne!(a, b, "different prompts should yield different completions"); + } + + #[tokio::test] + async fn tokens_stay_within_vocab() { + let engine = MockEngine::with_vocab_size(50); + let out = engine.generate(vec![7, 7, 7], 32, 0.0).await.unwrap(); + assert!(out.iter().all(|&t| t < 50), "ids must be within vocab range"); + } + + #[tokio::test] + async fn stream_matches_generate() { + let engine = MockEngine::new(); + let prompt = vec![10, 20, 30]; + let batched = engine.generate(prompt.clone(), 6, 0.0).await.unwrap(); + let mut streamed = Vec::new(); + let mut s = engine.generate_stream(prompt, 6, 0.0).await.unwrap(); + while let Some(tok) = s.next().await { + streamed.push(tok); + } + assert_eq!(batched, streamed, "streaming and batched output must agree"); + } +} diff --git a/src/backend/mod.rs b/src/backend/mod.rs index 68fc792..68e2dab 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -2,5 +2,4 @@ pub mod engine; pub mod llm_engine; pub mod mock; -pub use engine::Backend; -pub use llm_engine::LLMEngine; +pub use llm_engine::{engine, EngineHandle}; diff --git a/src/cli/chat.rs b/src/cli/chat.rs index d9bb260..844e92a 100644 --- a/src/cli/chat.rs +++ b/src/cli/chat.rs @@ -5,7 +5,7 @@ use rustyline_derive::{Completer, Helper, Highlighter, Validator}; use std::io::{self, Write}; use tokio_stream::StreamExt; -use crate::backend::{Backend, LLMEngine}; +use crate::backend::EngineHandle; #[derive(Clone)] struct PlaceholderHint { @@ -44,11 +44,8 @@ impl Hinter for PlaceholderHinter { } /// Interactive chat loop for puma run -pub async fn interactive_chat( - engine: &E, - model: &str, -) -> Result<(), io::Error> { - let mut conversation_history = Vec::new(); +pub async fn interactive_chat(engine: &EngineHandle, model: &str) -> Result<(), io::Error> { + let mut conversation_history: Vec<(String, String)> = Vec::new(); // Setup editor with placeholder hinter let helper = PlaceholderHinter { @@ -79,11 +76,15 @@ pub async fn interactive_chat( break; } - // Add user message to history - conversation_history.push(format!("User: {}", input)); + // Add user message to history (structured role/content turns) + conversation_history.push(("user".to_string(), input.clone())); - // Build prompt from conversation history - let prompt = conversation_history.join("\n") + "\nAssistant:"; + // Build prompt from conversation history via the shared formatter + let prompt = crate::utils::prompt::format_conversation( + conversation_history + .iter() + .map(|(role, content)| (role.as_str(), content.as_str())), + ); // Empty line before response println!(); @@ -103,7 +104,7 @@ pub async fn interactive_chat( println!("\n"); // Double newline after response // Add assistant response to history - conversation_history.push(format!("Assistant: {}", full_response.trim())); + conversation_history.push(("assistant".to_string(), full_response.trim().to_string())); } Err(e) => { eprintln!("Error: {}\n", e); diff --git a/src/cli/commands.rs b/src/cli/commands.rs index c4d003c..c1ab516 100644 --- a/src/cli/commands.rs +++ b/src/cli/commands.rs @@ -4,7 +4,7 @@ use prettytable::{format, row, Table}; use tokenizers::Tokenizer; -use crate::backend::llm_engine::LLMEngine; +use crate::backend::engine; use crate::backend::mock::MockEngine; use crate::cli::{chat, inspect, ls, rm}; use crate::downloader::{self, Provider}; @@ -20,6 +20,16 @@ pub struct Cli { command: Commands, } +impl Cli { + /// Whether this command should emit logs by default. + /// + /// Only the long-running server (`serve`) logs unprompted; one-shot and + /// interactive CLI commands stay quiet unless the user sets `RUST_LOG`. + pub fn wants_default_logging(&self) -> bool { + matches!(self.command, Commands::SERVE(_)) + } +} + #[derive(Subcommand)] #[allow(clippy::upper_case_acronyms)] enum Commands { @@ -223,11 +233,10 @@ pub async fn run(cli: Cli) { Ok(tok) => tok, Err(e) => { eprintln!( - "Warning: Could not load tokenizer from {}: {}. Using BPE tokenizer instead.", + "Warning: Could not load tokenizer from {}: {}.", tokenizer_path, e ); - use tokenizers::models::bpe::BPE; - Tokenizer::new(BPE::default()) + std::process::exit(1); } }; @@ -236,17 +245,14 @@ pub async fn run(cli: Cli) { // Real backend will use: registry.get_model(&args.model)?.metadata.cache.path let backend = MockEngine::new(); - // Create LLMEngine (integrates scheduler + memory management + tokenizer) - let engine = LLMEngine::new(backend, tokenizer, args.model.clone()); + // Create engine: cheap send-side handle + runner that owns the scheduler + let (handle, runner) = engine(backend, tokenizer, args.model.clone()); - // TODO: Spawn engine loop when real backend integration is done - // For now, engine.generate() sends events but backend is MockEngine - // tokio::spawn(async move { - // engine_copy.run().await; - // }); + // Spawn the runner's event loop; the handle submits work via events + tokio::spawn(runner.serve()); // Start interactive chat - if let Err(e) = chat::interactive_chat(&engine, &args.model).await { + if let Err(e) = chat::interactive_chat(&handle, &args.model).await { eprintln!("❌ Chat error: {}", e); std::process::exit(1); } diff --git a/src/cli/serve.rs b/src/cli/serve.rs index 5b44ca6..55aab77 100644 --- a/src/cli/serve.rs +++ b/src/cli/serve.rs @@ -1,8 +1,11 @@ use colored::Colorize; use std::sync::Arc; +use tokenizers::models::bpe::BPE; +use tokenizers::Tokenizer; use tracing::{debug, info}; use crate::api::routes::create_router; +use crate::backend::engine; use crate::backend::mock::MockEngine; use crate::registry::model_registry::ModelRegistry; @@ -30,16 +33,25 @@ pub async fn execute( info!("Starting PUMA to serve model: {}", model_name); // Initialize backend (MockEngine for now, replace with MLX later) - let engine = Arc::new(MockEngine::new()); - info!("Inference engine initialized"); + let backend = MockEngine::new(); debug!("Using MockEngine backend"); + // TODO: Load the model's real tokenizer; placeholder BPE for now + let tokenizer = Tokenizer::new(BPE::default()); + + // Create engine: cheap send-side handle + runner that owns the scheduler + let (handle, runner) = engine(backend, tokenizer, model_name.to_string()); + + // Spawn the runner's event loop; the handle submits work via events + tokio::spawn(runner.serve()); + info!("Inference engine initialized"); + // Initialize model registry let registry = Arc::new(ModelRegistry::new(None)); info!("Model registry loaded"); // Create router - let app = create_router(engine, registry); + let app = create_router(handle, registry); // Bind address let addr = format!("{}:{}", host, port); diff --git a/src/main.rs b/src/main.rs index 4ef25f1..14191fd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,20 +19,28 @@ use crate::cli::commands::{run, Cli}; use crate::utils::file; fn main() { - // Setup tracing subscriber for tower-http TraceLayer + let cli = Cli::parse(); + + // Setup tracing subscriber for tower-http TraceLayer. + // + // An explicit RUST_LOG always wins. Otherwise only the long-running server + // logs by default; one-shot/interactive CLI commands stay quiet so their + // output isn't buried under per-request INFO lines. + let default_filter = if cli.wants_default_logging() { + "info,hf_hub=warn,tower_http=info,rusqlite_migration=warn" + } else { + "off" + }; tracing_subscriber::fmt() .with_env_filter( - tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| { - "info,hf_hub=warn,tower_http=info,rusqlite_migration=warn".into() - }), + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| default_filter.into()), ) .init(); // Create the root folder if it doesn't exist. file::create_folder_if_not_exists(&file::root_home()).unwrap(); - let cli = Cli::parse(); - let runtime = Builder::new_multi_thread() .worker_threads(4) .enable_all() diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs index 7b713f9..8710b94 100644 --- a/src/scheduler/core.rs +++ b/src/scheduler/core.rs @@ -143,6 +143,30 @@ impl Scheduler { self.prefill_batch.len() > before || !self.decode_batch.is_empty() } + /// Drain the current prefill batch as inference work for the engine loop. + /// + /// Clears the prefill batch and returns `(seq_id, token_ids, max_tokens, + /// streaming)` per scheduled sequence. The engine loop runs the backend on + /// these and reports results back via [`send_chunk`](Self::send_chunk) and + /// [`complete_sequence`](Self::complete_sequence). The `streaming` flag tells + /// the loop whether to call the backend's streaming or single-shot path. + pub fn take_prefill_batch(&mut self) -> Vec<(SequenceId, Arc>, usize, bool)> { + let batch = std::mem::take(&mut self.prefill_batch); + batch + .into_iter() + .filter_map(|seq_id| match self.sequences.get(&seq_id) { + Some(SequenceState::Prefilling(s)) => { + let streaming = self + .response_channels + .get(&seq_id) + .is_some_and(|tx| tx.is_streaming()); + Some((seq_id, Arc::clone(&s.token_ids), s.max_tokens, streaming)) + } + _ => None, + }) + .collect() + } + // ===== Internal State Management (called by LLMEngine after inference) ===== pub fn append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { @@ -718,16 +742,21 @@ impl Scheduler { // ===== Response Channel Helpers ===== - /// Send a token to client (for streaming) - pub fn send_token(&mut self, seq_id: SequenceId, token: String) { + /// Send a decoded text chunk to a streaming client. + /// + /// The chunk is an incremental piece of detokenized output (a delta), not a + /// single token — see the engine's incremental decode. No-op for + /// non-streaming clients, which receive the full text via + /// [`complete_sequence`](Self::complete_sequence). + pub fn send_chunk(&mut self, seq_id: SequenceId, chunk: String) { if let Some(response_tx) = self.response_channels.get(&seq_id) { match response_tx { ResponseSender::Stream(tx) => { - let _ = tx.send(Ok(token)); + let _ = tx.send(Ok(chunk)); } ResponseSender::Single(_) => { - // Single response - can't stream individual tokens - // Will send complete result at the end + // Single response - streamed chunks are dropped; + // the full result is delivered at completion. } } } diff --git a/src/scheduler/events.rs b/src/scheduler/events.rs index 1589fc9..5a88c06 100644 --- a/src/scheduler/events.rs +++ b/src/scheduler/events.rs @@ -47,6 +47,13 @@ pub enum ResponseSender { Stream(tokio::sync::mpsc::UnboundedSender>), } +impl ResponseSender { + /// Whether this sequence expects token-by-token streaming. + pub fn is_streaming(&self) -> bool { + matches!(self, ResponseSender::Stream(_)) + } +} + /// Internal operations (token generation, state transitions, OOM handling) /// are NOT events - the scheduler handles them directly in its run loop. #[derive(Debug)] diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 73f5099..04e8ed2 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,2 +1,3 @@ pub mod file; pub mod format; +pub mod prompt; diff --git a/src/utils/prompt.rs b/src/utils/prompt.rs new file mode 100644 index 0000000..8432bc1 --- /dev/null +++ b/src/utils/prompt.rs @@ -0,0 +1,65 @@ +//! Prompt formatting — the seam between the chat API's role/content language +//! and the flat text the tokenizer consumes. +//! +//! Roles ("system"/"user"/"assistant") are an *API-level* concept. The +//! inference engine only ever sees token ids; it has no notion of turns. This +//! module serializes structured turns into one prompt string, which the engine +//! handle then tokenizes. +//! +//! Today it uses a simple `Role: content` convention with a trailing +//! `Assistant:` cue. When a real model is wired in, this is the single place to +//! swap in that model's chat template (e.g. `<|im_start|>` markers) — nothing +//! below the tokenizer changes. + +/// Format conversation turns into a single prompt string. +/// +/// Each turn is `(role, content)`. The result labels every turn as +/// `Role: content` (one per line) and ends with an empty `Assistant:` cue so +/// the model continues as the assistant. +pub fn format_conversation<'a, I>(turns: I) -> String +where + I: IntoIterator, +{ + let mut prompt = String::new(); + for (role, content) in turns { + prompt.push_str(&format!("{}: {}\n", display_role(role), content)); + } + prompt.push_str("Assistant:"); + prompt +} + +/// Capitalize a role label for display (`user` → `User`), defaulting unknown +/// roles to `Assistant` to match prior behavior. +fn display_role(role: &str) -> &'static str { + match role { + "system" => "System", + "user" => "User", + _ => "Assistant", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn formats_turns_with_assistant_cue() { + let prompt = format_conversation([ + ("system", "Be helpful."), + ("user", "Hello"), + ]); + assert_eq!(prompt, "System: Be helpful.\nUser: Hello\nAssistant:"); + } + + #[test] + fn unknown_role_defaults_to_assistant() { + let prompt = format_conversation([("tool", "result")]); + assert_eq!(prompt, "Assistant: result\nAssistant:"); + } + + #[test] + fn empty_conversation_is_just_the_cue() { + let prompt = format_conversation([]); + assert_eq!(prompt, "Assistant:"); + } +} From 0da5f1bd290c0c9db5c6fb5ce9675cfed38a8d72 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sat, 18 Jul 2026 22:17:28 +0100 Subject: [PATCH 27/33] add empty event Signed-off-by: kerthcet --- src/backend/llm_engine.rs | 42 +++---------- src/fsm/events.rs | 9 ++- src/fsm/states.rs | 10 +++- src/scheduler/core.rs | 123 ++++++++++++++++++++++++++++++++------ 4 files changed, 131 insertions(+), 53 deletions(-) diff --git a/src/backend/llm_engine.rs b/src/backend/llm_engine.rs index db73eef..82438a7 100644 --- a/src/backend/llm_engine.rs +++ b/src/backend/llm_engine.rs @@ -140,7 +140,6 @@ pub struct EngineRunner { backend: B, scheduler: Scheduler, tokenizer: Arc, - event_rx: mpsc::UnboundedReceiver, } impl EngineRunner { @@ -269,33 +268,10 @@ impl EngineRunner { tracing::info!("EngineRunner started"); loop { - // 1. Handle events (non-blocking, drain all) - while let Ok(event) = self.event_rx.try_recv() { - match event { - SchedulerEvent::AddRequest { - seq_id, - token_ids, - max_tokens, - response_tx, - } => { - self.scheduler - .add_request(seq_id, token_ids, max_tokens, response_tx); - } - - SchedulerEvent::CancelRequest { seq_id } => { - self.scheduler.cancel_request(seq_id); - } - - SchedulerEvent::GetBlocks { seq_id, response } => { - let result = self.scheduler.get_blocks(seq_id); - let _ = response.send(result); - } - - SchedulerEvent::GetStats { response } => { - let stats = self.scheduler.get_stats(); - let _ = response.send(stats); - } - } + // 1. Handle events (non-blocking, drain all). The scheduler owns the + // event channel; the loop drains it and dispatches each event. + while let Ok(event) = self.scheduler.event_rx.try_recv() { + self.scheduler.handle_event(event); } // 2. Schedule (always, not event-driven) @@ -341,12 +317,13 @@ pub fn engine( let allocator = Box::new(CpuAllocator::new(1024 * 1024 * 100)); let block_manager = BlockManager::new(allocator, 512); - // Create scheduler (max 32 batch size, 16 tokens per block) - let scheduler = Scheduler::new(block_manager, 32, 16); - - // Event channel + // Event channel: the handle produces events, the scheduler consumes them. let (event_tx, event_rx) = mpsc::unbounded_channel(); + // Create scheduler (max 32 batch size, 16 tokens per block); it owns the + // event receiver. + let scheduler = Scheduler::new(block_manager, event_rx, 32, 16); + let tokenizer = Arc::new(tokenizer); let handle = EngineHandle { @@ -360,7 +337,6 @@ pub fn engine( backend, scheduler, tokenizer, - event_rx, }; (handle, runner) diff --git a/src/fsm/events.rs b/src/fsm/events.rs index 47883b4..4e33445 100644 --- a/src/fsm/events.rs +++ b/src/fsm/events.rs @@ -1,5 +1,5 @@ use super::states::FinishReason; -use crate::block_manager::types::SequenceId; +use crate::block_manager::types::{SequenceId, TokenId}; /// FSM Events - pure data that triggers state transitions /// @@ -8,6 +8,13 @@ use crate::block_manager::types::SequenceId; /// SchedulerEvent instead. #[derive(Debug, Clone)] pub enum Event { + /// Create a new sequence: Empty → Waiting + Create { + seq_id: SequenceId, + token_ids: Vec, + max_tokens: usize, + }, + /// Allocate blocks and start prefilling Schedule { tokens_per_block: usize }, diff --git a/src/fsm/states.rs b/src/fsm/states.rs index 50426a5..f71ec08 100644 --- a/src/fsm/states.rs +++ b/src/fsm/states.rs @@ -8,8 +8,14 @@ use std::sync::Arc; /// ↓ ↓ /// └→ Preempted ←┘ -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub enum SequenceState { + /// Pre-birth placeholder: the default "from" state a new sequence + /// transitions out of via `Event::Create`. Never persisted in the + /// scheduler's sequence map. + #[default] + Empty, + /// Waiting in queue for scheduling Waiting(WaitingState), @@ -109,6 +115,7 @@ pub enum FinishReason { impl SequenceState { pub fn seq_id(&self) -> SequenceId { match self { + SequenceState::Empty => panic!("Empty state has no seq_id"), SequenceState::Waiting(s) => s.seq_id, SequenceState::Scheduling(s) => s.seq_id, SequenceState::Prefilling(s) => s.seq_id, @@ -148,6 +155,7 @@ impl SequenceState { pub fn variant_name(&self) -> &'static str { match self { + SequenceState::Empty => "Empty", SequenceState::Waiting(_) => "Waiting", SequenceState::Scheduling(_) => "Scheduling", SequenceState::Prefilling(_) => "Prefilling", diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs index 8710b94..cb8bba8 100644 --- a/src/scheduler/core.rs +++ b/src/scheduler/core.rs @@ -1,4 +1,4 @@ -use super::events::{ResponseSender, SchedulerStats}; +use super::events::{ResponseSender, SchedulerEvent, SchedulerEventReceiver, SchedulerStats}; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::*; use crate::fsm::{ @@ -41,6 +41,12 @@ pub struct Scheduler { /// Currently running sequences in decode phase decode_batch: Vec, + /// Receiver for external events (add/cancel/query) from clients. + /// + /// Public so the engine loop can drain it directly and feed each event to + /// [`handle_event`](Self::handle_event). + pub event_rx: SchedulerEventReceiver, + /// Configuration max_batch_size: usize, tokens_per_block: usize, @@ -49,6 +55,7 @@ pub struct Scheduler { impl Scheduler { pub fn new( block_manager: BlockManager, + event_rx: SchedulerEventReceiver, max_batch_size: usize, tokens_per_block: usize, ) -> Self { @@ -59,11 +66,43 @@ impl Scheduler { waiting_queue: VecDeque::new(), prefill_batch: Vec::new(), decode_batch: Vec::new(), + event_rx, max_batch_size, tokens_per_block, } } + /// Dispatch a single external event to its handler. + /// + /// One entry point for every `SchedulerEvent`: mutating events build the + /// corresponding FSM event and apply it; query events reply on their own + /// response channel. + pub fn handle_event(&mut self, event: SchedulerEvent) { + match event { + SchedulerEvent::AddRequest { + seq_id, + token_ids, + max_tokens, + response_tx, + } => self.add_request(seq_id, token_ids, max_tokens, response_tx), + + SchedulerEvent::CancelRequest { seq_id } => { + let event = Event::Abort { + reason: "User cancelled".to_string(), + }; + self.apply_abort(seq_id, event); + } + + SchedulerEvent::GetBlocks { seq_id, response } => { + let _ = response.send(self.get_blocks(seq_id)); + } + + SchedulerEvent::GetStats { response } => { + let _ = response.send(self.get_stats()); + } + } + } + // ===== Public API (called by LLMEngine) ===== pub fn add_request( @@ -83,40 +122,62 @@ impl Scheduler { // Check duplicate if self.sequences.contains_key(&seq_id) { warn!("Duplicate seq_id {:?}", seq_id); - self.send_error(response_tx, Error::InvalidTransition("Duplicate sequence ID")); + self.send_error( + response_tx, + Error::InvalidTransition("Duplicate sequence ID"), + ); return; } // Store response channel self.response_channels.insert(seq_id, response_tx); - // Create waiting state - tokens already provided by LLMEngine - let state = SequenceState::Waiting(WaitingState { + // Birth the sequence via an FSM transition: Empty → Waiting. + let event = Event::Create { seq_id, - token_ids: Arc::new(token_ids), + token_ids, max_tokens, - }); - - self.sequences.insert(seq_id, state); - self.waiting_queue.push_back(seq_id); + }; + match self.transition(SequenceState::Empty, event) { + Ok(state) => { + self.sequences.insert(seq_id, state); + self.waiting_queue.push_back(seq_id); + } + Err(e) => { + warn!("Failed to create sequence {:?}: {:?}", seq_id, e); + self.response_channels.remove(&seq_id); + } + } } + /// Cancel a request at the user's behest. pub fn cancel_request(&mut self, seq_id: SequenceId) { + self.abort_request(seq_id, "User cancelled".to_string()); + } + + /// Abort a sequence with a caller-supplied reason. + /// + /// The reason is passed through to the FSM `Abort` event, so callers (e.g. + /// the engine loop on a backend failure) describe *why* rather than + /// borrowing the generic cancel reason. + pub fn abort_request(&mut self, seq_id: SequenceId, reason: String) { + self.apply_abort(seq_id, Event::Abort { reason }); + } + + /// Apply an `Abort` FSM event to an existing sequence and evict it from the + /// queues/batches. Shared by the cancel/abort paths and the event loop. + fn apply_abort(&mut self, seq_id: SequenceId, event: Event) { let state = match self.sequences.remove(&seq_id) { Some(s) => s, None => { - warn!("Cannot cancel - sequence {:?} not found", seq_id); + warn!("Cannot abort - sequence {:?} not found", seq_id); return; } }; - let event = Event::Abort { - reason: "User cancelled".to_string(), - }; - match self.transition(state, event) { Ok(new_state) => { - info!("Cancelled sequence {:?}", seq_id); + info!("Aborted sequence {:?}", seq_id); self.sequences.insert(seq_id, new_state); // Remove from queues/batches @@ -125,7 +186,7 @@ impl Scheduler { self.decode_batch.retain(|&id| id != seq_id); } Err(e) => { - warn!("Failed to cancel {:?}: {:?}", seq_id, e); + warn!("Failed to abort {:?}: {:?}", seq_id, e); } } } @@ -474,6 +535,16 @@ impl Scheduler { use crate::fsm::Event; match (state, event) { + // Empty → Waiting (new sequence) + ( + SequenceState::Empty, + Event::Create { + seq_id, + token_ids, + max_tokens, + }, + ) => self.transition_create(seq_id, token_ids, max_tokens), + // Waiting → Prefilling (SequenceState::Waiting(s), Event::Schedule { .. }) => self.transition_schedule(s), @@ -512,6 +583,20 @@ impl Scheduler { // These private methods implement the actual transition logic. // Called by apply() after event dispatching. + /// Transition: Empty → Waiting (construct a new sequence's initial state) + fn transition_create( + &mut self, + seq_id: SequenceId, + token_ids: Vec, + max_tokens: usize, + ) -> Result { + Ok(SequenceState::Waiting(WaitingState { + seq_id, + token_ids: Arc::new(token_ids), + max_tokens, + })) + } + /// Transition: Waiting → Prefilling (allocate blocks for tokenized prompt) fn transition_schedule(&mut self, state: WaitingState) -> Result { let prompt_tokens = state.token_ids.len(); @@ -800,7 +885,8 @@ mod tests { fn create_test_scheduler() -> Scheduler { let allocator = Box::new(CpuAllocator::new(100_000)); let block_manager = BlockManager::new(allocator, 1024); - Scheduler::new(block_manager, 10, 16) + let (_event_tx, event_rx) = tokio::sync::mpsc::unbounded_channel(); + Scheduler::new(block_manager, event_rx, 10, 16) } fn create_dummy_response() -> ResponseSender { @@ -907,7 +993,8 @@ mod tests { // Small allocator - only 2 blocks let allocator = Box::new(CpuAllocator::new(2048)); let block_manager = BlockManager::new(allocator, 1024); - let mut scheduler = Scheduler::new(block_manager, 10, 16); + let (_event_tx, event_rx) = tokio::sync::mpsc::unbounded_channel(); + let mut scheduler = Scheduler::new(block_manager, event_rx, 10, 16); // First request needs 7 blocks (100 tokens / 16 tokens_per_block) scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); From 1a39cb032f894838d481aef2840b663f3f0116cb Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sat, 18 Jul 2026 23:54:14 +0100 Subject: [PATCH 28/33] temp save Signed-off-by: kerthcet --- src/backend/llm_engine.rs | 6 ++-- src/fsm/events.rs | 1 - src/scheduler/core.rs | 62 +++++++++++++++++---------------------- 3 files changed, 31 insertions(+), 38 deletions(-) diff --git a/src/backend/llm_engine.rs b/src/backend/llm_engine.rs index 82438a7..725494d 100644 --- a/src/backend/llm_engine.rs +++ b/src/backend/llm_engine.rs @@ -169,7 +169,8 @@ impl EngineRunner { Ok(tokens) => tokens, Err(e) => { tracing::error!("Backend inference failed for {:?}: {}", seq_id, e); - self.scheduler.cancel_request(seq_id); + self.scheduler + .abort_request(seq_id, format!("Backend inference failed: {}", e)); return; } }; @@ -210,7 +211,8 @@ impl EngineRunner { Ok(stream) => stream, Err(e) => { tracing::error!("Backend stream failed for {:?}: {}", seq_id, e); - self.scheduler.cancel_request(seq_id); + self.scheduler + .abort_request(seq_id, format!("Backend stream failed: {}", e)); return; } }; diff --git a/src/fsm/events.rs b/src/fsm/events.rs index 4e33445..a7d9d48 100644 --- a/src/fsm/events.rs +++ b/src/fsm/events.rs @@ -10,7 +10,6 @@ use crate::block_manager::types::{SequenceId, TokenId}; pub enum Event { /// Create a new sequence: Empty → Waiting Create { - seq_id: SequenceId, token_ids: Vec, max_tokens: usize, }, diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs index cb8bba8..432e694 100644 --- a/src/scheduler/core.rs +++ b/src/scheduler/core.rs @@ -87,10 +87,7 @@ impl Scheduler { } => self.add_request(seq_id, token_ids, max_tokens, response_tx), SchedulerEvent::CancelRequest { seq_id } => { - let event = Event::Abort { - reason: "User cancelled".to_string(), - }; - self.apply_abort(seq_id, event); + self.abort_request(seq_id, "User cancelled".to_string()) } SchedulerEvent::GetBlocks { seq_id, response } => { @@ -134,13 +131,12 @@ impl Scheduler { // Birth the sequence via an FSM transition: Empty → Waiting. let event = Event::Create { - seq_id, token_ids, max_tokens, }; - match self.transition(SequenceState::Empty, event) { - Ok(state) => { - self.sequences.insert(seq_id, state); + match self.transition(seq_id, SequenceState::Empty, event) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); self.waiting_queue.push_back(seq_id); } Err(e) => { @@ -150,23 +146,11 @@ impl Scheduler { } } - /// Cancel a request at the user's behest. - pub fn cancel_request(&mut self, seq_id: SequenceId) { - self.abort_request(seq_id, "User cancelled".to_string()); - } - /// Abort a sequence with a caller-supplied reason. /// - /// The reason is passed through to the FSM `Abort` event, so callers (e.g. - /// the engine loop on a backend failure) describe *why* rather than - /// borrowing the generic cancel reason. + /// The reason is passed through to the FSM `Abort` event, so callers (user + /// cancel, backend failure, …) describe *why* the sequence ended. pub fn abort_request(&mut self, seq_id: SequenceId, reason: String) { - self.apply_abort(seq_id, Event::Abort { reason }); - } - - /// Apply an `Abort` FSM event to an existing sequence and evict it from the - /// queues/batches. Shared by the cancel/abort paths and the event loop. - fn apply_abort(&mut self, seq_id: SequenceId, event: Event) { let state = match self.sequences.remove(&seq_id) { Some(s) => s, None => { @@ -175,7 +159,7 @@ impl Scheduler { } }; - match self.transition(state, event) { + match self.transition(seq_id, state, Event::Abort { reason }) { Ok(new_state) => { info!("Aborted sequence {:?}", seq_id); self.sequences.insert(seq_id, new_state); @@ -245,7 +229,7 @@ impl Scheduler { tokens_per_block: self.tokens_per_block, }; - match self.transition(state, event) { + match self.transition(seq_id, state, event) { Ok(new_state) => { self.sequences.insert(seq_id, new_state); } @@ -273,7 +257,7 @@ impl Scheduler { let backup = state.clone(); let event = Event::Complete { reason }; - match self.transition(state, event) { + match self.transition(seq_id, state, event) { Ok(new_state) => { info!("Completed sequence {:?}", seq_id); self.sequences.insert(seq_id, new_state); @@ -356,7 +340,7 @@ impl Scheduler { let backup = state.clone(); let event = Event::Preempt; - match self.transition(state, event) { + match self.transition(seq_id, state, event) { Ok(new_state) => { info!("Preempted sequence {:?}", seq_id); self.sequences.insert(seq_id, new_state); @@ -384,7 +368,7 @@ impl Scheduler { let backup = state.clone(); let event = Event::Resume; - match self.transition(state, event) { + match self.transition(seq_id, state, event) { Ok(new_state) => { debug!("Resumed sequence {:?}", seq_id); self.sequences.insert(seq_id, new_state); @@ -426,7 +410,7 @@ impl Scheduler { tokens_per_block: self.tokens_per_block, }; - match self.transition(state.clone(), event) { + match self.transition(seq_id, state.clone(), event) { Ok(new_state) => { self.sequences.insert(seq_id, new_state); self.prefill_batch.push(seq_id); @@ -517,21 +501,30 @@ impl Scheduler { /// /// # Arguments /// + /// * `seq_id` - Sequence the event applies to /// * `state` - Current sequence state (will be consumed) /// * `event` - FSM event to apply /// /// # Returns /// - /// New state on success, or `Error::InvalidTransition` for invalid state/event combinations + /// The new state on success (the caller stores it in the sequence map), or + /// the error on failure so callers can handle rollback/OOM policy. This + /// function computes state only; it does not touch the sequence map or the + /// scheduling queues. /// /// # Example /// /// ```rust,ignore /// let event = Event::Schedule { tokens_per_block: 16 }; - /// let new_state = self.transition(state, event)?; + /// let new_state = self.transition(seq_id, state, event)?; /// self.sequences.insert(seq_id, new_state); /// ``` - pub fn transition(&mut self, state: SequenceState, event: Event) -> Result { + pub fn transition( + &mut self, + seq_id: SequenceId, + state: SequenceState, + event: Event, + ) -> Result { use crate::fsm::Event; match (state, event) { @@ -539,7 +532,6 @@ impl Scheduler { ( SequenceState::Empty, Event::Create { - seq_id, token_ids, max_tokens, }, @@ -925,7 +917,7 @@ mod tests { let mut scheduler = create_test_scheduler(); scheduler.add_request(SequenceId(1), vec![0; 100], 200, create_dummy_response()); - scheduler.cancel_request(SequenceId(1)); + scheduler.abort_request(SequenceId(1), "User cancelled".to_string()); // Should transition to Aborted assert!(matches!( @@ -940,7 +932,7 @@ mod tests { let mut scheduler = create_test_scheduler(); // Should not panic - scheduler.cancel_request(SequenceId(999)); + scheduler.abort_request(SequenceId(999), "User cancelled".to_string()); assert_eq!(scheduler.sequences.len(), 0); } @@ -1207,7 +1199,7 @@ mod tests { scheduler.schedule_prefill(); assert_eq!(scheduler.prefill_batch.len(), 1); - scheduler.cancel_request(SequenceId(1)); + scheduler.abort_request(SequenceId(1), "User cancelled".to_string()); // Should remove from batch assert_eq!(scheduler.prefill_batch.len(), 0); From 6d6bcd4e969aaac92bcf69f780cb1b551eca622f Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 19 Jul 2026 01:55:01 +0100 Subject: [PATCH 29/33] add tests Signed-off-by: kerthcet --- src/block_manager/allocator.rs | 74 ++++ src/fsm/mod.rs | 2 - src/fsm/states.rs | 163 +++++++- src/fsm/transitions.rs | 373 ----------------- src/lib.rs | 1 + src/main.rs | 1 + src/scheduler/core.rs | 567 +++----------------------- src/sequence_manager/mod.rs | 715 +++++++++++++++++++++++++++++++++ src/system/system_info.rs | 79 ++++ 9 files changed, 1072 insertions(+), 903 deletions(-) delete mode 100644 src/fsm/transitions.rs create mode 100644 src/sequence_manager/mod.rs diff --git a/src/block_manager/allocator.rs b/src/block_manager/allocator.rs index 0db8aab..a9a7c35 100644 --- a/src/block_manager/allocator.rs +++ b/src/block_manager/allocator.rs @@ -77,3 +77,77 @@ impl MemoryAllocator for CpuAllocator { // TODO: Implement CudaAllocator later // pub struct CudaAllocator { ... } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_allocate_tracks_used_memory() { + let mut alloc = CpuAllocator::new(10_000); + assert_eq!(alloc.get_total_memory(), 10_000); + assert_eq!(alloc.get_available_memory(), 10_000); + + let addr = alloc.allocate(1024).unwrap(); + assert_eq!(addr.size, 1024); + assert!(!addr.ptr.is_null()); + assert_eq!(alloc.get_available_memory(), 10_000 - 1024); + + alloc.free(addr).unwrap(); + assert_eq!(alloc.get_available_memory(), 10_000); + } + + #[test] + fn test_allocate_out_of_memory() { + let mut alloc = CpuAllocator::new(1024); + + // First allocation fits exactly. + let addr = alloc.allocate(1024).unwrap(); + assert_eq!(alloc.get_available_memory(), 0); + + // Any further allocation exceeds capacity. + assert!(matches!(alloc.allocate(1), Err(Error::OutOfMemory))); + + alloc.free(addr).unwrap(); + } + + #[test] + fn test_allocate_exact_boundary() { + let mut alloc = CpuAllocator::new(2048); + let a = alloc.allocate(1024).unwrap(); + let b = alloc.allocate(1024).unwrap(); + assert_eq!(alloc.get_available_memory(), 0); + alloc.free(a).unwrap(); + alloc.free(b).unwrap(); + assert_eq!(alloc.get_available_memory(), 2048); + } + + #[test] + fn test_free_more_than_used_errors() { + let mut alloc = CpuAllocator::new(10_000); + let addr = alloc.allocate(512).unwrap(); + + // Fabricate an address claiming a larger size than is actually used. + let bogus = MemoryAddress { + ptr: addr.ptr, + size: 4096, + }; + assert!(matches!(alloc.free(bogus), Err(Error::FreeFailed(_)))); + + // The real allocation is still accounted for and can be freed. + assert_eq!(alloc.get_available_memory(), 10_000 - 512); + alloc.free(addr).unwrap(); + assert_eq!(alloc.get_available_memory(), 10_000); + } + + #[test] + fn test_reuse_after_free() { + let mut alloc = CpuAllocator::new(1024); + for _ in 0..3 { + let addr = alloc.allocate(1024).unwrap(); + assert_eq!(alloc.get_available_memory(), 0); + alloc.free(addr).unwrap(); + assert_eq!(alloc.get_available_memory(), 1024); + } + } +} diff --git a/src/fsm/mod.rs b/src/fsm/mod.rs index e945239..cbdadea 100644 --- a/src/fsm/mod.rs +++ b/src/fsm/mod.rs @@ -8,7 +8,6 @@ //! //! - `states`: State definitions (Waiting, Prefilling, Decoding, etc.) //! - `events`: Event types that trigger transitions -//! - `transitions`: (DEPRECATED - logic moved to Scheduler) //! //! # Usage //! @@ -38,7 +37,6 @@ pub mod events; pub mod states; -// pub mod transitions; // DEPRECATED: logic moved to Scheduler // Re-export commonly used types pub use events::Event; diff --git a/src/fsm/states.rs b/src/fsm/states.rs index f71ec08..17f5ceb 100644 --- a/src/fsm/states.rs +++ b/src/fsm/states.rs @@ -8,13 +8,12 @@ use std::sync::Arc; /// ↓ ↓ /// └→ Preempted ←┘ -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone)] pub enum SequenceState { - /// Pre-birth placeholder: the default "from" state a new sequence - /// transitions out of via `Event::Create`. Never persisted in the - /// scheduler's sequence map. - #[default] - Empty, + /// Pre-birth placeholder: the "from" state a new sequence transitions out + /// of via `Event::Create`. Carries the `seq_id` so every state variant can + /// answer `seq_id()`. Never persisted in the scheduler's sequence map. + Empty(SequenceId), /// Waiting in queue for scheduling Waiting(WaitingState), @@ -115,7 +114,7 @@ pub enum FinishReason { impl SequenceState { pub fn seq_id(&self) -> SequenceId { match self { - SequenceState::Empty => panic!("Empty state has no seq_id"), + SequenceState::Empty(seq_id) => *seq_id, SequenceState::Waiting(s) => s.seq_id, SequenceState::Scheduling(s) => s.seq_id, SequenceState::Prefilling(s) => s.seq_id, @@ -155,7 +154,7 @@ impl SequenceState { pub fn variant_name(&self) -> &'static str { match self { - SequenceState::Empty => "Empty", + SequenceState::Empty(_) => "Empty", SequenceState::Waiting(_) => "Waiting", SequenceState::Scheduling(_) => "Scheduling", SequenceState::Prefilling(_) => "Prefilling", @@ -166,3 +165,151 @@ impl SequenceState { } } } + +#[cfg(test)] +mod tests { + use super::*; + + fn waiting(seq_id: u64) -> SequenceState { + SequenceState::Waiting(WaitingState { + seq_id: SequenceId(seq_id), + token_ids: Arc::new(vec![1, 2, 3]), + max_tokens: 10, + }) + } + + fn prefilling(seq_id: u64, blocks: Vec) -> SequenceState { + SequenceState::Prefilling(PrefillingState { + seq_id: SequenceId(seq_id), + token_ids: Arc::new(vec![1, 2, 3]), + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: 3, + max_tokens: 10, + }) + } + + fn decoding(seq_id: u64, blocks: Vec) -> SequenceState { + SequenceState::Decoding(DecodingState { + seq_id: SequenceId(seq_id), + token_ids: Arc::new(vec![1, 2, 3]), + blocks: Arc::new(blocks), + num_tokens: 3, + max_tokens: 10, + }) + } + + #[test] + fn test_seq_id_for_every_variant() { + assert_eq!(SequenceState::Empty(SequenceId(7)).seq_id(), SequenceId(7)); + assert_eq!(waiting(1).seq_id(), SequenceId(1)); + assert_eq!(prefilling(2, vec![]).seq_id(), SequenceId(2)); + assert_eq!(decoding(3, vec![]).seq_id(), SequenceId(3)); + + let scheduling = SequenceState::Scheduling(SchedulingState { + seq_id: SequenceId(4), + prompt_tokens: 3, + max_tokens: 10, + blocks: Arc::new(vec![]), + }); + assert_eq!(scheduling.seq_id(), SequenceId(4)); + + let preempted = SequenceState::Preempted(PreemptedState { + seq_id: SequenceId(5), + token_ids: Arc::new(vec![1]), + num_tokens: 1, + max_tokens: 10, + }); + assert_eq!(preempted.seq_id(), SequenceId(5)); + + let finished = SequenceState::Finished(FinishedState { + seq_id: SequenceId(6), + finish_reason: FinishReason::Stop, + }); + assert_eq!(finished.seq_id(), SequenceId(6)); + + let aborted = SequenceState::Aborted(AbortedState { + seq_id: SequenceId(8), + reason: "boom".to_string(), + }); + assert_eq!(aborted.seq_id(), SequenceId(8)); + } + + #[test] + fn test_is_running() { + assert!(prefilling(1, vec![]).is_running()); + assert!(decoding(1, vec![]).is_running()); + assert!(!waiting(1).is_running()); + assert!(!SequenceState::Empty(SequenceId(1)).is_running()); + } + + #[test] + fn test_is_waiting() { + assert!(waiting(1).is_waiting()); + let scheduling = SequenceState::Scheduling(SchedulingState { + seq_id: SequenceId(1), + prompt_tokens: 3, + max_tokens: 10, + blocks: Arc::new(vec![]), + }); + assert!(scheduling.is_waiting()); + assert!(!decoding(1, vec![]).is_waiting()); + } + + #[test] + fn test_is_finished() { + let finished = SequenceState::Finished(FinishedState { + seq_id: SequenceId(1), + finish_reason: FinishReason::MaxTokens, + }); + let aborted = SequenceState::Aborted(AbortedState { + seq_id: SequenceId(1), + reason: "x".to_string(), + }); + assert!(finished.is_finished()); + assert!(aborted.is_finished()); + assert!(!decoding(1, vec![]).is_finished()); + } + + #[test] + fn test_blocks_exposed_only_for_resource_holding_states() { + // Prefilling/Decoding always expose their blocks. + assert_eq!( + prefilling(1, vec![BlockId(0), BlockId(1)]).blocks(), + Some(&[BlockId(0), BlockId(1)][..]) + ); + assert_eq!( + decoding(1, vec![BlockId(2)]).blocks(), + Some(&[BlockId(2)][..]) + ); + + // Scheduling exposes blocks only when it has partially allocated some. + let scheduling_with = SequenceState::Scheduling(SchedulingState { + seq_id: SequenceId(1), + prompt_tokens: 3, + max_tokens: 10, + blocks: Arc::new(vec![BlockId(9)]), + }); + assert_eq!(scheduling_with.blocks(), Some(&[BlockId(9)][..])); + + let scheduling_empty = SequenceState::Scheduling(SchedulingState { + seq_id: SequenceId(1), + prompt_tokens: 3, + max_tokens: 10, + blocks: Arc::new(vec![]), + }); + assert_eq!(scheduling_empty.blocks(), None); + + // States without KV blocks return None. + assert_eq!(waiting(1).blocks(), None); + assert_eq!(SequenceState::Empty(SequenceId(1)).blocks(), None); + } + + #[test] + fn test_variant_name() { + assert_eq!(SequenceState::Empty(SequenceId(1)).variant_name(), "Empty"); + assert_eq!(waiting(1).variant_name(), "Waiting"); + assert_eq!(prefilling(1, vec![]).variant_name(), "Prefilling"); + assert_eq!(decoding(1, vec![]).variant_name(), "Decoding"); + } +} diff --git a/src/fsm/transitions.rs b/src/fsm/transitions.rs deleted file mode 100644 index 1631658..0000000 --- a/src/fsm/transitions.rs +++ /dev/null @@ -1,373 +0,0 @@ -use super::events::Event; -use super::states::*; -use crate::block_manager::manager::BlockManager; -use crate::block_manager::types::*; -use std::sync::Arc; -use tracing::{debug, warn}; - -/// Apply an event to a state, producing a new state -/// -/// Pure function that handles all state transitions. -/// Pattern: (current_state, event) -> new_state -pub fn apply( - state: SequenceState, - event: Event, - block_manager: &mut BlockManager, -) -> Result { - match (state, event) { - // Waiting → Prefilling (Schedule) - (SequenceState::Waiting(s), Event::Schedule { tokens_per_block }) => { - apply_schedule(s, block_manager, tokens_per_block) - } - - // Prefilling → Decoding/Prefilling (AppendTokens) - ( - SequenceState::Prefilling(s), - Event::AppendTokens { - num_tokens, - tokens_per_block, - }, - ) => apply_append_tokens_prefilling(s, num_tokens, block_manager, tokens_per_block), - - // Decoding → Decoding/Finished (AppendTokens) - ( - SequenceState::Decoding(s), - Event::AppendTokens { - num_tokens, - tokens_per_block, - }, - ) => apply_append_tokens_decoding(s, num_tokens, block_manager, tokens_per_block), - - // Prefilling/Decoding → Finished (Complete) - (SequenceState::Prefilling(s), Event::Complete { reason }) => { - apply_complete_prefilling(s, reason, block_manager) - } - (SequenceState::Decoding(s), Event::Complete { reason }) => { - apply_complete_decoding(s, reason, block_manager) - } - - // Prefilling/Decoding → Preempted (Preempt) - (SequenceState::Prefilling(s), Event::Preempt) => { - apply_preempt_prefilling(s, block_manager) - } - (SequenceState::Decoding(s), Event::Preempt) => apply_preempt_decoding(s, block_manager), - - // Preempted → Waiting (Resume) - (SequenceState::Preempted(s), Event::Resume) => apply_resume(s), - - // Decoding → (Decoding, Decoding) (Fork) - (SequenceState::Decoding(s), Event::Fork { child_id }) => { - apply_fork(s, child_id, block_manager) - } - - // Any → Aborted (Abort) - (state, Event::Abort { reason }) => apply_abort(state, reason, block_manager), - - // Invalid transitions - (_state, _event) => Err(Error::InvalidTransition("Invalid state transition")), - } -} - -fn state_name(state: &SequenceState) -> &'static str { - match state { - SequenceState::Waiting(_) => "Waiting", - SequenceState::Scheduling(_) => "Scheduling", - SequenceState::Prefilling(_) => "Prefilling", - SequenceState::Decoding(_) => "Decoding", - SequenceState::Preempted(_) => "Preempted", - SequenceState::Finished(_) => "Finished", - SequenceState::Aborted(_) => "Aborted", - } -} - -// Helper functions for each transition - -fn apply_schedule( - state: WaitingState, - block_manager: &mut BlockManager, - tokens_per_block: usize, -) -> Result { - let blocks_needed = state.prompt_tokens.div_ceil(tokens_per_block); - - let mut blocks = Vec::new(); - for _ in 0..blocks_needed { - match block_manager.allocate() { - Ok(block_id) => blocks.push(block_id), - Err(Error::OutOfMemory) => { - // Cleanup on OOM - for block in blocks { - let _ = block_manager.free(block); - } - return Err(Error::OutOfMemory); - } - Err(e) => { - // Other error - cleanup and propagate - for block in blocks { - let _ = block_manager.free(block); - } - return Err(e); - } - } - } - - debug!( - "Scheduled seq {:?}: allocated {} blocks for {} tokens", - state.seq_id, - blocks.len(), - state.prompt_tokens - ); - - Ok(SequenceState::Prefilling(PrefillingState { - seq_id: state.seq_id, - blocks: Arc::new(blocks), - tokens_filled: 0, - tokens_total: state.prompt_tokens, - max_tokens: state.max_tokens, - })) -} - -fn apply_append_tokens_prefilling( - mut state: PrefillingState, - num_tokens: usize, - _block_manager: &mut BlockManager, - _tokens_per_block: usize, -) -> Result { - state.tokens_filled += num_tokens; - - debug!( - "Prefilling seq {:?}: {}/{} tokens", - state.seq_id, state.tokens_filled, state.tokens_total - ); - - if state.tokens_filled >= state.tokens_total { - debug!("Seq {:?} transitioning to Decoding", state.seq_id); - Ok(SequenceState::Decoding(DecodingState { - seq_id: state.seq_id, - blocks: state.blocks, - num_tokens: state.tokens_filled, - max_tokens: state.max_tokens, - })) - } else { - Ok(SequenceState::Prefilling(state)) - } -} - -fn apply_append_tokens_decoding( - mut state: DecodingState, - num_tokens: usize, - block_manager: &mut BlockManager, - tokens_per_block: usize, -) -> Result { - state.num_tokens += num_tokens; - - // Check if finished - if state.num_tokens >= state.max_tokens { - // Free blocks - for &block_id in state.blocks.iter() { - if let Err(e) = block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!( - "Seq {:?} finished: reached max_tokens ({})", - state.seq_id, state.max_tokens - ); - - return Ok(SequenceState::Finished(FinishedState { - seq_id: state.seq_id, - finish_reason: FinishReason::MaxTokens, - })); - } - - // Check if need more blocks - let blocks_needed = state.num_tokens.div_ceil(tokens_per_block); - let initial_block_count = state.blocks.len(); - - if state.blocks.len() < blocks_needed { - // Need to allocate more blocks - use Arc::make_mut for copy-on-write - let blocks = Arc::make_mut(&mut state.blocks); - - while blocks.len() < blocks_needed { - match block_manager.allocate() { - Ok(block_id) => { - blocks.push(block_id); - debug!( - "Seq {:?}: allocated block, total blocks: {}", - state.seq_id, - blocks.len() - ); - } - Err(e) => { - // Free any blocks we allocated in this call - warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); - for block_id in blocks.drain(initial_block_count..) { - let _ = block_manager.free(block_id); - } - return Err(e); - } - } - } - } - - Ok(SequenceState::Decoding(state)) -} - -fn apply_complete_prefilling( - state: PrefillingState, - reason: FinishReason, - block_manager: &mut BlockManager, -) -> Result { - // Free all blocks - for &block_id in state.blocks.iter() { - if let Err(e) = block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!( - "Completed seq {:?} during prefill: {:?}", - state.seq_id, reason - ); - - Ok(SequenceState::Finished(FinishedState { - seq_id: state.seq_id, - finish_reason: reason, - })) -} - -fn apply_complete_decoding( - state: DecodingState, - reason: FinishReason, - block_manager: &mut BlockManager, -) -> Result { - // Free all blocks - for &block_id in state.blocks.iter() { - if let Err(e) = block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!("Completed seq {:?}: {:?}", state.seq_id, reason); - - Ok(SequenceState::Finished(FinishedState { - seq_id: state.seq_id, - finish_reason: reason, - })) -} - -fn apply_preempt_prefilling( - state: PrefillingState, - block_manager: &mut BlockManager, -) -> Result { - // Free all blocks - for &block_id in state.blocks.iter() { - if let Err(e) = block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!("Preempted seq {:?} during prefill", state.seq_id); - - Ok(SequenceState::Preempted(PreemptedState { - seq_id: state.seq_id, - num_tokens: state.tokens_filled, - max_tokens: state.max_tokens, - })) -} - -fn apply_preempt_decoding( - state: DecodingState, - block_manager: &mut BlockManager, -) -> Result { - // Free all blocks - for &block_id in state.blocks.iter() { - if let Err(e) = block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!("Preempted seq {:?}", state.seq_id); - - Ok(SequenceState::Preempted(PreemptedState { - seq_id: state.seq_id, - num_tokens: state.num_tokens, - max_tokens: state.max_tokens, - })) -} - -fn apply_resume(state: PreemptedState) -> Result { - debug!("Resuming seq {:?}", state.seq_id); - - // Transition back to waiting for rescheduling - Ok(SequenceState::Waiting(WaitingState { - seq_id: state.seq_id, - prompt_tokens: state.num_tokens, - max_tokens: state.max_tokens, - })) -} - -fn apply_fork( - state: DecodingState, - child_id: SequenceId, - block_manager: &mut BlockManager, -) -> Result { - // Copy-on-write: increment ref counts - for &block_id in state.blocks.iter() { - block_manager.add_ref(block_id)?; - } - - let child = DecodingState { - seq_id: child_id, - blocks: state.blocks.clone(), - num_tokens: state.num_tokens, - max_tokens: state.max_tokens, - }; - - debug!("Forked seq {:?} → {:?}", state.seq_id, child.seq_id); - - // Fork returns parent, child needs to be handled separately by caller - // For now just return error - need to handle multi-state result - Err(Error::InvalidTransition( - "Fork not yet supported in unified apply", - )) -} - -fn apply_abort( - state: SequenceState, - reason: String, - block_manager: &mut BlockManager, -) -> Result { - let seq_id = state.seq_id(); - - // Free blocks if any - if let Some(blocks) = state.blocks() { - for &block_id in blocks { - if let Err(e) = block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, seq_id, e - ); - } - } - } - - warn!("Aborted seq {:?}: {}", seq_id, reason); - - Ok(SequenceState::Aborted(AbortedState { seq_id, reason })) -} diff --git a/src/lib.rs b/src/lib.rs index 5b5fab3..909efec 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,6 +8,7 @@ pub mod downloader; pub mod fsm; pub mod registry; pub mod scheduler; +pub mod sequence_manager; pub mod storage; pub mod system; pub mod utils; diff --git a/src/main.rs b/src/main.rs index 14191fd..a15b72d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,7 @@ mod downloader; mod fsm; mod registry; mod scheduler; +mod sequence_manager; mod storage; mod system; mod utils; diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs index 432e694..d9ffc16 100644 --- a/src/scheduler/core.rs +++ b/src/scheduler/core.rs @@ -1,10 +1,8 @@ use super::events::{ResponseSender, SchedulerEvent, SchedulerEventReceiver, SchedulerStats}; +use crate::sequence_manager::SequenceManager; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::*; -use crate::fsm::{ - AbortedState, DecodingState, Event, FinishReason, FinishedState, PreemptedState, - PrefillingState, SequenceState, WaitingState, -}; +use crate::fsm::{Event, FinishReason, SequenceState}; use std::collections::{HashMap, VecDeque}; use std::sync::Arc; use tracing::{debug, info, warn}; @@ -23,11 +21,8 @@ use tracing::{debug, info, warn}; /// - Scheduling policy (what to schedule when) /// - FSM transitions (internal methods with direct BlockManager access) pub struct Scheduler { - /// Block manager for memory allocation - block_manager: BlockManager, - - /// All sequences and their states - sequences: HashMap, + /// Owns sequence state + block memory and executes FSM transitions. + sequences: SequenceManager, /// Response channels to send results back to clients (supports streaming) response_channels: HashMap, @@ -60,8 +55,7 @@ impl Scheduler { tokens_per_block: usize, ) -> Self { Self { - block_manager, - sequences: HashMap::new(), + sequences: SequenceManager::new(block_manager, tokens_per_block), response_channels: HashMap::new(), waiting_queue: VecDeque::new(), prefill_batch: Vec::new(), @@ -117,7 +111,7 @@ impl Scheduler { ); // Check duplicate - if self.sequences.contains_key(&seq_id) { + if self.sequences.contains(seq_id) { warn!("Duplicate seq_id {:?}", seq_id); self.send_error( response_tx, @@ -129,14 +123,9 @@ impl Scheduler { // Store response channel self.response_channels.insert(seq_id, response_tx); - // Birth the sequence via an FSM transition: Empty → Waiting. - let event = Event::Create { - token_ids, - max_tokens, - }; - match self.transition(seq_id, SequenceState::Empty, event) { - Ok(new_state) => { - self.sequences.insert(seq_id, new_state); + // Birth the sequence (Empty → Waiting) inside the sequence manager. + match self.sequences.create(seq_id, token_ids, max_tokens) { + Ok(()) => { self.waiting_queue.push_back(seq_id); } Err(e) => { @@ -151,18 +140,9 @@ impl Scheduler { /// The reason is passed through to the FSM `Abort` event, so callers (user /// cancel, backend failure, …) describe *why* the sequence ended. pub fn abort_request(&mut self, seq_id: SequenceId, reason: String) { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Cannot abort - sequence {:?} not found", seq_id); - return; - } - }; - - match self.transition(seq_id, state, Event::Abort { reason }) { - Ok(new_state) => { + match self.sequences.advance(seq_id, Event::Abort { reason }) { + Ok(_) => { info!("Aborted sequence {:?}", seq_id); - self.sequences.insert(seq_id, new_state); // Remove from queues/batches self.waiting_queue.retain(|&id| id != seq_id); @@ -199,7 +179,7 @@ impl Scheduler { let batch = std::mem::take(&mut self.prefill_batch); batch .into_iter() - .filter_map(|seq_id| match self.sequences.get(&seq_id) { + .filter_map(|seq_id| match self.sequences.get(seq_id) { Some(SequenceState::Prefilling(s)) => { let streaming = self .response_channels @@ -215,52 +195,29 @@ impl Scheduler { // ===== Internal State Management (called by LLMEngine after inference) ===== pub fn append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Sequence {:?} not found", seq_id); - return; - } - }; - - let backup = state.clone(); let event = Event::AppendTokens { num_tokens, tokens_per_block: self.tokens_per_block, }; - match self.transition(seq_id, state, event) { - Ok(new_state) => { - self.sequences.insert(seq_id, new_state); - } + // On error the sequence manager restores the prior state; the scheduler + // only applies OOM policy. + match self.sequences.advance(seq_id, event) { + Ok(_) => {} Err(Error::OutOfMemory) => { warn!("OOM while appending tokens to {:?}", seq_id); - self.sequences.insert(seq_id, backup); self.handle_oom(); } Err(e) => { warn!("Failed to append tokens to {:?}: {:?}", seq_id, e); - self.sequences.insert(seq_id, backup); } } } pub fn complete_sequence(&mut self, seq_id: SequenceId, reason: FinishReason, result: String) { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Sequence {:?} not found", seq_id); - return; - } - }; - - let backup = state.clone(); - let event = Event::Complete { reason }; - - match self.transition(seq_id, state, event) { - Ok(new_state) => { + match self.sequences.advance(seq_id, Event::Complete { reason }) { + Ok(_) => { info!("Completed sequence {:?}", seq_id); - self.sequences.insert(seq_id, new_state); // Remove from batches self.prefill_batch.retain(|&id| id != seq_id); @@ -271,7 +228,6 @@ impl Scheduler { } Err(e) => { warn!("Failed to complete {:?}: {:?}", seq_id, e); - self.sequences.insert(seq_id, backup); // Send error back to client if let Some(response_tx) = self.response_channels.remove(&seq_id) { @@ -282,68 +238,13 @@ impl Scheduler { } pub fn fork_sequence(&mut self, parent_id: SequenceId, child_id: SequenceId) { - // Check duplicate child_id - if self.sequences.contains_key(&child_id) { - warn!("Child ID {:?} already exists", child_id); - return; - } - - let parent_state = match self.sequences.remove(&parent_id) { - Some(s) => s, - None => { - warn!("Parent sequence {:?} not found", parent_id); - return; - } - }; - - // TODO: Fork needs special handling - returns (parent, child) - // For now, use manual implementation - let backup = parent_state.clone(); - - if let SequenceState::Decoding(s) = parent_state { - // Copy-on-write: increment ref counts - for &block_id in s.blocks.iter() { - if let Err(e) = self.block_manager.add_ref(block_id) { - warn!("Failed to fork {:?}: {:?}", parent_id, e); - self.sequences.insert(parent_id, backup); - return; - } - } - - let child = DecodingState { - seq_id: child_id, - token_ids: Arc::clone(&s.token_ids), - blocks: Arc::clone(&s.blocks), - num_tokens: s.num_tokens, - max_tokens: s.max_tokens, - }; - - self.sequences.insert(parent_id, SequenceState::Decoding(s)); - self.sequences - .insert(child_id, SequenceState::Decoding(child)); - debug!("Forked sequence {:?} → {:?}", parent_id, child_id); - } else { - warn!("Cannot fork {:?} - not in Decoding state", parent_id); - self.sequences.insert(parent_id, backup); - } + let _ = self.sequences.fork(parent_id, child_id); } fn preempt_sequence(&mut self, seq_id: SequenceId) { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Sequence {:?} not found", seq_id); - return; - } - }; - - let backup = state.clone(); - let event = Event::Preempt; - - match self.transition(seq_id, state, event) { - Ok(new_state) => { + match self.sequences.advance(seq_id, Event::Preempt) { + Ok(_) => { info!("Preempted sequence {:?}", seq_id); - self.sequences.insert(seq_id, new_state); // Remove from batches self.prefill_batch.retain(|&id| id != seq_id); @@ -351,32 +252,18 @@ impl Scheduler { } Err(e) => { warn!("Failed to preempt {:?}: {:?}", seq_id, e); - self.sequences.insert(seq_id, backup); } } } fn resume_sequence(&mut self, seq_id: SequenceId) { - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Sequence {:?} not found", seq_id); - return; - } - }; - - let backup = state.clone(); - let event = Event::Resume; - - match self.transition(seq_id, state, event) { - Ok(new_state) => { + match self.sequences.advance(seq_id, Event::Resume) { + Ok(_) => { debug!("Resumed sequence {:?}", seq_id); - self.sequences.insert(seq_id, new_state); self.waiting_queue.push_back(seq_id); } Err(e) => { warn!("Failed to resume {:?}: {:?}", seq_id, e); - self.sequences.insert(seq_id, backup); } } } @@ -397,41 +284,27 @@ impl Scheduler { break; } - // Get and remove state - let state = match self.sequences.remove(&seq_id) { - Some(s) => s, - None => { - warn!("Sequence {:?} not found", seq_id); - continue; - } - }; - let event = Event::Schedule { tokens_per_block: self.tokens_per_block, }; - match self.transition(seq_id, state.clone(), event) { - Ok(new_state) => { - self.sequences.insert(seq_id, new_state); + // The sequence manager restores prior state on any error, so the + // scheduler only applies queueing policy. + match self.sequences.advance(seq_id, event) { + Ok(_) => { self.prefill_batch.push(seq_id); info!("Scheduled {:?} for prefill", seq_id); } Err(Error::OutOfMemory) => { - // Transition already cleaned up allocated blocks - // Restore state, put back in queue and stop - self.sequences.insert(seq_id, state); + // Put back at the front and stop scheduling this round. self.waiting_queue.push_front(seq_id); debug!("OOM - cannot schedule {:?}", seq_id); break; } Err(Error::InvalidTransition(_)) => { - // Not in Waiting state - restore and skip - self.sequences.insert(seq_id, state); warn!("Cannot schedule {:?} - not in Waiting state", seq_id); } Err(e) => { - // Other error - restore state - self.sequences.insert(seq_id, state); warn!("Failed to schedule {:?}: {:?}", seq_id, e); } } @@ -458,365 +331,19 @@ impl Scheduler { // ===== Query Methods ===== pub fn get_blocks(&self, seq_id: SequenceId) -> Result> { - match self.sequences.get(&seq_id) { - Some(state) => Ok(state.blocks().map(|b| b.to_vec()).unwrap_or_default()), - None => Err(Error::UnknownSequence(seq_id)), - } + self.sequences.blocks(seq_id) } pub fn get_stats(&self) -> SchedulerStats { - let num_running = self.sequences.values().filter(|s| s.is_running()).count(); - let num_waiting = self.waiting_queue.len(); - let num_preempted = self - .sequences - .values() - .filter(|s| matches!(s, SequenceState::Preempted(_))) - .count(); - SchedulerStats { num_sequences: self.sequences.len(), - num_running, - num_waiting, - num_preempted, - block_stats: self.block_manager.get_stats(), + num_running: self.sequences.num_running(), + num_waiting: self.waiting_queue.len(), + num_preempted: self.sequences.num_preempted(), + block_stats: self.sequences.block_stats(), } } - // ===== FSM Transitions ===== - - /// Transition sequence state using FSM event - single entry point for all transitions - /// - /// This method provides an event-based abstraction over the internal transition methods. - /// - /// # Design - /// - /// **Public API**: Event-based for maintainability - /// - Single entry point makes it easy to add cross-cutting concerns (logging, metrics) - /// - Event enum can be serialized for debugging/replay - /// - Clean interface for callers - /// - /// **Internal Implementation**: Type-safe methods - /// - Private `transition_*()` methods enforce correct state types at compile time - /// - Called by this method after event dispatching - /// - /// # Arguments - /// - /// * `seq_id` - Sequence the event applies to - /// * `state` - Current sequence state (will be consumed) - /// * `event` - FSM event to apply - /// - /// # Returns - /// - /// The new state on success (the caller stores it in the sequence map), or - /// the error on failure so callers can handle rollback/OOM policy. This - /// function computes state only; it does not touch the sequence map or the - /// scheduling queues. - /// - /// # Example - /// - /// ```rust,ignore - /// let event = Event::Schedule { tokens_per_block: 16 }; - /// let new_state = self.transition(seq_id, state, event)?; - /// self.sequences.insert(seq_id, new_state); - /// ``` - pub fn transition( - &mut self, - seq_id: SequenceId, - state: SequenceState, - event: Event, - ) -> Result { - use crate::fsm::Event; - - match (state, event) { - // Empty → Waiting (new sequence) - ( - SequenceState::Empty, - Event::Create { - token_ids, - max_tokens, - }, - ) => self.transition_create(seq_id, token_ids, max_tokens), - - // Waiting → Prefilling - (SequenceState::Waiting(s), Event::Schedule { .. }) => self.transition_schedule(s), - - // Prefilling → Prefilling/Decoding - (SequenceState::Prefilling(s), Event::AppendTokens { num_tokens, .. }) => { - self.transition_append_tokens_prefilling(s, num_tokens) - } - - // Decoding → Decoding/Finished - (SequenceState::Decoding(s), Event::AppendTokens { num_tokens, .. }) => { - self.transition_append_tokens_decoding(s, num_tokens) - } - - // Any running state → Finished - (state, Event::Complete { reason }) => self.transition_complete(state, reason), - - // Running → Preempted - ( - state @ (SequenceState::Prefilling(_) | SequenceState::Decoding(_)), - Event::Preempt, - ) => self.transition_preempt(state), - - // Preempted → Waiting - (SequenceState::Preempted(s), Event::Resume) => self.transition_resume(s), - - // Any → Aborted - (state, Event::Abort { reason }) => self.transition_abort(state, reason), - - // Invalid transitions - _ => Err(Error::InvalidTransition("Invalid state transition")), - } - } - - // ===== Internal Transition Methods ===== - // - // These private methods implement the actual transition logic. - // Called by apply() after event dispatching. - - /// Transition: Empty → Waiting (construct a new sequence's initial state) - fn transition_create( - &mut self, - seq_id: SequenceId, - token_ids: Vec, - max_tokens: usize, - ) -> Result { - Ok(SequenceState::Waiting(WaitingState { - seq_id, - token_ids: Arc::new(token_ids), - max_tokens, - })) - } - - /// Transition: Waiting → Prefilling (allocate blocks for tokenized prompt) - fn transition_schedule(&mut self, state: WaitingState) -> Result { - let prompt_tokens = state.token_ids.len(); - let blocks_needed = prompt_tokens.div_ceil(self.tokens_per_block); - - let mut blocks = Vec::new(); - for _ in 0..blocks_needed { - match self.block_manager.allocate() { - Ok(block_id) => blocks.push(block_id), - Err(Error::OutOfMemory) => { - // Cleanup on OOM - for block in blocks { - let _ = self.block_manager.free(block); - } - return Err(Error::OutOfMemory); - } - Err(e) => { - for block in blocks { - let _ = self.block_manager.free(block); - } - return Err(e); - } - } - } - - debug!( - "Scheduled seq {:?}: allocated {} blocks for {} tokens", - state.seq_id, - blocks.len(), - prompt_tokens - ); - - Ok(SequenceState::Prefilling(PrefillingState { - seq_id: state.seq_id, - token_ids: state.token_ids, - blocks: Arc::new(blocks), - tokens_filled: 0, - tokens_total: prompt_tokens, - max_tokens: state.max_tokens, - })) - } - - /// Transition: Prefilling → Decoding or Prefilling (append tokens) - fn transition_append_tokens_prefilling( - &mut self, - mut state: PrefillingState, - num_tokens: usize, - ) -> Result { - state.tokens_filled += num_tokens; - - debug!( - "Prefilling seq {:?}: {}/{} tokens", - state.seq_id, state.tokens_filled, state.tokens_total - ); - - if state.tokens_filled >= state.tokens_total { - debug!("Seq {:?} transitioning to Decoding", state.seq_id); - Ok(SequenceState::Decoding(DecodingState { - seq_id: state.seq_id, - token_ids: state.token_ids, - blocks: state.blocks, - num_tokens: state.tokens_filled, - max_tokens: state.max_tokens, - })) - } else { - Ok(SequenceState::Prefilling(state)) - } - } - - /// Transition: Decoding → Decoding or Finished (append tokens, maybe allocate more blocks) - fn transition_append_tokens_decoding( - &mut self, - mut state: DecodingState, - num_tokens: usize, - ) -> Result { - state.num_tokens += num_tokens; - - // Check if finished - if state.num_tokens >= state.max_tokens { - // Free blocks - for &block_id in state.blocks.iter() { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, state.seq_id, e - ); - } - } - - debug!( - "Seq {:?} finished: reached max_tokens ({})", - state.seq_id, state.max_tokens - ); - - return Ok(SequenceState::Finished(FinishedState { - seq_id: state.seq_id, - finish_reason: FinishReason::MaxTokens, - })); - } - - // Check if need more blocks - let blocks_needed = state.num_tokens.div_ceil(self.tokens_per_block); - let initial_block_count = state.blocks.len(); - - if state.blocks.len() < blocks_needed { - let blocks = Arc::make_mut(&mut state.blocks); - - while blocks.len() < blocks_needed { - match self.block_manager.allocate() { - Ok(block_id) => { - blocks.push(block_id); - debug!( - "Seq {:?}: allocated block, total blocks: {}", - state.seq_id, - blocks.len() - ); - } - Err(e) => { - warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); - for block_id in blocks.drain(initial_block_count..) { - let _ = self.block_manager.free(block_id); - } - return Err(e); - } - } - } - } - - Ok(SequenceState::Decoding(state)) - } - - /// Transition: Prefilling/Decoding → Finished - fn transition_complete( - &mut self, - state: SequenceState, - reason: FinishReason, - ) -> Result { - let seq_id = state.seq_id(); - - // Free blocks if any - if let Some(blocks) = state.blocks() { - for &block_id in blocks { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, seq_id, e - ); - } - } - } - - debug!("Completed seq {:?}: {:?}", seq_id, reason); - - Ok(SequenceState::Finished(FinishedState { - seq_id, - finish_reason: reason, - })) - } - - /// Transition: Prefilling/Decoding → Preempted (free blocks) - fn transition_preempt(&mut self, state: SequenceState) -> Result { - let seq_id = state.seq_id(); - let (token_ids, num_tokens, max_tokens) = match &state { - SequenceState::Prefilling(s) => { - (Arc::clone(&s.token_ids), s.tokens_filled, s.max_tokens) - } - SequenceState::Decoding(s) => (Arc::clone(&s.token_ids), s.num_tokens, s.max_tokens), - _ => { - return Err(Error::InvalidTransition( - "Can only preempt Prefilling/Decoding", - )) - } - }; - - // Free blocks - if let Some(blocks) = state.blocks() { - for &block_id in blocks { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, seq_id, e - ); - } - } - } - - debug!("Preempted seq {:?}", seq_id); - - Ok(SequenceState::Preempted(PreemptedState { - seq_id, - token_ids, - num_tokens, - max_tokens, - })) - } - - /// Transition: Preempted → Waiting - fn transition_resume(&mut self, state: PreemptedState) -> Result { - debug!("Resuming seq {:?}", state.seq_id); - - Ok(SequenceState::Waiting(WaitingState { - seq_id: state.seq_id, - token_ids: state.token_ids, - max_tokens: state.max_tokens, - })) - } - - /// Transition: Any → Aborted (free blocks, cleanup) - fn transition_abort(&mut self, state: SequenceState, reason: String) -> Result { - let seq_id = state.seq_id(); - - // Free blocks if any - if let Some(blocks) = state.blocks() { - for &block_id in blocks { - if let Err(e) = self.block_manager.free(block_id) { - warn!( - "Failed to free block {:?} for {:?}: {:?}", - block_id, seq_id, e - ); - } - } - } - - warn!("Aborted seq {:?}: {}", seq_id, reason); - - Ok(SequenceState::Aborted(AbortedState { seq_id, reason })) - } - // ===== Response Channel Helpers ===== /// Send a decoded text chunk to a streaming client. @@ -895,7 +422,7 @@ mod tests { assert_eq!(scheduler.sequences.len(), 1); assert_eq!(scheduler.waiting_queue.len(), 1); assert!(matches!( - scheduler.sequences.get(&SequenceId(1)), + scheduler.sequences.get(SequenceId(1)), Some(SequenceState::Waiting(_)) )); } @@ -921,7 +448,7 @@ mod tests { // Should transition to Aborted assert!(matches!( - scheduler.sequences.get(&SequenceId(1)), + scheduler.sequences.get(SequenceId(1)), Some(SequenceState::Aborted(_)) )); assert_eq!(scheduler.waiting_queue.len(), 0); @@ -945,7 +472,7 @@ mod tests { // Should move to prefilling assert!(matches!( - scheduler.sequences.get(&SequenceId(1)), + scheduler.sequences.get(SequenceId(1)), Some(SequenceState::Prefilling(_)) )); assert_eq!(scheduler.prefill_batch.len(), 1); @@ -996,7 +523,7 @@ mod tests { assert_eq!(scheduler.prefill_batch.len(), 0); assert_eq!(scheduler.waiting_queue.len(), 1); assert!(matches!( - scheduler.sequences.get(&SequenceId(1)), + scheduler.sequences.get(SequenceId(1)), Some(SequenceState::Waiting(_)) )); } @@ -1013,7 +540,7 @@ mod tests { // Should transition to Decoding assert!(matches!( - scheduler.sequences.get(&SequenceId(1)), + scheduler.sequences.get(SequenceId(1)), Some(SequenceState::Decoding(_)) )); } @@ -1029,7 +556,7 @@ mod tests { // Append more tokens scheduler.append_tokens(SequenceId(1), 10); - if let Some(SequenceState::Decoding(state)) = scheduler.sequences.get(&SequenceId(1)) { + if let Some(SequenceState::Decoding(state)) = scheduler.sequences.get(SequenceId(1)) { assert_eq!(state.num_tokens, 110); } else { panic!("Expected Decoding state"); @@ -1049,7 +576,7 @@ mod tests { // Should transition to Finished assert!(matches!( - scheduler.sequences.get(&SequenceId(1)), + scheduler.sequences.get(SequenceId(1)), Some(SequenceState::Finished(_)) )); } @@ -1066,7 +593,7 @@ mod tests { // Should be Finished assert!(matches!( - scheduler.sequences.get(&SequenceId(1)), + scheduler.sequences.get(SequenceId(1)), Some(SequenceState::Finished(_)) )); assert_eq!(scheduler.prefill_batch.len(), 0); @@ -1083,10 +610,10 @@ mod tests { scheduler.fork_sequence(SequenceId(1), SequenceId(2)); // Both should exist - assert!(scheduler.sequences.contains_key(&SequenceId(1))); - assert!(scheduler.sequences.contains_key(&SequenceId(2))); + assert!(scheduler.sequences.contains(SequenceId(1))); + assert!(scheduler.sequences.contains(SequenceId(2))); assert!(matches!( - scheduler.sequences.get(&SequenceId(2)), + scheduler.sequences.get(SequenceId(2)), Some(SequenceState::Decoding(_)) )); } @@ -1104,7 +631,7 @@ mod tests { // Original seq 2 should be unchanged (Waiting) assert!(matches!( - scheduler.sequences.get(&SequenceId(2)), + scheduler.sequences.get(SequenceId(2)), Some(SequenceState::Waiting(_)) )); } @@ -1121,7 +648,7 @@ mod tests { // Should be Preempted assert!(matches!( - scheduler.sequences.get(&SequenceId(1)), + scheduler.sequences.get(SequenceId(1)), Some(SequenceState::Preempted(_)) )); assert_eq!(scheduler.decode_batch.len(), 0); @@ -1140,7 +667,7 @@ mod tests { // Should be back to Waiting assert!(matches!( - scheduler.sequences.get(&SequenceId(1)), + scheduler.sequences.get(SequenceId(1)), Some(SequenceState::Waiting(_)) )); assert_eq!(scheduler.waiting_queue.len(), 1); diff --git a/src/sequence_manager/mod.rs b/src/sequence_manager/mod.rs new file mode 100644 index 0000000..9333f6a --- /dev/null +++ b/src/sequence_manager/mod.rs @@ -0,0 +1,715 @@ +//! SequenceManager - owns sequence state and the memory those sequences use. +//! +//! This is the *mechanism* layer: it holds the `sequences` map and the +//! `BlockManager`, and executes FSM transitions (which allocate/free blocks). +//! The `Scheduler` is the *policy* layer — it decides which events fire and +//! owns the queues/batches — and drives this manager without ever touching a +//! `BlockManager` or a `BlockId` directly. +//! +//! Keeping memory here (rather than in the scheduler) localizes future changes +//! like multi-GPU block pools or CPU/GPU swap: the scheduler keeps calling +//! `advance`/`create`/`can_allocate`, and the device details stay inside. + +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; +use crate::fsm::{ + AbortedState, DecodingState, Event, FinishReason, FinishedState, PreemptedState, + PrefillingState, SequenceState, WaitingState, +}; +use std::collections::HashMap; +use std::sync::Arc; +use tracing::{debug, warn}; + +/// Owns sequence state + block memory and executes FSM transitions. +pub struct SequenceManager { + /// Block manager for memory allocation + block_manager: BlockManager, + + /// All sequences and their states + sequences: HashMap, + + /// Tokens per block (used when computing block allocations) + tokens_per_block: usize, +} + +impl SequenceManager { + pub fn new(block_manager: BlockManager, tokens_per_block: usize) -> Self { + Self { + block_manager, + sequences: HashMap::new(), + tokens_per_block, + } + } + + // ===== Sequence map access (used by the scheduler for policy) ===== + + /// Whether a sequence with this id exists. + pub fn contains(&self, seq_id: SequenceId) -> bool { + self.sequences.contains_key(&seq_id) + } + + /// Look up a sequence's current state. + pub fn get(&self, seq_id: SequenceId) -> Option<&SequenceState> { + self.sequences.get(&seq_id) + } + + /// Total number of tracked sequences. + pub fn len(&self) -> usize { + self.sequences.len() + } + + pub fn is_empty(&self) -> bool { + self.sequences.is_empty() + } + + /// Count sequences currently running (prefill or decode). + pub fn num_running(&self) -> usize { + self.sequences.values().filter(|s| s.is_running()).count() + } + + /// Count sequences currently preempted. + pub fn num_preempted(&self) -> usize { + self.sequences + .values() + .filter(|s| matches!(s, SequenceState::Preempted(_))) + .count() + } + + /// Block ids held by a sequence (for KV-cache addressing). + pub fn blocks(&self, seq_id: SequenceId) -> Result> { + match self.sequences.get(&seq_id) { + Some(state) => Ok(state.blocks().map(|b| b.to_vec()).unwrap_or_default()), + None => Err(Error::UnknownSequence(seq_id)), + } + } + + /// Per-block-type memory statistics. + pub fn block_stats(&self) -> HashMap { + self.block_manager.get_stats() + } + + // ===== Transition entry points (used by the scheduler) ===== + + /// Create a new sequence in the Waiting state and store it. + /// + /// Returns `Err` if the id already exists. + pub fn create( + &mut self, + seq_id: SequenceId, + token_ids: Vec, + max_tokens: usize, + ) -> Result<()> { + if self.sequences.contains_key(&seq_id) { + return Err(Error::InvalidTransition("Duplicate sequence ID")); + } + let event = Event::Create { + token_ids, + max_tokens, + }; + let state = self.transition(SequenceState::Empty(seq_id), event)?; + self.sequences.insert(seq_id, state); + Ok(()) + } + + /// Apply an FSM event to an existing sequence, updating the map in place. + /// + /// On success the sequence's state is replaced with the new state. On + /// failure the original state is restored and the error is returned, so the + /// caller can apply rollback/OOM policy. + pub fn advance(&mut self, seq_id: SequenceId, event: Event) -> Result<&SequenceState> { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => return Err(Error::UnknownSequence(seq_id)), + }; + let backup = state.clone(); + + match self.transition(state, event) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + Ok(&self.sequences[&seq_id]) + } + Err(e) => { + // Restore prior state so the map is never left missing a seq. + self.sequences.insert(seq_id, backup); + Err(e) + } + } + } + + /// Fork a decoding sequence into a child via copy-on-write block sharing. + pub fn fork(&mut self, parent_id: SequenceId, child_id: SequenceId) -> Result<()> { + if self.sequences.contains_key(&child_id) { + warn!("Child ID {:?} already exists", child_id); + return Err(Error::InvalidTransition("Duplicate child ID")); + } + + let parent_state = match self.sequences.remove(&parent_id) { + Some(s) => s, + None => { + warn!("Parent sequence {:?} not found", parent_id); + return Err(Error::UnknownSequence(parent_id)); + } + }; + + let backup = parent_state.clone(); + + if let SequenceState::Decoding(s) = parent_state { + // Copy-on-write: increment ref counts on shared blocks. + for &block_id in s.blocks.iter() { + if let Err(e) = self.block_manager.add_ref(block_id) { + warn!("Failed to fork {:?}: {:?}", parent_id, e); + self.sequences.insert(parent_id, backup); + return Err(e); + } + } + + let child = DecodingState { + seq_id: child_id, + token_ids: Arc::clone(&s.token_ids), + blocks: Arc::clone(&s.blocks), + num_tokens: s.num_tokens, + max_tokens: s.max_tokens, + }; + + self.sequences.insert(parent_id, SequenceState::Decoding(s)); + self.sequences + .insert(child_id, SequenceState::Decoding(child)); + debug!("Forked sequence {:?} → {:?}", parent_id, child_id); + Ok(()) + } else { + warn!("Cannot fork {:?} - not in Decoding state", parent_id); + self.sequences.insert(parent_id, backup); + Err(Error::InvalidTransition("Can only fork Decoding sequences")) + } + } + + // ===== FSM Transitions ===== + + /// Compute the next state for `(state, event)`, allocating/freeing blocks as + /// needed. Pure with respect to the sequence map — does not read or write + /// `self.sequences`; the caller stores the result. + fn transition(&mut self, state: SequenceState, event: Event) -> Result { + let seq_id = state.seq_id(); + match (state, event) { + // Empty → Waiting (new sequence) + ( + SequenceState::Empty(_), + Event::Create { + token_ids, + max_tokens, + }, + ) => self.transition_create(seq_id, token_ids, max_tokens), + + // Waiting → Prefilling + (SequenceState::Waiting(s), Event::Schedule { .. }) => self.transition_schedule(s), + + // Prefilling → Prefilling/Decoding + (SequenceState::Prefilling(s), Event::AppendTokens { num_tokens, .. }) => { + self.transition_append_tokens_prefilling(s, num_tokens) + } + + // Decoding → Decoding/Finished + (SequenceState::Decoding(s), Event::AppendTokens { num_tokens, .. }) => { + self.transition_append_tokens_decoding(s, num_tokens) + } + + // Any running state → Finished + (state, Event::Complete { reason }) => self.transition_complete(state, reason), + + // Running → Preempted + ( + state @ (SequenceState::Prefilling(_) | SequenceState::Decoding(_)), + Event::Preempt, + ) => self.transition_preempt(state), + + // Preempted → Waiting + (SequenceState::Preempted(s), Event::Resume) => self.transition_resume(s), + + // Any → Aborted + (state, Event::Abort { reason }) => self.transition_abort(state, reason), + + // Invalid transitions + _ => Err(Error::InvalidTransition("Invalid state transition")), + } + } + + /// Transition: Empty → Waiting (construct a new sequence's initial state) + fn transition_create( + &mut self, + seq_id: SequenceId, + token_ids: Vec, + max_tokens: usize, + ) -> Result { + Ok(SequenceState::Waiting(WaitingState { + seq_id, + token_ids: Arc::new(token_ids), + max_tokens, + })) + } + + /// Transition: Waiting → Prefilling (allocate blocks for tokenized prompt) + fn transition_schedule(&mut self, state: WaitingState) -> Result { + let prompt_tokens = state.token_ids.len(); + let blocks_needed = prompt_tokens.div_ceil(self.tokens_per_block); + + let mut blocks = Vec::new(); + for _ in 0..blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => blocks.push(block_id), + Err(Error::OutOfMemory) => { + // Cleanup on OOM + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(Error::OutOfMemory); + } + Err(e) => { + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(e); + } + } + } + + debug!( + "Scheduled seq {:?}: allocated {} blocks for {} tokens", + state.seq_id, + blocks.len(), + prompt_tokens + ); + + Ok(SequenceState::Prefilling(PrefillingState { + seq_id: state.seq_id, + token_ids: state.token_ids, + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: prompt_tokens, + max_tokens: state.max_tokens, + })) + } + + /// Transition: Prefilling → Decoding or Prefilling (append tokens) + fn transition_append_tokens_prefilling( + &mut self, + mut state: PrefillingState, + num_tokens: usize, + ) -> Result { + state.tokens_filled += num_tokens; + + debug!( + "Prefilling seq {:?}: {}/{} tokens", + state.seq_id, state.tokens_filled, state.tokens_total + ); + + if state.tokens_filled >= state.tokens_total { + debug!("Seq {:?} transitioning to Decoding", state.seq_id); + Ok(SequenceState::Decoding(DecodingState { + seq_id: state.seq_id, + token_ids: state.token_ids, + blocks: state.blocks, + num_tokens: state.tokens_filled, + max_tokens: state.max_tokens, + })) + } else { + Ok(SequenceState::Prefilling(state)) + } + } + + /// Transition: Decoding → Decoding or Finished (append tokens, maybe allocate more blocks) + fn transition_append_tokens_decoding( + &mut self, + mut state: DecodingState, + num_tokens: usize, + ) -> Result { + state.num_tokens += num_tokens; + + // Check if finished + if state.num_tokens >= state.max_tokens { + // Free blocks + for &block_id in state.blocks.iter() { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!( + "Seq {:?} finished: reached max_tokens ({})", + state.seq_id, state.max_tokens + ); + + return Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: FinishReason::MaxTokens, + })); + } + + // Check if need more blocks + let blocks_needed = state.num_tokens.div_ceil(self.tokens_per_block); + let initial_block_count = state.blocks.len(); + + if state.blocks.len() < blocks_needed { + let blocks = Arc::make_mut(&mut state.blocks); + + while blocks.len() < blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => { + blocks.push(block_id); + debug!( + "Seq {:?}: allocated block, total blocks: {}", + state.seq_id, + blocks.len() + ); + } + Err(e) => { + warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); + for block_id in blocks.drain(initial_block_count..) { + let _ = self.block_manager.free(block_id); + } + return Err(e); + } + } + } + } + + Ok(SequenceState::Decoding(state)) + } + + /// Transition: Prefilling/Decoding → Finished + fn transition_complete( + &mut self, + state: SequenceState, + reason: FinishReason, + ) -> Result { + let seq_id = state.seq_id(); + + // Free blocks if any + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + debug!("Completed seq {:?}: {:?}", seq_id, reason); + + Ok(SequenceState::Finished(FinishedState { + seq_id, + finish_reason: reason, + })) + } + + /// Transition: Prefilling/Decoding → Preempted (free blocks) + fn transition_preempt(&mut self, state: SequenceState) -> Result { + let seq_id = state.seq_id(); + let (token_ids, num_tokens, max_tokens) = match &state { + SequenceState::Prefilling(s) => { + (Arc::clone(&s.token_ids), s.tokens_filled, s.max_tokens) + } + SequenceState::Decoding(s) => (Arc::clone(&s.token_ids), s.num_tokens, s.max_tokens), + _ => { + return Err(Error::InvalidTransition( + "Can only preempt Prefilling/Decoding", + )) + } + }; + + // Free blocks + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + debug!("Preempted seq {:?}", seq_id); + + Ok(SequenceState::Preempted(PreemptedState { + seq_id, + token_ids, + num_tokens, + max_tokens, + })) + } + + /// Transition: Preempted → Waiting + fn transition_resume(&mut self, state: PreemptedState) -> Result { + debug!("Resuming seq {:?}", state.seq_id); + + Ok(SequenceState::Waiting(WaitingState { + seq_id: state.seq_id, + token_ids: state.token_ids, + max_tokens: state.max_tokens, + })) + } + + /// Transition: Any → Aborted (free blocks, cleanup) + fn transition_abort(&mut self, state: SequenceState, reason: String) -> Result { + let seq_id = state.seq_id(); + + // Free blocks if any + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + warn!("Aborted seq {:?}: {}", seq_id, reason); + + Ok(SequenceState::Aborted(AbortedState { seq_id, reason })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + use crate::fsm::FinishReason; + + /// Build a manager with `max_blocks` worth of memory and the given + /// `tokens_per_block`. Block byte-size is fixed at 1024; capping allocator + /// memory is how we force OOM in tests. + fn manager(max_blocks: usize, tokens_per_block: usize) -> SequenceManager { + let block_bytes = 1024; + let allocator = Box::new(CpuAllocator::new(max_blocks * block_bytes)); + let block_manager = BlockManager::new(allocator, block_bytes); + SequenceManager::new(block_manager, tokens_per_block) + } + + fn schedule_event() -> Event { + Event::Schedule { + tokens_per_block: 4, + } + } + + fn append(num_tokens: usize) -> Event { + Event::AppendTokens { + num_tokens, + tokens_per_block: 4, + } + } + + #[test] + fn test_create_inserts_waiting_sequence() { + let mut mgr = manager(10, 4); + let id = SequenceId(1); + + mgr.create(id, vec![1, 2, 3], 20).unwrap(); + + assert!(mgr.contains(id)); + assert_eq!(mgr.len(), 1); + assert!(!mgr.is_empty()); + assert_eq!(mgr.get(id).unwrap().variant_name(), "Waiting"); + } + + #[test] + fn test_create_duplicate_id_errors() { + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3], 20).unwrap(); + + let err = mgr.create(id, vec![4, 5], 20).unwrap_err(); + assert!(matches!(err, Error::InvalidTransition(_))); + // Original sequence is untouched. + assert_eq!(mgr.len(), 1); + } + + #[test] + fn test_advance_unknown_sequence_errors() { + let mut mgr = manager(10, 4); + let err = mgr.advance(SequenceId(99), schedule_event()).unwrap_err(); + assert!(matches!(err, Error::UnknownSequence(SequenceId(99)))); + } + + #[test] + fn test_schedule_allocates_blocks() { + // 6 tokens, 4 tokens/block => 2 blocks needed. + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3, 4, 5, 6], 20).unwrap(); + + let state = mgr.advance(id, schedule_event()).unwrap(); + assert_eq!(state.variant_name(), "Prefilling"); + assert_eq!(mgr.blocks(id).unwrap().len(), 2); + } + + #[test] + fn test_schedule_oom_rolls_back_to_waiting() { + // Needs 2 blocks but only 1 is available -> OOM, no blocks leaked. + let mut mgr = manager(1, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3, 4, 5, 6], 20).unwrap(); + + let err = mgr.advance(id, schedule_event()).unwrap_err(); + assert!(matches!(err, Error::OutOfMemory)); + + // State restored to Waiting and no blocks are held. + assert_eq!(mgr.get(id).unwrap().variant_name(), "Waiting"); + assert!(mgr.blocks(id).unwrap().is_empty()); + + // The partially-allocated block was freed: a 1-block prompt now fits. + let id2 = SequenceId(2); + mgr.create(id2, vec![1, 2], 20).unwrap(); + assert!(mgr.advance(id2, schedule_event()).is_ok()); + } + + #[test] + fn test_prefill_to_decoding_transition() { + // 4 tokens, 4/block => 1 block, prefill completes in one append. + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(id, schedule_event()).unwrap(); + + // Partial prefill stays in Prefilling. + let state = mgr.advance(id, append(2)).unwrap(); + assert_eq!(state.variant_name(), "Prefilling"); + + // Reaching tokens_total flips to Decoding. + let state = mgr.advance(id, append(2)).unwrap(); + assert_eq!(state.variant_name(), "Decoding"); + assert_eq!(mgr.num_running(), 1); + } + + #[test] + fn test_decoding_reaches_max_tokens_finishes_and_frees() { + let mut mgr = manager(10, 4); + let id = SequenceId(1); + // prompt of 4 tokens, max_tokens 5. + mgr.create(id, vec![1, 2, 3, 4], 5).unwrap(); + mgr.advance(id, schedule_event()).unwrap(); + mgr.advance(id, append(4)).unwrap(); // -> Decoding, num_tokens=4 + + let state = mgr.advance(id, append(1)).unwrap(); // hits max_tokens=5 + assert_eq!(state.variant_name(), "Finished"); + // Blocks freed on finish. + assert!(mgr.blocks(id).unwrap().is_empty()); + assert_eq!(mgr.num_running(), 0); + } + + #[test] + fn test_preempt_frees_blocks_then_resume_returns_to_waiting() { + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(id, schedule_event()).unwrap(); + mgr.advance(id, append(4)).unwrap(); // Decoding + + let state = mgr.advance(id, Event::Preempt).unwrap(); + assert_eq!(state.variant_name(), "Preempted"); + assert_eq!(mgr.num_preempted(), 1); + assert!(mgr.blocks(id).unwrap().is_empty()); + + let state = mgr.advance(id, Event::Resume).unwrap(); + assert_eq!(state.variant_name(), "Waiting"); + assert_eq!(mgr.num_preempted(), 0); + } + + #[test] + fn test_complete_and_abort_finish_the_sequence() { + let mut mgr = manager(10, 4); + + let a = SequenceId(1); + mgr.create(a, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(a, schedule_event()).unwrap(); + let state = mgr + .advance( + a, + Event::Complete { + reason: FinishReason::Stop, + }, + ) + .unwrap(); + assert_eq!(state.variant_name(), "Finished"); + assert!(mgr.blocks(a).unwrap().is_empty()); + + let b = SequenceId(2); + mgr.create(b, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(b, schedule_event()).unwrap(); + let state = mgr + .advance( + b, + Event::Abort { + reason: "cancelled".to_string(), + }, + ) + .unwrap(); + assert_eq!(state.variant_name(), "Aborted"); + assert!(mgr.blocks(b).unwrap().is_empty()); + } + + #[test] + fn test_invalid_transition_preserves_state() { + // Resume on a Waiting sequence is invalid. + let mut mgr = manager(10, 4); + let id = SequenceId(1); + mgr.create(id, vec![1, 2, 3], 20).unwrap(); + + let err = mgr.advance(id, Event::Resume).unwrap_err(); + assert!(matches!(err, Error::InvalidTransition(_))); + // Sequence remains present and in its original state. + assert!(mgr.contains(id)); + assert_eq!(mgr.get(id).unwrap().variant_name(), "Waiting"); + } + + #[test] + fn test_fork_shares_blocks_of_decoding_parent() { + let mut mgr = manager(10, 4); + let parent = SequenceId(1); + let child = SequenceId(2); + mgr.create(parent, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(parent, schedule_event()).unwrap(); + mgr.advance(parent, append(4)).unwrap(); // Decoding + + let parent_blocks = mgr.blocks(parent).unwrap(); + mgr.fork(parent, child).unwrap(); + + // Child exists, is Decoding, and shares the parent's blocks. + assert_eq!(mgr.get(child).unwrap().variant_name(), "Decoding"); + assert_eq!(mgr.blocks(child).unwrap(), parent_blocks); + assert_eq!(mgr.num_running(), 2); + } + + #[test] + fn test_fork_non_decoding_parent_errors() { + let mut mgr = manager(10, 4); + let parent = SequenceId(1); + mgr.create(parent, vec![1, 2, 3], 20).unwrap(); // Waiting + + let err = mgr.fork(parent, SequenceId(2)).unwrap_err(); + assert!(matches!(err, Error::InvalidTransition(_))); + // Parent is restored and child was never inserted. + assert!(mgr.contains(parent)); + assert!(!mgr.contains(SequenceId(2))); + } + + #[test] + fn test_fork_duplicate_child_errors() { + let mut mgr = manager(10, 4); + let parent = SequenceId(1); + let child = SequenceId(2); + mgr.create(parent, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(parent, schedule_event()).unwrap(); + mgr.advance(parent, append(4)).unwrap(); + + mgr.create(child, vec![9], 20).unwrap(); // child id already taken + let err = mgr.fork(parent, child).unwrap_err(); + assert!(matches!(err, Error::InvalidTransition(_))); + } +} diff --git a/src/system/system_info.rs b/src/system/system_info.rs index ff62ebc..2cc2fcb 100644 --- a/src/system/system_info.rs +++ b/src/system/system_info.rs @@ -280,3 +280,82 @@ impl SystemInfo { println!(" Running Models: {}", self.running_models); } } + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + + #[test] + fn test_collect_populates_basic_fields() { + let info = SystemInfo::collect(); + + // Version comes from Cargo and must be non-empty. + assert_eq!(info.version, env!("CARGO_PKG_VERSION")); + // Any real host reports at least one CPU core. + assert!(info.cpu_cores >= 1); + // format_size never yields an empty string. + assert!(!info.total_memory.is_empty()); + assert!(!info.cache_size.is_empty()); + // running_models tracking is not implemented yet. + assert_eq!(info.running_models, 0); + } + + #[test] + fn test_serde_roundtrip() { + let info = SystemInfo { + version: "1.2.3".to_string(), + os: "TestOS".to_string(), + architecture: "x86_64".to_string(), + cpu_cores: 8, + total_memory: "16.00 GB".to_string(), + gpu_info: vec![GpuInfo { + name: "Test GPU".to_string(), + backend: "Metal".to_string(), + memory: Some("10 GPU cores".to_string()), + }], + cache_dir: "/tmp/cache".to_string(), + cache_size: "0 B".to_string(), + models_count: 2, + running_models: 0, + }; + + let json = serde_json::to_string(&info).unwrap(); + let back: SystemInfo = serde_json::from_str(&json).unwrap(); + + assert_eq!(back.version, info.version); + assert_eq!(back.cpu_cores, info.cpu_cores); + assert_eq!(back.gpu_info.len(), 1); + assert_eq!(back.gpu_info[0].backend, "Metal"); + assert_eq!(back.gpu_info[0].memory.as_deref(), Some("10 GPU cores")); + } + + #[test] + fn test_calculate_cache_size_missing_dir_is_zero() { + let missing = PathBuf::from("/nonexistent/puma/cache/path/xyz"); + assert_eq!(SystemInfo::calculate_cache_size(&missing), 0); + } + + #[test] + fn test_cache_size_sums_files_and_subdirs() { + let tmp = tempfile::tempdir().unwrap(); + let root = tmp.path().to_path_buf(); + + // A file at the top level and one nested in a subdirectory. + fs::write(root.join("a.bin"), vec![0u8; 4096]).unwrap(); + let sub = root.join("sub"); + fs::create_dir(&sub).unwrap(); + fs::write(sub.join("b.bin"), vec![0u8; 4096]).unwrap(); + + // Disk usage is block-based (blocks * 512); an empty dir contributes 0 + // files, so total should reflect both files' on-disk size (> 0) and + // recurse into the subdirectory. + let total = SystemInfo::calculate_cache_size(&root); + let sub_only = SystemInfo::dir_size(&sub); + assert!(sub_only > 0, "nested file should contribute size"); + assert!( + total >= sub_only, + "top-level total {total} should include nested {sub_only}" + ); + } +} From d2cb47800b7153f9ba088ebae27ae0cf809e6e06 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 19 Jul 2026 01:56:01 +0100 Subject: [PATCH 30/33] fix lint Signed-off-by: kerthcet --- src/api/chat.rs | 4 +++- src/backend/engine.rs | 2 +- src/backend/mock.rs | 5 ++++- src/cli/chat.rs | 3 ++- src/scheduler/core.rs | 2 +- src/utils/prompt.rs | 5 +---- 6 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/api/chat.rs b/src/api/chat.rs index a15e99e..c87f1f3 100644 --- a/src/api/chat.rs +++ b/src/api/chat.rs @@ -238,6 +238,8 @@ async fn chat_completions_stream( /// Format chat messages into a prompt via the shared prompt formatter. fn format_chat_messages(messages: &[ChatMessage]) -> String { crate::utils::prompt::format_conversation( - messages.iter().map(|m| (m.role.as_str(), m.content.as_str())), + messages + .iter() + .map(|m| (m.role.as_str(), m.content.as_str())), ) } diff --git a/src/backend/engine.rs b/src/backend/engine.rs index 93924f6..a3bb6f2 100644 --- a/src/backend/engine.rs +++ b/src/backend/engine.rs @@ -1,7 +1,7 @@ +use crate::block_manager::types::TokenId; use std::io; use std::pin::Pin; use tokio_stream::Stream; -use crate::block_manager::types::TokenId; /// Backend trait - low-level inference that works with token IDs /// diff --git a/src/backend/mock.rs b/src/backend/mock.rs index 16ecfad..3639aa7 100644 --- a/src/backend/mock.rs +++ b/src/backend/mock.rs @@ -142,7 +142,10 @@ mod tests { async fn tokens_stay_within_vocab() { let engine = MockEngine::with_vocab_size(50); let out = engine.generate(vec![7, 7, 7], 32, 0.0).await.unwrap(); - assert!(out.iter().all(|&t| t < 50), "ids must be within vocab range"); + assert!( + out.iter().all(|&t| t < 50), + "ids must be within vocab range" + ); } #[tokio::test] diff --git a/src/cli/chat.rs b/src/cli/chat.rs index 844e92a..087251e 100644 --- a/src/cli/chat.rs +++ b/src/cli/chat.rs @@ -104,7 +104,8 @@ pub async fn interactive_chat(engine: &EngineHandle, model: &str) -> Result<(), println!("\n"); // Double newline after response // Add assistant response to history - conversation_history.push(("assistant".to_string(), full_response.trim().to_string())); + conversation_history + .push(("assistant".to_string(), full_response.trim().to_string())); } Err(e) => { eprintln!("Error: {}\n", e); diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs index d9ffc16..4df64d6 100644 --- a/src/scheduler/core.rs +++ b/src/scheduler/core.rs @@ -1,8 +1,8 @@ use super::events::{ResponseSender, SchedulerEvent, SchedulerEventReceiver, SchedulerStats}; -use crate::sequence_manager::SequenceManager; use crate::block_manager::manager::BlockManager; use crate::block_manager::types::*; use crate::fsm::{Event, FinishReason, SequenceState}; +use crate::sequence_manager::SequenceManager; use std::collections::{HashMap, VecDeque}; use std::sync::Arc; use tracing::{debug, info, warn}; diff --git a/src/utils/prompt.rs b/src/utils/prompt.rs index 8432bc1..551706d 100644 --- a/src/utils/prompt.rs +++ b/src/utils/prompt.rs @@ -44,10 +44,7 @@ mod tests { #[test] fn formats_turns_with_assistant_cue() { - let prompt = format_conversation([ - ("system", "Be helpful."), - ("user", "Hello"), - ]); + let prompt = format_conversation([("system", "Be helpful."), ("user", "Hello")]); assert_eq!(prompt, "System: Be helpful.\nUser: Hello\nAssistant:"); } From 5c9b2714222d0470ae896898c60638d427228c63 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 19 Jul 2026 10:07:21 +0100 Subject: [PATCH 31/33] address comments Signed-off-by: kerthcet --- docs/fsm_architecture.md | 42 ++++++++++++-------- src/backend/llm_engine.rs | 77 +++++++++++++++++++++++++++++++++---- src/block_manager/types.rs | 3 ++ src/fsm/mod.rs | 32 +++++++-------- src/scheduler/core.rs | 63 +++++++++++++++++++++++++++--- src/sequence_manager/mod.rs | 58 ++++++++++++++++++++++++++++ 6 files changed, 229 insertions(+), 46 deletions(-) diff --git a/docs/fsm_architecture.md b/docs/fsm_architecture.md index 1d413b2..928d5e3 100644 --- a/docs/fsm_architecture.md +++ b/docs/fsm_architecture.md @@ -45,25 +45,33 @@ Events transform states: `Event + OldState → NewState` Invalid transitions return `Error::InvalidTransition`. -### FSM Integration with Scheduler +### FSM Integration with SequenceManager -The FSM logic is **integrated into the Scheduler** to solve the ownership boundary problem: +The FSM logic lives in the **`SequenceManager`** (the *mechanism* layer), which +owns `BlockManager`. The `Scheduler` (the *policy* layer) drives it with events +and never touches `BlockManager` or transition logic directly. -**Public API:** `Scheduler::transition(state, event)` - Event-based abstraction +**Public API:** `SequenceManager::advance(seq_id, event)` - Event-based abstraction - Single entry point for all state transitions - Easy to add logging, metrics, debugging -- Event replay capability for testing +- Backs up state and restores it on error (no block leaks) -**Internal Implementation:** Private `transition_*()` methods +**Internal Implementation:** Private `transition()` / `transition_*()` methods - Type-safe helpers that enforce correct state types - Direct access to `self.block_manager` - no parameter passing -- Called by `transition()` after event dispatching +- Called by `advance()` after looking up the current state ```rust -impl Scheduler { - pub fn transition(&mut self, state: SequenceState, event: Event) -> Result { +impl SequenceManager { + pub fn advance(&mut self, seq_id: SequenceId, event: Event) -> Result<&SequenceState> { + // Remove current state, run transition(), reinsert on success + // or restore the backup on error. + } + + fn transition(&mut self, state: SequenceState, event: Event) -> Result { + let seq_id = state.seq_id(); match (state, event) { - (SequenceState::Waiting(s), Event::Schedule{..}) => + (SequenceState::Waiting(s), Event::Schedule{..}) => self.transition_schedule(s), // ... dispatches to internal methods } @@ -100,21 +108,23 @@ impl Scheduler { └─ CompleteEvent → Finished ``` -## Scheduler Integration +## Scheduler / SequenceManager Integration -The Scheduler owns the FSM and manages all state transitions. +The `SequenceManager` owns the FSM and executes all state transitions; the +`Scheduler` owns the queues/batches and decides *which* events to fire. -**Implementation:** `src/scheduler/core.rs` +**Implementation:** `src/sequence_manager/mod.rs` (mechanism), +`src/scheduler/core.rs` (policy) **Key features:** - Synchronous API (LLMEngine handles async coordination) -- Maintains `HashMap` -- Owns `BlockManager` directly - no parameter passing needed +- `SequenceManager` maintains `HashMap` +- `SequenceManager` owns `BlockManager` directly - no parameter passing needed - Clones state before transitions (prevents loss on failure) - Restores backup on error (no block leaks) -- Event-based `apply()` method for maintainability +- Event-based `advance()` method for maintainability -**Tests:** `cargo test --lib` (120 tests passing) +**Tests:** `cargo test --lib` ## Arc Optimization diff --git a/src/backend/llm_engine.rs b/src/backend/llm_engine.rs index 725494d..5b42a6e 100644 --- a/src/backend/llm_engine.rs +++ b/src/backend/llm_engine.rs @@ -234,17 +234,21 @@ impl EngineRunner { completion_ids.push(token_id); let decoded = self.decode_tokens(&completion_ids).unwrap_or_default(); - // A trailing replacement char means an incomplete multi-byte - // character; wait for the next token before emitting more. - if decoded.ends_with('\u{FFFD}') { - continue; - } - if decoded.len() > sent_len { - let delta = decoded[sent_len..].to_string(); + if let Some(delta) = stream_delta(&decoded, sent_len) { sent_len = decoded.len(); - self.scheduler.send_chunk(seq_id, delta); + self.scheduler.send_chunk(seq_id, delta.to_string()); } } + + // Flush any tail held back by the incomplete-multibyte guard: if the + // stream ended while the last decode still had bytes beyond `sent_len` + // (e.g. a trailing `�` from a truncated multi-byte char), emit them now + // so the final output is never silently dropped. + let decoded = self.decode_tokens(&completion_ids).unwrap_or_default(); + if let Some(delta) = stream_flush(&decoded, sent_len) { + self.scheduler.send_chunk(seq_id, delta.to_string()); + } + let num_completion = completion_ids.len(); // Advance FSM state (prefill → decode → finished) so blocks are freed, @@ -305,6 +309,31 @@ impl EngineRunner { } } +/// Incremental streaming detokenization: given the full text decoded from the +/// completion prefix so far and how many bytes were already sent, return the +/// next delta to emit, or `None` to hold output back. +/// +/// A trailing replacement char (`�`) means the last multi-byte character is +/// still incomplete, so we wait for the next token rather than emit a partial +/// char. This mirrors the scheme vLLM and TGI use. +fn stream_delta(decoded: &str, sent_len: usize) -> Option<&str> { + if decoded.ends_with('\u{FFFD}') { + return None; + } + stream_flush(decoded, sent_len) +} + +/// Final flush: emit whatever bytes remain past `sent_len`, including any tail +/// the `�` guard in [`stream_delta`] held back when the stream ended. Returns +/// `None` when nothing is left to send. +fn stream_flush(decoded: &str, sent_len: usize) -> Option<&str> { + if decoded.len() > sent_len { + Some(&decoded[sent_len..]) + } else { + None + } +} + /// Construct a paired [`EngineHandle`] and [`EngineRunner`]. /// /// Spawn `runner.serve()` on a task and share the returned handle with the API / @@ -368,4 +397,36 @@ mod tests { let result = handle.generate("test-model", "Hello world", 100, 0.7).await; assert!(result.is_ok()); } + + #[test] + fn test_stream_delta_emits_complete_suffix() { + // Nothing sent yet -> emit the whole thing. + assert_eq!(stream_delta("hello", 0), Some("hello")); + // Some already sent -> emit only the new bytes. + assert_eq!(stream_delta("hello world", 5), Some(" world")); + // Nothing new -> hold. + assert_eq!(stream_delta("hello", 5), None); + } + + #[test] + fn test_stream_delta_holds_incomplete_multibyte() { + // Decode ended on a replacement char: the multi-byte char is not yet + // complete, so we must not emit anything this step. + assert_eq!(stream_delta("hi\u{FFFD}", 2), None); + } + + #[test] + fn test_stream_flush_emits_held_back_tail() { + // Simulates the bug scenario: the loop held back a trailing `�` and the + // stream then ended. The flush must still emit those bytes rather than + // drop them. + let decoded = "hi\u{FFFD}"; + assert_eq!(stream_delta(decoded, 2), None); // loop held it back + assert_eq!(stream_flush(decoded, 2), Some("\u{FFFD}")); // flush emits it + } + + #[test] + fn test_stream_flush_none_when_all_sent() { + assert_eq!(stream_flush("done", 4), None); + } } diff --git a/src/block_manager/types.rs b/src/block_manager/types.rs index 69a5c76..820ed1f 100644 --- a/src/block_manager/types.rs +++ b/src/block_manager/types.rs @@ -103,6 +103,9 @@ pub enum Error { #[error("Invalid state transition: {0}")] InvalidTransition(&'static str), + + #[error("Sequence aborted: {0}")] + Aborted(String), } pub type Result = std::result::Result; diff --git a/src/fsm/mod.rs b/src/fsm/mod.rs index cbdadea..cbe4b1a 100644 --- a/src/fsm/mod.rs +++ b/src/fsm/mod.rs @@ -1,8 +1,10 @@ //! FSM - Finite State Machine for sequence lifecycle //! -//! This module defines states and events for the sequence state machine. -//! The actual transition logic is implemented in `Scheduler::transition()` and -//! private `Scheduler::transition_*()` methods. +//! This module defines the *data* of the state machine: the states and the +//! events that trigger transitions. The actual transition logic lives in +//! [`SequenceManager`](crate::sequence_manager::SequenceManager) — its private +//! `transition()`/`transition_*()` methods compute the next state and +//! allocate/free blocks along the way. //! //! # Architecture //! @@ -11,26 +13,24 @@ //! //! # Usage //! -//! States and events are used by the Scheduler: +//! The `Scheduler` (policy) drives the `SequenceManager` (mechanism) with these +//! events; it never touches `BlockManager` or transition logic directly: //! //! ```rust,ignore -//! use crate::fsm::{Event, SequenceState}; -//! use crate::scheduler::Scheduler; +//! use crate::fsm::Event; //! -//! let mut scheduler = Scheduler::new(...); -//! let state = SequenceState::Waiting(WaitingState { ... }); -//! let event = Event::Schedule { tokens_per_block: 16 }; -//! -//! // Scheduler owns BlockManager and performs transitions -//! let new_state = scheduler.transition(state, event)?; +//! // seq_manager: SequenceManager, owned by the scheduler +//! seq_manager.create(seq_id, token_ids, max_tokens)?; +//! let new_state = seq_manager.advance(seq_id, Event::Schedule { tokens_per_block: 16 })?; //! ``` //! //! # Design Rationale //! -//! FSM logic is **inside Scheduler** rather than a separate `fsm::apply()` because: -//! - Scheduler owns `BlockManager` - no need to pass as parameter -//! - Direct access to scheduler context (tokens_per_block, etc.) -//! - Event-based `transition()` method provides single entry point for logging/metrics +//! Transition logic lives **inside `SequenceManager`** rather than a free +//! `fsm::apply()` because: +//! - The manager owns `BlockManager` - no need to pass it as a parameter +//! - Direct access to context (tokens_per_block, etc.) +//! - Event-based `advance()` gives a single entry point for logging/metrics //! - Type-safe internal `transition_*()` methods enforce correct state types //! //! See `docs/fsm_architecture.md` for detailed documentation. diff --git a/src/scheduler/core.rs b/src/scheduler/core.rs index 4df64d6..252b8c8 100644 --- a/src/scheduler/core.rs +++ b/src/scheduler/core.rs @@ -7,19 +7,20 @@ use std::collections::{HashMap, VecDeque}; use std::sync::Arc; use tracing::{debug, info, warn}; -/// Scheduler - memory and batch management (synchronous API) +/// Scheduler - batch and scheduling policy (synchronous API) /// /// Architecture /// - Pure synchronous methods (no async, no loop) /// - LLMEngine calls: add_request() → schedule() → process_outputs() -/// - Owns BlockManager and performs FSM transitions directly +/// - Drives a `SequenceManager` (which owns memory and executes FSM +/// transitions) via `create`/`advance`/`fork` /// /// Responsibilities: -/// - Owns BlockManager (memory allocation) -/// - Owns sequence state storage (FSM states) +/// - This is the *policy* layer; the `SequenceManager` is the *mechanism* +/// layer that owns `BlockManager` and the FSM transition logic /// - Manages batches (waiting/prefill/decode) /// - Scheduling policy (what to schedule when) -/// - FSM transitions (internal methods with direct BlockManager access) +/// - Owns client response channels and delivers results/errors pub struct Scheduler { /// Owns sequence state + block memory and executes FSM transitions. sequences: SequenceManager, @@ -140,7 +141,12 @@ impl Scheduler { /// The reason is passed through to the FSM `Abort` event, so callers (user /// cancel, backend failure, …) describe *why* the sequence ended. pub fn abort_request(&mut self, seq_id: SequenceId, reason: String) { - match self.sequences.advance(seq_id, Event::Abort { reason }) { + match self.sequences.advance( + seq_id, + Event::Abort { + reason: reason.clone(), + }, + ) { Ok(_) => { info!("Aborted sequence {:?}", seq_id); @@ -148,6 +154,13 @@ impl Scheduler { self.waiting_queue.retain(|&id| id != seq_id); self.prefill_batch.retain(|&id| id != seq_id); self.decode_batch.retain(|&id| id != seq_id); + + // Resolve the client channel so waiters unblock and the entry + // is not leaked. Streaming callers see the error and close; + // single callers get an `Err` instead of hanging forever. + if let Some(response_tx) = self.response_channels.remove(&seq_id) { + self.send_error(response_tx, Error::Aborted(reason)); + } } Err(e) => { warn!("Failed to abort {:?}: {:?}", seq_id, e); @@ -454,6 +467,44 @@ mod tests { assert_eq!(scheduler.waiting_queue.len(), 0); } + #[test] + fn test_abort_resolves_single_channel_and_removes_entry() { + let mut scheduler = create_test_scheduler(); + + // Keep the receiver so we can observe what the client sees. + let (tx, rx) = tokio::sync::oneshot::channel(); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, ResponseSender::Single(tx)); + assert_eq!(scheduler.response_channels.len(), 1); + + scheduler.abort_request(SequenceId(1), "User cancelled".to_string()); + + // The waiter is unblocked with an Aborted error (not left hanging)... + match rx.blocking_recv() { + Ok(Err(Error::Aborted(reason))) => assert_eq!(reason, "User cancelled"), + other => panic!("expected Ok(Err(Aborted)), got {:?}", other), + } + // ...and the channel entry is not leaked. + assert!(scheduler.response_channels.is_empty()); + } + + #[test] + fn test_abort_closes_streaming_channel() { + let mut scheduler = create_test_scheduler(); + + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + scheduler.add_request(SequenceId(1), vec![0; 100], 200, ResponseSender::Stream(tx)); + + scheduler.abort_request(SequenceId(1), "boom".to_string()); + + // Streaming client receives the error, then the channel closes. + match rx.blocking_recv() { + Some(Err(Error::Aborted(reason))) => assert_eq!(reason, "boom"), + other => panic!("expected Some(Err(Aborted)), got {:?}", other), + } + assert!(rx.blocking_recv().is_none(), "channel should be closed"); + assert!(scheduler.response_channels.is_empty()); + } + #[test] fn test_cancel_request_not_found() { let mut scheduler = create_test_scheduler(); diff --git a/src/sequence_manager/mod.rs b/src/sequence_manager/mod.rs index 9333f6a..dffb7f5 100644 --- a/src/sequence_manager/mod.rs +++ b/src/sequence_manager/mod.rs @@ -34,6 +34,7 @@ pub struct SequenceManager { impl SequenceManager { pub fn new(block_manager: BlockManager, tokens_per_block: usize) -> Self { + assert!(tokens_per_block > 0, "tokens_per_block must be > 0"); Self { block_manager, sequences: HashMap::new(), @@ -155,12 +156,21 @@ impl SequenceManager { if let SequenceState::Decoding(s) = parent_state { // Copy-on-write: increment ref counts on shared blocks. + let mut reffed = Vec::with_capacity(s.blocks.len()); for &block_id in s.blocks.iter() { if let Err(e) = self.block_manager.add_ref(block_id) { warn!("Failed to fork {:?}: {:?}", parent_id, e); + // Roll back the refs we already bumped so those blocks can + // still return to the free pool later (otherwise leak). + for &done in &reffed { + if let Err(e) = self.block_manager.free(done) { + warn!("Failed to roll back ref for {:?}: {:?}", done, e); + } + } self.sequences.insert(parent_id, backup); return Err(e); } + reffed.push(block_id); } let child = DecodingState { @@ -699,6 +709,54 @@ mod tests { assert!(!mgr.contains(SequenceId(2))); } + /// Sum of free blocks across all block types. + fn free_blocks(mgr: &SequenceManager) -> usize { + mgr.block_stats().values().map(|s| s.free_blocks).sum() + } + + #[test] + fn test_fork_rolls_back_refs_on_partial_failure() { + let mut mgr = manager(10, 4); + let parent = SequenceId(1); + mgr.create(parent, vec![1, 2, 3, 4], 20).unwrap(); + mgr.advance(parent, schedule_event()).unwrap(); + mgr.advance(parent, append(4)).unwrap(); // Decoding, holds 1 real block + + let real_block = mgr.blocks(parent).unwrap()[0]; + + // Inject a bogus second block so add_ref succeeds on the real block and + // then fails on the bogus one, exercising the mid-loop rollback path. + match mgr.sequences.get_mut(&parent) { + Some(SequenceState::Decoding(s)) => { + s.blocks = Arc::new(vec![real_block, BlockId(9999)]); + } + _ => panic!("parent should be Decoding"), + } + + let err = mgr.fork(parent, SequenceId(2)).unwrap_err(); + assert!(matches!(err, Error::InvalidBlockId(_))); + // Child was never created and the parent is restored. + assert!(!mgr.contains(SequenceId(2))); + assert!(mgr.contains(parent)); + + // The real block's ref was rolled back to 1 (parent only), so aborting + // the parent drops it to 0 and returns it to the free pool. Without the + // rollback its ref would be 2 and it would leak (never pooled). + let free_before = free_blocks(&mgr); + mgr.advance( + parent, + Event::Abort { + reason: "x".to_string(), + }, + ) + .unwrap(); + assert_eq!( + free_blocks(&mgr), + free_before + 1, + "real block should return to the free pool after abort" + ); + } + #[test] fn test_fork_duplicate_child_errors() { let mut mgr = manager(10, 4); From 667da3f6f7b6b251e45bb777242855b914f5bf41 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 19 Jul 2026 10:42:04 +0100 Subject: [PATCH 32/33] remove fork event Signed-off-by: kerthcet --- docs/fsm_architecture.md | 8 ++++++-- src/backend/llm_engine.rs | 2 ++ src/fsm/events.rs | 5 +---- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/docs/fsm_architecture.md b/docs/fsm_architecture.md index 928d5e3..769216b 100644 --- a/docs/fsm_architecture.md +++ b/docs/fsm_architecture.md @@ -35,9 +35,9 @@ Events transform states: `Event + OldState → NewState` **Implementation:** `src/fsm/events.rs` **Events:** +- `Create` - Empty → Waiting (birth a new sequence) - `Schedule` - Waiting → Prefilling (allocate blocks) - `AppendTokens` - Prefilling → Decoding or Decoding → Decoding/Finished -- `Fork` - Decoding → (Decoding, Decoding) for beam search - `Preempt` - Decoding → Preempted (free blocks on OOM) - `Resume` - Preempted → Waiting (re-queue) - `Complete` - * → Finished (free blocks) @@ -45,6 +45,10 @@ Events transform states: `Event + OldState → NewState` Invalid transitions return `Error::InvalidTransition`. +Forking (Decoding → (Decoding, Decoding) for beam search) is **not** an event: +it is a dedicated `SequenceManager::fork` method, since it produces a second +sequence rather than transforming one state into another. + ### FSM Integration with SequenceManager The FSM logic lives in the **`SequenceManager`** (the *mechanism* layer), which @@ -96,7 +100,7 @@ impl SequenceManager { │ │ ├─ Decoding (more tokens) │ │ └─ Finished (max_tokens) │ │ - │ ├─→ ForkEvent (beam search) + │ ├─→ fork() method (beam search, not an event) │ │ ├─ Parent: Decoding │ │ └─ Child: Decoding │ │ diff --git a/src/backend/llm_engine.rs b/src/backend/llm_engine.rs index 5b42a6e..21e71bc 100644 --- a/src/backend/llm_engine.rs +++ b/src/backend/llm_engine.rs @@ -232,6 +232,8 @@ impl EngineRunner { let mut sent_len = 0usize; // bytes of the decoded completion already sent while let Some(token_id) = stream.next().await { completion_ids.push(token_id); + // Advance decoding state and allocate KV blocks as needed. + self.scheduler.append_tokens(seq_id, 1); let decoded = self.decode_tokens(&completion_ids).unwrap_or_default(); if let Some(delta) = stream_delta(&decoded, sent_len) { diff --git a/src/fsm/events.rs b/src/fsm/events.rs index a7d9d48..69518cd 100644 --- a/src/fsm/events.rs +++ b/src/fsm/events.rs @@ -1,5 +1,5 @@ use super::states::FinishReason; -use crate::block_manager::types::{SequenceId, TokenId}; +use crate::block_manager::types::TokenId; /// FSM Events - pure data that triggers state transitions /// @@ -32,9 +32,6 @@ pub enum Event { /// Resume a preempted sequence Resume, - /// Fork a sequence (copy-on-write) - Fork { child_id: SequenceId }, - /// Abort a sequence with error Abort { reason: String }, } From 105e39909f2f9269beb047fe85cc68ca34d421d0 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 19 Jul 2026 10:46:51 +0100 Subject: [PATCH 33/33] updaet tdoc Signed-off-by: kerthcet --- docs/fsm_architecture.md | 51 ++++++++++++++++++++-------------------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/docs/fsm_architecture.md b/docs/fsm_architecture.md index 769216b..2ddccf5 100644 --- a/docs/fsm_architecture.md +++ b/docs/fsm_architecture.md @@ -1,6 +1,8 @@ # FSM Architecture for PUMA -Finite State Machine pattern for sequence lifecycle management, integrated into the Scheduler. +Finite State Machine pattern for sequence lifecycle management. The FSM states +and events are defined in `src/fsm/`; the transition logic lives in the +`SequenceManager` (`src/sequence_manager/mod.rs`), which the `Scheduler` drives. ## Core Concepts @@ -11,6 +13,7 @@ Each sequence is an independent state machine. Multiple sequences can be in diff **Implementation:** `src/fsm/states.rs` **States:** +- `Empty` - Pre-birth placeholder; the `from` state of `Create`, never persisted - `Waiting` - In queue, no resources allocated - `Scheduling` - Being scheduled (transitional) - `Prefilling` - Processing prompt @@ -85,6 +88,7 @@ impl SequenceManager { // Direct access to self.block_manager, self.tokens_per_block } } +``` ## State Diagram @@ -180,9 +184,9 @@ All state transitions are wrapped with backup/restore: PUMA uses a two-level event system (inspired by TokenSpeed): **Level 1: FSM Events (Internal)** - `src/fsm/events.rs` -- Created by Scheduler internally +- Created by the Scheduler internally - Include scheduler context (tokens_per_block, etc.) -- Used with `Scheduler::apply(state, event)` +- Applied via `SequenceManager::advance(seq_id, event)` - Examples: `Event::Schedule`, `Event::AppendTokens` **Level 2: Scheduler Events (External)** - `src/scheduler/events.rs` @@ -196,45 +200,40 @@ PUMA uses a two-level event system (inspired by TokenSpeed): - They don't know scheduler policies (tokens_per_block) - FSM events require scheduler context -See `ARCHITECTURE.md` for detailed explanation and flow diagrams. - ## Key Files - `src/fsm/states.rs` - State definitions and helper methods - `src/fsm/events.rs` - FSM event enum (internal transitions) -- `src/scheduler/core.rs` - Scheduler with integrated FSM (`apply()` method) +- `src/sequence_manager/mod.rs` - Owns memory + FSM transition logic (`advance()`) +- `src/scheduler/core.rs` - Scheduling policy that drives the SequenceManager - `src/scheduler/events.rs` - External scheduler events - `src/backend/llm_engine.rs` - Event coordinator -- `ARCHITECTURE.md` - Comprehensive architecture documentation ## Usage Example ```rust -// External component sends high-level event +// External component sends a high-level event scheduler_tx.send(SchedulerEvent::AddRequest { - seq_id: 1, - prompt_tokens: 10, + seq_id: SequenceId(1), + token_ids: vec![/* tokenized prompt */], max_tokens: 100, + response_tx, // channel the result is delivered on })?; -// Scheduler translates to FSM events internally +// The Scheduler translates it to an FSM event and drives the SequenceManager, +// which owns BlockManager and performs the transition. impl Scheduler { - pub fn schedule(&mut self) -> bool { - // Pop from waiting queue - let state = self.sequences.remove(&seq_id).unwrap(); - - // Create FSM event with scheduler context - let event = Event::Schedule { - tokens_per_block: self.tokens_per_block, - }; - - // Perform transition (direct access to self.block_manager) - match self.transition(state, event) { - Ok(new_state) => { - self.sequences.insert(seq_id, new_state); - self.prefill_batch.push(seq_id); + fn schedule_prefill(&mut self) { + while let Some(seq_id) = self.waiting_queue.pop_front() { + // Build the FSM event with scheduler context + let event = Event::Schedule { + tokens_per_block: self.tokens_per_block, + }; + + match self.sequences.advance(seq_id, event) { + Ok(_) => self.prefill_batch.push(seq_id), + Err(e) => { /* apply OOM / queueing policy */ } } - Err(e) => { /* handle error */ } } } }