feat: support event driven for sequences

[kerthcet]

What this PR does / why we need it

Which issue(s) this PR fixes

Fixes #

Special notes for your reviewer

Does this PR introduce a user-facing change?

[Copilot]

Pull request overview

This PR introduces an event-driven, FSM-based sequence lifecycle manager backed by a block allocator/manager, along with documentation and an example demonstrating basic usage. This appears to be foundational infrastructure for managing sequence scheduling, token append, forking, preemption, and completion in a structured way.

Changes:

• Added a block_manager subsystem (allocator + manager + types/errors) to allocate/refcount/free KV-cache-like blocks.
• Added a sequence_manager subsystem implementing an event-driven FSM with typed states and transition events, plus a thread-safe SequenceIdGenerator.
• Added developer-facing documentation (docs/fsm_architecture.md) and an executable example (examples/fsm_usage.rs), and introduced thiserror for error derivations.

Reviewed changes

Copilot reviewed 15 out of 16 changed files in this pull request and generated 12 comments.

Show a summary per file

File	Description
src/sequence_manager/states.rs	Defines FSM state types/variants and helper methods for sequence lifecycle tracking.
src/sequence_manager/mod.rs	Declares the sequence_manager module (and should expose a usable public API surface).
src/sequence_manager/id_generator.rs	Adds a thread-safe sequential SequenceId generator with unit tests.
src/sequence_manager/fsm_manager.rs	Implements the event loop and handlers for applying FSM transitions per sequence.
src/sequence_manager/fsm_events.rs	Defines the event/transition types that transform SequenceState with resource ownership semantics.
src/sequence_manager/events.rs	Defines the event protocol and stats payload used to drive/query the manager.
src/main.rs	Wires new modules into the binary crate.
src/lib.rs	Introduces a library surface re-exporting modules for external use (examples depend on this).
src/block_manager/types.rs	Adds block/sequence IDs, block stats, and thiserror-based error types.
src/block_manager/mod.rs	Declares the block_manager module (and should expose a usable public API surface).
src/block_manager/manager.rs	Implements block allocation, refcounting, free pooling, and stats reporting.
src/block_manager/allocator.rs	Defines the allocator trait and a simple CPU allocator implementation.
examples/fsm_usage.rs	Demonstrates end-to-end usage of the FSM sequence manager via async event sending.
docs/fsm_architecture.md	Documents the FSM approach, state diagram, and usage patterns for contributors.
Cargo.toml	Adds thiserror dependency.
Cargo.lock	Locks thiserror dependency version.

Comments suppressed due to low confidence (1)

src/block_manager/mod.rs:4

• This module is used from the example as puma::block_manager::{BlockManager, CpuAllocator}, but block_manager/mod.rs currently only declares submodules and doesn't re-export those types. As a result, the example (and any downstream crate using the same import style) will not compile.

pub mod allocator;
pub mod manager;
pub mod types;

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 15 out of 16 changed files in this pull request and generated 11 comments.

[kerthcet]

/kind feature

[Copilot]

Pull request overview

Copilot reviewed 22 out of 23 changed files in this pull request and generated 6 comments.

[Copilot]

Pull request overview

Copilot reviewed 22 out of 23 changed files in this pull request and generated 7 comments.

[Copilot]

Pull request overview

Copilot reviewed 35 out of 36 changed files in this pull request and generated 7 comments.

Comments suppressed due to low confidence (1)

src/downloader/mod.rs:34

• Downloader was changed from an async fn to a return-position impl Future, but existing implementors (e.g. src/downloader/huggingface.rs) still use async fn download_model and will no longer satisfy the trait, causing a compilation failure. Update the implementor signatures (return an async move { ... } future) or revert the trait back to async fn (optionally using async-trait).

pub trait Downloader {
    fn download_model(
        &self,
        name: &str,
    ) -> impl std::future::Future<Output = Result<(), DownloadError>> + Send;
}

[Copilot]

Pull request overview

Copilot reviewed 35 out of 36 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (4)

src/sequence_manager/mod.rs:161

• If add_ref fails partway through the loop, any refs already incremented are not rolled back. Restoring the parent sequence state does not undo those BlockManager ref_count increments, so blocks can leak (never reaching ref_count=0). Track which block IDs were successfully add_ref’d and decrement them before returning the error.
  src/backend/llm_engine.rs:303
• yield_now() in a tight infinite loop still results in a high-CPU busy-spin when there are no events and no work scheduled (the task is immediately re-polled). This can waste a full core when the engine is idle; consider sleeping briefly when idle or awaiting on the event channel.

            // 4. Small yield to prevent busy loop
            tokio::task::yield_now().await;

src/fsm/mod.rs:5

• The module docs say the transition logic lives in Scheduler::transition() / Scheduler::transition_*(), but in this PR the transition implementation is in SequenceManager::transition() / SequenceManager::transition_*() (and the scheduler delegates to the sequence manager). This comment will mislead contributors and readers of docs/fsm_architecture.md.

//! The actual transition logic is implemented in `Scheduler::transition()` and
//! private `Scheduler::transition_*()` methods.

src/backend/llm_engine.rs:97

• GenerateResponse.completion_tokens is always returned as 0, which makes the API’s usage fields incorrect (and total_tokens under-counted). Since the scheduler currently returns only the final text, you can compute the completion token count by re-encoding the completion text with the same tokenizer.

        Ok(GenerateResponse {
            text,
            prompt_tokens,
            completion_tokens: 0,
        })

[Copilot]

Pull request overview

Copilot reviewed 35 out of 36 changed files in this pull request and generated 4 comments.

Comments suppressed due to low confidence (4)

src/downloader/mod.rs:33

• The Downloader trait method signature was changed to return impl Future, but existing implementors (e.g. HuggingFaceDownloader) still define async fn download_model(...), which will not satisfy this trait and will break compilation. Either keep async fn in the trait, or update all implementors to the new signature in the same PR.

pub trait Downloader {
    fn download_model(
        &self,
        name: &str,
    ) -> impl std::future::Future<Output = Result<(), DownloadError>> + Send;

src/backend/llm_engine.rs:77

• EngineHandle::generate ignores the requested model argument, so callers can request one model but get responses from whatever model this handle was constructed with. This should validate the model name (or remove the parameter) to avoid serving the wrong model.

        _model: &str,

src/backend/llm_engine.rs:307

• EngineRunner::serve runs an infinite loop that always calls yield_now(), even when there is no work and no incoming events. This will spin the CPU at 100% in idle. Consider awaiting event_rx.recv() when idle (or using select! with a small sleep) to block efficiently.

            // 4. Small yield to prevent busy loop
            tokio::task::yield_now().await;

src/backend/llm_engine.rs:97

• EngineHandle::generate always returns completion_tokens: 0, which makes API usage accounting incorrect for both /v1/completions and /v1/chat/completions. At minimum, compute the completion token count from the returned text using the same tokenizer.

        Ok(GenerateResponse {
            text,
            prompt_tokens,
            completion_tokens: 0,
        })

[kerthcet]

/lgtm

[kerthcet]

/hold

[Copilot]

Pull request overview

Copilot reviewed 35 out of 36 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (4)

src/backend/llm_engine.rs:97

• EngineHandle::generate always returns completion_tokens: 0, which makes API usage.completion_tokens incorrect (and total_tokens undercounts). Since the scheduler/runner currently only returns the decoded text, the actual completion token count is being dropped.

        Ok(GenerateResponse {
            text,
            prompt_tokens,
            completion_tokens: 0,
        })

src/backend/llm_engine.rs:233

• Streaming path advances the FSM by 1 token per streamed completion token before marking prefill complete. This causes completion tokens to be counted as prefilling progress, delaying the Prefilling→Decoding transition and preventing KV block allocation/max_tokens accounting from lining up with the actual completion stream.

        let mut completion_ids: Vec<TokenId> = Vec::new();
        let mut sent_len = 0usize; // bytes of the decoded completion already sent
        while let Some(token_id) = stream.next().await {

src/backend/llm_engine.rs:260

• After streaming completes, the code calls append_tokens(seq_id, prompt_len) again. This double-counts the prompt tokens (they should only be applied once to finish prefill) and can push the sequence over max_tokens or trigger unexpected transitions/freeing.

        // Advance FSM state (prefill → decode → finished) so blocks are freed,
        // then close the stream channel.
        self.scheduler.append_tokens(seq_id, prompt_len);
        self.scheduler
            .complete_sequence(seq_id, FinishReason::Stop, String::new());

src/backend/llm_engine.rs:309

• The engine loop unconditionally yield_now()s every iteration, even when there are no events and no scheduled work. This results in a tight idle spin that can burn CPU while the server is otherwise idle.

            // 4. Small yield to prevent busy loop
            tokio::task::yield_now().await;

[kerthcet]

/unhold

[kerthcet]

/lgtm

[InftyAI-Agent]

Approved: PR has both lgtm and approved labels

[Copilot]

Pull request overview

Copilot reviewed 35 out of 36 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (6)

src/backend/llm_engine.rs:97

• GenerateResponse.completion_tokens is always returned as 0, so API usage.completion_tokens/total_tokens is incorrect for every non-streaming request.

        Ok(GenerateResponse {
            text,
            prompt_tokens,
            completion_tokens: 0,
        })

src/backend/llm_engine.rs:84

• The model parameter is currently ignored in EngineHandle::generate, so callers can request any model name and still get a response from whatever model this handle was constructed with. This makes the API contract ambiguous and can hide routing bugs.

        let token_ids = self.tokenize(prompt)?;
        let prompt_tokens = token_ids.len();

src/backend/llm_engine.rs:112

• The model parameter is currently ignored in EngineHandle::generate_stream, so the streaming API can return output for a different model than requested. This should be validated (or removed from the public signature) to keep model routing explicit.

        let token_ids = self.tokenize(prompt)?;

        let (response_tx, mut response_rx) = tokio::sync::mpsc::unbounded_channel();
        self.send_request(token_ids, max_tokens, ResponseSender::Stream(response_tx))?;

src/backend/llm_engine.rs:209

• In streaming mode, the scheduler remains in Prefilling when the backend starts emitting completion tokens, but append_tokens(seq_id, 1) is called for each completion token. This causes the first prompt_len completion tokens to be counted as prompt prefill progress, shifting the state machine and KV allocation accounting.

        let prompt_len = token_ids.len();
        let mut stream = match self
            .backend
            .generate_stream((*token_ids).clone(), max_tokens, 0.0)
            .await

src/backend/llm_engine.rs:260

• After streaming finishes, append_tokens(seq_id, prompt_len) is called again even though prompt prefill should already have been accounted for. This double-counts prompt tokens in the FSM and can trigger premature MaxTokens finishing / extra KV allocations.

        // Advance FSM state (prefill → decode → finished) so blocks are freed,
        // then close the stream channel.
        self.scheduler.append_tokens(seq_id, prompt_len);
        self.scheduler
            .complete_sequence(seq_id, FinishReason::Stop, String::new());

src/backend/llm_engine.rs:186

• forward_single completes the sequence without ever appending the generated completion token count into the scheduler/FSM. That means KV-block growth and max-token enforcement (for total tokens) never reflects the actual completion length, which will matter once the backend is real and OOM/preemption policies depend on accurate accounting.

        // Advance FSM state (prefill → decode → finished) so blocks are freed,
        // then deliver the result to the client channel.
        self.scheduler.append_tokens(seq_id, prompt_len);
        self.scheduler
            .complete_sequence(seq_id, FinishReason::Stop, text);