From 1a7423a92b71105a1d164810de2d714b6551448b Mon Sep 17 00:00:00 2001 From: kerthcet Date: Sun, 5 Jul 2026 16:59:51 +0100 Subject: [PATCH 01/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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)) - } -}