diff --git a/Cargo.lock b/Cargo.lock index 22e04fb..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,11 +1606,14 @@ dependencies = [ "reqwest", "rusqlite", "rusqlite_migration", + "rustyline", + "rustyline-derive", "serde", "serde_derive", "serde_json", "sysinfo", "tempfile", + "thiserror 2.0.18", "tokio", "tokio-stream", "tower 0.4.13", @@ -1573,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" @@ -1817,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" @@ -2419,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 622000e..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" @@ -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" @@ -39,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/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/docs/fsm_architecture.md b/docs/fsm_architecture.md new file mode 100644 index 0000000..2bc3a91 --- /dev/null +++ b/docs/fsm_architecture.md @@ -0,0 +1,438 @@ +# 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), + Scheduling(SchedulingState), + Prefilling(PrefillingState), + Decoding(DecodingState), + Preempted(PreemptedState), + Finished(FinishedState), + Aborted(AbortedState), +} + +// 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(); + + // Clone state before transition (prevents loss on failure) + let backup = state.clone(); + + // Apply event transformation + let event = AppendTokensEvent { + num_tokens, + block_manager: &mut self.block_manager, + tokens_per_block: self.tokens_per_block, + }; + + // Transform state + match event.apply(state) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + } + Err(e) => { + error!("Transition failed: {:?}", e); + // Restore original state to prevent sequence loss + self.sequences.insert(seq_id, backup); + } + } + } + // ... other events + } + } + } +} +``` + +## 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 | + +## Error Handling and State Recovery + +On transition failure, the state is cloned before applying the event: + +```rust +fn handle_append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { + let state = self.sequences.remove(&seq_id).unwrap(); + let backup = state.clone(); // Clone before transition + + let event = AppendTokensEvent { ... }; + + match event.apply(state) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + } + Err(Error::OutOfMemory) => { + // Restore state and trigger preemption + self.sequences.insert(seq_id, backup); + self.handle_oom(); + } + Err(e) => { + // Restore state on any error (prevents block leaks) + self.sequences.insert(seq_id, backup); + warn!("Transition failed: {:?}", e); + } + } +} +``` + +**Benefits:** +- No sequence loss on transition failure +- No block leaks (blocks remain owned by restored state) +- Can retry or handle errors without losing context +- Clone cost is negligible (~50-100 bytes per state) + +## Usage + +```rust +use puma::block_manager::{BlockManager, CpuAllocator}; +use puma::sequence_manager::{SequenceManager, SequenceEvent, SequenceIdGenerator}; + +// Setup +let allocator = Box::new(CpuAllocator::new(100_000_000)); +let block_manager = BlockManager::new(allocator, 4096); +let (seq_manager, event_tx) = SequenceManager::new(block_manager, 16); + +// ID generator for sequential IDs starting from 0 +let id_gen = SequenceIdGenerator::new(); +let seq_id = id_gen.next(); // SequenceId(0) + +// Run event loop +tokio::spawn(async move { seq_manager.run().await }); + +// Send events +event_tx.send(SequenceEvent::AddRequest { + seq_id, + prompt_tokens: 100, + max_tokens: 150, +}).unwrap(); + +event_tx.send(SequenceEvent::AppendTokens { + seq_id, + num_tokens: 100, +}).unwrap(); +``` diff --git a/src/block_manager/allocator.rs b/src/block_manager/allocator.rs new file mode 100644 index 0000000..0db8aab --- /dev/null +++ b/src/block_manager/allocator.rs @@ -0,0 +1,79 @@ +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::AllocationFailed(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<()> { + if addr.size > self.used_memory { + 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::FreeFailed(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..449236e --- /dev/null +++ b/src/block_manager/manager.rs @@ -0,0 +1,189 @@ +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, 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 + .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, + mem_addr, + ref_count: 1, + }, + ); + + Ok(block_id) + } + + pub fn allocate(&mut self) -> Result { + 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) + .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, + 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 + } +} + +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/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 new file mode 100644 index 0000000..7279f14 --- /dev/null +++ b/src/block_manager/mod.rs @@ -0,0 +1,8 @@ +pub mod allocator; +pub mod manager; +pub mod types; + +// Re-export for convenience + +#[cfg(test)] +mod manager_tests; diff --git a/src/block_manager/types.rs b/src/block_manager/types.rs new file mode 100644 index 0000000..3076802 --- /dev/null +++ b/src/block_manager/types.rs @@ -0,0 +1,105 @@ +/// 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 failed: {0}")] + AllocationFailed(String), + + #[error("Free failed: {0}")] + FreeFailed(String), + + #[error("Invalid state transition: {0}")] + InvalidTransition(&'static str), +} + +pub type Result = std::result::Result; diff --git a/src/cli/chat.rs b/src/cli/chat.rs new file mode 100644 index 0000000..6d7c5a9 --- /dev/null +++ b/src/cli/chat.rs @@ -0,0 +1,115 @@ +use rustyline::error::ReadlineError; +use rustyline::hint::{Hint, Hinter}; +use rustyline::{Context, Editor}; +use rustyline_derive::{Completer, Helper, Highlighter, Validator}; +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(io::Error::other)?; + 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) => { + break; + } + Err(err) => { + return Err(io::Error::other(err)); + } + }; + + // Exit commands + if input.is_empty() { + continue; + } + if input == "exit" { + break; + } + + // Add user message to history + conversation_history.push(format!("User: {}", input)); + + // 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) => { + 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; 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"); 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 diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..43546f1 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,12 @@ +#![allow(dead_code)] + +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..aa4feeb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,12 @@ +#![allow(dead_code)] + 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..7398b8e --- /dev/null +++ b/src/sequence_manager/events.rs @@ -0,0 +1,73 @@ +use crate::block_manager::types::*; +use std::collections::HashMap; +use tokio::sync::mpsc; + +/// Events that drive the sequence manager +#[derive(Debug)] +pub enum SequenceEvent { + // Request lifecycle + AddRequest { + seq_id: SequenceId, + prompt_tokens: usize, + max_tokens: usize, + }, + + // Token generation + AppendTokens { + seq_id: SequenceId, + num_tokens: usize, + }, + + // Sequence forking (for beam search, parallel sampling) + ForkSequence { + parent_id: SequenceId, + child_id: SequenceId, + }, + + // Completion + CompleteSequence { + seq_id: SequenceId, + }, + + // Preemption (when OOM) + PreemptSequence { + seq_id: SequenceId, + }, + + // Resume preempted sequence + ResumeSequence { + seq_id: SequenceId, + }, + + // Query + GetBlocks { + seq_id: SequenceId, + response: tokio::sync::oneshot::Sender>>, + }, + + // Stats + GetStats { + response: tokio::sync::oneshot::Sender, + }, +} + +#[derive(Debug, Clone)] +pub struct SequenceManagerStats { + pub num_sequences: usize, + pub num_running: usize, + pub num_waiting: usize, + pub num_preempted: usize, + pub block_stats: HashMap, +} + +pub type SequenceEventSender = mpsc::UnboundedSender; +pub type SequenceEventReceiver = mpsc::UnboundedReceiver; + +/// Create event channel for sequence manager +/// +/// TODO: Use bounded channel to prevent OOM when producers outpace the event loop. +/// Capacity should be calculated based on available memory and average event size. +/// Consider: capacity = (available_memory * 0.1) / sizeof(SequenceEvent) +pub fn create_event_channel() -> (SequenceEventSender, SequenceEventReceiver) { + mpsc::unbounded_channel() +} diff --git a/src/sequence_manager/fsm_events.rs b/src/sequence_manager/fsm_events.rs new file mode 100644 index 0000000..0c4904d --- /dev/null +++ b/src/sequence_manager/fsm_events.rs @@ -0,0 +1,740 @@ +use super::states::*; +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; +use std::sync::Arc; +use tracing::{debug, warn}; + +/// FSM Events - transform states with type-safe transitions +/// +/// Pattern: Event takes old state, returns new state +/// Invalid transitions return Error::invalid_transition +/// Schedule a waiting sequence (allocate blocks) +pub struct ScheduleEvent<'a> { + pub block_manager: &'a mut BlockManager, + pub tokens_per_block: usize, +} + +impl<'a> ScheduleEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Waiting(s) => self.apply_to_waiting(s), + _ => Err(Error::InvalidTransition( + "ScheduleEvent requires Waiting state", + )), + } + } + + fn apply_to_waiting(self, state: WaitingState) -> Result { + let blocks_needed = state.prompt_tokens.div_ceil(self.tokens_per_block); + + let mut blocks = Vec::new(); + for _ in 0..blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => blocks.push(block_id), + Err(Error::OutOfMemory) => { + // OOM during scheduling - free what we allocated + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(Error::OutOfMemory); + } + Err(e) => { + // Other error - cleanup and propagate + for block in blocks { + let _ = self.block_manager.free(block); + } + return Err(e); + } + } + } + + debug!( + "Scheduled seq {:?}: allocated {} blocks for {} tokens", + state.seq_id, + blocks.len(), + state.prompt_tokens + ); + + Ok(SequenceState::Prefilling(PrefillingState { + seq_id: state.seq_id, + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: state.prompt_tokens, + max_tokens: state.max_tokens, + })) + } +} + +/// Append tokens to a running sequence +pub struct AppendTokensEvent<'a> { + pub num_tokens: usize, + pub block_manager: &'a mut BlockManager, + pub tokens_per_block: usize, +} + +impl<'a> AppendTokensEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Prefilling(s) => self.apply_to_prefilling(s), + SequenceState::Decoding(s) => self.apply_to_decoding(s), + _ => Err(Error::InvalidTransition( + "AppendTokensEvent requires Prefilling or Decoding state", + )), + } + } + + fn apply_to_prefilling(self, mut state: PrefillingState) -> Result { + state.tokens_filled += self.num_tokens; + + debug!( + "Prefilling seq {:?}: {}/{} tokens", + state.seq_id, state.tokens_filled, state.tokens_total + ); + + if state.tokens_filled >= state.tokens_total { + // Transition to decode + debug!("Seq {:?} transitioning to Decoding", state.seq_id); + Ok(SequenceState::Decoding(DecodingState { + seq_id: state.seq_id, + blocks: state.blocks, + num_tokens: state.tokens_filled, + max_tokens: state.max_tokens, + })) + } else { + Ok(SequenceState::Prefilling(state)) + } + } + + fn apply_to_decoding(self, mut state: DecodingState) -> Result { + state.num_tokens += self.num_tokens; + + // Check if finished + if state.num_tokens >= state.max_tokens { + // Free blocks + for &block_id in state.blocks.iter() { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!( + "Seq {:?} finished: reached max_tokens ({})", + state.seq_id, state.max_tokens + ); + + return Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: FinishReason::MaxTokens, + })); + } + + // Check if need more blocks + let blocks_needed = state.num_tokens.div_ceil(self.tokens_per_block); + let initial_block_count = state.blocks.len(); + + if state.blocks.len() < blocks_needed { + // Need to allocate more blocks - use Arc::make_mut for copy-on-write + let blocks = Arc::make_mut(&mut state.blocks); + + while blocks.len() < blocks_needed { + match self.block_manager.allocate() { + Ok(block_id) => { + blocks.push(block_id); + debug!( + "Seq {:?}: allocated block, total blocks: {}", + state.seq_id, + blocks.len() + ); + } + Err(e) => { + // Free any blocks we allocated in this call + warn!("Failed to allocate block for {:?}: {:?}", state.seq_id, e); + for block_id in blocks.drain(initial_block_count..) { + let _ = self.block_manager.free(block_id); + } + return Err(e); + } + } + } + } + + Ok(SequenceState::Decoding(state)) + } +} + +/// Preempt a running sequence (free blocks) +pub struct PreemptEvent<'a> { + pub block_manager: &'a mut BlockManager, +} + +impl<'a> PreemptEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Decoding(s) => self.apply_to_decoding(s), + SequenceState::Prefilling(s) => self.apply_to_prefilling(s), + _ => Err(Error::InvalidTransition( + "PreemptEvent requires running state", + )), + } + } + + fn apply_to_decoding(self, state: DecodingState) -> Result { + // Free all blocks + for &block_id in state.blocks.iter() { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!("Preempted seq {:?}", state.seq_id); + + Ok(SequenceState::Preempted(PreemptedState { + seq_id: state.seq_id, + num_tokens: state.num_tokens, + max_tokens: state.max_tokens, + })) + } + + fn apply_to_prefilling(self, state: PrefillingState) -> Result { + // Free all blocks (best effort) + for &block_id in state.blocks.iter() { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!("Preempted seq {:?} during prefill", state.seq_id); + + Ok(SequenceState::Preempted(PreemptedState { + seq_id: state.seq_id, + num_tokens: state.tokens_filled, + max_tokens: state.max_tokens, + })) + } +} + +/// Resume a preempted sequence +pub struct ResumeEvent; + +impl ResumeEvent { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Preempted(s) => self.apply_to_preempted(s), + _ => Err(Error::InvalidTransition( + "ResumeEvent requires Preempted state", + )), + } + } + + fn apply_to_preempted(self, state: PreemptedState) -> Result { + debug!("Resuming seq {:?}", state.seq_id); + + // Transition back to waiting for rescheduling + Ok(SequenceState::Waiting(WaitingState { + seq_id: state.seq_id, + prompt_tokens: state.num_tokens, + max_tokens: state.max_tokens, + })) + } +} + +/// Fork a sequence (copy-on-write) +pub struct ForkEvent<'a> { + pub child_id: SequenceId, + pub block_manager: &'a mut BlockManager, +} + +impl<'a> ForkEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result<(SequenceState, SequenceState)> { + match state { + SequenceState::Decoding(s) => self.apply_to_decoding(s), + _ => Err(Error::InvalidTransition( + "ForkEvent requires Decoding state", + )), + } + } + + fn apply_to_decoding(self, state: DecodingState) -> Result<(SequenceState, SequenceState)> { + // Copy-on-write: increment ref counts + for &block_id in state.blocks.iter() { + self.block_manager.add_ref(block_id)?; + } + + let child = DecodingState { + seq_id: self.child_id, + blocks: state.blocks.clone(), + num_tokens: state.num_tokens, + max_tokens: state.max_tokens, + }; + + debug!("Forked seq {:?} → {:?}", state.seq_id, child.seq_id); + + Ok(( + SequenceState::Decoding(state), + SequenceState::Decoding(child), + )) + } +} + +/// Complete a sequence +pub struct CompleteEvent<'a> { + pub reason: FinishReason, + pub block_manager: &'a mut BlockManager, +} + +impl<'a> CompleteEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + match state { + SequenceState::Decoding(s) => self.apply_to_decoding(s), + SequenceState::Prefilling(s) => self.apply_to_prefilling(s), + _ => Err(Error::InvalidTransition( + "CompleteEvent requires running state", + )), + } + } + + fn apply_to_decoding(self, state: DecodingState) -> Result { + // Free all blocks + for &block_id in state.blocks.iter() { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!("Completed seq {:?}: {:?}", state.seq_id, self.reason); + + Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: self.reason, + })) + } + + fn apply_to_prefilling(self, state: PrefillingState) -> Result { + // Free all blocks + for &block_id in state.blocks.iter() { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, state.seq_id, e + ); + } + } + + debug!( + "Completed seq {:?} during prefill: {:?}", + state.seq_id, self.reason + ); + + Ok(SequenceState::Finished(FinishedState { + seq_id: state.seq_id, + finish_reason: self.reason, + })) + } +} + +/// Abort a sequence (error or cancellation) +pub struct AbortEvent<'a> { + pub reason: String, + pub block_manager: &'a mut BlockManager, +} + +impl<'a> AbortEvent<'a> { + pub fn apply(self, state: SequenceState) -> Result { + // Can abort from any state + let seq_id = state.seq_id(); + + // Free blocks if any + if let Some(blocks) = state.blocks() { + for &block_id in blocks { + if let Err(e) = self.block_manager.free(block_id) { + warn!( + "Failed to free block {:?} for {:?}: {:?}", + block_id, seq_id, e + ); + } + } + } + + warn!("Aborted seq {:?}: {}", seq_id, self.reason); + + Ok(SequenceState::Aborted(AbortedState { + seq_id, + reason: self.reason, + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + use crate::block_manager::manager::BlockManager; + + #[test] + fn test_schedule_event_waiting_to_prefilling() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let waiting = SequenceState::Waiting(WaitingState { + seq_id: SequenceId(1), + prompt_tokens: 100, + max_tokens: 150, + }); + + let event = ScheduleEvent { + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(waiting); + assert!(result.is_ok()); + + let new_state = result.unwrap(); + assert!(matches!(new_state, SequenceState::Prefilling(_))); + + if let SequenceState::Prefilling(state) = new_state { + assert_eq!(state.seq_id, SequenceId(1)); + assert_eq!(state.tokens_filled, 0); + assert_eq!(state.tokens_total, 100); + // Should allocate ceil(100/16) = 7 blocks + assert_eq!(state.blocks.len(), 7); + } + } + + #[test] + fn test_schedule_event_oom() { + // Small allocator - only 2 blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let waiting = SequenceState::Waiting(WaitingState { + seq_id: SequenceId(1), + prompt_tokens: 100, // Needs 7 blocks + max_tokens: 150, + }); + + let event = ScheduleEvent { + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(waiting); + assert!(matches!(result, Err(Error::OutOfMemory))); + + // Should have cleaned up - no blocks leaked + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_append_tokens_prefilling_to_decoding() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + // Allocate blocks + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let prefilling = SequenceState::Prefilling(PrefillingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + tokens_filled: 0, + tokens_total: 100, + max_tokens: 150, + }); + + let event = AppendTokensEvent { + num_tokens: 100, + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(prefilling).unwrap(); + + // Should transition to Decoding + assert!(matches!(result, SequenceState::Decoding(_))); + + if let SequenceState::Decoding(state) = result { + assert_eq!(state.num_tokens, 100); + assert_eq!(state.blocks.len(), 2); + } + } + + #[test] + fn test_append_tokens_decoding_needs_more_blocks() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![block_manager.allocate().unwrap()]; // 1 block + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 16, // 1 block worth + max_tokens: 150, + }); + + let event = AppendTokensEvent { + num_tokens: 20, // Now need 3 blocks total + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(decoding).unwrap(); + + if let SequenceState::Decoding(state) = result { + assert_eq!(state.num_tokens, 36); + // Should allocate 2 more blocks (total 3) + assert_eq!(state.blocks.len(), 3); + } + } + + #[test] + fn test_append_tokens_decoding_oom_cleanup() { + // Small allocator - only enough for initial blocks + let allocator = Box::new(CpuAllocator::new(2048)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + let initial_block_count = blocks.len(); + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 32, + max_tokens: 150, + }); + + // Try to append tokens that would need more blocks (will fail - OOM) + let event = AppendTokensEvent { + num_tokens: 50, // Would need 6 blocks total + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(decoding); + assert!(matches!(result, Err(Error::OutOfMemory))); + + // Original 2 blocks should still be allocated (not freed by error path) + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, initial_block_count); + } + + #[test] + fn test_append_tokens_reaches_max_tokens() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![block_manager.allocate().unwrap()]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 140, + max_tokens: 150, + }); + + let event = AppendTokensEvent { + num_tokens: 10, + block_manager: &mut block_manager, + tokens_per_block: 16, + }; + + let result = event.apply(decoding).unwrap(); + + // Should transition to Finished + assert!(matches!(result, SequenceState::Finished(_))); + + if let SequenceState::Finished(state) = result { + assert_eq!(state.seq_id, SequenceId(1)); + assert_eq!(state.finish_reason, FinishReason::MaxTokens); + } + + // Blocks should be freed + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_preempt_event() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 32, + max_tokens: 150, + }); + + let event = PreemptEvent { + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + + // Should transition to Preempted + assert!(matches!(result, SequenceState::Preempted(_))); + + if let SequenceState::Preempted(state) = result { + assert_eq!(state.num_tokens, 32); + } + + // Blocks should be freed + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_resume_event() { + let preempted = SequenceState::Preempted(PreemptedState { + seq_id: SequenceId(1), + num_tokens: 32, + max_tokens: 150, + }); + + let event = ResumeEvent; + let result = event.apply(preempted).unwrap(); + + // Should transition back to Waiting + assert!(matches!(result, SequenceState::Waiting(_))); + + if let SequenceState::Waiting(state) = result { + assert_eq!(state.prompt_tokens, 32); + assert_eq!(state.max_tokens, 150); + } + } + + #[test] + fn test_fork_event() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 32, + max_tokens: 150, + }); + + let event = ForkEvent { + child_id: SequenceId(2), + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + let (parent, child) = result; + + // Both should be in Decoding state + assert!(matches!(parent, SequenceState::Decoding(_))); + assert!(matches!(child, SequenceState::Decoding(_))); + + if let (SequenceState::Decoding(p), SequenceState::Decoding(c)) = (parent, child) { + assert_eq!(p.seq_id, SequenceId(1)); + assert_eq!(c.seq_id, SequenceId(2)); + assert_eq!(p.blocks, c.blocks); // Shared blocks + } + } + + #[test] + fn test_complete_event() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![block_manager.allocate().unwrap()]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 100, + max_tokens: 150, + }); + + let event = CompleteEvent { + reason: FinishReason::Stop, + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + + assert!(matches!(result, SequenceState::Finished(_))); + + if let SequenceState::Finished(state) = result { + assert_eq!(state.finish_reason, FinishReason::Stop); + } + + // Blocks should be freed + let stats = block_manager.get_stats(); + let block_stats = stats.get(&BlockType::StandardKV).unwrap(); + assert_eq!(block_stats.allocated_blocks, 0); + } + + #[test] + fn test_fork_duplicate_child_id() { + let allocator = Box::new(CpuAllocator::new(100_000)); + let mut block_manager = BlockManager::new(allocator, 1024); + + let blocks = vec![ + block_manager.allocate().unwrap(), + block_manager.allocate().unwrap(), + ]; + + let decoding = SequenceState::Decoding(DecodingState { + seq_id: SequenceId(1), + blocks: Arc::new(blocks), + num_tokens: 32, + max_tokens: 150, + }); + + // Fork with child_id = SequenceId(2) + let event = ForkEvent { + child_id: SequenceId(2), + block_manager: &mut block_manager, + }; + + let result = event.apply(decoding).unwrap(); + let (parent, child) = result; + + // Verify fork succeeded + if let SequenceState::Decoding(c) = &child { + assert_eq!(c.seq_id, SequenceId(2)); + } + + // Now try to fork again with the SAME child_id (should fail in manager) + // This test verifies the ForkEvent itself works, but the manager + // should check for duplicate child_id before calling this + drop(parent); + drop(child); + } +} diff --git a/src/sequence_manager/fsm_manager.rs b/src/sequence_manager/fsm_manager.rs new file mode 100644 index 0000000..762f93b --- /dev/null +++ b/src/sequence_manager/fsm_manager.rs @@ -0,0 +1,418 @@ +use super::events::*; +use super::fsm_events::*; +use super::states::*; +use crate::block_manager::manager::BlockManager; +use crate::block_manager::types::*; +use std::collections::{HashMap, VecDeque}; +use tracing::{debug, info, warn}; + +/// Sequence Manager with FSM-based state transitions +pub struct SequenceManager { + block_manager: BlockManager, + + // Each sequence is a state machine + sequences: HashMap, + + tokens_per_block: usize, + + // Event-driven + event_rx: SequenceEventReceiver, + + // Scheduling queues + waiting_queue: VecDeque, +} + +impl SequenceManager { + pub fn new( + block_manager: BlockManager, + tokens_per_block: usize, + ) -> (Self, SequenceEventSender) { + assert!(tokens_per_block > 0, "tokens_per_block must be > 0"); + + let (event_tx, event_rx) = create_event_channel(); + let manager = Self { + block_manager, + sequences: HashMap::new(), + tokens_per_block, + event_rx, + waiting_queue: VecDeque::new(), + }; + + (manager, event_tx) + } + + /// Main event loop + pub async fn run(mut self) { + info!("FSM SequenceManager event loop started"); + + while let Some(event) = self.event_rx.recv().await { + match event { + SequenceEvent::AddRequest { + seq_id, + prompt_tokens, + max_tokens, + } => { + self.handle_add_request(seq_id, prompt_tokens, max_tokens); + } + + SequenceEvent::AppendTokens { seq_id, num_tokens } => { + self.handle_append_tokens(seq_id, num_tokens); + } + + SequenceEvent::ForkSequence { + parent_id, + child_id, + } => { + self.handle_fork_sequence(parent_id, child_id); + } + + SequenceEvent::CompleteSequence { seq_id } => { + self.handle_complete_sequence(seq_id); + } + + SequenceEvent::PreemptSequence { seq_id } => { + self.handle_preempt_sequence(seq_id); + } + + SequenceEvent::ResumeSequence { seq_id } => { + self.handle_resume_sequence(seq_id); + } + + SequenceEvent::GetBlocks { seq_id, response } => { + let result = self.get_blocks(seq_id); + let _ = response.send(result); + } + + SequenceEvent::GetStats { response } => { + let stats = self.get_stats(); + let _ = response.send(stats); + } + } + } + + info!("FSM SequenceManager event loop stopped"); + } + + fn handle_add_request(&mut self, seq_id: SequenceId, prompt_tokens: usize, max_tokens: usize) { + debug!( + "Adding request: seq_id={:?}, prompt_tokens={}, max_tokens={}", + seq_id, prompt_tokens, max_tokens + ); + + // Check for duplicate seq_id + if self.sequences.contains_key(&seq_id) { + warn!( + "Duplicate AddRequest for {:?} - sequence already exists, ignoring", + seq_id + ); + return; + } + + let state = SequenceState::Waiting(WaitingState { + seq_id, + prompt_tokens, + max_tokens, + }); + + self.sequences.insert(seq_id, state); + self.waiting_queue.push_back(seq_id); + + // Try to schedule immediately + self.try_schedule(); + } + + fn handle_append_tokens(&mut self, seq_id: SequenceId, num_tokens: usize) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let backup = state.clone(); + let event = AppendTokensEvent { + num_tokens, + block_manager: &mut self.block_manager, + tokens_per_block: self.tokens_per_block, + }; + + match event.apply(state) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + } + Err(Error::OutOfMemory) => { + warn!("OOM while appending tokens to {:?}", seq_id); + self.sequences.insert(seq_id, backup); + self.handle_oom(); + } + Err(e) => { + warn!("Failed to append tokens to {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); + } + } + } + + fn handle_fork_sequence(&mut self, parent_id: SequenceId, child_id: SequenceId) { + // Check if child_id already exists + if self.sequences.contains_key(&child_id) { + warn!( + "Cannot fork {:?} → {:?}: child ID already exists", + parent_id, child_id + ); + return; + } + + let parent_state = match self.sequences.remove(&parent_id) { + Some(s) => s, + None => { + warn!("Parent sequence {:?} not found", parent_id); + return; + } + }; + + let backup = parent_state.clone(); + let event = ForkEvent { + child_id, + block_manager: &mut self.block_manager, + }; + + match event.apply(parent_state) { + Ok((parent_state, child_state)) => { + self.sequences.insert(parent_id, parent_state); + self.sequences.insert(child_id, child_state); + debug!("Forked sequence {:?} → {:?}", parent_id, child_id); + } + Err(e) => { + warn!("Failed to fork {:?}: {:?}", parent_id, e); + self.sequences.insert(parent_id, backup); + } + } + } + + fn handle_complete_sequence(&mut self, seq_id: SequenceId) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let backup = state.clone(); + let event = CompleteEvent { + reason: FinishReason::Stop, + block_manager: &mut self.block_manager, + }; + + match event.apply(state) { + Ok(new_state) => { + info!("Completed sequence {:?}", seq_id); + // Keep finished state for stats/debugging + self.sequences.insert(seq_id, new_state); + + // Try to schedule waiting sequences + self.try_schedule(); + } + Err(e) => { + warn!("Failed to complete {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); + } + } + } + + fn handle_preempt_sequence(&mut self, seq_id: SequenceId) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let backup = state.clone(); + let event = PreemptEvent { + block_manager: &mut self.block_manager, + }; + + match event.apply(state) { + Ok(new_state) => { + info!("Preempted sequence {:?}", seq_id); + self.sequences.insert(seq_id, new_state); + } + Err(e) => { + warn!("Failed to preempt {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); + } + } + } + + fn handle_resume_sequence(&mut self, seq_id: SequenceId) { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + warn!("Sequence {:?} not found", seq_id); + return; + } + }; + + let backup = state.clone(); + let event = ResumeEvent; + + match event.apply(state) { + Ok(new_state) => { + debug!("Resumed sequence {:?}", seq_id); + self.sequences.insert(seq_id, new_state); + self.waiting_queue.push_back(seq_id); + self.try_schedule(); + } + Err(e) => { + warn!("Failed to resume {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); + } + } + } + + fn try_schedule(&mut self) { + while let Some(&seq_id) = self.waiting_queue.front() { + let state = match self.sequences.remove(&seq_id) { + Some(s) => s, + None => { + self.waiting_queue.pop_front(); + continue; + } + }; + + // Only schedule if in Waiting state + if !matches!(state, SequenceState::Waiting(_)) { + self.sequences.insert(seq_id, state); + self.waiting_queue.pop_front(); + continue; + } + + let backup = state.clone(); + let event = ScheduleEvent { + block_manager: &mut self.block_manager, + tokens_per_block: self.tokens_per_block, + }; + + match event.apply(state) { + Ok(new_state) => { + self.sequences.insert(seq_id, new_state); + self.waiting_queue.pop_front(); + info!("Scheduled sequence {:?}", seq_id); + } + Err(Error::OutOfMemory) => { + // Can't schedule - restore state and leave in waiting queue + self.sequences.insert(seq_id, backup); + break; + } + Err(e) => { + warn!("Failed to schedule {:?}: {:?}", seq_id, e); + self.sequences.insert(seq_id, backup); + self.waiting_queue.pop_front(); + } + } + } + } + + fn handle_oom(&mut self) { + warn!("Handling OOM - preemption policy not yet implemented"); + // TODO: Implement preemption policy + // - Find lowest priority sequence + // - Preempt it + // - Retry scheduling + } + + fn get_blocks(&self, seq_id: SequenceId) -> Result> { + match self.sequences.get(&seq_id) { + None => Err(Error::UnknownSequence(seq_id)), + Some(state) => Ok(state.blocks().map(|b| b.to_vec()).unwrap_or_default()), + } + } + + fn get_stats(&self) -> SequenceManagerStats { + let num_running = self.sequences.values().filter(|s| s.is_running()).count(); + + let num_waiting = self.sequences.values().filter(|s| s.is_waiting()).count(); + + let num_preempted = self + .sequences + .values() + .filter(|s| matches!(s, SequenceState::Preempted(_))) + .count(); + + SequenceManagerStats { + num_sequences: self.sequences.len(), + num_running, + num_waiting, + num_preempted, + block_stats: self.block_manager.get_stats(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::block_manager::allocator::CpuAllocator; + + fn create_test_manager() -> (SequenceManager, SequenceEventSender) { + let allocator = Box::new(CpuAllocator::new(1024 * 1024)); // 1MB + let block_manager = BlockManager::new(allocator, 256); // 256 byte blocks + SequenceManager::new(block_manager, 16) // 16 tokens per block + } + + #[test] + fn test_duplicate_add_request_ignored() { + let (mut manager, _tx) = create_test_manager(); + + let seq_id = SequenceId(0); + + // Add first request + manager.handle_add_request(seq_id, 32, 128); + + // Verify it was added + assert!(manager.sequences.contains_key(&seq_id)); + let first_state = manager.sequences.get(&seq_id).cloned(); + + // Try to add duplicate - should be ignored + manager.handle_add_request(seq_id, 64, 256); // Different params + + // Verify the original state is unchanged + let current_state = manager.sequences.get(&seq_id); + assert_eq!(format!("{:?}", first_state), format!("{:?}", current_state)); + + // Verify original params (may have transitioned to Prefilling via try_schedule) + match manager.sequences.get(&seq_id) { + Some(SequenceState::Waiting(s)) => { + assert_eq!(s.prompt_tokens, 32); // Original value + assert_eq!(s.max_tokens, 128); // Original value + } + Some(SequenceState::Prefilling(s)) => { + // Scheduled automatically, check original params + assert_eq!(s.tokens_total, 32); // Original prompt_tokens + assert_eq!(s.max_tokens, 128); // Original value + } + _ => panic!( + "Expected Waiting or Prefilling state, got: {:?}", + manager.sequences.get(&seq_id) + ), + } + } + + #[test] + fn test_add_different_sequences() { + let (mut manager, _tx) = create_test_manager(); + + // Add multiple different sequences - all should succeed + manager.handle_add_request(SequenceId(0), 32, 128); + manager.handle_add_request(SequenceId(1), 64, 256); + manager.handle_add_request(SequenceId(2), 16, 64); + + assert_eq!(manager.sequences.len(), 3); + assert!(manager.sequences.contains_key(&SequenceId(0))); + assert!(manager.sequences.contains_key(&SequenceId(1))); + assert!(manager.sequences.contains_key(&SequenceId(2))); + } +} 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..138bdca --- /dev/null +++ b/src/sequence_manager/mod.rs @@ -0,0 +1,7 @@ +pub mod events; +pub mod fsm_events; +pub mod fsm_manager; +pub mod id_generator; +pub mod states; + +// Re-export public API for convenience diff --git a/src/sequence_manager/states.rs b/src/sequence_manager/states.rs new file mode 100644 index 0000000..2128f1a --- /dev/null +++ b/src/sequence_manager/states.rs @@ -0,0 +1,145 @@ +use crate::block_manager::types::*; +use std::sync::Arc; + +/// FSM States - each state owns its resources +/// +/// State transitions: +/// Waiting → Scheduling → Prefilling → Decoding → Finished +/// ↓ ↓ +/// └→ Preempted ←┘ + +#[derive(Debug, Clone)] +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, Clone)] +pub struct WaitingState { + pub seq_id: SequenceId, + pub prompt_tokens: usize, + pub max_tokens: usize, +} + +/// Scheduling state - in process of allocating blocks +#[derive(Debug, Clone)] +pub struct SchedulingState { + pub seq_id: SequenceId, + pub prompt_tokens: usize, + pub max_tokens: usize, + // Partially allocated blocks (if any) + pub blocks: Arc>, +} + +/// Prefilling state - processing prompt +#[derive(Debug, Clone)] +pub struct PrefillingState { + pub seq_id: SequenceId, + pub blocks: Arc>, // Arc for cheap cloning (O(1)) + pub tokens_filled: usize, + pub tokens_total: usize, + pub max_tokens: usize, +} + +/// Decoding state - generating tokens +#[derive(Debug, Clone)] +pub struct DecodingState { + pub seq_id: SequenceId, + pub blocks: Arc>, // Arc for cheap cloning (O(1)) + pub num_tokens: usize, + pub max_tokens: usize, +} + +/// Preempted state - blocks freed, waiting to resume +#[derive(Debug, Clone)] +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, Clone)] +pub struct FinishedState { + pub seq_id: SequenceId, + pub finish_reason: FinishReason, +} + +/// Aborted state - error or cancellation +#[derive(Debug, Clone)] +pub struct AbortedState { + pub seq_id: SequenceId, + pub reason: String, +} + +#[derive(Debug, Clone, PartialEq)] +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.as_ref()), + SequenceState::Decoding(s) => Some(s.blocks.as_ref()), + SequenceState::Scheduling(s) if !s.blocks.is_empty() => Some(s.blocks.as_ref()), + _ => None, + } + } +}