From 98eac6372f9641c62fdf903873b15b94f8b141b8 Mon Sep 17 00:00:00 2001 From: Streaky Date: Thu, 6 Aug 2026 02:55:02 +0100 Subject: [PATCH 1/7] Define Phase 8 scheduling baseline --- docs/outline.md | 78 ++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 77 insertions(+), 1 deletion(-) diff --git a/docs/outline.md b/docs/outline.md index d627a23..7d62f7b 100644 --- a/docs/outline.md +++ b/docs/outline.md @@ -1607,7 +1607,83 @@ Harden request ordering and make resource decisions explainable under sustained - report model residency, executor-slot occupancy, context placement, queue state and age, capacity reservations, transition costs, eviction and unload reasons, and scheduler decisions; - validate multi-model and mixed-context pressure, repeated load/unload cycles, cancellation storms, deadline expiry, and sustained operation against capacity, fairness, leak, and latency gates; - propagate one transport correlation identifier into canonical requests and scheduler/executor work while naming any distinct inference operation identifier explicitly; report response-header, first-event, and terminal durations separately instead of presenting header latency as whole-request duration; -- make routine diagnostics operationally non-interfering and machine-usable: encode JSON and SSE bodies structurally where possible, move record emission off the response-polling path through a bounded asynchronous sink, and expose overflow or loss rather than silently dropping records. If exact unredacted capture deliberately backpressures that sink, make the latency tradeoff operator-visible and enforce restricted sink access and retention alongside unmistakable credential and content warnings. +- make routine diagnostics operationally non-interfering and machine-usable: encode JSON and SSE bodies structurally where possible, move record emission off the response-polling path through a bounded asynchronous sink, expose overflow or loss rather than silently dropping records, and never let either privacy-safe or fully unredacted capture backpressure inference; fully unredacted capture must enforce restricted sink access and retention alongside unmistakable credential and content warnings. + +The initial Phase 8 scheduler is deliberately a replaceable fairness baseline, +not the final locality- or cost-aware planner. Rust should expose request +ordering as a policy over protocol-neutral runnable operations, with residency, +context placement, and transition estimates supplied as observations rather +than embedded as policy state. The baseline policy is priority-aware deficit +round robin with monotonic age promotion. It charges bounded work quanta, +preserves FIFO order among otherwise equivalent requests, and eventually +promotes every continuously runnable request. Class weights, deficit refill, +prefill width, and promotion intervals are versioned scheduler-policy +parameters recorded by status and proof artifacts, so measurements may tune +them without changing the execution-session or application APIs. Later +locality-aware planning may replace this policy behind the same boundary, but +may not weaken its explicit priority, age, deadline, tenant-isolation, or +starvation invariants. + +The baseline has three protocol-neutral request classes: `interactive`, +`standard`, and `batch`; `standard` is the default. Phase 8 does not expose +client-selected elevation through any external adapter. Adapters assign +`standard`, while controlled model-free and acceptance workloads may assign +other classes internally. The canonical request nevertheless records a typed +class and its trusted source so a post-v1 authorization policy could permit +client selection without changing scheduler or executor interfaces. No +transport field should be accepted and silently trusted before that policy +exists. Fairness accounting is defined per authenticated principal, with +request FIFO order inside an otherwise equivalent principal/class queue; +anonymous operation naturally has one principal. + +Scheduling requires resumable execution rather than a synchronous whole-request +`generate` call. A request-owned execution session retains sampler, +incremental UTF-8/stop frontier, usage, deadline, and unpublished successor +state. Tokenization is cancellation-aware preparation but is not interleaved. +Uncached prefill runs in bounded token chunks, and each decode quantum samples +and renders exactly one token. Activation and transfer are scheduler-owned +operations immediately preceding the native quantum that needs them. A session +holds a native slot only while activation, prefill, or decode work and its +native completion fence are outstanding; it may be suspended and later +reacquire a compatible slot at the resulting safe boundary. Cancellation +during native work requests the native abort, waits for the ownership fence, +and discards unpublished successor state. Logical context publication remains +one terminal transactional operation and never occurs at an intermediate +quantum. + +Scheduler diagnostics use distinct identifiers for distinct lifetimes: a +transport correlation ID names one adapter request, an inference-operation ID +names the protocol-neutral inference, and an execution-session ID names its +resumable scheduled state. Durable context IDs and immutable model epochs +remain separate. Every decision record includes these applicable IDs, the +principal and class, queue age and age promotion, quantum kind and charged +work, model and context placement, executor-slot occupancy, transition cost +and capacity reservations, cancellation/deadline state, and a machine-readable +selection or rejection reason. + +Diagnostic emission never backpressures inference in the Phase 8 baseline, +including in fully unredacted mode. Records move through a bounded asynchronous +sink whose configured capacity is published in status and proof output. On +overflow it drops a complete record, increments exact per-kind and total loss +counters, and emits a later loss summary when capacity returns; it never +silently truncates a record. The default stderr sink retains no in-process +history. Fully unredacted mode changes disclosure only, requires the existing +explicit warnings and controlled access, and leaves durable retention to an +operator-selected restricted sink rather than an implicit server buffer. + +The Phase 8 acceptance workload and thresholds are versioned proof inputs, not +constants inferred after a mixed-load result is known. The first vertical slice +must check in deterministic model-free fairness and cancellation fixtures plus +the real-GPU workload definition, measure isolated first-event and per-quantum +baselines on the declared model, executor, configuration, and GPU, and freeze +relative latency thresholds before running the sustained mixed workload. +Starvation is gated primarily by a maximum number of scheduler rounds while +wall-clock queue and first-event latency are reported separately. The mixed +artifact records class/principal/request mix, prompt and generation sizes, +arrival pattern, duration, cancellation and deadline injection, capacity +recovery, diagnostic loss, and all thresholds. Tuning those policy parameters +after observing a working system requires a new versioned fixture and evidence; +it does not require changing the resumable execution contract. The server already includes one narrow operator-diagnostic slice toward this phase: HTTP debug records correlate request metadata, terminal status, From e67905f480103cae0b4e84e0d96d15a541181deb Mon Sep 17 00:00:00 2001 From: Streaky Date: Thu, 6 Aug 2026 03:51:37 +0100 Subject: [PATCH 2/7] Implement Phase 8 workload scheduling --- AGENTS.md | 4 +- Cargo.lock | 1 + README.md | 43 +- compose.test.yaml | 24 + config/phase8-workload.json | 114 ++++ crates/cli/Cargo.toml | 1 + crates/cli/src/main.rs | 606 +++++++++++++++++- crates/executor/src/lib.rs | 5 + crates/server/src/lib.rs | 404 +++++++++--- crates/server/src/mapped.rs | 604 +++++++++++------- crates/server/src/residency.rs | 305 ++++++--- crates/server/src/scheduler.rs | 1093 ++++++++++++++++++++++++++++++++ docs/outline.md | 31 +- tools/phase8-report.sh | 25 + tools/report-phase8.py | 64 ++ 15 files changed, 2913 insertions(+), 411 deletions(-) create mode 100644 config/phase8-workload.json create mode 100644 crates/server/src/scheduler.rs create mode 100755 tools/phase8-report.sh create mode 100755 tools/report-phase8.py diff --git a/AGENTS.md b/AGENTS.md index 1c8cf43..1830a97 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -10,7 +10,7 @@ Keep this `AGENTS.md` up to date whenever development workflows, architecture, s ## Current state -Phases 1 through 7 are implemented. The real Gemma executor proof demonstrated exact checkpoint continuation. The logical context store owns shared token branches and transactional mappings. The physical manager owns tier capacity, representations, transfers, bindings, prepared transitions, eviction, observability, and committed device block tables. The server adds durable opaque external context IDs, atomic state recovery, model lifecycle APIs, bounded transition-cost scheduling, shared incremental streaming, canonical usage, authentication policy, OpenAI completion/chat adapters, and checked OpenAPI. Phase 5 adds transactional llama.cpp sequence mappings, reference-only activation, graph-reuse and byte-copy metrics, and a real-GPU staged-versus-mapped proof. Phase 6 adds the persistent mapped live executor, bounded lifecycle controls, and its real-GPU acceptance report. Phase 7 adds executor-reported operating points, capacity-admitted multi-model residency, transactional immutable model epochs, active-reference draining, pressure-driven idle LRU eviction, bounded local mapped-context spill/restore, and residency observability. +Phases 1 through 8 are implemented. The real Gemma executor proof demonstrated exact checkpoint continuation. The logical context store owns shared token branches and transactional mappings. The physical manager owns tier capacity, representations, transfers, bindings, prepared transitions, eviction, observability, and committed device block tables. The server adds durable opaque external context IDs, atomic state recovery, model lifecycle APIs, bounded transition-cost scheduling, shared incremental streaming, canonical usage, authentication policy, OpenAI completion/chat adapters, and checked OpenAPI. Phase 5 adds transactional llama.cpp sequence mappings, reference-only activation, graph-reuse and byte-copy metrics, and a real-GPU staged-versus-mapped proof. Phase 6 integrates one persistent mapped Gemma executor with bounded admission, request-owned incremental generation, cancellation/deadline propagation, graceful shutdown/restart, and a real-GPU acceptance matrix. Phase 7 adds capacity-admitted multi-model residency, transactional epoch lifecycle, active-reference draining, idle LRU pressure eviction, bounded native-state spill/restore, and real-GPU lifecycle evidence. Phase 8 adds resumable execution sessions, priority-aware deficit round-robin with monotonic age promotion, trusted scheduling metadata, bounded non-blocking scheduler diagnostics, and a versioned real-GPU mixed-workload proof. The repository currently contains: @@ -25,7 +25,7 @@ The repository currently contains: - `executor`: upstream and patch metadata; - `tools`: fetch, verification, integration-report, and coverage helpers. -The server dynamically admits and reuses multiple model epochs within configured device, host, and storage budgets. Each resident model currently owns one native execution slot, so requests for the same model serialize at that slot while distinct resident models can execute independently. HTTP transport diagnostics default off and support privacy-safe and fully unredacted levels through `--http-debug` or `CUSCO_HTTP_DEBUG`; full mode exposes headers, query values, credentials, and body content. Phase 8 workload scheduling and operational hardening is next; this diagnostic slice alone does not implement that phase. Broader compatibility and production hardening remain later work and must not be represented as implemented. +The server dynamically admits and reuses multiple model epochs within configured device, host, and storage budgets. Each resident model currently owns one native execution slot. The Phase 8 scheduler interleaves request-owned prefill and decode quanta with per-principal/class fairness within that slot, while distinct resident models can execute independently. Suspended sessions retain their model reference but release slot occupancy between native quanta; cancellation, deadline, and shutdown callbacks remain registered through each native ownership fence. HTTP transport diagnostics default off and support privacy-safe and fully unredacted levels through `--http-debug` or `CUSCO_HTTP_DEBUG`; full mode exposes headers, query values, credentials, and body content. Phase 9 compatibility, configuration, persistence, and production packaging is next. Broader model compatibility and later optimization remain prospective and must not be represented as implemented. The current unprefixed `/v1` completion and chat routes are minimal adapters, not the complete Phase 9 OpenAI compatibility profile. They do not yet preserve the full chat, tool, and stream-option semantics or emit OpenAI-native streaming chunks; unsupported compatibility input must not be represented as supported. diff --git a/Cargo.lock b/Cargo.lock index 7eb5db0..a438862 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -319,6 +319,7 @@ dependencies = [ "cusco-executor", "cusco-model-registry", "cusco-server", + "serde", "serde_json", "tokio", ] diff --git a/README.md b/README.md index e34c056..0a7a01d 100644 --- a/README.md +++ b/README.md @@ -6,19 +6,21 @@ The project separates responsibilities deliberately: Rust will manage logical co ## Current state -Cusco has completed Phases 1 through 7. The executor proof established exact +Cusco has completed Phases 1 through 8. The executor proof established exact checkpoint continuation for a hybrid/recurrent Gemma model, and the Rust layers now provide durable logical contexts, capacity-accounted physical state, transactional mapped activation, an authenticated HTTP API, immutable model -epochs, bounded live inference, and dynamic model residency. - -Phase 7 replaces the one-model process with a capacity-admitted residency -scheduler. Executor-reported operating points account model weights, context -capacity, and device/host placement; model loads, epoch reloads, retirement, -and unload use transactional publication and active-reference draining. -Pressure selects idle LRU victims, inactive mapped contexts can spill through -the native sequence-state ABI and restore exactly, and `/native/status` -reports configured budgets, resident epochs, and lifecycle/tier metrics. +epochs, bounded live inference, dynamic model residency, and resumable +priority-aware workload scheduling. + +Phase 8 schedules bounded prefill and one-token decode quanta with +priority-aware deficit round robin, per-principal fairness, FIFO ordering, and +monotonic age promotion. Request-owned execution sessions preserve unpublished +successor state across quanta, while bounded asynchronous scheduler diagnostics +attribute decisions without backpressuring inference. The checked-in versioned +mixed workload gates fairness, starvation, cancellation, deadlines, capacity +recovery, diagnostic loss, and relative first-event and per-quantum latency on +the declared real model and GPU. ## Run the real-model inference integration test @@ -86,6 +88,21 @@ inference. Machine-readable GPU, epoch, resident-set, capacity, lifecycle, and latency evidence is written to `results/phase7-server.json`. +## Run the Phase 8 scheduler report + +With the validation GGUF and NVIDIA runtime available, run: + +```sh +CUSCO_GPU_DEVICE_ID=0 tools/phase8-report.sh +``` + +The workflow runs the GPU-less 80%-per-file coverage gate and the versioned +real-GPU workload in `config/phase8-workload.json`. It records the isolated +baseline, mixed interactive/standard/batch results, cancellation and deadline +injections, post-workload capacity recovery, scheduler policy and counters, +fully drained attributed decision records, latency/starvation thresholds, build +and workload provenance, and selected GPU in `results/phase8-server.json`. + ## Run the production Compose service Place the initial model at `data/models/gemma-4-e2b-it.gguf`, then provide an @@ -142,10 +159,10 @@ implicitly fetch models. ## Future goals -Development now proceeds from dynamic residency toward: +Development now proceeds from measured workload scheduling toward: -- workload scheduling and operational hardening; -- broader protocol and model compatibility, recovery, and deployment behavior; +- compatibility, persistence, daemon configuration, and production packaging; +- broader model compatibility and later execution-policy optimization; - optional semantic context compaction once the underlying state system is proven reliable. Each stage is intended to remain gated by correctness and measurable capacity results. The full design and phased acceptance criteria are documented in `docs/outline.md`. diff --git a/compose.test.yaml b/compose.test.yaml index 5c5d75d..a742b7d 100644 --- a/compose.test.yaml +++ b/compose.test.yaml @@ -82,6 +82,30 @@ services: - cargo-target:/work/target deploy: { resources: { reservations: { devices: [{ driver: nvidia, device_ids: ["${CUSCO_GPU_DEVICE_ID:-1}"], capabilities: [gpu] }] } } } + phase8-proof: + build: *test-build + command: + - sh + - -c + - >- + umask 000 && chmod a+rwx /results && + nvidia-smi --query-gpu=uuid,name,compute_cap,driver_version + --format=csv,noheader > /results/phase8-gpu.csv && + cargo run -p cusco -- scheduler-proof + /models/gemma-4-e2b-it.gguf + --workload /work/config/phase8-workload.json + --output /results/phase8-server.json + environment: + CUSCO_REQUESTED_HOST_GPU: "${CUSCO_GPU_DEVICE_ID:-1}" + LD_LIBRARY_PATH: "/opt/llama-build/bin" + volumes: + - "${CUSCO_MODEL_DIR:-./models}:/models:ro" + - "${CUSCO_RESULT_DIR:-./results}:/results" + - cargo-registry:/root/.cargo/registry + - cargo-git:/root/.cargo/git + - cargo-target:/work/target + deploy: { resources: { reservations: { devices: [{ driver: nvidia, device_ids: ["${CUSCO_GPU_DEVICE_ID:-1}"], capabilities: [gpu] }] } } } + volumes: cargo-registry: cargo-git: diff --git a/config/phase8-workload.json b/config/phase8-workload.json new file mode 100644 index 0000000..bf9c623 --- /dev/null +++ b/config/phase8-workload.json @@ -0,0 +1,114 @@ +{ + "version": 2, + "name": "gemma-4-e2b-it-mixed-scheduler-v2", + "model_family": "gemma-4-e2b-it", + "policy": { + "version": 1, + "interactive_weight": 4, + "standard_weight": 2, + "batch_weight": 1, + "promotion_rounds": 8, + "deficit_refill": 1, + "prefill_tokens": 32, + "diagnostic_capacity": 2048 + }, + "baseline": { + "id": "isolated-standard", + "principal": "baseline", + "class": "standard", + "prompt": "Summarize why transactional state publication matters. ", + "prompt_repetitions": 4, + "max_tokens": 6, + "arrival_delay_ms": 0, + "expected": "complete" + }, + "mixed": [ + { + "id": "interactive-a-1", + "principal": "tenant-a", + "class": "interactive", + "prompt": "Give one concise scheduling invariant. ", + "prompt_repetitions": 1, + "max_tokens": 6, + "arrival_delay_ms": 0, + "expected": "complete" + }, + { + "id": "interactive-b-1", + "principal": "tenant-b", + "class": "interactive", + "prompt": "Name one cancellation safety rule. ", + "prompt_repetitions": 1, + "max_tokens": 6, + "arrival_delay_ms": 2, + "expected": "complete" + }, + { + "id": "standard-a-1", + "principal": "tenant-a", + "class": "standard", + "prompt": "Explain bounded fairness for an inference scheduler. ", + "prompt_repetitions": 4, + "max_tokens": 8, + "arrival_delay_ms": 0, + "expected": "complete" + }, + { + "id": "standard-c-cancel", + "principal": "tenant-c", + "class": "standard", + "prompt": "Generate a response that will be cancelled after its first token. ", + "prompt_repetitions": 3, + "max_tokens": 8, + "arrival_delay_ms": 1, + "expected": "cancel" + }, + { + "id": "standard-d-deadline", + "principal": "tenant-d", + "class": "standard", + "prompt": "Generate a response whose active deadline expires after its first token. ", + "prompt_repetitions": 3, + "max_tokens": 8, + "arrival_delay_ms": 1, + "expected": "deadline" + }, + { + "id": "batch-a-1", + "principal": "tenant-a", + "class": "batch", + "prompt": "Describe transactional model-state ownership and safe publication. ", + "prompt_repetitions": 12, + "max_tokens": 8, + "arrival_delay_ms": 0, + "expected": "complete" + }, + { + "id": "batch-b-1", + "principal": "tenant-b", + "class": "batch", + "prompt": "Describe checkpoint restoration, spill accounting, and capacity recovery. ", + "prompt_repetitions": 12, + "max_tokens": 8, + "arrival_delay_ms": 0, + "expected": "complete" + }, + { + "id": "batch-e-1", + "principal": "tenant-e", + "class": "batch", + "prompt": "Explain why scheduler priority and model residency must remain separate decisions. ", + "prompt_repetitions": 10, + "max_tokens": 8, + "arrival_delay_ms": 3, + "expected": "complete" + } + ], + "thresholds": { + "max_queue_age_rounds": 128, + "max_first_event_baseline_multiplier": 20, + "max_first_event_additive_ms": 5000, + "max_quantum_baseline_multiplier": 20, + "max_quantum_additive_ms": 1000 + } +} diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index 28138cc..7ef921b 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -10,5 +10,6 @@ clap.workspace = true cusco-executor = { path = "../executor" } cusco-model-registry = { path = "../model-registry" } cusco-server = { path = "../server" } +serde.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 6fe8237..b080033 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -2,8 +2,16 @@ use anyhow::{Context, Result, ensure}; use clap::{Parser, Subcommand, ValueEnum}; use cusco_executor::{Executor, logits_identical}; use cusco_model_registry::{GEMMA_URI, ModelRecord, fetch_hf, register_local}; -use serde_json::json; -use std::{fs, net::SocketAddr, path::PathBuf, sync::Arc, time::Instant}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use std::{ + fs, + net::SocketAddr, + path::PathBuf, + sync::{Arc, Barrier, Mutex}, + thread, + time::{Duration, Instant}, +}; #[derive(Parser)] struct Args { @@ -64,6 +72,22 @@ enum Command { #[arg(long, default_value = "/results/phase5.json")] output: PathBuf, }, + /// Run the versioned real-model Phase 8 scheduler acceptance workload. + SchedulerProof { + model: PathBuf, + #[arg(long, default_value = "/work/config/phase8-workload.json")] + workload: PathBuf, + #[arg(long, default_value = "/results/phase8-server.json")] + output: PathBuf, + #[arg(long, default_value_t = 4096)] + context: u32, + #[arg(long, default_value_t = 99)] + gpu_layers: i32, + #[arg(long, default_value_t = 8_589_934_592)] + device_bytes: usize, + #[arg(long, default_value_t = 17_179_869_184)] + host_bytes: usize, + }, Serve { model: PathBuf, #[arg(long, default_value = "gemma-4-e2b-it")] @@ -114,6 +138,20 @@ enum Command { active_time_ms: u64, #[arg(long, default_value_t = 8)] stream_buffer: usize, + #[arg(long, default_value_t = 32)] + scheduler_prefill_tokens: usize, + #[arg(long, default_value_t = 8)] + scheduler_interactive_weight: u32, + #[arg(long, default_value_t = 4)] + scheduler_standard_weight: u32, + #[arg(long, default_value_t = 1)] + scheduler_batch_weight: u32, + #[arg(long, default_value_t = 64)] + scheduler_promotion_rounds: u64, + #[arg(long, default_value_t = 1)] + scheduler_deficit_refill: u32, + #[arg(long, default_value_t = 1024)] + scheduler_diagnostic_capacity: usize, /// HTTP transport diagnostics: off, privacy-safe, or fully unredacted. #[arg( long, @@ -165,6 +203,23 @@ fn run(command: Command) -> Result<()> { prefix, output, } => mapped_proof(model, context, gpu_layers, &prefix, output)?, + Command::SchedulerProof { + model, + workload, + output, + context, + gpu_layers, + device_bytes, + host_bytes, + } => scheduler_proof( + model, + workload, + output, + context, + gpu_layers, + device_bytes, + host_bytes, + )?, Command::Serve { model, model_id, @@ -191,12 +246,19 @@ fn run(command: Command) -> Result<()> { wall_time_ms, active_time_ms, stream_buffer, + scheduler_prefill_tokens, + scheduler_interactive_weight, + scheduler_standard_weight, + scheduler_batch_weight, + scheduler_promotion_rounds, + scheduler_deficit_refill, + scheduler_diagnostic_capacity, http_debug, shutdown_grace_ms, } => { use cusco_server::{ AnonymousAdmin, AuthProvider, BearerAuth, ModelRecord, ResidencyConfig, - ResidentEngine, Server, ServerConfig, + ResidentEngine, SchedulerPolicyConfig, Server, ServerConfig, WorkloadScheduler, }; let anonymous = bearer_token.is_none(); let auth: Arc = match bearer_token { @@ -216,6 +278,19 @@ fn run(command: Command) -> Result<()> { }, spill_directory, )?; + let engine = WorkloadScheduler::new( + engine, + SchedulerPolicyConfig { + version: 1, + prefill_tokens: scheduler_prefill_tokens, + interactive_weight: scheduler_interactive_weight, + standard_weight: scheduler_standard_weight, + batch_weight: scheduler_batch_weight, + promotion_rounds: scheduler_promotion_rounds, + deficit_refill: scheduler_deficit_refill, + diagnostic_capacity: scheduler_diagnostic_capacity, + }, + )?; let server = Server::open(state, auth, engine)?; server.configure(ServerConfig { active_requests, @@ -269,6 +344,395 @@ fn run(command: Command) -> Result<()> { } Ok(()) } +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct SchedulerProofWorkload { + version: u32, + name: String, + model_family: String, + policy: cusco_server::SchedulerPolicyConfig, + baseline: SchedulerProofCase, + mixed: Vec, + thresholds: SchedulerProofThresholds, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct SchedulerProofCase { + id: String, + principal: String, + class: cusco_server::SchedulingClass, + prompt: String, + prompt_repetitions: usize, + max_tokens: usize, + arrival_delay_ms: u64, + expected: SchedulerProofOutcome, +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +enum SchedulerProofOutcome { + Complete, + Cancel, + Deadline, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct SchedulerProofThresholds { + max_queue_age_rounds: u64, + max_first_event_baseline_multiplier: u128, + max_first_event_additive_ms: u128, + max_quantum_baseline_multiplier: u128, + max_quantum_additive_ms: u128, +} + +#[derive(Debug, Serialize)] +struct SchedulerProofResult { + id: String, + principal: String, + class: cusco_server::SchedulingClass, + expected: SchedulerProofOutcome, + observed: String, + first_event_ms: Option, + total_ms: u128, + token_wait_ms: Vec, +} + +#[allow(clippy::too_many_arguments)] +fn scheduler_proof( + model_path: PathBuf, + workload_path: PathBuf, + output: PathBuf, + n_ctx: u32, + gpu_layers: i32, + device_bytes: usize, + host_bytes: usize, +) -> Result<()> { + use cusco_server::{MappedEngine, WorkloadScheduler}; + + let proof_started = Instant::now(); + let workload_bytes = fs::read(&workload_path) + .with_context(|| format!("read workload {}", workload_path.display()))?; + let workload: SchedulerProofWorkload = + serde_json::from_slice(&workload_bytes).context("parse scheduler workload")?; + ensure!( + matches!(workload.version, 1 | 2), + "unsupported scheduler workload version" + ); + ensure!(!workload.mixed.is_empty(), "mixed workload is empty"); + ensure!( + workload + .mixed + .iter() + .any(|case| case.expected == SchedulerProofOutcome::Cancel), + "mixed workload has no cancellation injection" + ); + ensure!( + workload + .mixed + .iter() + .any(|case| case.expected == SchedulerProofOutcome::Deadline), + "mixed workload has no deadline injection" + ); + + let model_size = if model_path.to_string_lossy().starts_with("mock://") { + 1 + } else { + fs::metadata(&model_path) + .with_context(|| format!("stat model {}", model_path.display()))? + .len() + }; + let engine = MappedEngine::open( + &workload.model_family, + &model_path, + n_ctx, + gpu_layers, + device_bytes, + host_bytes, + ) + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + let model = cusco_server::ModelRecord { + id: "phase8-proof-model".into(), + revision: "proof".into(), + path: model_path.clone(), + sha256: "verified-by-proof-wrapper".into(), + aliases: Vec::new(), + family: workload.model_family.clone(), + size_bytes: model_size, + epoch: 1, + }; + let diagnostics = Arc::new(Mutex::new(Vec::::new())); + let diagnostic_rows = diagnostics.clone(); + let scheduler = WorkloadScheduler::new_with_diagnostics( + engine.clone(), + workload.policy, + Some(move |line: &str| { + if let Ok(value) = serde_json::from_str(line) { + diagnostic_rows.lock().expect("diagnostic lock").push(value); + } + }), + ) + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + + let baseline = run_scheduler_proof_case( + scheduler.clone(), + model.clone(), + workload.baseline.clone(), + None, + ); + ensure!( + baseline.observed == "complete", + "isolated baseline did not complete" + ); + let baseline_first = baseline + .first_event_ms + .context("isolated baseline emitted no token")?; + let baseline_quantum = percentile(&baseline.token_wait_ms, 95).max(1); + + let barrier = Arc::new(Barrier::new(workload.mixed.len() + 1)); + let mut workers = Vec::with_capacity(workload.mixed.len()); + for case in workload.mixed.clone() { + let scheduler = scheduler.clone(); + let model = model.clone(); + let barrier = barrier.clone(); + workers.push(thread::spawn(move || { + barrier.wait(); + if case.arrival_delay_ms > 0 { + thread::sleep(Duration::from_millis(case.arrival_delay_ms)); + } + run_scheduler_proof_case(scheduler, model, case, None) + })); + } + barrier.wait(); + let mixed = workers + .into_iter() + .map(|worker| worker.join().expect("scheduler proof worker panicked")) + .collect::>(); + + ensure!( + wait_for_scheduler_idle(&scheduler, Duration::from_secs(5)), + "scheduler did not become idle after the mixed workload" + ); + let before_recovery_status = scheduler.status(); + let before_recovery_metrics = engine.metrics(); + let recovery_case = SchedulerProofCase { + id: "capacity-recovery".into(), + principal: "recovery".into(), + class: cusco_server::SchedulingClass::Interactive, + prompt: "Capacity recovered.".into(), + prompt_repetitions: 1, + max_tokens: 2, + arrival_delay_ms: 0, + expected: SchedulerProofOutcome::Complete, + }; + let recovery = run_scheduler_proof_case(scheduler.clone(), model, recovery_case, None); + ensure!( + wait_for_scheduler_idle(&scheduler, Duration::from_secs(5)), + "scheduler did not become idle after the recovery request" + ); + ensure!( + scheduler.flush_diagnostics(Duration::from_secs(5)), + "scheduler diagnostic sink did not drain" + ); + let status = scheduler.status(); + let after_recovery_metrics = engine.metrics(); + let diagnostic_rows = diagnostics.lock().expect("diagnostic lock").clone(); + let decisions = diagnostic_rows + .iter() + .filter(|row| row["type"] == "scheduler_decision") + .cloned() + .collect::>(); + let max_queue_age_rounds = decisions + .iter() + .filter_map(|row| row["queue_age_rounds"].as_u64()) + .max() + .unwrap_or(0); + let mixed_first = mixed + .iter() + .filter_map(|result| result.first_event_ms) + .collect::>(); + let mixed_quantums = mixed + .iter() + .flat_map(|result| result.token_wait_ms.iter().copied()) + .collect::>(); + let p95_first_event_ms = percentile(&mixed_first, 95); + let p95_token_wait_ms = percentile(&mixed_quantums, 95); + let first_event_limit_ms = baseline_first + .saturating_mul(workload.thresholds.max_first_event_baseline_multiplier) + .saturating_add(workload.thresholds.max_first_event_additive_ms); + let quantum_limit_ms = baseline_quantum + .saturating_mul(workload.thresholds.max_quantum_baseline_multiplier) + .saturating_add(workload.thresholds.max_quantum_additive_ms); + let outcomes_match = mixed.iter().all(|result| { + result.observed + == match result.expected { + SchedulerProofOutcome::Complete => "complete", + SchedulerProofOutcome::Cancel => "cancelled", + SchedulerProofOutcome::Deadline => "deadline", + } + }); + let capacity_recovered = recovery.observed == "complete" + && before_recovery_status.metrics.runnable == 0 + && before_recovery_status.metrics.waiting_for_consumer == 0 + && status.metrics.runnable == 0 + && status.metrics.waiting_for_consumer == 0 + && after_recovery_metrics.requests == before_recovery_metrics.requests.saturating_add(1); + let gates = json!({ + "outcomes_match": outcomes_match, + "capacity_recovered": capacity_recovered, + "diagnostics_lossless": status.metrics.diagnostic_records_lost == 0 + && status.metrics.diagnostic_records == status.metrics.diagnostic_records_delivered, + "starvation_round_bound": max_queue_age_rounds <= workload.thresholds.max_queue_age_rounds, + "first_event_latency_bound": p95_first_event_ms <= first_event_limit_ms, + "quantum_latency_bound": p95_token_wait_ms <= quantum_limit_ms, + }); + let passed = gates + .as_object() + .expect("gates object") + .values() + .all(|value| value == &Value::Bool(true)); + let artifact = json!({ + "phase": "8", + "passed": passed, + "workload": workload, + "provenance": { + "model_path": model_path, + "model_bytes": model_size, + "context": n_ctx, + "gpu_layers": gpu_layers, + "device_bytes": device_bytes, + "host_bytes": host_bytes, + }, + "baseline": baseline, + "mixed": mixed, + "capacity_recovery": recovery, + "capacity_recovery_evidence": { + "before": { + "scheduler": before_recovery_status, + "mapped_metrics": before_recovery_metrics, + }, + "after": { + "scheduler": status.clone(), + "mapped_metrics": after_recovery_metrics, + }, + }, + "measurements": { + "max_queue_age_rounds": max_queue_age_rounds, + "p95_first_event_ms": p95_first_event_ms, + "p95_token_wait_ms": p95_token_wait_ms, + "first_event_limit_ms": first_event_limit_ms, + "quantum_limit_ms": quantum_limit_ms, + "elapsed_ms": proof_started.elapsed().as_millis(), + }, + "scheduler": status, + "mapped_metrics": engine.metrics(), + "diagnostics": diagnostic_rows, + "gates": gates, + }); + if let Some(parent) = output.parent() { + fs::create_dir_all(parent)?; + } + fs::write(&output, serde_json::to_vec_pretty(&artifact)?)?; + ensure!(passed, "Phase 8 scheduler proof gates failed"); + println!("{}", output.display()); + Ok(()) +} + +fn run_scheduler_proof_case( + scheduler: Arc, + model: cusco_server::ModelRecord, + case: SchedulerProofCase, + start_barrier: Option>, +) -> SchedulerProofResult { + use cusco_server::{ + EngineRequest, Error as ServerError, FrontierControl, InferenceEngine, RequestControl, + SchedulingMetadata, + }; + + if let Some(barrier) = start_barrier { + barrier.wait(); + } + let control = Arc::new(RequestControl::new()); + let sink_control = control.clone(); + let expected = case.expected; + let started = Instant::now(); + let mut last_token = started; + let mut first_event_ms = None; + let mut token_wait_ms = Vec::new(); + let prompt = case.prompt.repeat(case.prompt_repetitions); + let result = scheduler.generate( + EngineRequest { + model, + prompt, + max_tokens: case.max_tokens, + prior_tokens: Vec::new(), + control, + scheduling: SchedulingMetadata { + class: case.class, + source: cusco_server::PrioritySource::ControlledWorkload, + principal: case.principal.clone(), + correlation_id: format!("phase8-transport-{}", case.id), + inference_id: format!("phase8-inference-{}", case.id), + }, + prefill_chunk_tokens: scheduler.status().policy.prefill_tokens, + }, + &mut |_, _, _| { + let now = Instant::now(); + first_event_ms.get_or_insert_with(|| now.duration_since(started).as_millis()); + token_wait_ms.push(now.duration_since(last_token).as_millis()); + last_token = now; + match expected { + SchedulerProofOutcome::Complete => {} + SchedulerProofOutcome::Cancel => sink_control.cancel(), + SchedulerProofOutcome::Deadline => sink_control.expire(), + } + Ok(FrontierControl::Continue) + }, + ); + let observed = match result { + Ok(_) => "complete", + Err(ServerError::Cancelled) => "cancelled", + Err(ServerError::Deadline) => "deadline", + Err(_) => "failed", + } + .to_owned(); + SchedulerProofResult { + id: case.id, + principal: case.principal, + class: case.class, + expected, + observed, + + first_event_ms, + total_ms: started.elapsed().as_millis(), + token_wait_ms, + } +} +fn wait_for_scheduler_idle(scheduler: &cusco_server::WorkloadScheduler, timeout: Duration) -> bool { + let started = Instant::now(); + loop { + let metrics = scheduler.status().metrics; + if metrics.runnable == 0 && metrics.waiting_for_consumer == 0 { + return true; + } + if started.elapsed() >= timeout { + return false; + } + thread::sleep(Duration::from_millis(1)); + } +} + +fn percentile(values: &[u128], percentile: usize) -> u128 { + if values.is_empty() { + return 0; + } + let mut sorted = values.to_vec(); + sorted.sort_unstable(); + let index = (sorted.len() - 1).saturating_mul(percentile) / 100; + sorted[index] +} + fn mapped_proof( model: PathBuf, n_ctx: u32, @@ -509,6 +973,37 @@ mod tests { assert_eq!(http_debug, HttpDebugLevelArg::Full); } + #[test] + fn scheduler_proof_cli_parses_versioned_workload_inputs() { + let args = Args::try_parse_from([ + "cusco", + "scheduler-proof", + "model.gguf", + "--workload", + "workload.json", + "--output", + "artifact.json", + "--context", + "2048", + ]) + .unwrap(); + let Command::SchedulerProof { + model, + workload, + output, + context, + .. + } = args.command + else { + panic!("scheduler-proof command expected") + }; + assert_eq!(model, PathBuf::from("model.gguf")); + assert_eq!(workload, PathBuf::from("workload.json")); + assert_eq!(output, PathBuf::from("artifact.json")); + assert_eq!(context, 2048); + assert_eq!(percentile(&[40, 10, 30, 20], 95), 30); + } + #[test] fn http_debug_defaults_off_and_declares_environment_source() { use clap::CommandFactory; @@ -601,6 +1096,97 @@ mod tests { .unwrap() > 0 ); + let workload = root.join("scheduler-workload.json"); + fs::write( + &workload, + serde_json::to_vec(&json!({ + "version": 1, + "name": "model-free", + "model_family": "gemma-4-e2b-it", + "policy": { + "version": 1, + "interactive_weight": 4, + "standard_weight": 2, + "batch_weight": 1, + "promotion_rounds": 8, + "deficit_refill": 1, + "prefill_tokens": 8, + "diagnostic_capacity": 1024 + }, + "baseline": { + "id": "baseline", + "principal": "baseline", + "class": "standard", + "prompt": "baseline", + "prompt_repetitions": 1, + "max_tokens": 2, + "arrival_delay_ms": 0, + "expected": "complete" + }, + "mixed": [ + { + "id": "interactive", + "principal": "one", + "class": "interactive", + "prompt": "interactive", + "prompt_repetitions": 1, + "max_tokens": 2, + "arrival_delay_ms": 0, + "expected": "complete" + }, + { + "id": "cancel", + "principal": "two", + "class": "standard", + "prompt": "cancel", + "prompt_repetitions": 1, + "max_tokens": 2, + "arrival_delay_ms": 0, + "expected": "cancel" + }, + { + "id": "deadline", + "principal": "three", + "class": "batch", + "prompt": "deadline", + "prompt_repetitions": 1, + "max_tokens": 2, + "arrival_delay_ms": 0, + "expected": "deadline" + } + ], + "thresholds": { + "max_queue_age_rounds": 100, + "max_first_event_baseline_multiplier": 100, + "max_first_event_additive_ms": 1000, + "max_quantum_baseline_multiplier": 100, + "max_quantum_additive_ms": 1000 + } + })) + .unwrap(), + ) + .unwrap(); + let scheduler_output = root.join("scheduler.json"); + run(Command::SchedulerProof { + model: PathBuf::from("mock://deterministic"), + workload, + output: scheduler_output.clone(), + context: 128, + gpu_layers: 0, + device_bytes: 1 << 20, + host_bytes: 1 << 20, + }) + .unwrap(); + let scheduler_artifact: Value = + serde_json::from_slice(&fs::read(scheduler_output).unwrap()).unwrap(); + assert_eq!(scheduler_artifact["passed"], true); + assert_eq!(scheduler_artifact["mixed"].as_array().unwrap().len(), 3); + assert!( + scheduler_artifact["measurements"]["max_queue_age_rounds"] + .as_u64() + .unwrap() + > 0 + ); let public: SocketAddr = "0.0.0.0:8080".parse().unwrap(); assert!( run(Command::Serve { @@ -631,6 +1217,13 @@ mod tests { stream_buffer: 1, shutdown_grace_ms: 100, http_debug: HttpDebugLevelArg::Off, + scheduler_prefill_tokens: 32, + scheduler_interactive_weight: 8, + scheduler_standard_weight: 4, + scheduler_batch_weight: 1, + scheduler_promotion_rounds: 64, + scheduler_deficit_refill: 1, + scheduler_diagnostic_capacity: 1024, }) .is_err() ); @@ -663,6 +1256,13 @@ mod tests { stream_buffer: 1, shutdown_grace_ms: 100, http_debug: HttpDebugLevelArg::Off, + scheduler_prefill_tokens: 32, + scheduler_interactive_weight: 8, + scheduler_standard_weight: 4, + scheduler_batch_weight: 1, + scheduler_promotion_rounds: 64, + scheduler_deficit_refill: 1, + scheduler_diagnostic_capacity: 1024, }) .unwrap_err(); assert!(unsupported.to_string().contains("unsupported model family")); diff --git a/crates/executor/src/lib.rs b/crates/executor/src/lib.rs index 11401a1..cb90c8c 100644 --- a/crates/executor/src/lib.rs +++ b/crates/executor/src/lib.rs @@ -137,6 +137,11 @@ pub struct GreedySampler { raw: NonNull, } +// SAFETY: a sampler is request-owned and all access still requires an exclusive +// borrow of the executor that created it. Moving a suspended request between +// scheduler threads does not permit concurrent native sampler access. +unsafe impl Send for GreedySampler {} + impl Executor { pub fn open(path: &str, n_ctx: u32, gpu_layers: i32) -> Result { let path = CString::new(path).map_err(|_| Error::InvalidPath)?; diff --git a/crates/server/src/lib.rs b/crates/server/src/lib.rs index f85ed7c..22adc88 100644 --- a/crates/server/src/lib.rs +++ b/crates/server/src/lib.rs @@ -32,6 +32,7 @@ use thiserror::Error; use uuid::Uuid; mod generation; mod mapped; +mod scheduler; pub use generation::{ FinishReason, FrontierControl, GenerationFrontier, MAX_STOP_BYTES, MAX_STOP_SEQUENCES, @@ -39,6 +40,7 @@ pub use generation::{ }; pub use mapped::{ExecutionProfile, MappedEngine, MappedMetrics}; pub use residency::{ResidencyConfig, ResidencyMetrics, ResidentEngine, ResidentModelStatus}; +pub use scheduler::{SchedulerMetrics, SchedulerStatus, WorkloadScheduler}; #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(transparent)] @@ -110,6 +112,44 @@ pub enum StreamEvent { message: String, }, } +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum SchedulingClass { + Interactive, + #[default] + Standard, + Batch, +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum PrioritySource { + #[default] + AdapterDefault, + ControlledWorkload, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct SchedulingMetadata { + pub class: SchedulingClass, + pub source: PrioritySource, + pub principal: String, + pub correlation_id: String, + pub inference_id: String, +} + +impl Default for SchedulingMetadata { + fn default() -> Self { + Self { + class: SchedulingClass::Standard, + source: PrioritySource::AdapterDefault, + principal: "anonymous-admin".into(), + correlation_id: String::new(), + inference_id: String::new(), + } + } +} + #[derive(Clone, Debug, Serialize, Deserialize)] pub struct InferRequest { pub model: String, @@ -124,8 +164,8 @@ pub struct InferRequest { pub stop: Vec, #[serde(default)] pub raw_continuation: bool, - #[serde(default)] - pub priority: i32, + #[serde(skip, default)] + pub scheduling: SchedulingMetadata, } fn default_tokens() -> usize { 16 @@ -207,6 +247,52 @@ impl Default for ServerConfig { } } } + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Serialize)] +pub struct SchedulerPolicyConfig { + pub version: u32, + pub interactive_weight: u32, + pub standard_weight: u32, + pub batch_weight: u32, + pub deficit_refill: u32, + pub prefill_tokens: usize, + pub promotion_rounds: u64, + pub diagnostic_capacity: usize, +} + +impl Default for SchedulerPolicyConfig { + fn default() -> Self { + Self { + version: 1, + interactive_weight: 4, + standard_weight: 2, + batch_weight: 1, + deficit_refill: 1, + prefill_tokens: 32, + promotion_rounds: 64, + diagnostic_capacity: 1024, + } + } +} + +impl SchedulerPolicyConfig { + fn validate(self) -> Result { + if self.version == 0 + || self.interactive_weight == 0 + || self.standard_weight == 0 + || self.batch_weight == 0 + || self.deficit_refill == 0 + || self.prefill_tokens == 0 + || self.promotion_rounds == 0 + || self.diagnostic_capacity == 0 + { + return Err(Error::State( + "scheduler policy parameters must be nonzero".into(), + )); + } + Ok(self) + } +} const HTTP_DEBUG_BODY_LIMIT: usize = 64 << 10; #[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq, Serialize)] #[serde(rename_all = "snake_case")] @@ -473,7 +559,7 @@ pub struct RequestControl { abort: Mutex>>, } impl RequestControl { - fn new() -> Self { + pub fn new() -> Self { Self { state: AtomicU8::new(CONTROL_RUNNING), abort: Mutex::new(None), @@ -488,7 +574,7 @@ impl RequestControl { } } - fn cancel(&self) { + pub fn cancel(&self) { if self .state .compare_exchange( @@ -503,7 +589,7 @@ impl RequestControl { } } - fn expire(&self) { + pub fn expire(&self) { if self .state .compare_exchange( @@ -597,12 +683,15 @@ impl AuthProvider for BearerAuth { } } -pub struct EngineRequest<'a> { - pub model: &'a ModelRecord, - pub prompt: &'a str, +#[derive(Clone)] +pub struct EngineRequest { + pub model: ModelRecord, + pub prompt: String, pub max_tokens: usize, - pub prior_tokens: &'a [i32], - pub control: &'a RequestControl, + pub prior_tokens: Vec, + pub control: Arc, + pub scheduling: SchedulingMetadata, + pub prefill_chunk_tokens: usize, } pub type TokenSink<'a> = dyn FnMut(i32, &[u8], bool) -> Result + 'a; @@ -616,12 +705,80 @@ pub struct EngineOutput { pub prefill: PrefillMetrics, } +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum QuantumKind { + Preparation, + Prefill, + Decode, + Publication, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct QuantumObservation { + pub kind: QuantumKind, + pub charged_tokens: usize, + pub context_placement: String, + pub executor_slot_occupied: bool, + pub transition_cost_bytes: u64, + pub capacity_reserved_bytes: u64, +} + +impl QuantumObservation { + fn model_free(kind: QuantumKind, charged_tokens: usize) -> Self { + Self { + kind, + charged_tokens, + context_placement: "model_free".into(), + executor_slot_occupied: false, + transition_cost_bytes: 0, + capacity_reserved_bytes: 0, + } + } +} + +pub enum SessionStep { + Progress(QuantumObservation), + Token { + id: i32, + piece: Vec, + terminal_or_control: bool, + observation: QuantumObservation, + }, + Finished(EngineOutput), +} + +pub trait ExecutionSession: Send { + fn step(&mut self) -> Result; + fn finish(&mut self) -> Result; +} + pub trait InferenceEngine: Send + Sync { + fn start_session(&self, request: EngineRequest) -> Result, Error>; + fn generate( &self, - request: EngineRequest<'_>, + request: EngineRequest, sink: &mut TokenSink<'_>, - ) -> Result; + ) -> Result { + let mut session = self.start_session(request)?; + loop { + match session.step()? { + SessionStep::Progress(_) => {} + SessionStep::Token { + id, + piece, + terminal_or_control, + .. + } => { + if sink(id, &piece, terminal_or_control)? == FrontierControl::Stop { + return session.finish(); + } + } + SessionStep::Finished(output) => return Ok(output), + } + } + } fn prepare_model(&self, _model: &ModelRecord) -> Result<(), Error> { Ok(()) @@ -639,36 +796,51 @@ pub trait InferenceEngine: Send + Sync { None } } + #[derive(Default)] pub struct DeterministicEngine; -impl InferenceEngine for DeterministicEngine { - fn generate( - &self, - request: EngineRequest<'_>, - sink: &mut TokenSink<'_>, - ) -> Result { - request.control.check()?; - let input_tokens = request.prompt.split_whitespace().count(); - for (index, piece) in request - .prompt - .split_whitespace() - .rev() - .cycle() - .take(request.max_tokens) - .enumerate() - { - request.control.check()?; - let piece = if index == 0 { - piece.to_owned() - } else { - format!(" {piece}") - }; - if sink(-(index as i32) - 1, piece.as_bytes(), false)? == FrontierControl::Stop { - break; - } - } + +struct DeterministicSession { + request: EngineRequest, + pieces: Vec, + index: usize, + prepared: bool, +} + +impl ExecutionSession for DeterministicSession { + fn step(&mut self) -> Result { + self.request.control.check()?; + if !self.prepared { + self.prepared = true; + return Ok(SessionStep::Progress(QuantumObservation::model_free( + QuantumKind::Preparation, + self.pieces.len(), + ))); + } + if self.index >= self.request.max_tokens || self.pieces.is_empty() { + return self.finish().map(SessionStep::Finished); + } + let piece = &self.pieces[self.index % self.pieces.len()]; + let piece = if self.index == 0 { + piece.clone() + } else { + format!(" {piece}") + }; + let id = -(self.index as i32) - 1; + self.index += 1; + Ok(SessionStep::Token { + id, + piece: piece.into_bytes(), + terminal_or_control: false, + observation: QuantumObservation::model_free(QuantumKind::Decode, 1), + }) + } + + fn finish(&mut self) -> Result { + self.request.control.check()?; + let input_tokens = self.pieces.len(); Ok(EngineOutput { - successor_tokens: request.prior_tokens.to_vec(), + successor_tokens: self.request.prior_tokens.clone(), input_tokens, cached_tokens: 0, evaluated_tokens: input_tokens, @@ -681,6 +853,24 @@ impl InferenceEngine for DeterministicEngine { } } +impl InferenceEngine for DeterministicEngine { + fn start_session(&self, request: EngineRequest) -> Result, Error> { + request.control.check()?; + let pieces = request + .prompt + .split_whitespace() + .rev() + .map(str::to_owned) + .collect(); + Ok(Box::new(DeterministicSession { + request, + pieces, + index: 0, + prepared: false, + })) + } +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct SlotCandidate { pub slot: usize, @@ -1429,11 +1619,13 @@ impl Server { let mut delta_index = 0; let generated = self.engine.generate( EngineRequest { - model: &model, - prompt: &req.prompt, + model: model.clone(), + prompt: req.prompt.clone(), max_tokens: req.max_tokens, - prior_tokens, - control: &control, + prior_tokens: prior_tokens.to_vec(), + control: control.clone(), + scheduling: req.scheduling.clone(), + prefill_chunk_tokens: 32, }, &mut |_id, piece, terminal_or_control| { if deadline.is_some_and(|deadline| Instant::now() >= deadline) { @@ -1822,7 +2014,7 @@ async fn completion( _permit: permit, }: PrequeueJson, ) -> Result { - auth(&s, &headers, Scope::Inference)?; + let request_context = auth(&s, &headers, Scope::Inference)?; s.model(&r.model)?; drop(permit); infer_response( @@ -1836,6 +2028,7 @@ async fn completion( r.raw_continuation, r.deadline_ms, retained_bytes, + request_context.principal, ) .await } @@ -1848,7 +2041,7 @@ async fn chat( _permit: permit, }: PrequeueJson, ) -> Result { - auth(&s, &headers, Scope::Inference)?; + let request_context = auth(&s, &headers, Scope::Inference)?; s.model(&r.model)?; drop(permit); infer_response( @@ -1866,6 +2059,7 @@ async fn chat( false, r.deadline_ms, retained_bytes, + request_context.principal, ) .await } @@ -1922,9 +2116,11 @@ async fn infer_response( raw_continuation: bool, deadline_ms: Option, retained_bytes: usize, + principal: String, ) -> Result { server.model(&model)?; let id = Uuid::new_v4().to_string(); + let correlation_id = id.clone(); let config = server.config(); let wall_limit = Duration::from_millis( deadline_ms @@ -1956,7 +2152,13 @@ async fn infer_response( ), stop, raw_continuation, - priority: 0, + scheduling: SchedulingMetadata { + class: SchedulingClass::Standard, + source: PrioritySource::AdapterDefault, + principal, + correlation_id: correlation_id.clone(), + inference_id: Uuid::new_v4().to_string(), + }, }; if streaming { let (started, receiver) = server.infer_stream_reserved(id, request, admission).await?; @@ -1977,7 +2179,12 @@ async fn infer_response( let rows = first .chain(rest) .map(|event| Ok::<_, Infallible>(Event::default().json_data(event).unwrap())); - Ok(Sse::new(rows).into_response()) + let mut response = Sse::new(rows).into_response(); + response.headers_mut().insert( + "x-request-id", + HeaderValue::from_str(&correlation_id).expect("UUID is a valid header value"), + ); + Ok(response) } else { let mut disconnect = DisconnectGuard::new(control); let result = tokio::task::spawn_blocking(move || { @@ -1988,7 +2195,12 @@ async fn infer_response( .map_err(state_err)?; disconnect.disarm(); let (response, _) = result?; - Ok(Json(json!({"id":response.id,"object":"text_completion","choices":[{"text":response.text}],"usage":response.usage})).into_response()) + let mut response = Json(json!({"id":response.id,"object":"text_completion","choices":[{"text":response.text}],"usage":response.usage})).into_response(); + response.headers_mut().insert( + "x-request-id", + HeaderValue::from_str(&correlation_id).expect("UUID is a valid header value"), + ); + Ok(response) } } async fn list_models(State(s): State, headers: HeaderMap) -> Result, Error> { @@ -2303,11 +2515,7 @@ mod tests { struct FailingEngine; impl InferenceEngine for FailingEngine { - fn generate( - &self, - _: EngineRequest<'_>, - _: &mut TokenSink<'_>, - ) -> Result { + fn start_session(&self, _: EngineRequest) -> Result, Error> { Err(Error::State("generation failed".into())) } } @@ -2315,31 +2523,65 @@ mod tests { struct SlowEngine { started: Arc, } - impl InferenceEngine for SlowEngine { - fn generate( - &self, - request: EngineRequest<'_>, - sink: &mut TokenSink<'_>, - ) -> Result { - for index in 0..request.max_tokens { - sink(index as i32, b"x", false)?; - if index == 0 { - self.started.notify_one(); - } - std::thread::sleep(Duration::from_millis(20)); - } - Ok(EngineOutput { - successor_tokens: (0..request.max_tokens as i32).collect(), - input_tokens: request.prompt.split_whitespace().count(), + + struct SlowSession { + request: EngineRequest, + started: Arc, + index: usize, + } + + impl SlowSession { + fn output(&self) -> EngineOutput { + let input_tokens = self.request.prompt.split_whitespace().count(); + EngineOutput { + successor_tokens: (0..self.request.max_tokens as i32).collect(), + input_tokens, cached_tokens: 0, - evaluated_tokens: request.prompt.split_whitespace().count(), + evaluated_tokens: input_tokens, prefill: PrefillMetrics { - total_tokens: request.prompt.split_whitespace().count(), - uncached_tokens: request.prompt.split_whitespace().count(), + total_tokens: input_tokens, + uncached_tokens: input_tokens, ..PrefillMetrics::default() }, + } + } + } + + impl ExecutionSession for SlowSession { + fn step(&mut self) -> Result { + self.request.control.check()?; + if self.index >= self.request.max_tokens { + return Ok(SessionStep::Finished(self.output())); + } + let id = self.index as i32; + self.index += 1; + if id == 0 { + self.started.notify_one(); + } + std::thread::sleep(Duration::from_millis(20)); + Ok(SessionStep::Token { + id, + piece: b"x".to_vec(), + terminal_or_control: false, + observation: QuantumObservation::model_free(QuantumKind::Decode, 1), }) } + + fn finish(&mut self) -> Result { + Ok(self.output()) + } + } + impl InferenceEngine for SlowEngine { + fn start_session( + &self, + request: EngineRequest, + ) -> Result, Error> { + Ok(Box::new(SlowSession { + request, + started: self.started.clone(), + index: 0, + })) + } } #[test] fn persistence_branching_and_ids_survive_restart() { @@ -2369,7 +2611,7 @@ mod tests { max_tokens: 2, context_id: None, deadline_ms: Some(1000), - priority: 2, + scheduling: SchedulingMetadata::default(), stop: vec![], raw_continuation: false, }; @@ -2388,7 +2630,7 @@ mod tests { max_tokens: 1, context_id: Some(before.id.clone()), deadline_ms: None, - priority: 0, + scheduling: SchedulingMetadata::default(), stop: vec![], raw_continuation: false, }, @@ -2405,7 +2647,7 @@ mod tests { max_tokens: 1, context_id: Some(before.id), deadline_ms: Some(0), - priority: 0, + scheduling: SchedulingMetadata::default(), stop: vec![], raw_continuation: false, } @@ -2422,7 +2664,7 @@ mod tests { max_tokens: 1, context_id: None, deadline_ms: None, - priority: 0, + scheduling: SchedulingMetadata::default(), stop: vec![], raw_continuation: false, }, @@ -2446,7 +2688,7 @@ mod tests { max_tokens: 1, context_id: None, deadline_ms: None, - priority: 0, + scheduling: SchedulingMetadata::default(), stop: vec![], raw_continuation: false, } @@ -2950,7 +3192,7 @@ mod tests { deadline_ms: None, stop: vec![], raw_continuation: false, - priority: 0, + scheduling: SchedulingMetadata::default(), }, admission, ) @@ -2997,7 +3239,7 @@ mod tests { deadline_ms: None, stop: vec![], raw_continuation: false, - priority: 0, + scheduling: SchedulingMetadata::default(), }, admission, ) @@ -3056,7 +3298,7 @@ mod tests { deadline_ms: None, stop: vec![], raw_continuation: false, - priority: 0, + scheduling: SchedulingMetadata::default(), }, ) }) @@ -3076,7 +3318,7 @@ mod tests { deadline_ms: Some(5), stop: vec![], raw_continuation: false, - priority: 0, + scheduling: SchedulingMetadata::default(), }, ), Err(Error::Deadline) diff --git a/crates/server/src/mapped.rs b/crates/server/src/mapped.rs index 2155b5a..e79885b 100644 --- a/crates/server/src/mapped.rs +++ b/crates/server/src/mapped.rs @@ -1,16 +1,17 @@ use crate::{ - EngineOutput, EngineRequest, Error, FrontierControl, InferenceEngine, PrefillMetrics, TokenSink, + EngineOutput, EngineRequest, Error, ExecutionSession, InferenceEngine, PrefillMetrics, + QuantumKind, QuantumObservation, SessionStep, }; use cusco_context_store::{ AdapterEpoch, ComponentMask, ContextStore, EvaluatedPrefixId, LogicalContextId, ModelEpoch, PersistentTokenSequence, }; -use cusco_executor::{Decode, Executor, MappingId, MappingState, OperatingPoint}; +use cusco_executor::{Decode, Executor, GreedySampler, MappingId, MappingState, OperatingPoint}; use cusco_physical_manager::{ Capacity, Component, PhysicalManager, PhysicalRepresentationId, Tier, }; use parking_lot::Mutex; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use std::{ collections::HashMap, fs, @@ -139,7 +140,7 @@ impl ExecutionProfile { } } -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)] pub struct MappedMetrics { pub requests: u64, pub cache_hits: u64, @@ -181,7 +182,7 @@ pub struct MappedEngine { context_capacity: usize, spill_dir: Option, spill_capacity: usize, - state: Mutex, + state: Arc>, } impl MappedEngine { @@ -275,7 +276,7 @@ impl MappedEngine { model_path, spill_dir, spill_capacity, - state: Mutex::new(MappedState { + state: Arc::new(Mutex::new(MappedState { executor, logical: ContextStore::default(), physical: PhysicalManager::new(Capacity { @@ -286,7 +287,7 @@ impl MappedEngine { model_epoch, metrics: MappedMetrics::default(), spill_bytes: 0, - }), + })), })) } @@ -368,49 +369,69 @@ impl MappedEngine { } } -impl InferenceEngine for MappedEngine { - fn generate( - &self, - request: EngineRequest<'_>, - sink: &mut TokenSink<'_>, - ) -> Result { - request.control.check()?; - if request.model.path.to_str() != Some(self.model_path()) { - return Err(Error::State( - "Phase 6B admits only the process-owned model".into(), - )); - } - let prompt_started = Instant::now(); +struct MappedSession { + profile: ExecutionProfile, + state: Arc>, + request: EngineRequest, + prompt_started: Instant, + stage: MappedStage, + tokens: Vec, + logical_context: Option, + active_mapping: Option, + parent: Option, + next: Option, + sampler: Option, + input_tokens: usize, + cached: usize, + evaluated_tokens: usize, + evaluated: usize, + generated: usize, + prefill: PrefillMetrics, + finished: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum MappedStage { + Preparation, + Prefill, + Decode, + Publication, + Finished, +} + +impl MappedSession { + fn prepare(&mut self) -> Result { + self.request.control.check()?; + let tokenization_started = Instant::now(); let mut state = self.state.lock(); state.executor.reset_cancellation(); let cancellation = state.executor.cancellation_handle(); - let _abort = request.control.register_abort(Arc::new(move || { + let _abort = self.request.control.register_abort(Arc::new(move || { cancellation.cancel(); })); - request.control.check()?; - let tokenization_started = Instant::now(); let prompt = state .executor - .tokenize(request.prompt) + .tokenize(&self.request.prompt) .map_err(state_error)?; - request.control.check()?; - let tokenization_ns = elapsed_ns(tokenization_started); - let input_tokens = prompt.len(); - let total = request + self.request.control.check()?; + self.prefill.tokenization_ns = elapsed_ns(tokenization_started); + self.input_tokens = prompt.len(); + let total = self + .request .prior_tokens .len() .checked_add(prompt.len()) - .and_then(|value| value.checked_add(request.max_tokens)) + .and_then(|value| value.checked_add(self.request.max_tokens)) .ok_or_else(|| Error::State("request token count overflow".into()))?; - if total > self.context_capacity { + if total > self.profile.context_limit { return Err(Error::State( "request exceeds model context capacity".into(), )); } - let mut tokens = request.prior_tokens.to_vec(); - tokens.extend_from_slice(&prompt); - let sequence = PersistentTokenSequence::default().append(&tokens); + self.tokens = self.request.prior_tokens.clone(); + self.tokens.extend_from_slice(&prompt); + let sequence = PersistentTokenSequence::default().append(&self.tokens); let model_epoch = state.model_epoch; let logical_context = state.logical.create(sequence, model_epoch, ADAPTER_EPOCH); let prefix_lookup_started = Instant::now(); @@ -418,131 +439,300 @@ impl InferenceEngine for MappedEngine { .logical .longest_valid_prefix(logical_context) .map_err(state_error)?; - let prefix_lookup_ns = elapsed_ns(prefix_lookup_started); - let cached = prefix.as_ref().map_or(0, |mapping| mapping.represented_end); - let evaluated_tokens = tokens.len().saturating_sub(cached); - let mapping_activation_started = Instant::now(); - activate_prefix(&mut state, logical_context, prefix.as_deref())?; - request.control.check()?; - - let mut active_mapping = state.executor.mapping_metrics().active; - if cached == 0 { - active_mapping = fork_and_activate(&mut state.executor, active_mapping)?; - } - let mapping_activation_ns = elapsed_ns(mapping_activation_started); - let uncached_prefill_started = Instant::now(); - let mut evaluated = cached; - let mut parent = prefix.as_ref().map(|mapping| mapping.id); - let mut next = prefix + self.prefill.prefix_lookup_ns = elapsed_ns(prefix_lookup_started); + self.cached = prefix.as_ref().map_or(0, |mapping| mapping.represented_end); + self.evaluated_tokens = self.tokens.len().saturating_sub(self.cached); + self.evaluated = self.cached; + self.parent = prefix.as_ref().map(|mapping| mapping.id); + self.next = prefix .as_ref() .and_then(|mapping| state.resident.get(&mapping.id)) .map(|resident| resident.continuation.clone()); - while evaluated < tokens.len() { - request.control.check()?; - let end = tokens - .len() - .min(((evaluated / self.profile.block_size) + 1) * self.profile.block_size); - next = Some( - state - .executor - .decode(&tokens[evaluated..end]) - .map_err(state_error)?, - ); - request.control.check()?; - state.metrics.decoded_tokens += (end - evaluated) as u64; - evaluated = end; - if evaluated % self.profile.block_size == 0 { - parent = Some(publish_block( - &mut state, - &self.profile, - logical_context, - evaluated, - parent, - active_mapping, - next.as_ref().expect("decode result exists"), - )?); - active_mapping = fork_and_activate(&mut state.executor, active_mapping)?; + + let activation_started = Instant::now(); + let transfer_before = state.physical.metrics().transfer_bytes; + activate_prefix(&mut state, logical_context, prefix.as_deref())?; + let source = state.executor.mapping_metrics().active; + self.active_mapping = Some(fork_and_activate(&mut state.executor, source)?); + self.prefill.mapping_activation_ns = elapsed_ns(activation_started); + let transfer_after = state.physical.metrics().transfer_bytes; + self.logical_context = Some(logical_context); + self.prefill.total_tokens = self.tokens.len(); + self.prefill.cached_tokens = self.cached; + self.prefill.uncached_tokens = self.evaluated_tokens; + self.stage = MappedStage::Prefill; + Ok(SessionStep::Progress(QuantumObservation { + kind: QuantumKind::Preparation, + charged_tokens: 1, + context_placement: "device".into(), + executor_slot_occupied: true, + transition_cost_bytes: transfer_after.saturating_sub(transfer_before), + capacity_reserved_bytes: 0, + })) + } + + fn prefill(&mut self) -> Result { + if self.evaluated >= self.tokens.len() { + if self.request.max_tokens > 0 && self.next.is_none() { + return Err(Error::State( + "an exact cached prefix cannot supply uncached logits".into(), + )); } + let mut state = self.state.lock(); + self.sampler = Some(state.executor.greedy_sampler().map_err(state_error)?); + self.stage = MappedStage::Decode; + return Ok(SessionStep::Progress(QuantumObservation::model_free( + QuantumKind::Prefill, + 1, + ))); } - let uncached_prefill_ns = elapsed_ns(uncached_prefill_started); - if request.max_tokens > 0 && next.is_none() { - return Err(Error::State( - "an exact cached prefix cannot supply uncached logits".into(), - )); + + let started = Instant::now(); + let mut state = self.state.lock(); + self.request.control.check()?; + state.executor.reset_cancellation(); + let cancellation = state.executor.cancellation_handle(); + let _abort = self.request.control.register_abort(Arc::new(move || { + cancellation.cancel(); + })); + self.request.control.check()?; + self.activate_owned(&mut state)?; + let next_block = ((self.evaluated / self.profile.block_size) + 1) + .saturating_mul(self.profile.block_size); + let end = self.tokens.len().min(next_block).min( + self.evaluated + .saturating_add(self.request.prefill_chunk_tokens), + ); + let charged = end.saturating_sub(self.evaluated); + self.next = Some( + state + .executor + .decode(&self.tokens[self.evaluated..end]) + .map_err(state_error)?, + ); + self.request.control.check()?; + state.metrics.decoded_tokens += charged as u64; + self.evaluated = end; + if self.evaluated % self.profile.block_size == 0 { + self.parent = Some(publish_block( + &mut state, + &self.profile, + self.logical_context.expect("prepared context exists"), + self.evaluated, + self.parent, + self.active_mapping.expect("prepared mapping exists"), + self.next.as_ref().expect("decode result exists"), + )?); + let source = self.active_mapping.expect("published mapping exists"); + self.active_mapping = Some(fork_and_activate(&mut state.executor, source)?); } - let mut prefill = PrefillMetrics { - total_tokens: tokens.len(), - cached_tokens: cached, - uncached_tokens: evaluated_tokens, - tokenization_ns, - prefix_lookup_ns, - mapping_activation_ns, - uncached_prefill_ns, - total_ns: elapsed_ns(prompt_started), - ..PrefillMetrics::default() - }; - let mut sampler = state.executor.greedy_sampler().map_err(state_error)?; + self.prefill.uncached_prefill_ns = self + .prefill + .uncached_prefill_ns + .saturating_add(elapsed_ns(started)); + Ok(SessionStep::Progress(QuantumObservation { + kind: QuantumKind::Prefill, + charged_tokens: charged.max(1), + context_placement: "device".into(), + executor_slot_occupied: true, + transition_cost_bytes: 0, + capacity_reserved_bytes: 0, + })) + } + + fn decode(&mut self) -> Result { + if self.generated >= self.request.max_tokens { + self.stage = MappedStage::Publication; + return self.publish().map(SessionStep::Finished); + } + let mut state = self.state.lock(); + self.request.control.check()?; + state.executor.reset_cancellation(); + let cancellation = state.executor.cancellation_handle(); + let _abort = self.request.control.register_abort(Arc::new(move || { + cancellation.cancel(); + })); + self.request.control.check()?; + self.activate_owned(&mut state)?; + let sampled = self + .sampler + .as_mut() + .expect("decode owns sampler") + .sample(&mut state.executor) + .map_err(state_error)?; + let terminal_or_control = self.profile.terminal_tokens.contains(&sampled); let mut piece = Vec::with_capacity(32); - for index in 0..request.max_tokens { - request.control.check()?; - let sampled = sampler.sample(&mut state.executor).map_err(state_error)?; - let terminal_or_control = self.profile.terminal_tokens.contains(&sampled); - if terminal_or_control { - piece.clear(); - } else { - state - .executor - .render_token(sampled, &mut piece) - .map_err(state_error)?; - } - request.control.check()?; - tokens.push(sampled); + if !terminal_or_control { state - .logical - .append(logical_context, &[sampled]) + .executor + .render_token(sampled, &mut piece) .map_err(state_error)?; - let control = sink(sampled, &piece, terminal_or_control)?; - if control == FrontierControl::Stop - || terminal_or_control - || index + 1 == request.max_tokens - { - break; - } - next = Some(state.executor.decode(&[sampled]).map_err(state_error)?); - request.control.check()?; + } + self.request.control.check()?; + self.tokens.push(sampled); + state + .logical + .append( + self.logical_context.expect("prepared context exists"), + &[sampled], + ) + .map_err(state_error)?; + self.generated += 1; + if terminal_or_control || self.generated == self.request.max_tokens { + self.stage = MappedStage::Publication; + } else { + self.next = Some(state.executor.decode(&[sampled]).map_err(state_error)?); + self.request.control.check()?; state.metrics.decoded_tokens += 1; - evaluated += 1; - if evaluated % self.profile.block_size == 0 { - parent = Some(publish_block( + self.evaluated += 1; + if self.evaluated % self.profile.block_size == 0 { + self.parent = Some(publish_block( &mut state, &self.profile, - logical_context, - evaluated, - parent, - active_mapping, - next.as_ref().expect("decode result exists"), + self.logical_context.expect("prepared context exists"), + self.evaluated, + self.parent, + self.active_mapping.expect("prepared mapping exists"), + self.next.as_ref().expect("decode result exists"), )?); - active_mapping = fork_and_activate(&mut state.executor, active_mapping)?; + let source = self.active_mapping.expect("published mapping exists"); + self.active_mapping = Some(fork_and_activate(&mut state.executor, source)?); + } + } + Ok(SessionStep::Token { + id: sampled, + piece, + terminal_or_control, + observation: QuantumObservation { + kind: QuantumKind::Decode, + charged_tokens: 1, + context_placement: "device".into(), + executor_slot_occupied: true, + transition_cost_bytes: 0, + capacity_reserved_bytes: 0, + }, + }) + } + + fn activate_owned(&self, state: &mut MappedState) -> Result<(), Error> { + state + .executor + .activate_mapping(self.active_mapping.expect("prepared mapping exists")) + .map_err(state_error) + } + + fn publish(&mut self) -> Result { + if let Some(output) = &self.finished { + return Ok(output.clone()); + } + self.request.control.check()?; + self.sampler = None; + let mut state = self.state.lock(); + if let Some(active) = self.active_mapping.take() { + let fallback = self + .parent + .and_then(|id| state.resident.get(&id)) + .and_then(|resident| resident.native) + .unwrap_or(MappingId(0)); + if active != fallback { + state + .executor + .activate_mapping(fallback) + .map_err(state_error)?; + state.executor.remove_mapping(active).map_err(state_error)?; } } - request.control.check()?; let native_metrics = state.executor.mapping_metrics(); state.metrics.requests += 1; - state.metrics.cached_tokens += cached as u64; - state.metrics.cache_hits += u64::from(cached != 0); + state.metrics.cached_tokens += self.cached as u64; + state.metrics.cache_hits += u64::from(self.cached != 0); state.metrics.reference_switches = native_metrics.reference_switches; state.metrics.activation_bytes_copied = native_metrics.activation_bytes_copied; let physical_metrics = state.physical.metrics(); - prefill.transfer_bytes = physical_metrics.transfer_bytes; - prefill.device_bytes = physical_metrics.device_total; - prefill.host_bytes = physical_metrics.host_used; - Ok(EngineOutput { - successor_tokens: tokens, - input_tokens, - cached_tokens: cached, - evaluated_tokens, - prefill, - }) + self.prefill.transfer_bytes = physical_metrics.transfer_bytes; + self.prefill.device_bytes = physical_metrics.device_total; + self.prefill.host_bytes = physical_metrics.host_used; + self.prefill.total_ns = elapsed_ns(self.prompt_started); + let output = EngineOutput { + successor_tokens: self.tokens.clone(), + input_tokens: self.input_tokens, + cached_tokens: self.cached, + evaluated_tokens: self.evaluated_tokens, + prefill: self.prefill, + }; + self.finished = Some(output.clone()); + self.stage = MappedStage::Finished; + Ok(output) + } +} + +impl ExecutionSession for MappedSession { + fn step(&mut self) -> Result { + match self.stage { + MappedStage::Preparation => self.prepare(), + MappedStage::Prefill => self.prefill(), + MappedStage::Decode => self.decode(), + MappedStage::Publication => self.publish().map(SessionStep::Finished), + MappedStage::Finished => Ok(SessionStep::Finished( + self.finished.clone().expect("finished output exists"), + )), + } + } + + fn finish(&mut self) -> Result { + self.publish() + } +} + +impl Drop for MappedSession { + fn drop(&mut self) { + self.sampler = None; + if let Some(active) = self.active_mapping.take() { + let mut state = self.state.lock(); + let fallback = self + .parent + .and_then(|id| state.resident.get(&id)) + .and_then(|resident| resident.native) + .unwrap_or(MappingId(0)); + if active != fallback { + let _ = state.executor.activate_mapping(fallback); + let _ = state.executor.remove_mapping(active); + } + } + } +} + +impl InferenceEngine for MappedEngine { + fn start_session(&self, request: EngineRequest) -> Result, Error> { + request.control.check()?; + if request.model.path.to_str() != Some(self.model_path()) { + return Err(Error::State( + "Phase 8 admits only a resident process-owned model".into(), + )); + } + let context_limit = self.context_capacity.min(self.profile.context_limit); + let mut profile = self.profile.clone(); + profile.context_limit = context_limit; + Ok(Box::new(MappedSession { + profile, + state: self.state.clone(), + request, + prompt_started: Instant::now(), + stage: MappedStage::Preparation, + tokens: Vec::new(), + logical_context: None, + active_mapping: None, + parent: None, + next: None, + sampler: None, + input_tokens: 0, + cached: 0, + evaluated_tokens: 0, + evaluated: 0, + generated: 0, + prefill: PrefillMetrics::default(), + finished: None, + })) } fn demote_inactive(&self) -> Result { @@ -844,8 +1034,8 @@ fn state_error(error: impl std::fmt::Display) -> Error { #[cfg(test)] mod tests { use super::*; - use crate::ModelRecord; - use std::path::PathBuf; + use crate::{FrontierControl, ModelRecord, RequestControl, SchedulingMetadata}; + use std::{path::PathBuf, sync::Arc}; fn model() -> ModelRecord { ModelRecord { @@ -860,6 +1050,23 @@ mod tests { } } + fn test_request( + model: &ModelRecord, + prompt: impl Into, + max_tokens: usize, + prior_tokens: &[i32], + ) -> EngineRequest { + EngineRequest { + model: model.clone(), + prompt: prompt.into(), + max_tokens, + prior_tokens: prior_tokens.to_vec(), + control: Arc::new(RequestControl::new()), + scheduling: SchedulingMetadata::default(), + prefill_chunk_tokens: 32, + } + } + #[derive(Debug)] struct CollectedOutput { pieces: Vec>, @@ -869,10 +1076,10 @@ mod tests { } trait GenerateCollected { - fn generate_collected(&self, request: EngineRequest<'_>) -> Result; + fn generate_collected(&self, request: EngineRequest) -> Result; } impl GenerateCollected for MappedEngine { - fn generate_collected(&self, request: EngineRequest<'_>) -> Result { + fn generate_collected(&self, request: EngineRequest) -> Result { let mut pieces = Vec::new(); let output = self.generate(request, &mut |_, piece, _| { pieces.push(piece.to_vec()); @@ -930,22 +1137,10 @@ mod tests { .unwrap(); let model = model(); let first = engine - .generate_collected(EngineRequest { - model: &model, - prompt: &"a".repeat(40), - max_tokens: 2, - prior_tokens: &[], - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, &"a".repeat(40), 2, &[])) .unwrap(); let second = engine - .generate_collected(EngineRequest { - model: &model, - prompt: "z", - max_tokens: 1, - prior_tokens: &first.successor_tokens, - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, "z", 1, &first.successor_tokens)) .unwrap(); assert!(!second.pieces.is_empty()); let metrics = engine.metrics(); @@ -1001,13 +1196,10 @@ mod tests { .unwrap(); let model = model(); let first = engine - .generate_collected(EngineRequest { - model: &model, - prompt: &"a".repeat(40), - max_tokens: 2, - prior_tokens: &[], - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, &"a".repeat(40), 2, &[])) + .unwrap(); + engine + .generate_collected(test_request(&model, &"b".repeat(40), 1, &[])) .unwrap(); let control = MappedEngine::open( "gemma-4-e2b-it", @@ -1019,34 +1211,21 @@ mod tests { ) .unwrap(); let control_first = control - .generate_collected(EngineRequest { - model: &model, - prompt: &"a".repeat(40), - max_tokens: 2, - prior_tokens: &[], - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, &"a".repeat(40), 2, &[])) .unwrap(); let baseline = control - .generate_collected(EngineRequest { - model: &model, - prompt: "z", - max_tokens: 2, - prior_tokens: &control_first.successor_tokens, - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request( + &model, + "z", + 2, + &control_first.successor_tokens, + )) .unwrap(); assert_eq!(engine.spill_inactive_mappings().unwrap(), 1); assert_eq!(fs::read_dir(&spill_dir).unwrap().count(), 1); let resumed = engine - .generate_collected(EngineRequest { - model: &model, - prompt: "z", - max_tokens: 2, - prior_tokens: &first.successor_tokens, - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, "z", 2, &first.successor_tokens)) .unwrap(); assert_eq!(resumed.cached_tokens, 32); assert_eq!(resumed.pieces, baseline.pieces); @@ -1068,13 +1247,7 @@ mod tests { .unwrap(); let model = model(); let error = engine - .generate_collected(EngineRequest { - model: &model, - prompt: &"x".repeat(65), - max_tokens: 1, - prior_tokens: &[], - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, &"x".repeat(65), 1, &[])) .unwrap_err(); assert!(error.to_string().contains("context capacity")); assert_eq!(engine.metrics(), MappedMetrics::default()); @@ -1093,30 +1266,17 @@ mod tests { .unwrap(); let model = model(); let first = engine - .generate_collected(EngineRequest { - model: &model, - prompt: &"a".repeat(40), - max_tokens: 1, - prior_tokens: &[], - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, &"a".repeat(40), 1, &[])) .unwrap(); - let failed = engine.generate_collected(EngineRequest { - model: &model, - prompt: &"b".repeat(25), - max_tokens: 1, - prior_tokens: &first.successor_tokens, - control: &crate::RequestControl::new(), - }); + let failed = engine.generate_collected(test_request( + &model, + &"b".repeat(25), + 1, + &first.successor_tokens, + )); assert!(failed.unwrap_err().to_string().contains("cannot publish")); let resumed = engine - .generate_collected(EngineRequest { - model: &model, - prompt: "c", - max_tokens: 1, - prior_tokens: &first.successor_tokens, - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, "c", 1, &first.successor_tokens)) .unwrap(); assert_eq!(resumed.pieces.len(), 1); assert_eq!(engine.metrics().requests, 2); @@ -1136,22 +1296,10 @@ mod tests { .unwrap(); let model = model(); let primed = engine - .generate_collected(EngineRequest { - model: &model, - prompt: &"p".repeat(32), - max_tokens: 0, - prior_tokens: &[], - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, &"p".repeat(32), 0, &[])) .unwrap(); let resumed = engine - .generate_collected(EngineRequest { - model: &model, - prompt: "", - max_tokens: 1, - prior_tokens: &primed.successor_tokens, - control: &crate::RequestControl::new(), - }) + .generate_collected(test_request(&model, "", 1, &primed.successor_tokens)) .unwrap(); assert_eq!(resumed.pieces.len(), 1); assert_eq!(resumed.cached_tokens, 32); @@ -1189,7 +1337,7 @@ mod tests { max_tokens: 1, context_id: None, deadline_ms: None, - priority: 0, + scheduling: SchedulingMetadata::default(), stop: vec![], raw_continuation: false, }, @@ -1207,7 +1355,7 @@ mod tests { max_tokens: 1, context_id: Some(durable.id), deadline_ms: None, - priority: 0, + scheduling: SchedulingMetadata::default(), stop: vec![], raw_continuation: false, }, diff --git a/crates/server/src/residency.rs b/crates/server/src/residency.rs index fd4c137..8ef05c9 100644 --- a/crates/server/src/residency.rs +++ b/crates/server/src/residency.rs @@ -1,5 +1,5 @@ use crate::{ - EngineOutput, EngineRequest, Error, InferenceEngine, MappedEngine, ModelRecord, TokenSink, + EngineRequest, Error, ExecutionSession, InferenceEngine, MappedEngine, ModelRecord, SessionStep, }; use cusco_context_store::ModelEpoch; use cusco_executor::OperatingPoint; @@ -54,6 +54,7 @@ pub struct ResidencyMetrics { pub host_bytes: u64, pub storage_bytes: u64, pub active_slots: usize, + pub sessions: usize, } #[derive(Clone, Debug, Eq, PartialEq, Serialize)] @@ -62,6 +63,7 @@ pub struct ResidentModelStatus { pub revision: String, pub epoch: u64, pub active_slots: usize, + pub sessions: usize, pub retiring: bool, pub context_resident: bool, pub last_used: u64, @@ -112,6 +114,7 @@ struct ResidentModel { engine: Arc, point: OperatingPoint, active: AtomicUsize, + sessions: AtomicUsize, retiring: AtomicBool, context_resident: AtomicBool, last_used: AtomicU64, @@ -124,6 +127,7 @@ impl ResidentModel { revision: self.record.revision.clone(), epoch: self.record.epoch, active_slots: self.active.load(Ordering::Acquire), + sessions: self.sessions.load(Ordering::Acquire), retiring: self.retiring.load(Ordering::Acquire), context_resident: self.context_resident.load(Ordering::Acquire), last_used: self.last_used.load(Ordering::Acquire), @@ -141,7 +145,7 @@ struct ResidencyState { pub struct ResidentEngine { config: ResidencyConfig, - state: Mutex, + state: Arc>, loader: Arc, } @@ -158,7 +162,7 @@ impl ResidentEngine { fs::create_dir_all(&spill_dir).map_err(super::state_err)?; Ok(Arc::new(Self { config: config.validate()?, - state: Mutex::new(ResidencyState::default()), + state: Arc::new(Mutex::new(ResidencyState::default())), loader: Arc::new(NativeLoader { spill_dir }), })) } @@ -167,7 +171,7 @@ impl ResidentEngine { fn with_loader(config: ResidencyConfig, loader: Arc) -> Arc { Arc::new(Self { config: config.validate().unwrap(), - state: Mutex::new(ResidencyState::default()), + state: Arc::new(Mutex::new(ResidencyState::default())), loader, }) } @@ -204,7 +208,11 @@ impl ResidentEngine { competent: !self.config.require_competent || self.config.gpu_layers > 0, }) } - fn fits(&self, point: OperatingPoint, retained: impl Iterator) -> bool { + fn fits_with( + config: ResidencyConfig, + point: OperatingPoint, + retained: impl Iterator, + ) -> bool { let (device, host, storage) = retained.fold( (point.device_bytes, point.host_bytes, point.model_bytes), |(device, host, storage), item| { @@ -215,9 +223,13 @@ impl ResidentEngine { ) }, ); - device <= self.config.device_bytes - && host <= self.config.host_bytes - && storage <= self.config.storage_bytes + device <= config.device_bytes + && host <= config.host_bytes + && storage <= config.storage_bytes + } + + fn fits(&self, point: OperatingPoint, retained: impl Iterator) -> bool { + Self::fits_with(self.config, point, retained) } fn effective_point(entry: &ResidentModel) -> OperatingPoint { @@ -259,7 +271,9 @@ impl ResidentEngine { .entries .values() .filter(|entry| { - entry.active.load(Ordering::Acquire) == 0 && entry.record.id != model.id + entry.active.load(Ordering::Acquire) == 0 + && entry.sessions.load(Ordering::Acquire) == 0 + && entry.record.id != model.id }) .cloned() .collect::>(); @@ -359,6 +373,7 @@ impl ResidentEngine { engine, point, active: AtomicUsize::new(0), + sessions: AtomicUsize::new(0), retiring: AtomicBool::new(false), context_resident: AtomicBool::new(true), last_used: AtomicU64::new(last_used), @@ -390,6 +405,11 @@ impl ResidentEngine { .values() .map(|entry| entry.active.load(Ordering::Acquire)) .sum(); + state.metrics.sessions = state + .entries + .values() + .map(|entry| entry.sessions.load(Ordering::Acquire)) + .sum(); } fn retire_other_epochs(&self, model: &ModelRecord) { @@ -402,15 +422,22 @@ impl ResidentEngine { .collect::>(); for (key, entry) in keys { entry.retiring.store(true, Ordering::Release); - if entry.active.load(Ordering::Acquire) == 0 && state.entries.remove(&key).is_some() { + if entry.active.load(Ordering::Acquire) == 0 + && entry.sessions.load(Ordering::Acquire) == 0 + && state.entries.remove(&key).is_some() + { state.metrics.unloads += 1; state.metrics.reloads += 1; } } Self::refresh_metrics(&mut state); } - fn acquire(&self, target: &Arc) -> Result { - let mut state = self.state.lock(); + fn acquire_entry( + config: ResidencyConfig, + shared: &Arc>, + target: &Arc, + ) -> Result { + let mut state = shared.lock(); let key = (target.record.id.clone(), target.record.epoch); let Some(current) = state.entries.get(&key) else { return Ok(false); @@ -425,6 +452,7 @@ impl ResidentEngine { .filter(|entry| { !Arc::ptr_eq(entry, target) && entry.active.load(Ordering::Acquire) == 0 + && entry.sessions.load(Ordering::Acquire) == 0 && entry.context_resident.load(Ordering::Acquire) }) .cloned() @@ -436,7 +464,7 @@ impl ResidentEngine { .values() .filter(|entry| !Arc::ptr_eq(entry, target)) .map(|entry| Self::effective_point(entry)); - if self.fits(target.point, retained) { + if Self::fits_with(config, target.point, retained) { target.context_resident.store(true, Ordering::Release); state.metrics.context_reloads += 1; break; @@ -463,10 +491,17 @@ impl ResidentEngine { Ok(true) } - fn release(&self, entry: &Arc) { + fn acquire(&self, target: &Arc) -> Result { + Self::acquire_entry(self.config, &self.state, target) + } + + fn release_entry(shared: &Arc>, entry: &Arc) { entry.active.fetch_sub(1, Ordering::AcqRel); - let mut state = self.state.lock(); - if entry.retiring.load(Ordering::Acquire) && entry.active.load(Ordering::Acquire) == 0 { + let mut state = shared.lock(); + if entry.retiring.load(Ordering::Acquire) + && entry.active.load(Ordering::Acquire) == 0 + && entry.sessions.load(Ordering::Acquire) == 0 + { let key = (entry.record.id.clone(), entry.record.epoch); if state.entries.remove(&key).is_some() { state.metrics.unloads += 1; @@ -475,6 +510,10 @@ impl ResidentEngine { Self::refresh_metrics(&mut state); } + fn release(&self, entry: &Arc) { + Self::release_entry(&self.state, entry); + } + pub fn metrics(&self) -> ResidencyMetrics { let mut state = self.state.lock(); Self::refresh_metrics(&mut state); @@ -494,6 +533,72 @@ impl ResidentEngine { } } +struct ResidentSession { + config: ResidencyConfig, + entry: Arc, + state: Arc>, + inner: Box, +} + +struct ResidentQuantumLease { + entry: Arc, + state: Arc>, +} + +impl Drop for ResidentQuantumLease { + fn drop(&mut self) { + ResidentEngine::release_entry(&self.state, &self.entry); + } +} + +impl ResidentSession { + fn acquire(&self) -> Result { + if !ResidentEngine::acquire_entry(self.config, &self.state, &self.entry)? { + return Err(Error::State( + "resident model disappeared while its session was suspended".into(), + )); + } + Ok(ResidentQuantumLease { + entry: self.entry.clone(), + state: self.state.clone(), + }) + } +} + +impl ExecutionSession for ResidentSession { + fn step(&mut self) -> Result { + let lease = self.acquire()?; + let result = self.inner.step(); + drop(lease); + result + } + + fn finish(&mut self) -> Result { + let lease = self.acquire()?; + let result = self.inner.finish(); + drop(lease); + result + } +} + +impl Drop for ResidentSession { + fn drop(&mut self) { + let previous = self.entry.sessions.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "resident session count underflow"); + let mut state = self.state.lock(); + if self.entry.retiring.load(Ordering::Acquire) + && self.entry.active.load(Ordering::Acquire) == 0 + && self.entry.sessions.load(Ordering::Acquire) == 0 + { + let key = (self.entry.record.id.clone(), self.entry.record.epoch); + if state.entries.remove(&key).is_some() { + state.metrics.unloads += 1; + } + } + ResidentEngine::refresh_metrics(&mut state); + } +} + impl InferenceEngine for ResidentEngine { fn prepare_model(&self, model: &ModelRecord) -> Result<(), Error> { self.load(model, None)?; @@ -509,7 +614,9 @@ impl InferenceEngine for ResidentEngine { let key = (id.to_owned(), epoch); if let Some(entry) = state.entries.get(&key).cloned() { entry.retiring.store(true, Ordering::Release); - if entry.active.load(Ordering::Acquire) == 0 { + if entry.active.load(Ordering::Acquire) == 0 + && entry.sessions.load(Ordering::Acquire) == 0 + { state.entries.remove(&key); state.metrics.unloads += 1; } @@ -525,27 +632,37 @@ impl InferenceEngine for ResidentEngine { })) } - fn generate( - &self, - request: EngineRequest<'_>, - sink: &mut TokenSink<'_>, - ) -> Result { + fn start_session(&self, request: EngineRequest) -> Result, Error> { let entry = loop { - let entry = self.load(request.model, Some(request.control))?; + let entry = self.load(&request.model, Some(&request.control))?; if self.acquire(&entry)? { break entry; } }; - let result = entry.engine.generate(request, sink); - self.release(&entry); - result + let started = entry.engine.start_session(request); + match started { + Ok(inner) => { + entry.sessions.fetch_add(1, Ordering::AcqRel); + self.release(&entry); + Ok(Box::new(ResidentSession { + config: self.config, + entry, + state: self.state.clone(), + inner, + })) + } + Err(error) => { + self.release(&entry); + Err(error) + } + } } } #[cfg(test)] mod tests { use super::*; - use crate::{DeterministicEngine, FrontierControl, RequestControl}; + use crate::{DeterministicEngine, FrontierControl, RequestControl, SchedulingMetadata}; use std::{path::PathBuf, sync::Barrier, thread}; struct FixtureLoader { @@ -559,27 +676,50 @@ mod tests { release: Arc, } + struct BlockingSession { + entered: Arc, + release: Arc, + inner: Box, + blocked: bool, + } + + impl ExecutionSession for BlockingSession { + fn step(&mut self) -> Result { + if !self.blocked { + self.entered.wait(); + self.release.wait(); + self.blocked = true; + } + self.inner.step() + } + + fn finish(&mut self) -> Result { + self.inner.finish() + } + } + impl InferenceEngine for BlockingEngine { - fn generate( + fn start_session( &self, - request: EngineRequest<'_>, - sink: &mut TokenSink<'_>, - ) -> Result { - self.entered.wait(); - self.release.wait(); - DeterministicEngine.generate(request, sink) + request: EngineRequest, + ) -> Result, Error> { + Ok(Box::new(BlockingSession { + entered: self.entered.clone(), + release: self.release.clone(), + inner: DeterministicEngine.start_session(request)?, + blocked: false, + })) } } struct DemotableEngine; impl InferenceEngine for DemotableEngine { - fn generate( + fn start_session( &self, - request: EngineRequest<'_>, - sink: &mut TokenSink<'_>, - ) -> Result { - DeterministicEngine.generate(request, sink) + request: EngineRequest, + ) -> Result, Error> { + DeterministicEngine.start_session(request) } fn demote_inactive(&self) -> Result { @@ -587,6 +727,23 @@ mod tests { } } + fn test_request( + model: &ModelRecord, + prompt: impl Into, + max_tokens: usize, + control: Arc, + ) -> EngineRequest { + EngineRequest { + model: model.clone(), + prompt: prompt.into(), + max_tokens, + prior_tokens: vec![], + control, + scheduling: SchedulingMetadata::default(), + prefill_chunk_tokens: 32, + } + } + struct DemotableLoader { point: OperatingPoint, } @@ -663,6 +820,32 @@ mod tests { } } + #[test] + fn suspended_session_releases_its_native_slot_without_becoming_evictable() { + let loader = Arc::new(FixtureLoader { + point: point(10), + failures: Mutex::new(vec![]), + blocker: None, + }); + let engine = ResidentEngine::with_loader(config(21), loader); + let model = model("suspended", "r", 1, 10); + let mut session = engine + .start_session(test_request( + &model, + "yield", + 2, + Arc::new(RequestControl::new()), + )) + .unwrap(); + + assert_eq!(engine.status()[0].active_slots, 0); + assert_eq!(engine.status()[0].sessions, 1); + let _ = session.step().unwrap(); + assert_eq!(engine.status()[0].active_slots, 0); + assert_eq!(engine.metrics().active_slots, 0); + drop(session); + assert_eq!(engine.status()[0].sessions, 0); + } #[test] fn failed_reload_preserves_the_published_epoch() { let loader = Arc::new(FixtureLoader { @@ -706,19 +889,10 @@ mod tests { engine.prepare_model(&first).unwrap(); engine.prepare_model(&second).unwrap(); - let control = RequestControl::new(); + let control = Arc::new(RequestControl::new()); let mut sink = |_: i32, _: &[u8], _: bool| Ok(FrontierControl::Continue); engine - .generate( - EngineRequest { - model: &first, - prompt: "touch", - max_tokens: 1, - prior_tokens: &[], - control: &control, - }, - &mut sink, - ) + .generate(test_request(&first, "touch", 1, control), &mut sink) .unwrap(); engine.prepare_model(&third).unwrap(); @@ -763,19 +937,10 @@ mod tests { let worker_engine = engine.clone(); let worker = thread::spawn(move || { - let control = RequestControl::new(); + let control = Arc::new(RequestControl::new()); let mut sink = |_: i32, _: &[u8], _: bool| Ok(FrontierControl::Continue); worker_engine - .generate( - EngineRequest { - model: &old, - prompt: "old epoch", - max_tokens: 1, - prior_tokens: &[], - control: &control, - }, - &mut sink, - ) + .generate(test_request(&old, "old epoch", 1, control), &mut sink) .unwrap(); }); entered.wait(); @@ -852,12 +1017,8 @@ mod tests { limits.gpu_layers = 0; let engine = ResidentEngine::with_loader(limits, Arc::new(DemotableLoader { point })); - engine - .prepare_model(&model("first", "r", 1, 10)) - .unwrap(); - engine - .prepare_model(&model("second", "r", 2, 10)) - .unwrap(); + engine.prepare_model(&model("first", "r", 1, 10)).unwrap(); + engine.prepare_model(&model("second", "r", 2, 10)).unwrap(); assert_eq!(engine.status().len(), 2); assert_eq!(engine.metrics().context_spills, 1); @@ -873,18 +1034,12 @@ mod tests { blocker: None, }); let engine = ResidentEngine::with_loader(config(100), loader); - let control = RequestControl::new(); + let control = Arc::new(RequestControl::new()); control.cancel(); let mut sink = |_: i32, _: &[u8], _: bool| Ok(FrontierControl::Continue); let error = engine .generate( - EngineRequest { - model: &model("cancelled", "r", 1, 10), - prompt: "cancel", - max_tokens: 1, - prior_tokens: &[], - control: &control, - }, + test_request(&model("cancelled", "r", 1, 10), "cancel", 1, control), &mut sink, ) .unwrap_err(); diff --git a/crates/server/src/scheduler.rs b/crates/server/src/scheduler.rs new file mode 100644 index 0000000..7c84f5e --- /dev/null +++ b/crates/server/src/scheduler.rs @@ -0,0 +1,1093 @@ +use crate::{ + EngineOutput, EngineRequest, Error, ExecutionSession, InferenceEngine, PrioritySource, + QuantumKind, QuantumObservation, SchedulerPolicyConfig, SchedulingClass, SessionStep, +}; +use parking_lot::Mutex; +use serde::Serialize; +use serde_json::{Value, json}; +use std::{ + collections::{HashMap, VecDeque}, + io::Write, + sync::{ + Arc, + atomic::{AtomicU64, AtomicUsize, Ordering}, + mpsc::{self, Receiver, Sender, SyncSender, TrySendError}, + }, + thread, + time::{Duration, Instant}, +}; +use uuid::Uuid; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)] +pub struct SchedulerMetrics { + pub admitted: u64, + pub completed: u64, + pub cancelled: u64, + pub deadlines: u64, + pub failed: u64, + pub rounds: u64, + pub runnable: usize, + pub waiting_for_consumer: usize, + pub diagnostic_records: u64, + pub diagnostic_records_lost: u64, + pub diagnostic_records_delivered: u64, +} + +#[derive(Clone, Debug, Serialize)] +pub struct SchedulerStatus { + pub policy: SchedulerPolicyConfig, + pub metrics: SchedulerMetrics, + pub diagnostic_loss_by_kind: HashMap, +} + +#[derive(Clone, Debug, Serialize)] +struct SchedulerDecision { + #[serde(rename = "type")] + record_type: &'static str, + transport_correlation_id: String, + inference_operation_id: String, + execution_session_id: String, + principal: String, + class: SchedulingClass, + priority_source: PrioritySource, + model_id: String, + model_epoch: u64, + queue_age_ms: u128, + queue_age_rounds: u64, + age_promotions: u8, + quantum_kind: QuantumKind, + charged_tokens: usize, + quantum_duration_ns: u128, + context_placement: String, + executor_slot_occupied: bool, + transition_cost_bytes: u64, + capacity_reserved_bytes: u64, + cancelled: bool, + deadline_expired: bool, + reason: &'static str, +} + +#[derive(Clone, Debug, Serialize)] +struct DiagnosticLossSummary { + #[serde(rename = "type")] + record_type: &'static str, + total: u64, + by_kind: HashMap, +} + +#[derive(Default)] +struct DiagnosticLoss { + pending_total: u64, + pending_by_kind: HashMap, + all_by_kind: HashMap, +} + +struct SchedulerDiagnostics { + sender: SyncSender, + emitted: AtomicU64, + delivered: Arc, + lost: AtomicU64, + loss: Mutex, +} + +impl SchedulerDiagnostics { + fn stderr(capacity: usize) -> Arc { + Self::with_sink(capacity, |line| { + let mut stderr = std::io::stderr().lock(); + let _ = writeln!(stderr, "{line}"); + }) + } + + fn with_sink(capacity: usize, sink: impl Fn(&str) + Send + 'static) -> Arc { + let delivered = Arc::new(AtomicU64::new(0)); + let worker_delivered = delivered.clone(); + let (sender, receiver) = mpsc::sync_channel::(capacity); + thread::Builder::new() + .name("cusco-scheduler-diagnostics".into()) + .spawn(move || { + while let Ok(line) = receiver.recv() { + sink(&line); + worker_delivered.fetch_add(1, Ordering::Release); + } + }) + .expect("scheduler diagnostic worker starts"); + Arc::new(Self { + sender, + emitted: AtomicU64::new(0), + lost: AtomicU64::new(0), + delivered, + loss: Mutex::new(DiagnosticLoss::default()), + }) + } + + fn emit(&self, kind: &str, record: &T) { + let Ok(line) = serde_json::to_string(record) else { + self.record_loss("serialization_error"); + return; + }; + self.flush_loss_summary(); + match self.sender.try_send(line) { + Ok(()) => { + self.emitted.fetch_add(1, Ordering::Relaxed); + } + Err(TrySendError::Full(_)) | Err(TrySendError::Disconnected(_)) => { + self.record_loss(kind); + } + } + } + + fn flush_loss_summary(&self) { + let summary = { + let loss = self.loss.lock(); + if loss.pending_total == 0 { + return; + } + DiagnosticLossSummary { + record_type: "scheduler_diagnostic_loss", + total: loss.pending_total, + by_kind: loss.pending_by_kind.clone(), + } + }; + let Ok(line) = serde_json::to_string(&summary) else { + return; + }; + if self.sender.try_send(line).is_ok() { + self.emitted.fetch_add(1, Ordering::Relaxed); + let mut loss = self.loss.lock(); + loss.pending_total = 0; + loss.pending_by_kind.clear(); + } + } + + fn record_loss(&self, kind: &str) { + self.lost.fetch_add(1, Ordering::Relaxed); + let mut loss = self.loss.lock(); + loss.pending_total = loss.pending_total.saturating_add(1); + *loss.pending_by_kind.entry(kind.into()).or_default() += 1; + *loss.all_by_kind.entry(kind.into()).or_default() += 1; + } + + fn loss_by_kind(&self) -> HashMap { + self.loss.lock().all_by_kind.clone() + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +struct FlowKey { + principal: String, + class: SchedulingClass, +} + +struct PolicyEntry { + id: u64, + flow: FlowKey, + admitted_round: u64, + admitted_at: Instant, +} + +struct Selection { + id: u64, + queue_age_rounds: u64, + queue_age_ms: u128, + promotions: u8, + charge_scale: i64, +} + +struct FairPolicy { + config: SchedulerPolicyConfig, + runnable: VecDeque, + deficits: HashMap, + round: u64, + cursor: usize, +} + +impl FairPolicy { + fn new(config: SchedulerPolicyConfig) -> Self { + Self { + config, + runnable: VecDeque::new(), + deficits: HashMap::new(), + round: 0, + cursor: 0, + } + } + + fn enqueue(&mut self, id: u64, flow: FlowKey, admitted_round: u64, admitted_at: Instant) { + self.runnable.push_back(PolicyEntry { + id, + flow, + admitted_round, + admitted_at, + }); + } + + fn effective_class(&self, entry: &PolicyEntry) -> (SchedulingClass, u8) { + let promotions = ((self.round.saturating_sub(entry.admitted_round)) + / self.config.promotion_rounds) + .min(2) as u8; + let class = match (entry.flow.class, promotions) { + (SchedulingClass::Batch, 1) => SchedulingClass::Standard, + (SchedulingClass::Batch, 2) | (SchedulingClass::Standard, 1..) => { + SchedulingClass::Interactive + } + (class, _) => class, + }; + (class, promotions) + } + + fn weight(&self, class: SchedulingClass) -> i64 { + let weight = match class { + SchedulingClass::Interactive => self.config.interactive_weight, + SchedulingClass::Standard => self.config.standard_weight, + SchedulingClass::Batch => self.config.batch_weight, + }; + i64::from(weight.saturating_mul(self.config.deficit_refill)) + } + + fn select(&mut self) -> Option { + if self.runnable.is_empty() { + return None; + } + self.round = self.round.saturating_add(1); + let mut heads = Vec::<(usize, FlowKey, SchedulingClass, u8, i64)>::new(); + for (index, entry) in self.runnable.iter().enumerate() { + if heads.iter().any(|(_, flow, _, _, _)| flow == &entry.flow) { + continue; + } + let (class, promotions) = self.effective_class(entry); + let credit = self.weight(class); + *self.deficits.entry(entry.flow.clone()).or_default() = self + .deficits + .get(&entry.flow) + .copied() + .unwrap_or_default() + .saturating_add(credit) + .min(credit.saturating_mul(8)); + heads.push((index, entry.flow.clone(), class, promotions, credit)); + } + let total_weight = heads + .iter() + .map(|(_, _, _, _, weight)| *weight) + .sum::() + .max(1); + let max_deficit = heads + .iter() + .map(|(_, flow, _, _, _)| self.deficits.get(flow).copied().unwrap_or_default()) + .max() + .unwrap_or_default(); + let eligible = heads + .iter() + .filter(|(_, flow, _, _, _)| { + self.deficits.get(flow).copied().unwrap_or_default() == max_deficit + }) + .collect::>(); + let selected = eligible[self.cursor % eligible.len()]; + self.cursor = (self.cursor + 1) % eligible.len(); + let entry = self + .runnable + .remove(selected.0) + .expect("selected entry exists"); + Some(Selection { + id: entry.id, + queue_age_rounds: self.round.saturating_sub(entry.admitted_round), + queue_age_ms: entry.admitted_at.elapsed().as_millis(), + promotions: selected.3, + charge_scale: total_weight, + }) + } + + fn charge(&mut self, flow: &FlowKey, tokens: usize, scale: i64) { + *self.deficits.entry(flow.clone()).or_default() -= + (tokens.max(1) as i64).saturating_mul(scale); + } + + fn remove(&mut self, id: u64) -> bool { + let previous = self.runnable.len(); + self.runnable.retain(|entry| entry.id != id); + self.runnable.len() != previous + } +} + +enum ClientMessage { + Token { + id: i32, + piece: Vec, + terminal_or_control: bool, + }, + Finished(EngineOutput), + Failed(Error), +} + +enum Command { + Submit { + id: u64, + request: EngineRequest, + response: Sender, + }, + Continue(u64), + Stop(u64), + Cancel(u64), + Shutdown, +} + +struct Job { + request: Option, + session: Option>, + control: Arc, + response: Sender, + flow: FlowKey, + scheduling: crate::SchedulingMetadata, + model_id: String, + model_epoch: u64, + execution_session_id: String, + admitted_round: u64, + admitted_at: Instant, +} + +struct SchedulerCounters { + admitted: AtomicU64, + completed: AtomicU64, + cancelled: AtomicU64, + deadlines: AtomicU64, + failed: AtomicU64, + rounds: AtomicU64, + runnable: AtomicUsize, + waiting: AtomicUsize, +} + +impl Default for SchedulerCounters { + fn default() -> Self { + Self { + admitted: AtomicU64::new(0), + completed: AtomicU64::new(0), + cancelled: AtomicU64::new(0), + deadlines: AtomicU64::new(0), + failed: AtomicU64::new(0), + rounds: AtomicU64::new(0), + runnable: AtomicUsize::new(0), + waiting: AtomicUsize::new(0), + } + } +} + +pub struct WorkloadScheduler { + inner: Arc, + commands: Sender, + next_id: AtomicU64, + policy: SchedulerPolicyConfig, + counters: Arc, + diagnostics: Arc, + worker: Mutex>>, +} + +impl WorkloadScheduler { + pub fn new( + inner: Arc, + policy: SchedulerPolicyConfig, + ) -> Result, Error> { + Self::new_with_diagnostics(inner, policy, None::) + } + pub fn new_with_diagnostics( + inner: Arc, + policy: SchedulerPolicyConfig, + sink: Option, + ) -> Result, Error> + where + F: Fn(&str) + Send + 'static, + { + let policy = policy.validate()?; + let diagnostics = match sink { + Some(sink) => SchedulerDiagnostics::with_sink(policy.diagnostic_capacity, sink), + None => SchedulerDiagnostics::stderr(policy.diagnostic_capacity), + }; + let counters = Arc::new(SchedulerCounters::default()); + let (commands, receiver) = mpsc::channel(); + let worker_inner = inner.clone(); + let worker_counters = counters.clone(); + let worker_diagnostics = diagnostics.clone(); + let worker = thread::Builder::new() + .name("cusco-workload-scheduler".into()) + .spawn(move || { + run_scheduler( + worker_inner, + policy, + receiver, + worker_counters, + worker_diagnostics, + ) + }) + .map_err(|error| Error::State(error.to_string()))?; + Ok(Arc::new(Self { + inner, + commands, + next_id: AtomicU64::new(0), + policy, + counters, + diagnostics, + worker: Mutex::new(Some(worker)), + })) + } + + pub fn status(&self) -> SchedulerStatus { + SchedulerStatus { + policy: self.policy, + metrics: SchedulerMetrics { + admitted: self.counters.admitted.load(Ordering::Acquire), + completed: self.counters.completed.load(Ordering::Acquire), + cancelled: self.counters.cancelled.load(Ordering::Acquire), + deadlines: self.counters.deadlines.load(Ordering::Acquire), + failed: self.counters.failed.load(Ordering::Acquire), + rounds: self.counters.rounds.load(Ordering::Acquire), + runnable: self.counters.runnable.load(Ordering::Acquire), + waiting_for_consumer: self.counters.waiting.load(Ordering::Acquire), + diagnostic_records: self.diagnostics.emitted.load(Ordering::Acquire), + diagnostic_records_lost: self.diagnostics.lost.load(Ordering::Acquire), + diagnostic_records_delivered: self.diagnostics.delivered.load(Ordering::Acquire), + }, + diagnostic_loss_by_kind: self.diagnostics.loss_by_kind(), + } + } + + pub fn flush_diagnostics(&self, timeout: Duration) -> bool { + let target = self.diagnostics.emitted.load(Ordering::Acquire); + let started = Instant::now(); + while self.diagnostics.delivered.load(Ordering::Acquire) < target { + if started.elapsed() >= timeout { + return false; + } + thread::sleep(Duration::from_millis(1)); + } + true + } +} + +impl Drop for WorkloadScheduler { + fn drop(&mut self) { + let _ = self.commands.send(Command::Shutdown); + if let Some(worker) = self.worker.lock().take() { + let _ = worker.join(); + } + } +} + +struct SchedulerClientSession { + id: u64, + control: Arc, + commands: Sender, + receiver: Receiver, + waiting_ack: bool, + finished: Option, +} + +impl SchedulerClientSession { + fn receive(&mut self) -> Result { + match self.receiver.recv().map_err(|_| Error::Cancelled)? { + ClientMessage::Token { + id, + piece, + terminal_or_control, + } => { + self.waiting_ack = true; + Ok(SessionStep::Token { + id, + piece, + terminal_or_control, + observation: QuantumObservation::model_free(QuantumKind::Decode, 1), + }) + } + ClientMessage::Finished(output) => { + self.finished = Some(output.clone()); + Ok(SessionStep::Finished(output)) + } + ClientMessage::Failed(error) => Err(error), + } + } +} + +impl ExecutionSession for SchedulerClientSession { + fn step(&mut self) -> Result { + if let Some(output) = &self.finished { + return Ok(SessionStep::Finished(output.clone())); + } + if self.waiting_ack { + self.commands + .send(Command::Continue(self.id)) + .map_err(|_| Error::Cancelled)?; + self.waiting_ack = false; + } + self.receive() + } + + fn finish(&mut self) -> Result { + if let Some(output) = &self.finished { + return Ok(output.clone()); + } + if self.waiting_ack { + self.commands + .send(Command::Stop(self.id)) + .map_err(|_| Error::Cancelled)?; + self.waiting_ack = false; + } else { + self.commands + .send(Command::Cancel(self.id)) + .map_err(|_| Error::Cancelled)?; + } + match self.receive()? { + SessionStep::Finished(output) => Ok(output), + _ => Err(Error::State("scheduler returned work after stop".into())), + } + } +} + +impl Drop for SchedulerClientSession { + fn drop(&mut self) { + if self.finished.is_none() { + self.control.cancel(); + let _ = self.commands.send(Command::Cancel(self.id)); + } + } +} + +impl InferenceEngine for WorkloadScheduler { + fn start_session( + &self, + mut request: EngineRequest, + ) -> Result, Error> { + request.control.check()?; + if request.scheduling.correlation_id.is_empty() { + request.scheduling.correlation_id = Uuid::new_v4().to_string(); + } + if request.scheduling.inference_id.is_empty() { + request.scheduling.inference_id = Uuid::new_v4().to_string(); + } + request.prefill_chunk_tokens = self.policy.prefill_tokens; + let control = request.control.clone(); + let id = self + .next_id + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |id| id.checked_add(1)) + .map_err(|_| Error::State("scheduler session id space exhausted".into()))?; + let (response, receiver) = mpsc::channel(); + self.commands + .send(Command::Submit { + id, + request, + response, + }) + .map_err(|_| Error::ShuttingDown)?; + Ok(Box::new(SchedulerClientSession { + id, + control, + commands: self.commands.clone(), + receiver, + waiting_ack: false, + finished: None, + })) + } + fn prepare_model(&self, model: &crate::ModelRecord) -> Result<(), Error> { + self.inner.prepare_model(model) + } + + fn commit_model(&self, model: &crate::ModelRecord, replaced_epoch: Option) { + self.inner.commit_model(model, replaced_epoch); + } + + fn retire_model(&self, id: &str, epoch: u64) { + self.inner.retire_model(id, epoch); + } + + fn demote_inactive(&self) -> Result { + self.inner.demote_inactive() + } + + fn residency_status(&self) -> Option { + let mut status = self.inner.residency_status().unwrap_or_else(|| json!({})); + if let Some(object) = status.as_object_mut() { + object.insert( + "scheduler".into(), + serde_json::to_value(self.status()).expect("scheduler status serializes"), + ); + } + Some(status) + } +} + +fn run_scheduler( + inner: Arc, + config: SchedulerPolicyConfig, + commands: Receiver, + counters: Arc, + diagnostics: Arc, +) { + let mut policy = FairPolicy::new(config); + let mut jobs = HashMap::::new(); + loop { + while let Ok(command) = commands.try_recv() { + if matches!(command, Command::Shutdown) { + cancel_all_jobs(&mut jobs); + return; + } + handle_command(command, &mut jobs, &mut policy, &counters); + } + counters + .runnable + .store(policy.runnable.len(), Ordering::Release); + if let Some(selection) = policy.select() { + counters.rounds.store(policy.round, Ordering::Release); + execute_quantum( + selection, + &inner, + &mut jobs, + &mut policy, + &counters, + &diagnostics, + ); + continue; + } + match commands.recv() { + Ok(Command::Shutdown) | Err(_) => { + cancel_all_jobs(&mut jobs); + break; + } + Ok(command) => handle_command(command, &mut jobs, &mut policy, &counters), + } + } +} + +fn cancel_all_jobs(jobs: &mut HashMap) { + for (_, job) in jobs.drain() { + job.control.cancel(); + let _ = job.response.send(ClientMessage::Failed(Error::Cancelled)); + } +} + +fn handle_command( + command: Command, + jobs: &mut HashMap, + policy: &mut FairPolicy, + counters: &SchedulerCounters, +) { + match command { + Command::Submit { + id, + request, + response, + } => { + let flow = FlowKey { + principal: request.scheduling.principal.clone(), + class: request.scheduling.class, + }; + let admitted_at = Instant::now(); + let admitted_round = policy.round; + let scheduling = request.scheduling.clone(); + let model_id = request.model.id.clone(); + let model_epoch = request.model.epoch; + let control = request.control.clone(); + jobs.insert( + id, + Job { + request: Some(request), + session: None, + control, + response, + flow: flow.clone(), + scheduling, + model_id, + model_epoch, + execution_session_id: Uuid::new_v4().to_string(), + admitted_round, + admitted_at, + }, + ); + policy.enqueue(id, flow, admitted_round, admitted_at); + counters.admitted.fetch_add(1, Ordering::Relaxed); + } + Command::Continue(id) => { + if let Some(job) = jobs.get(&id) { + policy.enqueue(id, job.flow.clone(), job.admitted_round, job.admitted_at); + counters.waiting.fetch_sub(1, Ordering::AcqRel); + } + } + Command::Stop(id) => { + counters.waiting.fetch_sub(1, Ordering::AcqRel); + if let Some(mut job) = jobs.remove(&id) { + let result = job + .session + .as_mut() + .ok_or_else(|| Error::State("scheduler session was not started".into())) + .and_then(|session| session.finish()); + send_terminal(job, result, counters); + } + } + Command::Cancel(id) => { + let was_runnable = policy.remove(id); + if let Some(job) = jobs.remove(&id) { + if job.session.is_some() && !was_runnable { + counters.waiting.fetch_sub(1, Ordering::AcqRel); + } + job.control.cancel(); + let _ = job.response.send(ClientMessage::Failed(Error::Cancelled)); + counters.cancelled.fetch_add(1, Ordering::Relaxed); + } + } + Command::Shutdown => unreachable!("shutdown is handled by the scheduler loop"), + } +} + +fn execute_quantum( + selection: Selection, + inner: &Arc, + jobs: &mut HashMap, + policy: &mut FairPolicy, + counters: &SchedulerCounters, + diagnostics: &SchedulerDiagnostics, +) { + let Some(job) = jobs.get_mut(&selection.id) else { + return; + }; + let quantum_started = Instant::now(); + let result = if job.session.is_none() { + let request = job.request.take().expect("unstarted job retains request"); + inner.start_session(request).map(|session| { + job.session = Some(session); + SessionStep::Progress(QuantumObservation::model_free(QuantumKind::Preparation, 1)) + }) + } else { + job.session.as_mut().expect("session exists").step() + }; + + let (observation, wait_for_consumer, terminal) = match result { + Ok(SessionStep::Progress(observation)) => (observation, false, None), + Ok(SessionStep::Token { + id, + piece, + terminal_or_control, + observation, + }) => { + if job + .response + .send(ClientMessage::Token { + id, + piece, + terminal_or_control, + }) + .is_err() + { + (observation, false, Some(Err(Error::Cancelled))) + } else { + counters.waiting.fetch_add(1, Ordering::AcqRel); + (observation, true, None) + } + } + Ok(SessionStep::Finished(output)) => ( + QuantumObservation::model_free(QuantumKind::Publication, 1), + false, + Some(Ok(output)), + ), + Err(error) => ( + QuantumObservation::model_free(QuantumKind::Preparation, 1), + false, + Some(Err(error)), + ), + }; + let quantum_duration_ns = quantum_started.elapsed().as_nanos(); + policy.charge( + &job.flow, + observation.charged_tokens, + selection.charge_scale, + ); + diagnostics.emit( + "scheduler_decision", + &SchedulerDecision { + record_type: "scheduler_decision", + transport_correlation_id: job.scheduling.correlation_id.clone(), + inference_operation_id: job.scheduling.inference_id.clone(), + execution_session_id: job.execution_session_id.clone(), + principal: job.scheduling.principal.clone(), + class: job.scheduling.class, + priority_source: job.scheduling.source, + model_id: job.model_id.clone(), + model_epoch: job.model_epoch, + queue_age_ms: selection.queue_age_ms, + queue_age_rounds: selection.queue_age_rounds, + age_promotions: selection.promotions, + quantum_kind: observation.kind, + charged_tokens: observation.charged_tokens, + quantum_duration_ns, + context_placement: observation.context_placement, + executor_slot_occupied: observation.executor_slot_occupied, + transition_cost_bytes: observation.transition_cost_bytes, + capacity_reserved_bytes: observation.capacity_reserved_bytes, + cancelled: terminal + .as_ref() + .is_some_and(|result| matches!(result, Err(Error::Cancelled))), + deadline_expired: terminal + .as_ref() + .is_some_and(|result| matches!(result, Err(Error::Deadline))), + reason: if selection.promotions > 0 { + "age_promoted" + } else { + "class_ready" + }, + }, + ); + + if let Some(result) = terminal { + let job = jobs.remove(&selection.id).expect("terminal job exists"); + send_terminal(job, result, counters); + } else if !wait_for_consumer { + let job = jobs.get(&selection.id).expect("runnable job exists"); + policy.enqueue( + selection.id, + job.flow.clone(), + job.admitted_round, + job.admitted_at, + ); + } +} + +fn send_terminal(job: Job, result: Result, counters: &SchedulerCounters) { + match result { + Ok(output) => { + let _ = job.response.send(ClientMessage::Finished(output)); + counters.completed.fetch_add(1, Ordering::Relaxed); + } + Err(Error::Cancelled) => { + let _ = job.response.send(ClientMessage::Failed(Error::Cancelled)); + counters.cancelled.fetch_add(1, Ordering::Relaxed); + } + Err(Error::Deadline) => { + let _ = job.response.send(ClientMessage::Failed(Error::Deadline)); + counters.deadlines.fetch_add(1, Ordering::Relaxed); + } + Err(error) => { + let _ = job.response.send(ClientMessage::Failed(error)); + counters.failed.fetch_add(1, Ordering::Relaxed); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{DeterministicEngine, ModelRecord, RequestControl}; + use parking_lot::Mutex; + use std::{path::PathBuf, time::Duration}; + + fn request(class: SchedulingClass, principal: &str, words: usize) -> EngineRequest { + EngineRequest { + model: ModelRecord { + id: "m".into(), + revision: "r".into(), + path: PathBuf::from("m.gguf"), + sha256: String::new(), + aliases: vec![], + family: "gemma-4-e2b-it".into(), + size_bytes: 1, + epoch: 1, + }, + prompt: std::iter::repeat_n("x", words) + .collect::>() + .join(" "), + max_tokens: 2, + prior_tokens: vec![], + control: Arc::new(RequestControl::new()), + scheduling: crate::SchedulingMetadata { + class, + source: PrioritySource::ControlledWorkload, + principal: principal.into(), + correlation_id: Uuid::new_v4().to_string(), + inference_id: Uuid::new_v4().to_string(), + }, + prefill_chunk_tokens: 1, + } + } + + #[test] + fn policy_preserves_flow_fifo_and_promotes_waiting_batch_work() { + let mut config = SchedulerPolicyConfig::default(); + config.promotion_rounds = 2; + let mut policy = FairPolicy::new(config); + let now = Instant::now(); + policy.enqueue( + 1, + FlowKey { + principal: "p".into(), + class: SchedulingClass::Batch, + }, + 0, + now, + ); + policy.enqueue( + 2, + FlowKey { + principal: "p".into(), + class: SchedulingClass::Batch, + }, + 0, + now, + ); + for id in 10..14 { + policy.enqueue( + id, + FlowKey { + principal: format!("i{id}"), + class: SchedulingClass::Interactive, + }, + 0, + now, + ); + } + let selected = (0..6) + .filter_map(|_| policy.select().map(|s| s.id)) + .collect::>(); + assert!(selected.contains(&1)); + if selected.contains(&2) { + assert!( + selected.iter().position(|id| *id == 1) < selected.iter().position(|id| *id == 2) + ); + } + } + + #[test] + fn weighted_deficits_preserve_class_shares_without_starvation() { + let config = SchedulerPolicyConfig { + interactive_weight: 4, + standard_weight: 2, + batch_weight: 1, + promotion_rounds: 10_000, + ..SchedulerPolicyConfig::default() + }; + let mut policy = FairPolicy::new(config); + let now = Instant::now(); + let flows = [ + FlowKey { + principal: "interactive".into(), + class: SchedulingClass::Interactive, + }, + FlowKey { + principal: "standard".into(), + class: SchedulingClass::Standard, + }, + FlowKey { + principal: "batch".into(), + class: SchedulingClass::Batch, + }, + ]; + for (id, flow) in flows.iter().enumerate() { + policy.enqueue(id as u64, flow.clone(), 0, now); + } + let mut counts = [0usize; 3]; + for _ in 0..70 { + let selected = policy.select().unwrap(); + let index = selected.id as usize; + counts[index] += 1; + policy.charge(&flows[index], 1, selected.charge_scale); + policy.enqueue(selected.id, flows[index].clone(), 0, now); + } + assert_eq!(counts, [40, 20, 10]); + } + + #[test] + fn dropping_a_scheduled_session_cancels_its_shared_control_and_waiting_count() { + let scheduler = WorkloadScheduler::new( + Arc::new(DeterministicEngine), + SchedulerPolicyConfig::default(), + ) + .unwrap(); + let request = request(SchedulingClass::Standard, "principal", 1); + let control = request.control.clone(); + let mut session = scheduler.start_session(request).unwrap(); + assert!(matches!(session.step().unwrap(), SessionStep::Token { .. })); + assert_eq!(scheduler.status().metrics.waiting_for_consumer, 1); + drop(session); + assert!(matches!(control.check(), Err(Error::Cancelled))); + for _ in 0..100 { + if scheduler.status().metrics.waiting_for_consumer == 0 { + break; + } + thread::sleep(Duration::from_millis(2)); + } + assert_eq!(scheduler.status().metrics.waiting_for_consumer, 0); + } + + #[test] + fn dropping_scheduler_cancels_and_joins_outstanding_work() { + let scheduler = WorkloadScheduler::new( + Arc::new(DeterministicEngine), + SchedulerPolicyConfig::default(), + ) + .unwrap(); + let request = request(SchedulingClass::Standard, "principal", 2); + let control = request.control.clone(); + let mut session = scheduler.start_session(request).unwrap(); + drop(scheduler); + assert!(matches!(control.check(), Err(Error::Cancelled))); + assert!(matches!(session.step(), Err(Error::Cancelled))); + } + + #[test] + fn scheduler_interleaves_principals_and_emits_attributed_records() { + let records = Arc::new(Mutex::new(Vec::new())); + let captured = records.clone(); + let scheduler = WorkloadScheduler::new_with_diagnostics( + Arc::new(DeterministicEngine), + SchedulerPolicyConfig::default(), + Some(move |line: &str| captured.lock().push(line.to_owned())), + ) + .unwrap(); + let mut first = scheduler + .start_session(request(SchedulingClass::Standard, "one", 2)) + .unwrap(); + let mut second = scheduler + .start_session(request(SchedulingClass::Standard, "two", 2)) + .unwrap(); + assert!(matches!(first.step().unwrap(), SessionStep::Token { .. })); + assert!(matches!(second.step().unwrap(), SessionStep::Token { .. })); + first.finish().unwrap(); + second.finish().unwrap(); + assert!(scheduler.flush_diagnostics(Duration::from_secs(1))); + let parsed = records + .lock() + .iter() + .map(|line| serde_json::from_str::(line).unwrap()) + .collect::>(); + assert!(parsed.iter().any(|record| record["principal"] == "one")); + assert!(parsed.iter().any(|record| record["principal"] == "two")); + assert!( + parsed + .iter() + .all(|record| record["execution_session_id"].is_string()) + ); + } + + #[test] + fn diagnostic_overflow_is_counted_without_blocking_scheduler() { + let scheduler = WorkloadScheduler::new_with_diagnostics( + Arc::new(DeterministicEngine), + SchedulerPolicyConfig { + diagnostic_capacity: 1, + ..SchedulerPolicyConfig::default() + }, + Some(|_: &str| thread::sleep(Duration::from_millis(20))), + ) + .unwrap(); + let mut sessions = (0..8) + .map(|index| { + scheduler + .start_session(request(SchedulingClass::Standard, &format!("p{index}"), 1)) + .unwrap() + }) + .collect::>(); + for session in &mut sessions { + let _ = session.step(); + let _ = session.finish(); + } + assert!(scheduler.status().metrics.diagnostic_records_lost > 0); + assert!( + scheduler + .status() + .diagnostic_loss_by_kind + .contains_key("scheduler_decision") + ); + } +} diff --git a/docs/outline.md b/docs/outline.md index 7d62f7b..4c2bc91 100644 --- a/docs/outline.md +++ b/docs/outline.md @@ -13,8 +13,8 @@ The first project decision is therefore an architectural boundary, not a cache p The document uses three status categories: - **Historical design hypothesis:** the decisive executor question and early recommendations explain why implementation began with checkpoint equivalence rather than a broad server. -- **Implemented current state:** Phases 1 through 5 are complete. The real Gemma proof demonstrated exact checkpoint continuation; the logical and physical managers implement transactional shared state and tiering; the minimal server implements durable contexts, scheduling, lifecycle APIs, streaming events, authentication policy, compatibility adapters, and checked OpenAPI; and the mapped-execution proof demonstrated transactional sequence mappings and reference-only activation. -- **Remaining prospective contract:** Phase 6 begins the next architectural cutover by integrating the persistent mapped executor with the live server. Phases 7 through 10 remain planned work and are requirements, not claims of implementation. +- **Implemented current state:** Phases 1 through 8 are complete. The real Gemma proof demonstrated exact checkpoint continuation; the logical and physical managers implement transactional shared state and tiering; the live server implements persistent mapped execution, bounded lifecycle, dynamic multi-model residency, resumable priority-aware workload scheduling, non-blocking attributed diagnostics, authentication, compatibility adapters, and checked OpenAPI; and the real-GPU acceptance artifacts cover mapped activation, live lifecycle, residency, and mixed-workload fairness. +- **Remaining prospective contract:** Phases 9 and 10 remain planned work and are requirements, not claims of implementation. ## Decisive technical hypothesis @@ -1685,13 +1685,26 @@ recovery, diagnostic loss, and all thresholds. Tuning those policy parameters after observing a working system requires a new versioned fixture and evidence; it does not require changing the resumable execution contract. -The server already includes one narrow operator-diagnostic slice toward this -phase: HTTP debug records correlate request metadata, terminal status, -duration, and streaming response chunks. Diagnostics default off and may be -selected by CLI or environment as privacy-safe records with omitted headers, -redacted JSON strings, and bounded bodies, or as fully unredacted URI, header, -and body records for controlled diagnosis. This transport trace is not -scheduler decision attribution and does not satisfy the Phase 8 exit gate. +The implemented Phase 8 contract uses a protocol-neutral `ExecutionSession` +boundary and a replaceable priority-aware deficit-round-robin scheduler with +per-principal/class flows, FIFO equivalence ordering, monotonic age promotion, +bounded prefill quanta, and one-token decode quanta. Trusted scheduling +metadata carries distinct transport, inference-operation, and execution-session +identifiers. Native slot occupancy is scoped to each activation, prefill, or +decode quantum, and native abort registration remains live through that +quantum's completion fence. Every scheduler decision records policy attribution +and resource observations through a bounded asynchronous sink with exact loss +counters; HTTP debug tracing remains a separate opt-in transport diagnostic. + +`config/phase8-workload.json` freezes the versioned real-GPU workload, policy, +and acceptance thresholds. `tools/phase8-report.sh` runs the per-file coverage +gate and records model, build, workload, GPU, isolated-baseline, mixed-load, +cancellation, deadline, capacity-recovery, starvation, latency, and diagnostic +evidence in `results/phase8-server.json`; proof snapshots wait for every accepted +diagnostic record to reach the sink. Deterministic model-free fixtures cover +weighted class shares, per-flow FIFO, age promotion, cancellation, suspended +session slot release, diagnostic overflow, scheduler shutdown, and the proof +path. ### Phase 9: compatibility, persistence, and production packaging diff --git a/tools/phase8-report.sh b/tools/phase8-report.sh new file mode 100755 index 0000000..7643096 --- /dev/null +++ b/tools/phase8-report.sh @@ -0,0 +1,25 @@ +#!/bin/sh +set -eu + +compose() { + docker compose -f compose.test.yaml "$@" +} + +result_dir=${CUSCO_RESULT_DIR:-./results} +model_dir=${CUSCO_MODEL_DIR:-./models} +service=phase8-proof + +mkdir -p "$result_dir" +rm -f \ + "$result_dir/phase8-coverage.json" \ + "$result_dir/phase8-gpu.csv" \ + "$result_dir/phase8-server.json" + +compose run --build --rm test +printf '%s\n' '{"passed":true,"command":"docker compose -f compose.test.yaml run --build --rm test","per_file_line_floor_percent":80}' \ + > "$result_dir/phase8-coverage.json" +compose run --build --rm --no-deps --entrypoint chmod "$service" a+rwx /results +compose run --build --rm "$service" +python3 tools/report-phase8.py \ + "$result_dir" \ + "$model_dir/gemma-4-e2b-it.gguf" diff --git a/tools/report-phase8.py b/tools/report-phase8.py new file mode 100755 index 0000000..74a946a --- /dev/null +++ b/tools/report-phase8.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +import hashlib +import json +import pathlib +import sys + +RESULT_DIR = pathlib.Path(sys.argv[1] if len(sys.argv) > 1 else "results") +MODEL = pathlib.Path(sys.argv[2] if len(sys.argv) > 2 else "models/gemma-4-e2b-it.gguf") +ARTIFACT = RESULT_DIR / "phase8-server.json" +WORKLOAD = pathlib.Path("config/phase8-workload.json") + + +def digest(path): + value = hashlib.sha256() + with path.open("rb") as source: + while chunk := source.read(1024 * 1024): + value.update(chunk) + return value.hexdigest() + + +artifact = json.loads(ARTIFACT.read_text(encoding="utf-8")) +coverage = json.loads((RESULT_DIR / "phase8-coverage.json").read_text(encoding="utf-8")) +workload_bytes = WORKLOAD.read_bytes() +header = pathlib.Path("native/include/cusco_executor.h").read_text(encoding="utf-8") +abi_line = next( + line for line in header.splitlines() if line.startswith("#define CUSCO_EXECUTOR_ABI_VERSION ") +) +artifact["model"] = {"path": str(MODEL), "sha256": digest(MODEL)} +artifact["workload_provenance"] = { + "path": str(WORKLOAD), + "sha256": hashlib.sha256(workload_bytes).hexdigest(), +} +artifact["build"] = { + "llama_cpp_tag": pathlib.Path("llama.cpp-version.txt").read_text(encoding="utf-8").strip(), + "executor_abi": int(abi_line.split()[-1].removesuffix("u")), +} +artifact["gpu"] = (RESULT_DIR / "phase8-gpu.csv").read_text(encoding="utf-8").strip() +artifact["coverage"] = coverage +artifact["checks"] = { + "CLI scheduler gates passed": artifact.get("passed") is True, + "coverage gate passed": coverage.get("passed") is True, + "GPU provenance recorded": bool(artifact["gpu"]), + "mixed workload includes all classes": { + row.get("class") for row in artifact.get("mixed", []) + } == {"interactive", "standard", "batch"}, + "multiple principals competed": len( + {row.get("principal") for row in artifact.get("mixed", [])} + ) >= 3, + "cancellation was observed": any( + row.get("observed") == "cancelled" for row in artifact.get("mixed", []) + ), + "deadline expiry was observed": any( + row.get("observed") == "deadline" for row in artifact.get("mixed", []) + ), + "capacity recovered": artifact.get("capacity_recovery", {}).get("observed") == "complete", + "diagnostic records were lossless": artifact.get("scheduler", {}) + .get("metrics", {}) + .get("diagnostic_records_lost") == 0, +} +ARTIFACT.write_text(json.dumps(artifact, indent=2) + "\n", encoding="utf-8") +failed = [name for name, passed in artifact["checks"].items() if not passed] +if failed: + raise SystemExit("Phase 8 proof failed: " + ", ".join(failed)) +print(f"Phase 8 proof: PASS. Artifact: `{ARTIFACT}`") From 9cc28e26ec259b1842300c1c84ed462d91d3bc07 Mon Sep 17 00:00:00 2001 From: Streaky Date: Thu, 6 Aug 2026 20:46:19 +0100 Subject: [PATCH 3/7] fix scheduler deficit lifecycle --- crates/server/src/scheduler.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/crates/server/src/scheduler.rs b/crates/server/src/scheduler.rs index 7c84f5e..eda55ce 100644 --- a/crates/server/src/scheduler.rs +++ b/crates/server/src/scheduler.rs @@ -301,6 +301,14 @@ impl FairPolicy { (tokens.max(1) as i64).saturating_mul(scale); } + fn clear_if_idle(&mut self, flow: &FlowKey, jobs: &HashMap) { + if !self.runnable.iter().any(|entry| &entry.flow == flow) + && !jobs.values().any(|job| &job.flow == flow) + { + self.deficits.remove(flow); + } + } + fn remove(&mut self, id: u64) -> bool { let previous = self.runnable.len(); self.runnable.retain(|entry| entry.id != id); @@ -699,7 +707,6 @@ fn handle_command( }, ); policy.enqueue(id, flow, admitted_round, admitted_at); - counters.admitted.fetch_add(1, Ordering::Relaxed); } Command::Continue(id) => { if let Some(job) = jobs.get(&id) { @@ -710,23 +717,27 @@ fn handle_command( Command::Stop(id) => { counters.waiting.fetch_sub(1, Ordering::AcqRel); if let Some(mut job) = jobs.remove(&id) { + let flow = job.flow.clone(); let result = job .session .as_mut() .ok_or_else(|| Error::State("scheduler session was not started".into())) .and_then(|session| session.finish()); send_terminal(job, result, counters); + policy.clear_if_idle(&flow, jobs); } } Command::Cancel(id) => { let was_runnable = policy.remove(id); if let Some(job) = jobs.remove(&id) { + let flow = job.flow.clone(); if job.session.is_some() && !was_runnable { counters.waiting.fetch_sub(1, Ordering::AcqRel); } job.control.cancel(); let _ = job.response.send(ClientMessage::Failed(Error::Cancelled)); counters.cancelled.fetch_add(1, Ordering::Relaxed); + policy.clear_if_idle(&flow, jobs); } } Command::Shutdown => unreachable!("shutdown is handled by the scheduler loop"), @@ -830,10 +841,11 @@ fn execute_quantum( }, }, ); - if let Some(result) = terminal { let job = jobs.remove(&selection.id).expect("terminal job exists"); + let flow = job.flow.clone(); send_terminal(job, result, counters); + policy.clear_if_idle(&flow, jobs); } else if !wait_for_consumer { let job = jobs.get(&selection.id).expect("runnable job exists"); policy.enqueue( From 9ddc6de3b81f8248adecde6fe3d863a59e7531b3 Mon Sep 17 00:00:00 2001 From: Streaky Date: Thu, 6 Aug 2026 20:46:21 +0100 Subject: [PATCH 4/7] validate mapped prefill chunk size --- crates/server/src/mapped.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/server/src/mapped.rs b/crates/server/src/mapped.rs index e79885b..369ccf0 100644 --- a/crates/server/src/mapped.rs +++ b/crates/server/src/mapped.rs @@ -705,6 +705,9 @@ impl Drop for MappedSession { impl InferenceEngine for MappedEngine { fn start_session(&self, request: EngineRequest) -> Result, Error> { request.control.check()?; + if request.prefill_chunk_tokens == 0 { + return Err(Error::State("prefill chunk size must be nonzero".into())); + } if request.model.path.to_str() != Some(self.model_path()) { return Err(Error::State( "Phase 8 admits only a resident process-owned model".into(), From fbdf5b96aa913c8bd25460779ea56671eaa0c682 Mon Sep 17 00:00:00 2001 From: Streaky Date: Thu, 6 Aug 2026 20:46:26 +0100 Subject: [PATCH 5/7] exclude first token from quantum waits --- crates/cli/src/main.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index b080033..4dd538e 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -679,8 +679,11 @@ fn run_scheduler_proof_case( }, &mut |_, _, _| { let now = Instant::now(); - first_event_ms.get_or_insert_with(|| now.duration_since(started).as_millis()); - token_wait_ms.push(now.duration_since(last_token).as_millis()); + if first_event_ms.is_some() { + token_wait_ms.push(now.duration_since(last_token).as_millis()); + } else { + first_event_ms = Some(now.duration_since(started).as_millis()); + } last_token = now; match expected { SchedulerProofOutcome::Complete => {} From 23b6842ef43e77d9cd7fab37deb1bb971ed85d57 Mon Sep 17 00:00:00 2001 From: Streaky Date: Thu, 6 Aug 2026 20:46:28 +0100 Subject: [PATCH 6/7] persist phase8 wrapper gate status --- tools/report-phase8.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/report-phase8.py b/tools/report-phase8.py index 74a946a..13a9c2e 100755 --- a/tools/report-phase8.py +++ b/tools/report-phase8.py @@ -57,6 +57,7 @@ def digest(path): .get("metrics", {}) .get("diagnostic_records_lost") == 0, } +artifact["passed"] = all(artifact["checks"].values()) ARTIFACT.write_text(json.dumps(artifact, indent=2) + "\n", encoding="utf-8") failed = [name for name, passed in artifact["checks"].items() if not passed] if failed: From ab8dd6a10c2afdd4377cc8dc5b8c1019540113fb Mon Sep 17 00:00:00 2001 From: Streaky Date: Thu, 6 Aug 2026 20:46:30 +0100 Subject: [PATCH 7/7] reject unsupported scheduler policy versions --- crates/server/src/lib.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/server/src/lib.rs b/crates/server/src/lib.rs index 22adc88..fb393c1 100644 --- a/crates/server/src/lib.rs +++ b/crates/server/src/lib.rs @@ -277,8 +277,13 @@ impl Default for SchedulerPolicyConfig { impl SchedulerPolicyConfig { fn validate(self) -> Result { - if self.version == 0 - || self.interactive_weight == 0 + if self.version != 1 { + return Err(Error::State(format!( + "unsupported scheduler policy version {}", + self.version + ))); + } + if self.interactive_weight == 0 || self.standard_weight == 0 || self.batch_weight == 0 || self.deficit_refill == 0