From 852258542531f4d2318b5263d25258d6ee7b30b8 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 17:32:26 -0700 Subject: [PATCH 01/24] fix(router): stop charging decode for old prompts Signed-off-by: PeaBrane --- .../configuration/groups/kv_router_args.py | 14 +++ components/src/dynamo/router/README.md | 5 +- components/src/dynamo/router/__main__.py | 1 + docs/components/router/README.md | 1 + docs/components/router/router-guide.md | 9 ++ lib/bench/kv_router/active_sequences_bench.rs | 1 + lib/bindings/c/src/lib.rs | 6 + lib/bindings/python/rust/llm/entrypoint.rs | 4 +- lib/bindings/python/tests/test_replay.py | 1 + lib/kv-router/src/protocols.rs | 34 ++++++ lib/kv-router/src/scheduling/config.rs | 46 +++++++- lib/kv-router/src/scheduling/local.rs | 105 ++++++++++++++++-- lib/kv-router/src/scheduling/policy.rs | 1 + lib/kv-router/src/scheduling/queue.rs | 43 ++++++- lib/kv-router/src/scheduling/types.rs | 1 + lib/kv-router/src/sequences/multi_worker.rs | 75 ++++++++++++- lib/kv-router/src/sequences/single.rs | 85 +++++++++++++- lib/llm/src/kv_router.rs | 16 ++- lib/llm/src/kv_router/prefill_router.rs | 37 +++++- lib/llm/src/kv_router/scheduler.rs | 4 +- lib/llm/src/kv_router/sequence.rs | 6 + lib/mocker/src/replay/router/offline.rs | 2 + lib/mocker/src/replay/router/online.rs | 1 + 23 files changed, 465 insertions(+), 33 deletions(-) diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index da5747dd8cca..8a97729ce860 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -26,6 +26,7 @@ "router_track_active_blocks", "router_track_output_blocks", "router_assume_kv_reuse", + "router_track_prefill_tokens", "router_snapshot_threshold", "router_reset_states", "router_ttl_secs", @@ -51,6 +52,7 @@ class KvRouterConfigBase(ConfigBase): router_track_active_blocks: bool router_track_output_blocks: bool router_assume_kv_reuse: bool + router_track_prefill_tokens: bool router_snapshot_threshold: int router_reset_states: bool router_ttl_secs: float @@ -173,6 +175,18 @@ def add_arguments(self, parser) -> None: ), obsolete_flag="--assume-kv-reuse", ) + add_negatable_bool_argument( + g, + flag_name="--router-track-prefill-tokens", + env_var="DYN_ROUTER_TRACK_PREFILL_TOKENS", + default=True, + dest="router_track_prefill_tokens", + help=( + "KV Router: Include prompt-side prefill tokens in active load accounting. " + "Use --no-router-track-prefill-tokens to ignore prompt tokens in router " + "prefill-token load, queue pressure, and active_prefill_tokens metrics." + ), + ) add_argument( g, flag_name="--router-snapshot-threshold", diff --git a/components/src/dynamo/router/README.md b/components/src/dynamo/router/README.md index a738380c28db..2ddacd3d49d8 100644 --- a/components/src/dynamo/router/README.md +++ b/components/src/dynamo/router/README.md @@ -29,7 +29,7 @@ python -m dynamo.router \ - `--endpoint`: Full endpoint path for workers in the format `namespace.component.endpoint` (e.g., `dynamo.prefill.generate`) **Router Configuration:** -All router options use the `--router-*` prefix (e.g., `--router-block-size`, `--router-kv-overlap-score-weight`, `--router-temperature`, `--router-kv-events` / `--no-router-kv-events`, `--router-replica-sync`, `--router-snapshot-threshold`, `--router-reset-states`, `--router-track-active-blocks` / `--no-router-track-active-blocks`). Legacy names without the prefix (e.g., `--block-size`, `--kv-events`) are still accepted but deprecated. For detailed descriptions, see the [Router Guide](/docs/components/router/router-guide.md). +All router options use the `--router-*` prefix (e.g., `--router-block-size`, `--router-kv-overlap-score-weight`, `--router-temperature`, `--router-kv-events` / `--no-router-kv-events`, `--router-replica-sync`, `--router-snapshot-threshold`, `--router-reset-states`, `--router-track-active-blocks` / `--no-router-track-active-blocks`, `--router-track-prefill-tokens` / `--no-router-track-prefill-tokens`). Legacy names without the prefix (e.g., `--block-size`, `--kv-events`) are still accepted but deprecated. For detailed descriptions, see the [Router Guide](/docs/components/router/router-guide.md). ## Architecture @@ -74,6 +74,9 @@ python -m dynamo.vllm --model MODEL_NAME --block-size 64 --disaggregation-mode p > **Why `--no-router-track-active-blocks` for prefill routing?** > Active block tracking is used for load balancing across decode (generation) phases. For prefill-only routing, decode load is not relevant, so disabling this reduces overhead and simplifies the router state. > +> **When should I use `--no-router-track-prefill-tokens`?** +> Use it on decode-only routers that should ignore already-completed prompt work. This keeps `active_prefill_tokens`, queue pressure, and load estimates focused on decode-side work after a prefill-to-decode handoff. +> > **Why `--router-block-size` is required for standalone routers:** > Unlike the frontend router which can infer block size from the ModelDeploymentCard (MDC) during worker registration, standalone routers cannot access the MDC and must have the block size explicitly specified. This is a work in progress to enable automatic inference. diff --git a/components/src/dynamo/router/__main__.py b/components/src/dynamo/router/__main__.py index 4adf13077347..c7fa95efcbe1 100644 --- a/components/src/dynamo/router/__main__.py +++ b/components/src/dynamo/router/__main__.py @@ -171,6 +171,7 @@ async def worker(runtime: DistributedRuntime): f"router_track_active_blocks={config.router_track_active_blocks}, " f"router_track_output_blocks={config.router_track_output_blocks}, " f"router_assume_kv_reuse={config.router_assume_kv_reuse}, " + f"router_track_prefill_tokens={config.router_track_prefill_tokens}, " f"router_ttl_secs={config.router_ttl_secs}, " f"router_max_tree_size={config.router_max_tree_size}, " f"router_prune_target_ratio={config.router_prune_target_ratio}" diff --git a/docs/components/router/README.md b/docs/components/router/README.md index 79488ae77864..670b2a020760 100644 --- a/docs/components/router/README.md +++ b/docs/components/router/README.md @@ -23,6 +23,7 @@ For Kubernetes, set `DYN_ROUTER_MODE=kv` on the Frontend service. Workers automa | `--no-router-kv-events` | enabled | Fall back to approximate routing (no event consumption from workers) | | `--router-queue-threshold` | `4.0` | Backpressure queue threshold; enables priority scheduling via `nvext.agent_hints.priority` | | `--router-queue-policy` | `fcfs` | Queue scheduling policy: `fcfs` (tail TTFT), `wspt` (avg TTFT), or `lcfs` (comparison-only reverse ordering) | +| `--no-router-track-prefill-tokens` | disabled | Ignore prompt-side prefill tokens in router load accounting; useful for decode-only routing paths | ### Standalone Router diff --git a/docs/components/router/router-guide.md b/docs/components/router/router-guide.md index 95f0f713e8fd..3e2dfc292fc2 100644 --- a/docs/components/router/router-guide.md +++ b/docs/components/router/router-guide.md @@ -256,6 +256,8 @@ The main KV-aware routing arguments (frontend uses the same `--router-*` flag na - `--no-router-assume-kv-reuse`: When tracking active blocks, disables the assumption of KV cache reuse. By default (`router_assume_kv_reuse=true`), the router computes actual block hashes for sequence tracking to deduplicate blocks and optimize load balancing. When disabled via this flag, the router generates random hashes for sequence blocks, treating each request's blocks as unique. This is useful in disaggregated setups where prefill transfers blocks to decode workers that may already have those blocks cached, but the engine cannot coordinate transfers to avoid duplication. Without this flag, the router's load balancing heuristics would undercount decode blocks when duplicates exist. +- `--no-router-track-prefill-tokens`: Disables prompt-side prefill token accounting in the router's active load model. By default (`router_track_prefill_tokens=true`), the router counts uncached prompt tokens toward `active_prefill_tokens`, queue pressure, and potential prefill-token load. Disable this for decode-only routing paths where prompt processing has already happened elsewhere and the decode router should ignore transferred prompt load. In normal live disaggregated serving, the decode-stage override applies this behavior automatically. + - `--router-replica-sync`: Disabled by default. Enables NATS-based synchronization of local routing decisions between router replicas. When enabled, routers share their active sequence information and local predictions of block usage, improving routing consistency across instances. Note that this does not sync the radix tree or cached KV block states themselves - in JetStream mode those are synchronized through JetStream events; in local indexer mode (default) each router queries workers directly. ### KV Indexer / Approx KV Indexer @@ -280,6 +282,8 @@ Use `--no-router-kv-events` when you are not confident that your backend engine Use `--no-router-assume-kv-reuse` in disaggregated setups where the decode worker does not reuse transferred KV cache blocks. By default the router assumes KV blocks transferred from prefill to decode will be deduplicated on the decode side, but vLLM and SGLang decode workers currently do not support this — only TensorRT-LLM does. Without this flag, the router undercounts decode blocks when duplicates exist, leading to inaccurate load estimates. +Use `--no-router-track-prefill-tokens` when a router is serving decode-only traffic and prompt processing has already completed elsewhere. This keeps decode routing decisions focused on decode-side load instead of briefly charging prompt tokens to the decode worker after handoff. The built-in live disaggregated decode path applies the equivalent per-request override automatically. + Use `--router-track-output-blocks` **(experimental)** when your workload is output-heavy and you want the router to account for output-side KV cache growth in load balancing. This is useful in two scenarios: (1) workloads with long output sequences and little multi-turn reuse, where output blocks dominate the KV cache footprint; (2) agentic schedulers (e.g. NAT or other LLM routers) that can accurately predict the expected output sequence length per request. When enabled, the router adds placeholder blocks as tokens are generated. If you additionally pass `nvext.agent_hints.osl` (expected output sequence length in tokens) per request, the router applies fractional decay to output blocks — each output block's weight starts at 1.0 and decays linearly toward 0.0 as generation approaches the expected OSL. This lets the router predict that a request nearing completion will soon free its blocks, effectively modeling the future load trajectory rather than just the current snapshot. Without `osl`, output blocks are added at full weight with no decay. The flag requires `--router-track-active-blocks` (the default). The `--router-queue-threshold` (default: 4.0) controls when incoming requests are held in a priority queue. The router holds requests while all workers exceed the given fraction of `max_num_batched_tokens`, releasing them as capacity frees up. This defers the routing decision so it is made with the freshest load metrics, rather than dispatching into an already-saturated system. It also enables priority scheduling via `nvext.agent_hints.priority`. Set to None to disable queueing entirely. @@ -310,6 +314,11 @@ The prefill router is automatically created when: - **Seamlessly integrated** into the request pipeline between preprocessing and decode routing - **Falls back gracefully** to decode-only mode if prefill fails or no prefill workers are available +**Key characteristics of the decode routing stage in disaggregated mode:** +- **Disables overlap scoring** (`overlap_score_weight=0`) because decode routing should not chase prefix reuse +- **Disables KV reuse assumption** (`assume_kv_reuse=false`) unless the backend can truly deduplicate transferred blocks +- **Disables prefill-token tracking** (`track_prefill_tokens=false`) so decode-side load reflects decode work rather than already-completed prompt work + ### Setup Example When both workers are registered, requests are automatically routed. diff --git a/lib/bench/kv_router/active_sequences_bench.rs b/lib/bench/kv_router/active_sequences_bench.rs index 1016e312ef69..1933f47dcbf2 100644 --- a/lib/bench/kv_router/active_sequences_bench.rs +++ b/lib/bench/kv_router/active_sequences_bench.rs @@ -478,6 +478,7 @@ async fn apply_entry( token_sequence: Some(block_hashes), isl, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: Some(output_length as u32), worker, lora_name: None, diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index bf8d095661d5..5540c80ac04c 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -488,6 +488,8 @@ impl RouterHandles { let config_override = if is_disaggregated { Some(RouterConfigOverride { overlap_score_weight: Some(0.0), + assume_kv_reuse: Some(false), + track_prefill_tokens: Some(false), ..Default::default() }) } else { @@ -564,6 +566,9 @@ fn kv_router_config_from_env() -> KvRouterConfig { if let Some(v) = env_bool("DYN_ROUTER_TRACK_OUTPUT_BLOCKS") { cfg.router_track_output_blocks = v; } + if let Some(v) = env_bool("DYN_ROUTER_TRACK_PREFILL_TOKENS") { + cfg.router_track_prefill_tokens = v; + } if let Some(v) = env_f64("DYN_ROUTER_QUEUE_THRESHOLD") { cfg.router_queue_threshold = Some(v); } @@ -575,6 +580,7 @@ fn kv_router_config_from_env() -> KvRouterConfig { router_replica_sync = cfg.router_replica_sync, router_track_active_blocks = cfg.router_track_active_blocks, router_track_output_blocks = cfg.router_track_output_blocks, + router_track_prefill_tokens = cfg.router_track_prefill_tokens, router_queue_threshold = ?cfg.router_queue_threshold, "KvRouterConfig initialized (DYN_* env overrides applied)" ); diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 519168733b28..47f0ee5a2e9b 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -58,7 +58,7 @@ impl KvRouterConfig { #[pymethods] impl KvRouterConfig { #[new] - #[pyo3(signature = (overlap_score_weight=1.0, router_temperature=0.0, use_kv_events=true, durable_kv_events=false, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_snapshot_threshold=1000000, router_reset_states=false, router_ttl_secs=120.0, router_max_tree_size=1048576, router_prune_target_ratio=0.8, router_queue_threshold=Some(4.0), router_event_threads=4, router_enable_cache_control=false, min_initial_workers=1, router_queue_policy="fcfs", remote_indexer_component=None))] + #[pyo3(signature = (overlap_score_weight=1.0, router_temperature=0.0, use_kv_events=true, durable_kv_events=false, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_snapshot_threshold=1000000, router_reset_states=false, router_ttl_secs=120.0, router_max_tree_size=1048576, router_prune_target_ratio=0.8, router_queue_threshold=Some(4.0), router_event_threads=4, router_enable_cache_control=false, min_initial_workers=1, router_queue_policy="fcfs", remote_indexer_component=None))] #[allow(clippy::too_many_arguments)] fn new( overlap_score_weight: f64, @@ -69,6 +69,7 @@ impl KvRouterConfig { router_track_active_blocks: bool, router_track_output_blocks: bool, router_assume_kv_reuse: bool, + router_track_prefill_tokens: bool, router_snapshot_threshold: Option, router_reset_states: bool, router_ttl_secs: f64, @@ -91,6 +92,7 @@ impl KvRouterConfig { router_track_active_blocks, router_track_output_blocks, router_assume_kv_reuse, + router_track_prefill_tokens, router_snapshot_threshold, router_reset_states, router_ttl_secs, diff --git a/lib/bindings/python/tests/test_replay.py b/lib/bindings/python/tests/test_replay.py index 064827676257..31d61441d9d9 100644 --- a/lib/bindings/python/tests/test_replay.py +++ b/lib/bindings/python/tests/test_replay.py @@ -75,6 +75,7 @@ def _router_config_payload(): "router_track_active_blocks": True, "router_track_output_blocks": False, "router_assume_kv_reuse": True, + "router_track_prefill_tokens": True, "router_snapshot_threshold": 1000000, "router_reset_states": False, "router_ttl_secs": 120.0, diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 2a40d5208c65..2fad5af529f3 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -8,6 +8,10 @@ use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use xxhash_rust::xxh3; +const fn default_track_prefill_tokens() -> bool { + true +} + /// The event subject that workers publish KV cache events on. pub const KV_EVENT_SUBJECT: &str = "kv-events"; @@ -363,6 +367,8 @@ pub enum ActiveSequenceEventData { token_sequence: Option>, isl: usize, overlap: u32, + #[serde(default = "default_track_prefill_tokens")] + track_prefill_tokens: bool, expected_output_tokens: Option, }, Free, @@ -1020,4 +1026,32 @@ mod tests { } if request_id == "req-123" )); } + + #[test] + fn test_active_sequence_add_request_serialization_preserves_track_prefill_tokens() { + let event = ActiveSequenceEvent { + request_id: "req-123".to_string(), + worker: WorkerWithDpRank::new(7, 0), + data: ActiveSequenceEventData::AddRequest { + token_sequence: Some(vec![11, 22]), + isl: 128, + overlap: 1, + track_prefill_tokens: false, + expected_output_tokens: Some(32), + }, + router_id: 9, + lora_name: None, + }; + + let serialized = serde_json::to_string(&event).unwrap(); + let deserialized: ActiveSequenceEvent = serde_json::from_str(&serialized).unwrap(); + + match deserialized.data { + ActiveSequenceEventData::AddRequest { + track_prefill_tokens, + .. + } => assert!(!track_prefill_tokens), + _ => panic!("expected add request event"), + } + } } diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index 4c566d425982..f84104fdd4c0 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -15,6 +15,10 @@ const fn default_min_initial_workers() -> usize { 1 } +const fn default_track_prefill_tokens() -> bool { + true +} + #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum RouterQueuePolicy { @@ -61,6 +65,9 @@ pub struct RouterConfigOverride { #[builder(default)] pub assume_kv_reuse: Option, + + #[builder(default)] + pub track_prefill_tokens: Option, } /// KV Router configuration parameters @@ -96,6 +103,12 @@ pub struct KvRouterConfig { /// When false, generates random hashes (assuming no KV cache reuse). pub router_assume_kv_reuse: bool, + /// Whether to include prompt-side prefill tokens in active load accounting (default: true). + /// When false, prompt tokens are excluded from active prefill token tracking, queue pressure, + /// and potential prefill-token load calculations. + #[serde(default = "default_track_prefill_tokens")] + pub router_track_prefill_tokens: bool, + /// Threshold for triggering snapshots. If None, no snapshots will be performed. #[validate(range(min = 1))] pub router_snapshot_threshold: Option, @@ -169,6 +182,7 @@ impl Default for KvRouterConfig { router_track_active_blocks: true, router_track_output_blocks: false, router_assume_kv_reuse: true, + router_track_prefill_tokens: default_track_prefill_tokens(), router_snapshot_threshold: Some(1000000), router_reset_states: false, router_ttl_secs: 120.0, @@ -206,6 +220,18 @@ fn validate_kv_router_config(config: &KvRouterConfig) -> Result<(), ValidationEr } impl KvRouterConfig { + pub fn assume_kv_reuse(&self, config_override: Option<&RouterConfigOverride>) -> bool { + config_override + .and_then(|cfg| cfg.assume_kv_reuse) + .unwrap_or(self.router_assume_kv_reuse) + } + + pub fn track_prefill_tokens(&self, config_override: Option<&RouterConfigOverride>) -> bool { + config_override + .and_then(|cfg| cfg.track_prefill_tokens) + .unwrap_or(self.router_track_prefill_tokens) + } + /// Compute sequence hashes for active block tracking based on configuration. /// /// Returns: @@ -228,9 +254,7 @@ impl KvRouterConfig { return Some(Vec::new()); } - let assume_kv_reuse = config_override - .and_then(|cfg| cfg.assume_kv_reuse) - .unwrap_or(self.router_assume_kv_reuse); + let assume_kv_reuse = self.assume_kv_reuse(config_override); if assume_kv_reuse { let block_hashes = compute_block_hash_for_seq(tokens, block_size, None, lora_name); @@ -280,6 +304,11 @@ mod tests { assert_eq!(KvRouterConfig::default().min_initial_workers, 1); } + #[test] + fn kv_router_config_defaults_to_tracking_prefill_tokens() { + assert!(KvRouterConfig::default().router_track_prefill_tokens); + } + #[test] fn kv_router_config_rejects_zero_initial_workers() { let cfg = KvRouterConfig { @@ -288,4 +317,15 @@ mod tests { }; assert!(cfg.validate().is_err()); } + + #[test] + fn router_config_override_serde_round_trip_preserves_track_prefill_tokens() { + let serialized = serde_json::to_string(&RouterConfigOverride { + track_prefill_tokens: Some(false), + ..Default::default() + }) + .unwrap(); + let deserialized: RouterConfigOverride = serde_json::from_str(&serialized).unwrap(); + assert_eq!(deserialized.track_prefill_tokens, Some(false)); + } } diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 2304e5a7d27f..4da0291cb6db 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -30,6 +30,7 @@ where request_tx: mpsc::Sender, slots: Arc>, queue: Arc>, + track_prefill_tokens_default: bool, worker_type: &'static str, } @@ -48,6 +49,7 @@ where block_size: u32, selector: Sel, policy: S, + track_prefill_tokens_default: bool, cancellation_token: CancellationToken, worker_type: &'static str, monitor_worker_configs: bool, @@ -135,6 +137,7 @@ where request_tx, slots, queue, + track_prefill_tokens_default, worker_type, } } @@ -154,6 +157,9 @@ where allowed_worker_ids: Option>, ) -> Result { let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); + let track_prefill_tokens = router_config_override + .and_then(|cfg| cfg.track_prefill_tokens) + .unwrap_or(self.track_prefill_tokens_default); let request = SchedulingRequest { maybe_request_id, token_seq, @@ -161,6 +167,7 @@ where overlaps, decode_blocks: HashMap::new(), prefill_tokens: HashMap::new(), + track_prefill_tokens, router_config_override: router_config_override.cloned(), update_states, lora_name, @@ -224,10 +231,16 @@ where token_seq: Option>, isl_tokens: usize, overlaps: OverlapScores, + track_prefill_tokens: bool, ) -> Vec { - let (decode_blocks, prefill_tokens) = - self.slots - .potential_blocks_and_tokens(token_seq.as_deref(), isl_tokens, overlaps); + let (decode_blocks, prefill_tokens) = self + .slots + .potential_blocks_and_tokens_with_prefill_tracking( + token_seq.as_deref(), + isl_tokens, + overlaps, + track_prefill_tokens, + ); let mut workers: HashSet = HashSet::new(); workers.extend(decode_blocks.keys().copied()); @@ -300,6 +313,7 @@ mod tests { 64, DefaultWorkerSelector::new(None, "test"), FcfsPolicy, + true, cancel_token.clone(), "test", monitor_worker_configs, @@ -344,6 +358,48 @@ mod tests { cancel_token.cancel(); } + #[tokio::test] + async fn test_schedule_override_can_disable_prefill_tracking() { + let mut workers = HashMap::new(); + workers.insert( + 0, + SimpleWorkerConfig { + max_num_batched_tokens: Some(64), + ..Default::default() + }, + ); + let (scheduler, slots, _cfg_tx, cancel_token) = make_scheduler(workers, None, true); + + scheduler + .schedule( + Some("req-1".to_string()), + 64, + Some(vec![1, 2, 3, 4]), + OverlapScores::default(), + Some(&crate::config::RouterConfigOverride { + track_prefill_tokens: Some(false), + ..Default::default() + }), + true, + None, + 0.0, + None, + None, + ) + .await + .unwrap(); + + assert_eq!( + slots + .active_tokens() + .get(&WorkerWithDpRank::new(0, 0)) + .copied(), + Some(0) + ); + + cancel_token.cancel(); + } + #[tokio::test] async fn test_mark_prefill_completed_drains_pending_queue() { let mut workers = HashMap::new(); @@ -474,7 +530,7 @@ mod tests { .collect(); expected.sort_by_key(|load| (load.worker_id, load.dp_rank)); - let mut actual = scheduler.get_potential_loads(Some(token_seq), 128, overlaps); + let mut actual = scheduler.get_potential_loads(Some(token_seq), 128, overlaps, true); actual.sort_by_key(|load| (load.worker_id, load.dp_rank)); assert_eq!(actual.len(), expected.len()); @@ -500,7 +556,7 @@ mod tests { make_scheduler(HashMap::new(), None, false); scheduler.register_workers(&HashSet::from([42])); - let loads = scheduler.get_potential_loads(None, 64, OverlapScores::default()); + let loads = scheduler.get_potential_loads(None, 64, OverlapScores::default(), true); assert_eq!(loads.len(), 1); assert_eq!(loads[0].worker_id, 42); @@ -517,7 +573,7 @@ mod tests { assert_eq!( scheduler - .get_potential_loads(None, 64, OverlapScores::default()) + .get_potential_loads(None, 64, OverlapScores::default(), true) .len(), 1 ); @@ -536,7 +592,7 @@ mod tests { tokio::time::timeout(Duration::from_secs(1), async { loop { if scheduler - .get_potential_loads(None, 64, OverlapScores::default()) + .get_potential_loads(None, 64, OverlapScores::default(), true) .len() == 3 { @@ -550,4 +606,39 @@ mod tests { cancel_token.cancel(); } + + #[tokio::test] + async fn test_get_potential_loads_can_ignore_prefill_tokens() { + let mut workers = HashMap::new(); + workers.insert( + 0, + SimpleWorkerConfig { + max_num_batched_tokens: Some(256), + ..Default::default() + }, + ); + let (scheduler, _slots, _cfg_tx, cancel_token) = make_scheduler(workers, None, true); + + scheduler + .schedule( + Some("req-1".to_string()), + 64, + Some(vec![11, 22]), + OverlapScores::default(), + None, + true, + None, + 0.0, + None, + None, + ) + .await + .unwrap(); + + let loads = scheduler.get_potential_loads(None, 64, OverlapScores::default(), false); + assert_eq!(loads.len(), 1); + assert_eq!(loads[0].potential_prefill_tokens, 64); + + cancel_token.cancel(); + } } diff --git a/lib/kv-router/src/scheduling/policy.rs b/lib/kv-router/src/scheduling/policy.rs index 33548ee75c22..8a613da17469 100644 --- a/lib/kv-router/src/scheduling/policy.rs +++ b/lib/kv-router/src/scheduling/policy.rs @@ -135,6 +135,7 @@ mod tests { overlaps, decode_blocks: HashMap::new(), prefill_tokens: HashMap::new(), + track_prefill_tokens: true, router_config_override: None, update_states: false, lora_name: None, diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 663a7842f921..93334f6e1061 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -191,11 +191,14 @@ impl< /// Run the full scheduling pipeline for a single request: /// compute potential load -> select worker -> respond -> book via add_request. async fn schedule(&self, mut request: SchedulingRequest) { - let (decode_blocks, prefill_tokens) = self.slots.potential_blocks_and_tokens( - request.token_seq.as_deref(), - request.isl_tokens, - request.overlaps.clone(), - ); + let (decode_blocks, prefill_tokens) = self + .slots + .potential_blocks_and_tokens_with_prefill_tracking( + request.token_seq.as_deref(), + request.isl_tokens, + request.overlaps.clone(), + request.track_prefill_tokens, + ); request.decode_blocks = decode_blocks; request.prefill_tokens = prefill_tokens; @@ -235,6 +238,7 @@ impl< token_sequence: request.token_seq, isl: request.isl_tokens, overlap: selection.overlap_blocks, + track_prefill_tokens: request.track_prefill_tokens, expected_output_tokens: request.expected_output_tokens, worker: selection.worker, lora_name: request.lora_name.clone(), @@ -376,6 +380,7 @@ mod tests { overlaps: OverlapScores::default(), decode_blocks: HashMap::new(), prefill_tokens: HashMap::new(), + track_prefill_tokens: true, router_config_override: None, update_states: true, lora_name: None, @@ -695,6 +700,7 @@ mod tests { overlaps: OverlapScores::default(), decode_blocks: HashMap::new(), prefill_tokens: HashMap::new(), + track_prefill_tokens: true, router_config_override: None, update_states: true, lora_name: None, @@ -719,4 +725,31 @@ mod tests { .unwrap(); slots.free(&"filter-0".to_string()).await.unwrap(); } + + #[tokio::test(flavor = "multi_thread")] + async fn test_queue_busy_check_ignores_untracked_prefill_tokens() { + let (queue, slots) = make_queue(1, 16, 256, Some(0.0)); + + let (mut req1, rx1) = make_request("req-1", 256); + req1.track_prefill_tokens = false; + queue.enqueue(req1).await; + let _resp1 = rx1.await.unwrap().unwrap(); + assert_eq!( + slots + .active_tokens() + .get(&WorkerWithDpRank::new(0, 0)) + .copied(), + Some(0) + ); + + let (req2, rx2) = make_request("req-2", 256); + queue.enqueue(req2).await; + let _resp2 = rx2.await.unwrap().unwrap(); + assert_eq!(queue.pending_count(), 0); + + let _ = slots.mark_prefill_completed(&"req-1".to_string()).await; + let _ = slots.free(&"req-1".to_string()).await; + let _ = slots.mark_prefill_completed(&"req-2".to_string()).await; + let _ = slots.free(&"req-2".to_string()).await; + } } diff --git a/lib/kv-router/src/scheduling/types.rs b/lib/kv-router/src/scheduling/types.rs index 9fe246601bc3..ce6bc2114fe7 100644 --- a/lib/kv-router/src/scheduling/types.rs +++ b/lib/kv-router/src/scheduling/types.rs @@ -42,6 +42,7 @@ pub struct SchedulingRequest { pub overlaps: OverlapScores, pub decode_blocks: HashMap, pub prefill_tokens: HashMap, + pub track_prefill_tokens: bool, pub router_config_override: Option, pub update_states: bool, pub lora_name: Option, diff --git a/lib/kv-router/src/sequences/multi_worker.rs b/lib/kv-router/src/sequences/multi_worker.rs index c00d1de013e3..c6122cadaec4 100644 --- a/lib/kv-router/src/sequences/multi_worker.rs +++ b/lib/kv-router/src/sequences/multi_worker.rs @@ -94,6 +94,7 @@ pub struct SequenceRequest { pub token_sequence: Option>, pub isl: usize, pub overlap: u32, + pub track_prefill_tokens: bool, pub expected_output_tokens: Option, pub worker: WorkerWithDpRank, pub lora_name: Option, @@ -218,6 +219,7 @@ impl ActiveSequencesMultiWorker

{ token_sequence, isl, overlap, + track_prefill_tokens, expected_output_tokens, } => { self.request_to_worker @@ -230,12 +232,13 @@ impl ActiveSequencesMultiWorker

{ let table = self.workers.read(); if let Some(&idx) = table.index.get(&event.worker) { - table.slots[idx].1.write().add_request( + table.slots[idx].1.write().add_request_with_prefill_tracking( event.request_id.clone(), token_sequence.clone(), *isl, *overlap, *expected_output_tokens, + *track_prefill_tokens, ); } else { tracing::warn!( @@ -370,6 +373,7 @@ impl ActiveSequencesMultiWorker

{ token_sequence, isl, overlap, + track_prefill_tokens, expected_output_tokens, worker, lora_name, @@ -394,6 +398,7 @@ impl ActiveSequencesMultiWorker

{ token_sequence: token_sequence.clone(), isl, overlap, + track_prefill_tokens, expected_output_tokens, }, router_id: self.router_id, @@ -415,12 +420,13 @@ impl ActiveSequencesMultiWorker

{ .get(&worker) .ok_or(SequenceError::WorkerNotFound { worker })?; let mut seq = table.slots[idx].1.write(); - seq.add_request( + seq.add_request_with_prefill_tracking( request_id, token_sequence, isl, overlap, expected_output_tokens, + track_prefill_tokens, ) }; @@ -634,6 +640,19 @@ impl ActiveSequencesMultiWorker

{ ) -> ( HashMap, HashMap, + ) { + self.potential_blocks_and_tokens_with_prefill_tracking(token_sequence, isl, overlaps, true) + } + + pub fn potential_blocks_and_tokens_with_prefill_tracking( + &self, + token_sequence: Option<&[SequenceHash]>, + isl: usize, + overlaps: OverlapScores, + track_prefill_tokens: bool, + ) -> ( + HashMap, + HashMap, ) { #[cfg(feature = "bench")] let start = tokio::time::Instant::now(); @@ -649,9 +668,14 @@ impl ActiveSequencesMultiWorker

{ for (worker, lock) in &table.slots { let overlap = *overlaps.scores.get(worker).unwrap_or(&0); - let (blocks, tokens) = - lock.read() - .potential_blocks_and_tokens(token_sequence, isl, overlap); + let (blocks, tokens) = lock + .read() + .potential_blocks_and_tokens_with_prefill_tracking( + token_sequence, + isl, + overlap, + track_prefill_tokens, + ); potential_blocks.insert(*worker, blocks); potential_tokens.insert(*worker, tokens); } @@ -770,3 +794,44 @@ impl ActiveSequencesMultiWorker

{ }); } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::test_utils::NoopSequencePublisher; + + fn make_sequences() -> ActiveSequencesMultiWorker { + ActiveSequencesMultiWorker::new( + NoopSequencePublisher, + 4, + HashMap::from([(1_u64, (0_u32, 1_u32))]), + false, + 0, + "test", + ) + } + + #[tokio::test] + async fn add_request_can_skip_prefill_token_tracking() { + let sequences = make_sequences(); + let worker = WorkerWithDpRank::new(1, 0); + + sequences + .add_request(SequenceRequest { + request_id: "req-1".to_string(), + token_sequence: Some(vec![1, 2, 3]), + isl: 12, + overlap: 0, + track_prefill_tokens: false, + expected_output_tokens: None, + worker, + lora_name: None, + }) + .await + .unwrap(); + + assert_eq!(sequences.active_tokens().get(&worker).copied(), Some(0)); + } +} diff --git a/lib/kv-router/src/sequences/single.rs b/lib/kv-router/src/sequences/single.rs index 99116eb298b0..f27cadaefa24 100644 --- a/lib/kv-router/src/sequences/single.rs +++ b/lib/kv-router/src/sequences/single.rs @@ -143,6 +143,27 @@ impl ActiveSequences { isl: usize, overlap: u32, expected_output_tokens: Option, + ) -> HashSet { + self.add_request_with_prefill_tracking( + request_id, + token_sequence, + isl, + overlap, + expected_output_tokens, + true, + ) + } + + /// Add a new request with optional prompt-token load accounting. + /// Returns the set of expired request IDs that were removed during cleanup. + pub fn add_request_with_prefill_tracking( + &mut self, + request_id: RequestId, + token_sequence: Option>, + isl: usize, + overlap: u32, + expected_output_tokens: Option, + track_prefill_tokens: bool, ) -> HashSet { // Check for double-add and log error, returning early if self.active_seqs.contains_key(&request_id) { @@ -153,7 +174,11 @@ impl ActiveSequences { // Lazily check and clean up expired requests, capturing removed IDs let removed_requests = self.force_expiry(); - let prefill_tokens = self.new_tokens(isl, overlap); + let prefill_tokens = if track_prefill_tokens { + self.new_tokens(isl, overlap) + } else { + 0 + }; self.prefill_tokens .insert(request_id.clone(), prefill_tokens); self.active_tokens += prefill_tokens; @@ -208,13 +233,27 @@ impl ActiveSequences { token_sequence: Option<&[SequenceHash]>, isl: usize, overlap: u32, + ) -> (usize, usize) { + self.potential_blocks_and_tokens_with_prefill_tracking(token_sequence, isl, overlap, true) + } + + pub fn potential_blocks_and_tokens_with_prefill_tracking( + &self, + token_sequence: Option<&[SequenceHash]>, + isl: usize, + overlap: u32, + track_prefill_tokens: bool, ) -> (usize, usize) { let potential_blocks = if let Some(token_seq) = token_sequence { self.new_blocks(token_seq) + self.active_blocks() } else { self.active_blocks() }; - let potential_tokens = self.new_tokens(isl, overlap) + self.active_tokens; + let potential_tokens = if track_prefill_tokens { + self.new_tokens(isl, overlap) + self.active_tokens + } else { + self.active_tokens + }; (potential_blocks, potential_tokens) } @@ -424,6 +463,48 @@ mod tests { assert_eq!(seq_manager.active_tokens(), 0); } + #[test] + fn test_add_request_without_prefill_tracking_keeps_active_tokens_zero() { + let mut seq_manager = ActiveSequences::new(4); + + seq_manager.add_request_with_prefill_tracking( + "r1".to_string(), + Some(vec![1, 2, 3]), + 12, + 0, + None, + false, + ); + + assert_eq!(seq_manager.active_tokens(), 0); + seq_manager.mark_prefill_completed(&"r1".to_string()); + assert_eq!(seq_manager.active_tokens(), 0); + seq_manager.free(&"r1".to_string()); + assert_eq!(seq_manager.active_blocks(), 0); + } + + #[test] + fn test_potential_blocks_and_tokens_without_prefill_tracking_ignores_prompt_load() { + let mut seq_manager = ActiveSequences::new(4); + seq_manager.add_request_with_prefill_tracking( + "r1".to_string(), + Some(vec![1, 2, 3]), + 12, + 0, + None, + false, + ); + + let (blocks, tokens) = seq_manager.potential_blocks_and_tokens_with_prefill_tracking( + Some(&[1, 2, 3, 4]), + 16, + 0, + false, + ); + assert_eq!(blocks, 4); + assert_eq!(tokens, 0); + } + #[tokio::test(start_paused = true)] async fn test_force_expiry() { let block_size = 4; diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 537e0729a82f..18ef5089c293 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -505,6 +505,9 @@ where router_config_override, lora_name.as_deref(), ); + let track_prefill_tokens = self + .kv_router_config + .track_prefill_tokens(router_config_override); if let Err(e) = self .scheduler @@ -513,6 +516,7 @@ where token_sequence: maybe_seq_hashes, isl: isl_tokens, overlap: overlap_blocks, + track_prefill_tokens, expected_output_tokens, worker, lora_name, @@ -584,10 +588,16 @@ where router_config_override, lora_name, ); + let track_prefill_tokens = self + .kv_router_config + .track_prefill_tokens(router_config_override); - Ok(self - .scheduler - .get_potential_loads(maybe_seq_hashes, isl_tokens, overlap_scores)) + Ok(self.scheduler.get_potential_loads( + maybe_seq_hashes, + isl_tokens, + overlap_scores, + track_prefill_tokens, + )) } /// Dump all events from the indexer diff --git a/lib/llm/src/kv_router/prefill_router.rs b/lib/llm/src/kv_router/prefill_router.rs index 561e25005275..6c891e56e4f5 100644 --- a/lib/llm/src/kv_router/prefill_router.rs +++ b/lib/llm/src/kv_router/prefill_router.rs @@ -59,6 +59,17 @@ enum PrefillOutcome { Completed(PrefillResult), } +fn build_decode_router_override( + existing_override: Option, +) -> RouterConfigOverride { + RouterConfigOverride { + overlap_score_weight: Some(0.0), + assume_kv_reuse: Some(false), + track_prefill_tokens: Some(false), + ..existing_override.unwrap_or_default() + } +} + /// The inner router used by PrefillRouter #[derive(Clone)] enum InnerPrefillRouter { @@ -717,12 +728,10 @@ impl // - overlap_score_weight = 0 (no KV cache overlap scoring for decode) // - assume_kv_reuse = false (generate random hashes since decode workers // may already have blocks cached from prefill transfer) + // - track_prefill_tokens = false (decode router should ignore prompt-side load) let existing_override = decode_req.router_config_override.take(); - decode_req.router_config_override = Some(RouterConfigOverride { - overlap_score_weight: Some(0.0), - assume_kv_reuse: Some(false), - ..existing_override.unwrap_or_default() - }); + decode_req.router_config_override = + Some(build_decode_router_override(existing_override)); // Map the modified request through with preserved context let decode_request = context.map(|_| decode_req); @@ -739,3 +748,21 @@ impl } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn decode_router_override_disables_overlap_and_prefill_tracking() { + let override_config = build_decode_router_override(Some(RouterConfigOverride { + router_temperature: Some(0.7), + ..Default::default() + })); + + assert_eq!(override_config.overlap_score_weight, Some(0.0)); + assert_eq!(override_config.assume_kv_reuse, Some(false)); + assert_eq!(override_config.track_prefill_tokens, Some(false)); + assert_eq!(override_config.router_temperature, Some(0.7)); + } +} diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs index 6b72f53e20e8..f38413b0a5dd 100644 --- a/lib/llm/src/kv_router/scheduler.rs +++ b/lib/llm/src/kv_router/scheduler.rs @@ -81,6 +81,7 @@ where block_size, selector, policy, + kv_router_config.router_track_prefill_tokens, component.drt().child_token(), worker_type, watch_worker_configs, @@ -180,9 +181,10 @@ where token_seq: Option>, isl_tokens: usize, overlaps: OverlapScores, + track_prefill_tokens: bool, ) -> Vec { self.inner - .get_potential_loads(token_seq, isl_tokens, overlaps) + .get_potential_loads(token_seq, isl_tokens, overlaps, track_prefill_tokens) } pub fn get_active_lora_counts(&self) -> HashMap { diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index ebf540645933..1620a01e497f 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -223,6 +223,7 @@ mod tests { token_sequence: Some(vec![0, 1, 2]), isl: 12, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::new(0, 0), lora_name: None, @@ -235,6 +236,7 @@ mod tests { token_sequence: Some(vec![3, 4]), isl: 8, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::new(0, 1), lora_name: None, @@ -247,6 +249,7 @@ mod tests { token_sequence: Some(vec![0, 1, 2, 3]), isl: 16, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::new(1, 0), lora_name: None, @@ -373,6 +376,7 @@ mod tests { token_sequence: None, isl: 12, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::from_worker_id(0), lora_name: None, @@ -385,6 +389,7 @@ mod tests { token_sequence: None, isl: 8, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::from_worker_id(1), lora_name: None, @@ -397,6 +402,7 @@ mod tests { token_sequence: None, isl: 16, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::from_worker_id(2), lora_name: None, diff --git a/lib/mocker/src/replay/router/offline.rs b/lib/mocker/src/replay/router/offline.rs index 59cb17f9b084..0c52c96c690d 100644 --- a/lib/mocker/src/replay/router/offline.rs +++ b/lib/mocker/src/replay/router/offline.rs @@ -77,6 +77,7 @@ impl PendingRequest { overlaps: self.overlaps.clone(), decode_blocks, prefill_tokens, + track_prefill_tokens: true, router_config_override: None, update_states: true, lora_name: None, @@ -271,6 +272,7 @@ impl OfflineReplayRouter { token_sequence: request.token_seq, isl: request.isl_tokens, overlap: selection.overlap_blocks, + track_prefill_tokens: true, expected_output_tokens: request.expected_output_tokens, worker: selection.worker, lora_name: None, diff --git a/lib/mocker/src/replay/router/online.rs b/lib/mocker/src/replay/router/online.rs index b5f7eb6671fc..8ba7793bcf41 100644 --- a/lib/mocker/src/replay/router/online.rs +++ b/lib/mocker/src/replay/router/online.rs @@ -138,6 +138,7 @@ impl KvReplayRouter { args.block_size as u32, selector, policy, + config.router_track_prefill_tokens, CancellationToken::new(), "replay", false, From 42dd884abb4fd8a31efd8296c9ae31d2eda68395 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 17:42:11 -0700 Subject: [PATCH 02/24] refactor replay timing without moving the clock Signed-off-by: PeaBrane --- lib/mocker/src/replay/offline/mod.rs | 1 + lib/mocker/src/replay/offline/multi.rs | 282 ++++++++++++------ .../src/replay/offline/runtime_utils.rs | 200 +++++++++++++ 3 files changed, 387 insertions(+), 96 deletions(-) create mode 100644 lib/mocker/src/replay/offline/runtime_utils.rs diff --git a/lib/mocker/src/replay/offline/mod.rs b/lib/mocker/src/replay/offline/mod.rs index 225780a0af03..a84cfd1b0c6f 100644 --- a/lib/mocker/src/replay/offline/mod.rs +++ b/lib/mocker/src/replay/offline/mod.rs @@ -10,6 +10,7 @@ use dynamo_kv_router::config::KvRouterConfig; pub(crate) mod core; pub(crate) mod events; pub(crate) mod multi; +pub(crate) mod runtime_utils; pub(crate) mod single; pub(crate) mod state; diff --git a/lib/mocker/src/replay/offline/multi.rs b/lib/mocker/src/replay/offline/multi.rs index 847133d556db..41d939707d82 100644 --- a/lib/mocker/src/replay/offline/multi.rs +++ b/lib/mocker/src/replay/offline/multi.rs @@ -1,8 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use super::events::{SimulationEvent, SimulationEventKind}; +use super::events::SimulationEvent; use super::normalize_trace_requests; +use super::runtime_utils::{ + WorkerCompletionPayload, next_timestamp as choose_next_timestamp, pop_next_concurrency_ready, + pop_next_trace_ready, pop_ready_worker_completion, push_worker_completion, +}; #[cfg(test)] use super::state::OfflineWorkerSnapshot; use super::state::OfflineWorkerState; @@ -295,21 +299,7 @@ impl OfflineRuntime { (ReplayMode::Concurrency { .. }, AdmissionSource::Requests(_)) => None, }; - match (next_arrival_ms, next_event_ms) { - (Some(arrival_ms), Some(event_ms)) => Some(arrival_ms.min(event_ms)), - (Some(arrival_ms), None) => Some(arrival_ms), - (None, Some(event_ms)) => Some(event_ms), - (None, None) => None, - } - } - - fn push_event(&mut self, at_ms: f64, kind: SimulationEventKind) { - self.events.push(SimulationEvent { - at_ms, - seq_no: self.next_event_seq, - kind, - }); - self.next_event_seq += 1; + choose_next_timestamp(next_arrival_ms, next_event_ms) } fn apply_completed_requests(&mut self, worker_idx: usize, completed_requests: usize) { @@ -395,25 +385,13 @@ impl OfflineRuntime { fn apply_worker_completions(&mut self) -> anyhow::Result { let mut changed = false; - loop { - let Some(event) = self.events.peek() else { - break; - }; - if event.at_ms != self.now_ms { - break; - } - if !matches!(event.kind, SimulationEventKind::WorkerCompletion { .. }) { - break; - } - - let event = self.events.pop().expect("event must exist after peek"); - let SimulationEventKind::WorkerCompletion { - worker_idx, - completed_requests, - output_signals, - kv_events, - } = event.kind; - + while let Some(WorkerCompletionPayload { + worker_idx, + completed_requests, + output_signals, + kv_events, + }) = pop_ready_worker_completion(&mut self.events, self.now_ms) + { self.workers[worker_idx].mark_idle(); self.process_completed_pass(worker_idx, completed_requests, output_signals, kv_events)?; changed = true; @@ -424,39 +402,33 @@ impl OfflineRuntime { fn release_trace_arrivals(&mut self) -> anyhow::Result { let mut released_any = false; - let mut ready_requests = Vec::new(); - - match &mut self.admission { - AdmissionSource::Requests(pending) => { - while pending - .front() - .and_then(|request| request.arrival_timestamp_ms) - .is_some_and(|arrival_ms| arrival_ms <= self.now_ms) - { - let request = pending - .pop_front() - .expect("front request must exist when arrival is ready"); - let arrival_ms = request - .arrival_timestamp_ms - .expect("trace replay requests must have an arrival timestamp"); - ready_requests.push((request, arrival_ms, None)); - } - } - AdmissionSource::Workload(driver) => { - ready_requests.extend(driver.pop_ready(self.now_ms, usize::MAX).into_iter().map( - |ready| { - ( - ready.request, - ready.scheduled_ready_at_ms, - ready.replay_hashes, - ) - }, - )); + if matches!(self.admission, AdmissionSource::Requests(_)) { + loop { + let next_ready = match &mut self.admission { + AdmissionSource::Requests(pending) => { + pop_next_trace_ready(pending, self.now_ms) + } + AdmissionSource::Workload(_) => unreachable!(), + }; + let Some((request, arrival_ms)) = next_ready else { + break; + }; + self.assign_request(request, arrival_ms, None)?; + released_any = true; } + return Ok(released_any); } - for (request, arrival_ms, replay_hashes) in ready_requests { - self.assign_request(request, arrival_ms, replay_hashes)?; + let ready_requests = match &mut self.admission { + AdmissionSource::Requests(_) => unreachable!(), + AdmissionSource::Workload(driver) => driver.pop_ready(self.now_ms, usize::MAX), + }; + for ready in ready_requests { + self.assign_request( + ready.request, + ready.scheduled_ready_at_ms, + ready.replay_hashes, + )?; released_any = true; } @@ -465,33 +437,38 @@ impl OfflineRuntime { fn top_off_concurrency(&mut self, max_in_flight: usize) -> anyhow::Result { let mut released_any = false; + if matches!(self.admission, AdmissionSource::Requests(_)) { + loop { + let cluster_in_flight = self.cluster_in_flight(); + let next_ready = match &mut self.admission { + AdmissionSource::Requests(pending) => pop_next_concurrency_ready( + pending, + self.now_ms, + cluster_in_flight, + max_in_flight, + ), + AdmissionSource::Workload(_) => unreachable!(), + }; + let Some((request, arrival_ms)) = next_ready else { + break; + }; + self.assign_request(request, arrival_ms, None)?; + released_any = true; + } + return Ok(released_any); + } + let available = max_in_flight.saturating_sub(self.cluster_in_flight()); if available == 0 { return Ok(false); } - let mut ready_requests = Vec::new(); - match &mut self.admission { - AdmissionSource::Requests(pending) => { - for _ in 0..available { - let Some(request) = pending.pop_front() else { - break; - }; - ready_requests.push((request, None)); - } - } - AdmissionSource::Workload(driver) => { - ready_requests.extend( - driver - .pop_ready(self.now_ms, available) - .into_iter() - .map(|ready| (ready.request, ready.replay_hashes)), - ); - } - } - - for (request, replay_hashes) in ready_requests { - self.assign_request(request, self.now_ms, replay_hashes)?; + let ready_requests = match &mut self.admission { + AdmissionSource::Requests(_) => unreachable!(), + AdmissionSource::Workload(driver) => driver.pop_ready(self.now_ms, available), + }; + for ready in ready_requests { + self.assign_request(ready.request, self.now_ms, ready.replay_hashes)?; released_any = true; } @@ -531,9 +508,11 @@ impl OfflineRuntime { } self.workers[worker_idx].mark_busy(); - self.push_event( + push_worker_completion( + &mut self.events, + &mut self.next_event_seq, executed.end_ms, - SimulationEventKind::WorkerCompletion { + WorkerCompletionPayload { worker_idx, completed_requests: executed.completed_requests, output_signals: executed.output_signals, @@ -1483,16 +1462,16 @@ mod tests { let request_1 = collector.snapshot(Uuid::from_u128(1)).unwrap(); let request_2 = collector.snapshot(Uuid::from_u128(2)).unwrap(); let request_3 = collector.snapshot(Uuid::from_u128(3)).unwrap(); + let first_unblock_ms = request_1 + .first_token_ms + .unwrap() + .min(request_2.first_token_ms.unwrap()); assert!(stats.max_router_pending > 0); assert!(request_3.first_admit_ms.unwrap() > request_3.arrival_time_ms); - assert!( - request_3.first_admit_ms.unwrap() - < request_1 - .last_token_ms - .unwrap() - .min(request_2.last_token_ms.unwrap()) - ); + assert_eq!(request_3.first_admit_ms.unwrap(), first_unblock_ms); + assert!(request_3.first_admit_ms.unwrap() < request_1.last_token_ms.unwrap()); + assert!(request_3.first_admit_ms.unwrap() < request_2.last_token_ms.unwrap()); } #[test] @@ -1561,6 +1540,72 @@ mod tests { ); } + #[test] + fn test_multi_worker_trace_kv_router_fcfs_and_lcfs_admit_queued_requests_in_opposite_timestamp_order() + { + let requests = vec![ + DirectRequest { + tokens: vec![10; 64], + max_output_tokens: 8, + uuid: Some(Uuid::from_u128(10)), + dp_rank: 0, + arrival_timestamp_ms: Some(0.0), + }, + DirectRequest { + tokens: vec![20; 128], + max_output_tokens: 8, + uuid: Some(Uuid::from_u128(20)), + dp_rank: 0, + arrival_timestamp_ms: Some(0.0), + }, + DirectRequest { + tokens: vec![30; 64], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(30)), + dp_rank: 0, + arrival_timestamp_ms: Some(0.1), + }, + DirectRequest { + tokens: vec![40; 64], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(40)), + dp_rank: 0, + arrival_timestamp_ms: Some(0.2), + }, + ]; + + let (fcfs_collector, fcfs_stats) = run_trace_multi_collect_with_stats( + &queueing_router_args(RouterQueuePolicy::Fcfs), + requests.clone(), + 2, + ReplayRouterMode::KvRouter, + ); + let (lcfs_collector, lcfs_stats) = run_trace_multi_collect_with_stats( + &queueing_router_args(RouterQueuePolicy::Lcfs), + requests, + 2, + ReplayRouterMode::KvRouter, + ); + + let fcfs_request_30 = fcfs_collector.snapshot(Uuid::from_u128(30)).unwrap(); + let fcfs_request_40 = fcfs_collector.snapshot(Uuid::from_u128(40)).unwrap(); + let lcfs_request_30 = lcfs_collector.snapshot(Uuid::from_u128(30)).unwrap(); + let lcfs_request_40 = lcfs_collector.snapshot(Uuid::from_u128(40)).unwrap(); + + assert!(fcfs_stats.max_router_pending > 0); + assert!(lcfs_stats.max_router_pending > 0); + assert_eq!( + &fcfs_stats.dispatch_order[2..4], + &[Uuid::from_u128(30), Uuid::from_u128(40)] + ); + assert_eq!( + &lcfs_stats.dispatch_order[2..4], + &[Uuid::from_u128(40), Uuid::from_u128(30)] + ); + assert!(fcfs_request_30.first_admit_ms.unwrap() < fcfs_request_40.first_admit_ms.unwrap()); + assert!(lcfs_request_40.first_admit_ms.unwrap() < lcfs_request_30.first_admit_ms.unwrap()); + } + #[test] fn test_multi_worker_concurrency_kv_router_respects_max_in_flight() { let args = queueing_router_args(RouterQueuePolicy::Fcfs); @@ -1605,6 +1650,51 @@ mod tests { assert!(stats.max_router_pending > 0); } + #[test] + fn test_multi_worker_concurrency_kv_router_records_backfill_timing() { + let args = queueing_router_args(RouterQueuePolicy::Fcfs); + let (collector, stats) = run_concurrency_multi_collect_with_stats( + &args, + vec![ + DirectRequest { + tokens: vec![1; 64], + max_output_tokens: 2, + uuid: Some(Uuid::from_u128(11)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + DirectRequest { + tokens: vec![2; 64], + max_output_tokens: 4, + uuid: Some(Uuid::from_u128(22)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + DirectRequest { + tokens: vec![3; 64], + max_output_tokens: 2, + uuid: Some(Uuid::from_u128(33)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + ], + 2, + 2, + ReplayRouterMode::KvRouter, + ); + + let request_1 = collector.snapshot(Uuid::from_u128(11)).unwrap(); + let request_2 = collector.snapshot(Uuid::from_u128(22)).unwrap(); + let request_3 = collector.snapshot(Uuid::from_u128(33)).unwrap(); + + assert_eq!(request_1.arrival_time_ms, 0.0); + assert_eq!(request_2.arrival_time_ms, 0.0); + assert_eq!(request_3.arrival_time_ms, request_1.last_token_ms.unwrap()); + assert!(request_3.arrival_time_ms < request_2.last_token_ms.unwrap()); + assert_eq!(request_3.first_admit_ms.unwrap(), request_3.arrival_time_ms); + assert_eq!(stats.max_in_flight_seen, 2); + } + #[test] fn test_multi_worker_trace_single_worker_round_robin_matches_single_runtime() { let args = replay_args(true, true); diff --git a/lib/mocker/src/replay/offline/runtime_utils.rs b/lib/mocker/src/replay/offline/runtime_utils.rs new file mode 100644 index 000000000000..1050c34dab05 --- /dev/null +++ b/lib/mocker/src/replay/offline/runtime_utils.rs @@ -0,0 +1,200 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{BinaryHeap, VecDeque}; + +use dynamo_kv_router::protocols::RouterEvent; + +use super::events::{SimulationEvent, SimulationEventKind}; +use crate::common::protocols::{DirectRequest, OutputSignal}; + +#[derive(Debug)] +pub(super) struct WorkerCompletionPayload { + pub worker_idx: usize, + pub completed_requests: usize, + pub output_signals: Vec, + pub kv_events: Vec, +} + +pub(super) fn next_timestamp( + next_arrival_ms: Option, + next_event_ms: Option, +) -> Option { + match (next_arrival_ms, next_event_ms) { + (Some(arrival_ms), Some(event_ms)) => Some(arrival_ms.min(event_ms)), + (Some(arrival_ms), None) => Some(arrival_ms), + (None, Some(event_ms)) => Some(event_ms), + (None, None) => None, + } +} + +pub(super) fn pop_next_trace_ready( + pending: &mut VecDeque, + now_ms: f64, +) -> Option<(DirectRequest, f64)> { + let arrival_ms = pending + .front() + .and_then(|request| request.arrival_timestamp_ms) + .filter(|arrival_ms| *arrival_ms <= now_ms)?; + let request = pending + .pop_front() + .expect("front request must exist when arrival is ready"); + Some((request, arrival_ms)) +} + +pub(super) fn pop_next_concurrency_ready( + pending: &mut VecDeque, + now_ms: f64, + cluster_in_flight: usize, + max_in_flight: usize, +) -> Option<(DirectRequest, f64)> { + if cluster_in_flight >= max_in_flight { + return None; + } + let request = pending.pop_front()?; + Some((request, now_ms)) +} + +pub(super) fn push_worker_completion( + events: &mut BinaryHeap, + next_event_seq: &mut u64, + at_ms: f64, + payload: WorkerCompletionPayload, +) { + events.push(SimulationEvent { + at_ms, + seq_no: *next_event_seq, + kind: SimulationEventKind::WorkerCompletion { + worker_idx: payload.worker_idx, + completed_requests: payload.completed_requests, + output_signals: payload.output_signals, + kv_events: payload.kv_events, + }, + }); + *next_event_seq += 1; +} + +pub(super) fn pop_ready_worker_completion( + events: &mut BinaryHeap, + now_ms: f64, +) -> Option { + let event = events.peek()?; + if event.at_ms != now_ms { + return None; + } + let event = events.pop().expect("event must exist after peek"); + let SimulationEventKind::WorkerCompletion { + worker_idx, + completed_requests, + output_signals, + kv_events, + } = event.kind; + Some(WorkerCompletionPayload { + worker_idx, + completed_requests, + output_signals, + kv_events, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use uuid::Uuid; + + fn direct_request(uuid: u128, arrival_timestamp_ms: Option) -> DirectRequest { + DirectRequest { + tokens: vec![1; 8], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(uuid)), + dp_rank: 0, + arrival_timestamp_ms, + } + } + + #[test] + fn test_next_timestamp_matches_current_choice_logic() { + assert_eq!(next_timestamp(Some(1.0), Some(2.0)), Some(1.0)); + assert_eq!(next_timestamp(Some(2.0), Some(1.0)), Some(1.0)); + assert_eq!(next_timestamp(Some(3.0), None), Some(3.0)); + assert_eq!(next_timestamp(None, Some(4.0)), Some(4.0)); + assert_eq!(next_timestamp(None, None), None); + } + + #[test] + fn test_pop_next_trace_ready_releases_only_arrivals_at_or_before_now() { + let mut pending = VecDeque::from(vec![ + direct_request(1, Some(1.0)), + direct_request(2, Some(1.1)), + direct_request(3, Some(2.0)), + ]); + + let (request_1, arrival_1) = pop_next_trace_ready(&mut pending, 1.0).unwrap(); + assert_eq!(request_1.uuid, Some(Uuid::from_u128(1))); + assert_eq!(arrival_1, 1.0); + + assert!(pop_next_trace_ready(&mut pending, 1.0).is_none()); + + let (request_2, arrival_2) = pop_next_trace_ready(&mut pending, 1.1).unwrap(); + assert_eq!(request_2.uuid, Some(Uuid::from_u128(2))); + assert_eq!(arrival_2, 1.1); + assert_eq!(pending.len(), 1); + } + + #[test] + fn test_pop_next_concurrency_ready_stops_at_max_in_flight() { + let mut pending = VecDeque::from(vec![direct_request(1, None), direct_request(2, None)]); + + assert!(pop_next_concurrency_ready(&mut pending, 5.0, 2, 2).is_none()); + + let (request, arrival_ms) = pop_next_concurrency_ready(&mut pending, 5.0, 1, 2).unwrap(); + assert_eq!(request.uuid, Some(Uuid::from_u128(1))); + assert_eq!(arrival_ms, 5.0); + assert_eq!(pending.len(), 1); + } + + #[test] + fn test_worker_completion_helpers_preserve_same_time_sequence_ordering() { + let mut events = BinaryHeap::new(); + let mut next_event_seq = 0; + + push_worker_completion( + &mut events, + &mut next_event_seq, + 10.0, + WorkerCompletionPayload { + worker_idx: 7, + completed_requests: 1, + output_signals: vec![OutputSignal { + uuid: Uuid::from_u128(7), + completed: true, + }], + kv_events: Vec::new(), + }, + ); + push_worker_completion( + &mut events, + &mut next_event_seq, + 10.0, + WorkerCompletionPayload { + worker_idx: 8, + completed_requests: 2, + output_signals: vec![OutputSignal { + uuid: Uuid::from_u128(8), + completed: false, + }], + kv_events: Vec::new(), + }, + ); + + assert!(pop_ready_worker_completion(&mut events, 9.0).is_none()); + + let first = pop_ready_worker_completion(&mut events, 10.0).unwrap(); + let second = pop_ready_worker_completion(&mut events, 10.0).unwrap(); + assert_eq!(first.worker_idx, 7); + assert_eq!(first.completed_requests, 1); + assert_eq!(second.worker_idx, 8); + assert_eq!(second.completed_requests, 2); + assert!(events.is_empty()); + } +} From b4d281c87d59541cf2d81957149201badb4a5b56 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 18:49:20 -0700 Subject: [PATCH 03/24] feat: teach replay to disagg without drama Signed-off-by: PeaBrane --- docs/benchmarks/mocker-trace-replay.md | 24 +- docs/mocker/mocker.md | 33 +- lib/bindings/python/rust/llm/replay.rs | 417 +++++--- lib/bindings/python/src/dynamo/replay/api.py | 16 + lib/bindings/python/src/dynamo/replay/main.py | 63 +- lib/bindings/python/tests/test_replay.py | 125 ++- lib/mocker/src/replay/entrypoints.rs | 119 ++- lib/mocker/src/replay/mod.rs | 47 +- lib/mocker/src/replay/offline/README.md | 41 +- lib/mocker/src/replay/offline/disagg.rs | 984 ++++++++++++++++++ lib/mocker/src/replay/offline/events.rs | 8 + lib/mocker/src/replay/offline/mod.rs | 37 + lib/mocker/src/replay/offline/multi.rs | 5 +- .../src/replay/offline/runtime_utils.rs | 11 +- lib/mocker/src/replay/offline/state.rs | 4 + lib/mocker/src/replay/router/offline.rs | 7 +- lib/mocker/src/replay/router/shared.rs | 2 +- lib/mocker/src/replay/validate.rs | 67 ++ lib/mocker/src/scheduler/mod.rs | 7 + lib/mocker/src/scheduler/sglang/core.rs | 4 + lib/mocker/src/scheduler/vllm/core.rs | 4 + 21 files changed, 1853 insertions(+), 172 deletions(-) create mode 100644 lib/mocker/src/replay/offline/disagg.rs diff --git a/docs/benchmarks/mocker-trace-replay.md b/docs/benchmarks/mocker-trace-replay.md index b6ee91753c3d..6defdabb57cc 100644 --- a/docs/benchmarks/mocker-trace-replay.md +++ b/docs/benchmarks/mocker-trace-replay.md @@ -117,6 +117,8 @@ The dedicated replay CLI exposes: - `--replay-mode offline|online` - `--router-mode round_robin|kv_router` - `--num-workers` +- `--num-prefill-workers` +- `--num-decode-workers` - `--replay-concurrency` - `--arrival-interval-ms` - `--arrival-speedup-ratio` @@ -125,6 +127,8 @@ The dedicated replay CLI exposes: - `--num-prefix-groups` - `--inter-turn-delay-ms` - `--extra-engine-args` (JSON string) +- `--prefill-engine-args` (JSON string) +- `--decode-engine-args` (JSON string) - `--router-config` (JSON string) - `--report-json` @@ -164,6 +168,14 @@ as `block_size`, `engine_type`, `dp_size`, `speedup_ratio`, and `decode_speedup_ `--extra-engine-args`, not as top-level replay CLI flags. Unspecified fields fall back to the same defaults used by `MockEngineArgs::default()` and `KvRouterConfig::default()`. +Offline disagg replay uses staged engine args instead of `--extra-engine-args`: + +- `--prefill-engine-args` for the prefill worker config +- `--decode-engine-args` for the decode worker config +- `--num-prefill-workers` and `--num-decode-workers` for pool sizes + +The staged JSON accepts `worker_type: "prefill" | "decode"` for CLI convenience. + ### Synthetic Replay Synthetic replay bypasses trace loading and generates in-memory requests with fixed input/output @@ -320,9 +332,11 @@ If `--report-json` is not provided, `python -m dynamo.replay` writes a timestamp Shared replay constraints: -- aggregated mode - `extra_engine_args.engine_type` must be `vllm` or `sglang` -- `extra_engine_args.dp_size` must be `1` +- aggregated replay requires the existing aggregated args path +- disagg replay requires both `prefill_engine_args` and `decode_engine_args` +- disagg replay requires `router_mode=kv_router` +- replay `dp_size` must be `1` Additional offline constraints: @@ -330,6 +344,7 @@ Additional offline constraints: - single-worker offline replay is still a dedicated fast path for `vllm`, but it now supports both flat request replay and workload-driven multi-turn replay - `sglang` still goes through the shared multi-worker replay runtime even when `num_workers=1` +- offline disagg replay is a separate two-stage runtime with prefill and decode worker pools Additional online constraints: @@ -343,12 +358,13 @@ If you violate those constraints, replay fails immediately with a validation err either a trace file, or all of `--input-tokens`, `--output-tokens`, and `--request-count` - `--replay-concurrency` works with both trace replay and synthetic replay - mocker compute-speed knobs such as `speedup_ratio` still affect simulated timing when passed via - `--extra-engine-args` + the engine-args JSON for the chosen replay mode - `--arrival-speedup-ratio` affects trace timestamps, not worker compute speed - `--arrival-interval-ms` only applies to synthetic replay - `--turns-per-session`, `--shared-prefix-ratio`, `--num-prefix-groups`, and `--inter-turn-delay-ms` only apply to synthetic replay -- `--extra-engine-args` and `--router-config` are JSON strings on the standalone replay CLI +- `--extra-engine-args`, `--prefill-engine-args`, `--decode-engine-args`, and `--router-config` + are JSON strings on the standalone replay CLI - offline replay does not need planner runtime setup, router registration, or external event transport - the replay block size should match the trace block size, because token synthesis expands `hash_ids` using the configured block size diff --git a/docs/mocker/mocker.md b/docs/mocker/mocker.md index 20f2822cbae5..75b5881bc9b5 100644 --- a/docs/mocker/mocker.md +++ b/docs/mocker/mocker.md @@ -126,10 +126,12 @@ python -m dynamo.mocker \ The mocker supports replaying Mooncake-style traces through the dedicated replay CLI, which exposes `offline|online`, `round_robin|kv_router`, `arrival_speedup_ratio`, closed-loop concurrency -admission, and synthetic workload generation directly: +admission, synthetic workload generation, and offline disaggregated prefill/decode replay directly: -The replay CLI defaults to `--replay-mode offline` and `--router-mode round_robin`. Engine settings -such as `block_size`, `engine_type`, and compute speedups still belong in `--extra-engine-args`. +The replay CLI defaults to `--replay-mode offline` and `--router-mode round_robin`. Aggregated +replay uses `--extra-engine-args`. Offline disagg replay instead uses +`--prefill-engine-args` plus `--decode-engine-args`, together with +`--num-prefill-workers` and `--num-decode-workers`. ```bash python -m dynamo.replay /path/to/mooncake_trace.jsonl \ @@ -197,6 +199,31 @@ Replay supports aggregated `vllm` and `sglang` engine configs. Internally replay `block_size`; for `sglang`, `sglang.page_size` is still accepted as a compatibility alias as long as it matches `block_size` when both are provided. +Offline replay also supports disaggregated `kv_router` mode. In that mode: + +- `--prefill-engine-args` must describe a prefill worker +- `--decode-engine-args` must describe a decode worker +- `--router-mode` must be `kv_router` +- only offline replay is supported + +Example: + +```bash +python -m dynamo.replay \ + --input-tokens 4096 \ + --output-tokens 256 \ + --request-count 100 \ + --replay-mode offline \ + --router-mode kv_router \ + --replay-concurrency 32 \ + --num-prefill-workers 2 \ + --num-decode-workers 6 \ + --prefill-engine-args '{"worker_type":"prefill","block_size":512}' \ + --decode-engine-args '{"worker_type":"decode","block_size":512}' \ + --router-config '{"router_queue_policy":"wspt"}' \ + --report-json /tmp/replay-report.json +``` + ## Performance Modeling Setup By default, the mocker uses hardcoded polynomial formulas to estimate prefill and decode timing. For more realistic simulations, pass `--planner-profile-data` with either: diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index 636b055d85ec..943b23e8919d 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -289,67 +289,113 @@ impl MockEngineArgs { } #[pyfunction] -#[pyo3(signature = (trace_file, extra_engine_args=None, router_config=None, num_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0))] +#[pyo3(signature = (trace_file, extra_engine_args=None, prefill_engine_args=None, decode_engine_args=None, router_config=None, num_workers=1, num_prefill_workers=1, num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0))] #[allow(clippy::too_many_arguments)] pub fn run_mocker_trace_replay( py: Python<'_>, trace_file: PathBuf, extra_engine_args: Option, + prefill_engine_args: Option, + decode_engine_args: Option, router_config: Option, num_workers: usize, + num_prefill_workers: usize, + num_decode_workers: usize, replay_concurrency: Option, replay_mode: &str, router_mode: &str, arrival_speedup_ratio: f64, ) -> PyResult { - let args = load_replay_mocker_args(py, extra_engine_args)?; + let args_selection = load_replay_args_selection( + py, + extra_engine_args, + prefill_engine_args, + decode_engine_args, + num_workers, + num_prefill_workers, + num_decode_workers, + )?; let router_config = load_replay_router_config(router_config); let replay_mode = replay_mode.to_owned(); let router_mode = parse_replay_router_mode(router_mode)?; let report = py.allow_threads(move || { let replay_concurrency = parse_replay_concurrency(replay_concurrency)?; - match (replay_mode.as_str(), replay_concurrency) { - ("offline", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_file_with_router_mode( - args, - router_config.clone(), - &trace_file, - max_in_flight, - num_workers, - router_mode, - ) - } - ("offline", None) => dynamo_mocker::replay::simulate_trace_file_with_router_mode( - args, - router_config.clone(), - &trace_file, - num_workers, - arrival_speedup_ratio, - router_mode, - ), - ("online", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_live_file_with_router_mode( - args, - router_config.clone(), - &trace_file, - max_in_flight, - num_workers, - router_mode, - ) + match args_selection { + ReplayArgsSelection::Aggregated(args) => { + match (replay_mode.as_str(), replay_concurrency) { + ("offline", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_file_with_router_mode( + *args, + router_config.clone(), + &trace_file, + max_in_flight, + num_workers, + router_mode, + ) + } + ("offline", None) => { + dynamo_mocker::replay::simulate_trace_file_with_router_mode( + *args, + router_config.clone(), + &trace_file, + num_workers, + arrival_speedup_ratio, + router_mode, + ) + } + ("online", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_live_file_with_router_mode( + *args, + router_config.clone(), + &trace_file, + max_in_flight, + num_workers, + router_mode, + ) + } + ("online", None) => { + dynamo_mocker::replay::simulate_trace_live_file_with_router_mode( + *args, + router_config.clone(), + &trace_file, + num_workers, + arrival_speedup_ratio, + router_mode, + ) + } + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + } } - ("online", None) => dynamo_mocker::replay::simulate_trace_live_file_with_router_mode( - args, - router_config.clone(), - &trace_file, - num_workers, - arrival_speedup_ratio, - router_mode, - ), - (other, _) => anyhow::bail!( - "replay_mode must be either 'offline' or 'online', got '{}'", - other - ), + ReplayArgsSelection::Disagg(config) => match (replay_mode.as_str(), replay_concurrency) + { + ("offline", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_file_disagg_with_router_mode( + *config, + router_config.clone(), + &trace_file, + max_in_flight, + router_mode, + ) + } + ("offline", None) => { + dynamo_mocker::replay::simulate_trace_file_disagg_with_router_mode( + *config, + router_config.clone(), + &trace_file, + arrival_speedup_ratio, + router_mode, + ) + } + ("online", _) => anyhow::bail!("disagg replay only supports replay_mode='offline'"), + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + }, } }); let report = report.map_err(to_pyerr)?; @@ -359,7 +405,7 @@ pub fn run_mocker_trace_replay( } #[pyfunction] -#[pyo3(signature = (input_tokens, output_tokens, request_count, extra_engine_args=None, router_config=None, num_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0, arrival_interval_ms=1.0, turns_per_session=1, shared_prefix_ratio=0.0, num_prefix_groups=0, inter_turn_delay_ms=0.0))] +#[pyo3(signature = (input_tokens, output_tokens, request_count, extra_engine_args=None, prefill_engine_args=None, decode_engine_args=None, router_config=None, num_workers=1, num_prefill_workers=1, num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0, arrival_interval_ms=1.0, turns_per_session=1, shared_prefix_ratio=0.0, num_prefix_groups=0, inter_turn_delay_ms=0.0))] #[allow(clippy::too_many_arguments)] pub fn run_mocker_synthetic_trace_replay( py: Python<'_>, @@ -367,8 +413,12 @@ pub fn run_mocker_synthetic_trace_replay( output_tokens: usize, request_count: usize, extra_engine_args: Option, + prefill_engine_args: Option, + decode_engine_args: Option, router_config: Option, num_workers: usize, + num_prefill_workers: usize, + num_decode_workers: usize, replay_concurrency: Option, replay_mode: &str, router_mode: &str, @@ -379,10 +429,22 @@ pub fn run_mocker_synthetic_trace_replay( num_prefix_groups: usize, inter_turn_delay_ms: f64, ) -> PyResult { - let args = load_replay_mocker_args(py, extra_engine_args)?; + let args_selection = load_replay_args_selection( + py, + extra_engine_args, + prefill_engine_args, + decode_engine_args, + num_workers, + num_prefill_workers, + num_decode_workers, + )?; let router_config = load_replay_router_config(router_config); let replay_mode = replay_mode.to_owned(); let router_mode = parse_replay_router_mode(router_mode)?; + let block_size = match &args_selection { + ReplayArgsSelection::Aggregated(args) => args.block_size.max(1), + ReplayArgsSelection::Disagg(config) => config.prefill_args.block_size.max(1), + }; let report = py.allow_threads(move || { let replay_concurrency = parse_replay_concurrency(replay_concurrency)?; let use_workload = turns_per_session > 1 @@ -392,7 +454,7 @@ pub fn run_mocker_synthetic_trace_replay( if use_workload { let mut trace = build_synthetic_workload( - args.block_size.max(1), + block_size, input_tokens, output_tokens, request_count, @@ -406,103 +468,159 @@ pub fn run_mocker_synthetic_trace_replay( trace = trace.speed_up_timing(arrival_speedup_ratio)?; } - return match (replay_mode.as_str(), replay_concurrency) { + return match args_selection { + ReplayArgsSelection::Aggregated(args) => match (replay_mode.as_str(), replay_concurrency) + { + ("offline", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_workload_with_router_mode( + *args, + router_config.clone(), + trace, + max_in_flight, + num_workers, + router_mode, + ) + } + ("offline", None) => { + dynamo_mocker::replay::simulate_trace_workload_with_router_mode( + *args, + router_config.clone(), + trace, + num_workers, + router_mode, + ) + } + ("online", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_live_workload_with_router_mode( + *args, + router_config.clone(), + trace, + max_in_flight, + num_workers, + router_mode, + ) + } + ("online", None) => { + dynamo_mocker::replay::simulate_trace_live_workload_with_router_mode( + *args, + router_config.clone(), + trace, + num_workers, + router_mode, + ) + } + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + }, + ReplayArgsSelection::Disagg(config) => { + match (replay_mode.as_str(), replay_concurrency) { + ("offline", Some(max_in_flight)) => dynamo_mocker::replay::simulate_concurrency_workload_disagg_with_router_mode( + *config, + router_config.clone(), + trace, + max_in_flight, + router_mode, + ), + ("offline", None) => dynamo_mocker::replay::simulate_trace_workload_disagg_with_router_mode( + *config, + router_config.clone(), + trace, + router_mode, + ), + ("online", _) => anyhow::bail!( + "disagg replay only supports replay_mode='offline'" + ), + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + } + } + }; + } + + let requests = build_synthetic_requests( + input_tokens, + output_tokens, + request_count, + arrival_interval_ms, + replay_concurrency.is_none(), + )?; + + match args_selection { + ReplayArgsSelection::Aggregated(args) => match (replay_mode.as_str(), replay_concurrency) + { ("offline", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_workload_with_router_mode( - args, + dynamo_mocker::replay::simulate_concurrency_requests_with_router_mode( + *args, router_config.clone(), - trace, + requests, max_in_flight, num_workers, router_mode, ) } - ("offline", None) => { - dynamo_mocker::replay::simulate_trace_workload_with_router_mode( - args, + ("offline", None) => dynamo_mocker::replay::simulate_trace_requests_with_router_mode( + *args, + router_config.clone(), + requests, + num_workers, + arrival_speedup_ratio, + router_mode, + ), + ("online", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_live_requests_with_router_mode( + *args, router_config.clone(), - trace, + requests, + max_in_flight, num_workers, router_mode, ) } - ("online", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_live_workload_with_router_mode( - args, + ("online", None) => { + dynamo_mocker::replay::simulate_trace_live_requests_with_router_mode( + *args, router_config.clone(), - trace, - max_in_flight, + requests, num_workers, + arrival_speedup_ratio, router_mode, ) } - ("online", None) => { - dynamo_mocker::replay::simulate_trace_live_workload_with_router_mode( - args, + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + }, + ReplayArgsSelection::Disagg(config) => match (replay_mode.as_str(), replay_concurrency) + { + ("offline", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_requests_disagg_with_router_mode( + *config, router_config.clone(), - trace, - num_workers, + requests, + max_in_flight, router_mode, ) } + ("offline", None) => { + dynamo_mocker::replay::simulate_trace_requests_disagg_with_router_mode( + *config, + router_config.clone(), + requests, + arrival_speedup_ratio, + router_mode, + ) + } + ("online", _) => anyhow::bail!("disagg replay only supports replay_mode='offline'"), (other, _) => anyhow::bail!( "replay_mode must be either 'offline' or 'online', got '{}'", other ), - }; - } - - let requests = build_synthetic_requests( - input_tokens, - output_tokens, - request_count, - arrival_interval_ms, - replay_concurrency.is_none(), - )?; - - match (replay_mode.as_str(), replay_concurrency) { - ("offline", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_requests_with_router_mode( - args, - router_config.clone(), - requests, - max_in_flight, - num_workers, - router_mode, - ) - } - ("offline", None) => dynamo_mocker::replay::simulate_trace_requests_with_router_mode( - args, - router_config.clone(), - requests, - num_workers, - arrival_speedup_ratio, - router_mode, - ), - ("online", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_live_requests_with_router_mode( - args, - router_config.clone(), - requests, - max_in_flight, - num_workers, - router_mode, - ) - } - ("online", None) => { - dynamo_mocker::replay::simulate_trace_live_requests_with_router_mode( - args, - router_config.clone(), - requests, - num_workers, - arrival_speedup_ratio, - router_mode, - ) - } - (other, _) => anyhow::bail!( - "replay_mode must be either 'offline' or 'online', got '{}'", - other - ), + }, } }); let report = report.map_err(to_pyerr)?; @@ -511,15 +629,68 @@ pub fn run_mocker_synthetic_trace_replay( .map(|obj| obj.unbind()) } -fn load_replay_mocker_args( +enum ReplayArgsSelection { + Aggregated(Box), + Disagg(Box), +} + +fn load_replay_args_selection( py: Python<'_>, extra_engine_args: Option, -) -> PyResult { - let mut args = match extra_engine_args { - Some(extra_args) => extra_args.inner(), - None => RsMockEngineArgs::default(), - }; + prefill_engine_args: Option, + decode_engine_args: Option, + num_workers: usize, + num_prefill_workers: usize, + num_decode_workers: usize, +) -> PyResult { + let aggregated_args = load_optional_replay_mocker_args(py, extra_engine_args)?; + let prefill_args = load_optional_replay_mocker_args(py, prefill_engine_args)?; + let decode_args = load_optional_replay_mocker_args(py, decode_engine_args)?; + + if aggregated_args.is_some() && (prefill_args.is_some() || decode_args.is_some()) { + return Err(PyException::new_err( + "extra_engine_args cannot be combined with prefill_engine_args/decode_engine_args", + )); + } + + match (aggregated_args, prefill_args, decode_args) { + (Some(args), None, None) => Ok(ReplayArgsSelection::Aggregated(Box::new(args))), + (None, None, None) => Ok(ReplayArgsSelection::Aggregated(Box::default())), + (None, Some(prefill_args), Some(decode_args)) => { + if num_workers != 1 { + return Err(PyException::new_err( + "num_workers is only used for aggregated replay; use num_prefill_workers and num_decode_workers for disagg replay", + )); + } + Ok(ReplayArgsSelection::Disagg(Box::new( + dynamo_mocker::replay::OfflineDisaggReplayConfig { + prefill_args, + decode_args, + num_prefill_workers, + num_decode_workers, + }, + ))) + } + (None, Some(_), None) | (None, None, Some(_)) => Err(PyException::new_err( + "prefill_engine_args and decode_engine_args must be provided together", + )), + _ => unreachable!(), + } +} +fn load_optional_replay_mocker_args( + py: Python<'_>, + extra_engine_args: Option, +) -> PyResult> { + extra_engine_args + .map(|extra_args| materialize_replay_mocker_args(py, extra_args.inner())) + .transpose() +} + +fn materialize_replay_mocker_args( + py: Python<'_>, + mut args: RsMockEngineArgs, +) -> PyResult { if let Some(ref backend_name) = args.aic_backend.clone() { let backend = backend_name.clone(); let system = args.aic_system.as_deref().unwrap_or("h200_sxm").to_string(); diff --git a/lib/bindings/python/src/dynamo/replay/api.py b/lib/bindings/python/src/dynamo/replay/api.py index 25d3a3207bb1..049a1afe6f8d 100644 --- a/lib/bindings/python/src/dynamo/replay/api.py +++ b/lib/bindings/python/src/dynamo/replay/api.py @@ -11,8 +11,12 @@ def run_trace_replay( trace_file, *, extra_engine_args=None, + prefill_engine_args=None, + decode_engine_args=None, router_config=None, num_workers=1, + num_prefill_workers=1, + num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", @@ -21,8 +25,12 @@ def run_trace_replay( return _run_mocker_trace_replay( trace_file, extra_engine_args=extra_engine_args, + prefill_engine_args=prefill_engine_args, + decode_engine_args=decode_engine_args, router_config=router_config, num_workers=num_workers, + num_prefill_workers=num_prefill_workers, + num_decode_workers=num_decode_workers, replay_concurrency=replay_concurrency, replay_mode=replay_mode, router_mode=router_mode, @@ -36,8 +44,12 @@ def run_synthetic_trace_replay( request_count, *, extra_engine_args=None, + prefill_engine_args=None, + decode_engine_args=None, router_config=None, num_workers=1, + num_prefill_workers=1, + num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", @@ -53,8 +65,12 @@ def run_synthetic_trace_replay( output_tokens, request_count, extra_engine_args=extra_engine_args, + prefill_engine_args=prefill_engine_args, + decode_engine_args=decode_engine_args, router_config=router_config, num_workers=num_workers, + num_prefill_workers=num_prefill_workers, + num_decode_workers=num_decode_workers, replay_concurrency=replay_concurrency, replay_mode=replay_mode, router_mode=router_mode, diff --git a/lib/bindings/python/src/dynamo/replay/main.py b/lib/bindings/python/src/dynamo/replay/main.py index 5cb254917e80..645f2ea70923 100644 --- a/lib/bindings/python/src/dynamo/replay/main.py +++ b/lib/bindings/python/src/dynamo/replay/main.py @@ -18,10 +18,42 @@ from dynamo.replay.reporting import format_report_table, write_report_json +def _load_engine_args(raw_args: str | None): + if raw_args is None: + return None + + raw = json.loads(raw_args) + worker_type = raw.pop("worker_type", None) + if worker_type is not None: + if "is_prefill" in raw or "is_decode" in raw: + raise ValueError( + "worker_type cannot be combined with is_prefill or is_decode" + ) + if worker_type == "prefill": + raw["is_prefill"] = True + elif worker_type == "decode": + raw["is_decode"] = True + elif worker_type != "aggregated": + raise ValueError( + "worker_type must be one of 'aggregated', 'prefill', or 'decode'" + ) + if "planner_profile_data" in raw: + profile_data_result = resolve_planner_profile_data( + Path(raw["planner_profile_data"]) + ) + if profile_data_result.npz_path is not None: + raw["planner_profile_data"] = str(profile_data_result.npz_path) + else: + del raw["planner_profile_data"] + return MockEngineArgs.from_json(json.dumps(raw)) + + def main(argv: Sequence[str] | None = None) -> int: parser = argparse.ArgumentParser(prog="python -m dynamo.replay") parser.add_argument("trace_file", nargs="?") parser.add_argument("--extra-engine-args") + parser.add_argument("--prefill-engine-args") + parser.add_argument("--decode-engine-args") parser.add_argument("--router-config") parser.add_argument("--input-tokens", type=int) parser.add_argument("--output-tokens", type=int) @@ -36,6 +68,8 @@ def main(argv: Sequence[str] | None = None) -> int: parser.add_argument("--num-prefix-groups", type=int, default=0) parser.add_argument("--inter-turn-delay-ms", type=float, default=0.0) parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--num-prefill-workers", type=int, default=1) + parser.add_argument("--num-decode-workers", type=int, default=1) parser.add_argument("--replay-concurrency", type=int) parser.add_argument( "--replay-mode", @@ -74,24 +108,9 @@ def main(argv: Sequence[str] | None = None) -> int: "synthetic replay requires --input-tokens, --output-tokens, and --request-count" ) - # Resolve planner_profile_data directory -> NPZ before passing to Rust. - # Rust only accepts NPZ files; resolve_planner_profile_data handles conversion. - profile_data_result = None - if args.extra_engine_args is not None: - raw = json.loads(args.extra_engine_args) - if "planner_profile_data" in raw: - profile_data_result = resolve_planner_profile_data( - Path(raw["planner_profile_data"]) - ) - if profile_data_result.npz_path is not None: - raw["planner_profile_data"] = str(profile_data_result.npz_path) - else: - del raw["planner_profile_data"] - extra_engine_args = MockEngineArgs.from_json(json.dumps(raw)) - else: - extra_engine_args = MockEngineArgs.from_json(args.extra_engine_args) - else: - extra_engine_args = None + extra_engine_args = _load_engine_args(args.extra_engine_args) + prefill_engine_args = _load_engine_args(args.prefill_engine_args) + decode_engine_args = _load_engine_args(args.decode_engine_args) router_config = ( KvRouterConfig.from_json(args.router_config) if args.router_config is not None @@ -102,8 +121,12 @@ def main(argv: Sequence[str] | None = None) -> int: report = run_trace_replay( args.trace_file, extra_engine_args=extra_engine_args, + prefill_engine_args=prefill_engine_args, + decode_engine_args=decode_engine_args, router_config=router_config, num_workers=args.num_workers, + num_prefill_workers=args.num_prefill_workers, + num_decode_workers=args.num_decode_workers, replay_concurrency=args.replay_concurrency, replay_mode=args.replay_mode, router_mode=args.router_mode, @@ -115,8 +138,12 @@ def main(argv: Sequence[str] | None = None) -> int: args.output_tokens, args.request_count, extra_engine_args=extra_engine_args, + prefill_engine_args=prefill_engine_args, + decode_engine_args=decode_engine_args, router_config=router_config, num_workers=args.num_workers, + num_prefill_workers=args.num_prefill_workers, + num_decode_workers=args.num_decode_workers, replay_concurrency=args.replay_concurrency, replay_mode=args.replay_mode, router_mode=args.router_mode, diff --git a/lib/bindings/python/tests/test_replay.py b/lib/bindings/python/tests/test_replay.py index 603952a5e4ed..cb491d63df4f 100644 --- a/lib/bindings/python/tests/test_replay.py +++ b/lib/bindings/python/tests/test_replay.py @@ -5,6 +5,7 @@ import os import subprocess import sys +from pathlib import Path import pytest @@ -195,6 +196,14 @@ def _sglang_args(): return MockEngineArgs.from_json(json.dumps(_sglang_args_payload())) +def _prefill_args(): + return MockEngineArgs(block_size=64, speedup_ratio=1000.0, worker_type="prefill") + + +def _decode_args(): + return MockEngineArgs(block_size=64, speedup_ratio=1000.0, worker_type="decode") + + def _write_router_config(tmp_path): config_path = tmp_path / "router_config.json" config_path.write_text( @@ -230,8 +239,12 @@ def _assert_basic_report_metrics(report): def _replay_cli_env() -> dict[str, str]: + repo_root = Path(__file__).resolve().parents[4] env = os.environ.copy() - pythonpath_entries = ["lib/bindings/python/src", "components/src"] + pythonpath_entries = [ + str(repo_root / "lib/bindings/python/src"), + str(repo_root / "components/src"), + ] existing_pythonpath = env.get("PYTHONPATH") if existing_pythonpath: pythonpath_entries.append(existing_pythonpath) @@ -554,6 +567,82 @@ def test_run_trace_replay_accepts_partial_extra_engine_args_json(tmp_path, repla ) +def test_run_trace_replay_supports_disagg_offline(tmp_path): + trace_path = _write_trace_and_args(tmp_path) + + report = run_trace_replay( + trace_path, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode="offline", + router_mode="kv_router", + ) + + _assert_basic_report_counts( + report, + num_requests=2, + input_tokens=64, + output_tokens=2, + ) + _assert_basic_report_metrics(report) + + +def test_run_synthetic_trace_replay_disagg_preserves_expected_output_tokens(): + report = run_synthetic_trace_replay( + 128, + 7, + 6, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode="offline", + router_mode="kv_router", + ) + + _assert_basic_report_counts( + report, + num_requests=6, + input_tokens=128, + output_tokens=7, + ) + _assert_basic_report_metrics(report) + + +def test_run_trace_replay_rejects_partial_disagg_args(tmp_path): + trace_path = _write_trace_and_args(tmp_path) + + with pytest.raises(Exception, match="must be provided together"): + run_trace_replay( + trace_path, + prefill_engine_args=_prefill_args(), + replay_mode="offline", + router_mode="kv_router", + ) + + +def test_run_trace_replay_rejects_online_disagg(tmp_path): + trace_path = _write_trace_and_args(tmp_path) + + with pytest.raises( + Exception, match="disagg replay only supports replay_mode='offline'" + ): + run_trace_replay( + trace_path, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode="online", + router_mode="kv_router", + ) + + def test_format_report_table_matches_aiperf_shape(): report = { "mean_ttft_ms": 18.26, @@ -799,6 +888,40 @@ def test_replay_cli_subprocess_trace_smoke(tmp_path): _assert_basic_report_metrics(report) +@pytest.mark.timeout(30) +def test_replay_cli_subprocess_trace_disagg_smoke(tmp_path): + trace_path = _write_cli_smoke_trace(tmp_path) + report_path = tmp_path / "trace_disagg_report.json" + + completed = _run_replay_cli( + tmp_path, + str(trace_path), + "--replay-mode", + "offline", + "--router-mode", + "kv_router", + "--num-prefill-workers", + "2", + "--num-decode-workers", + "2", + "--report-json", + str(report_path), + "--prefill-engine-args", + '{"block_size":64,"speedup_ratio":1000.0,"worker_type":"prefill"}', + "--decode-engine-args", + '{"block_size":64,"speedup_ratio":1000.0,"worker_type":"decode"}', + ) + + report = _assert_replay_cli_outputs(completed, report_path) + _assert_basic_report_counts( + report, + num_requests=10, + input_tokens=250, + output_tokens=25, + ) + _assert_basic_report_metrics(report) + + @pytest.mark.timeout(30) def test_replay_cli_subprocess_multiturn_trace_smoke(tmp_path): trace_path = _write_multiturn_trace(tmp_path) diff --git a/lib/mocker/src/replay/entrypoints.rs b/lib/mocker/src/replay/entrypoints.rs index a25befff3349..0032fc2b8794 100644 --- a/lib/mocker/src/replay/entrypoints.rs +++ b/lib/mocker/src/replay/entrypoints.rs @@ -9,10 +9,11 @@ use dynamo_kv_router::config::KvRouterConfig; use super::online; use super::validate::{ - validate_offline_concurrency_args, validate_offline_replay_args, + validate_offline_concurrency_args, validate_offline_disagg_concurrency_args, + validate_offline_disagg_replay_args, validate_offline_replay_args, validate_online_concurrency_args, validate_online_replay_args, }; -use super::{ReplayRouterMode, TraceSimulationReport}; +use super::{OfflineDisaggReplayConfig, ReplayRouterMode, TraceSimulationReport}; use crate::common::protocols::{DirectRequest, MockEngineArgs}; use crate::loadgen::Trace; @@ -56,6 +57,24 @@ pub fn simulate_trace_file_with_router_mode( Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } +pub fn simulate_trace_file_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace_path: &Path, + arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_replay_args(&config, router_mode)?; + let trace = Trace::from_mooncake(trace_path, config.prefill_args.block_size)? + .normalize_session_starts()? + .speed_up_timing(arrival_speedup_ratio)?; + let started_at = Instant::now(); + let report = + crate::replay::offline::simulate_trace_workload_disagg(config, router_config, trace)?; + Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) +} + pub fn simulate_trace_live_file( args: MockEngineArgs, trace_path: &Path, @@ -130,6 +149,29 @@ pub fn simulate_trace_requests_with_router_mode( Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } +pub fn simulate_trace_requests_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_replay_args(&config, router_mode)?; + if requests.is_empty() { + bail!("trace replay requires at least one request"); + } + + let started_at = Instant::now(); + let report = crate::replay::offline::simulate_trace_disagg( + config, + router_config, + requests, + arrival_speedup_ratio, + )?; + Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) +} + pub fn simulate_trace_live_requests( args: MockEngineArgs, requests: Vec, @@ -209,6 +251,27 @@ pub fn simulate_concurrency_file_with_router_mode( Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } +pub fn simulate_concurrency_file_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace_path: &Path, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_concurrency_args(&config, max_in_flight, router_mode)?; + let trace = Trace::from_mooncake(trace_path, config.prefill_args.block_size)?; + let started_at = Instant::now(); + let report = simulate_concurrency_workload_disagg_with_router_mode( + config, + router_config, + trace, + max_in_flight, + router_mode, + )?; + Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) +} + pub fn simulate_concurrency_live_file( args: MockEngineArgs, trace_path: &Path, @@ -326,6 +389,27 @@ pub fn simulate_concurrency_requests_with_router_mode( ) } +pub fn simulate_concurrency_requests_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_concurrency_args(&config, max_in_flight, router_mode)?; + if requests.is_empty() { + bail!("concurrency replay requires at least one request"); + } + + crate::replay::offline::simulate_concurrency_disagg( + config, + router_config, + requests, + max_in_flight, + ) +} + pub fn simulate_trace_workload( args: MockEngineArgs, trace: Trace, @@ -360,6 +444,20 @@ pub fn simulate_trace_workload_with_router_mode( Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } +pub fn simulate_trace_workload_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_replay_args(&config, router_mode)?; + let started_at = Instant::now(); + let report = + crate::replay::offline::simulate_trace_workload_disagg(config, router_config, trace)?; + Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) +} + pub fn simulate_trace_live_workload( args: MockEngineArgs, trace: Trace, @@ -422,6 +520,23 @@ pub fn simulate_concurrency_workload_with_router_mode( ) } +pub fn simulate_concurrency_workload_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_concurrency_args(&config, max_in_flight, router_mode)?; + crate::replay::offline::simulate_concurrency_workload_disagg( + config, + router_config, + trace, + max_in_flight, + ) +} + pub fn simulate_concurrency_live_workload( args: MockEngineArgs, trace: Trace, diff --git a/lib/mocker/src/replay/mod.rs b/lib/mocker/src/replay/mod.rs index 30d4fef6dedf..b64a7ad3aecd 100644 --- a/lib/mocker/src/replay/mod.rs +++ b/lib/mocker/src/replay/mod.rs @@ -10,7 +10,7 @@ mod validate; use std::collections::VecDeque; -use crate::common::protocols::DirectRequest; +use crate::common::protocols::{DirectRequest, MockEngineArgs}; pub(crate) use collector::TraceCollector; #[cfg(test)] @@ -25,19 +25,42 @@ pub enum ReplayRouterMode { KvRouter, } +#[derive(Clone, Debug)] +pub struct OfflineDisaggReplayConfig { + pub prefill_args: MockEngineArgs, + pub decode_args: MockEngineArgs, + pub num_prefill_workers: usize, + pub num_decode_workers: usize, +} + +impl OfflineDisaggReplayConfig { + pub fn normalized(self) -> anyhow::Result { + Ok(Self { + prefill_args: self.prefill_args.normalized()?, + decode_args: self.decode_args.normalized()?, + num_prefill_workers: self.num_prefill_workers, + num_decode_workers: self.num_decode_workers, + }) + } +} + pub use entrypoints::{ - simulate_concurrency_file, simulate_concurrency_file_with_router_mode, - simulate_concurrency_live_file, simulate_concurrency_live_file_with_router_mode, - simulate_concurrency_live_requests, simulate_concurrency_live_requests_with_router_mode, - simulate_concurrency_live_workload, simulate_concurrency_live_workload_with_router_mode, - simulate_concurrency_requests, simulate_concurrency_requests_with_router_mode, - simulate_concurrency_workload, simulate_concurrency_workload_with_router_mode, - simulate_trace_file, simulate_trace_file_with_router_mode, simulate_trace_live_file, - simulate_trace_live_file_with_router_mode, simulate_trace_live_requests, - simulate_trace_live_requests_with_router_mode, simulate_trace_live_workload, - simulate_trace_live_workload_with_router_mode, simulate_trace_requests, + simulate_concurrency_file, simulate_concurrency_file_disagg_with_router_mode, + simulate_concurrency_file_with_router_mode, simulate_concurrency_live_file, + simulate_concurrency_live_file_with_router_mode, simulate_concurrency_live_requests, + simulate_concurrency_live_requests_with_router_mode, simulate_concurrency_live_workload, + simulate_concurrency_live_workload_with_router_mode, simulate_concurrency_requests, + simulate_concurrency_requests_disagg_with_router_mode, + simulate_concurrency_requests_with_router_mode, simulate_concurrency_workload, + simulate_concurrency_workload_disagg_with_router_mode, + simulate_concurrency_workload_with_router_mode, simulate_trace_file, + simulate_trace_file_disagg_with_router_mode, simulate_trace_file_with_router_mode, + simulate_trace_live_file, simulate_trace_live_file_with_router_mode, + simulate_trace_live_requests, simulate_trace_live_requests_with_router_mode, + simulate_trace_live_workload, simulate_trace_live_workload_with_router_mode, + simulate_trace_requests, simulate_trace_requests_disagg_with_router_mode, simulate_trace_requests_with_router_mode, simulate_trace_workload, - simulate_trace_workload_with_router_mode, + simulate_trace_workload_disagg_with_router_mode, simulate_trace_workload_with_router_mode, }; pub(crate) fn normalize_trace_requests( diff --git a/lib/mocker/src/replay/offline/README.md b/lib/mocker/src/replay/offline/README.md index 2a9c2a73090c..d446c491b452 100644 --- a/lib/mocker/src/replay/offline/README.md +++ b/lib/mocker/src/replay/offline/README.md @@ -15,10 +15,11 @@ The public replay entrypoints live one level up in `lib/mocker/src/replay/entryp Offline replay starts in `lib/mocker/src/replay/offline/mod.rs`. -`offline/mod.rs` chooses between two implementations: +`offline/mod.rs` chooses between three implementations: - `lib/mocker/src/replay/offline/single.rs` for the special case `num_workers == 1` with the vLLM engine - `lib/mocker/src/replay/offline/multi.rs` for everything else, including multi-worker replay and `kv_router` replay +- `lib/mocker/src/replay/offline/disagg.rs` for offline disaggregated prefill/decode replay ## File Map @@ -28,6 +29,8 @@ Offline replay starts in `lib/mocker/src/replay/offline/mod.rs`. Minimal replay loop for one vLLM worker. - `lib/mocker/src/replay/offline/multi.rs` General offline cluster simulator for multi-worker replay and KV-router replay. +- `lib/mocker/src/replay/offline/disagg.rs` + Offline two-stage replay harness with separate prefill and decode pools. - `lib/mocker/src/replay/offline/state.rs` Per-worker wrapper around `EngineCore`, including optional KV event capture. - `lib/mocker/src/replay/offline/events.rs` @@ -114,7 +117,7 @@ So offline replay is not a toy simulator. It reuses the real per-pass mocker sch ## Completion Event Queue -The multi-worker harness uses `SimulationEvent` from `lib/mocker/src/replay/offline/events.rs` as a min-time priority queue implemented with `BinaryHeap`. +The multi-worker and disagg harnesses use `SimulationEvent` from `lib/mocker/src/replay/offline/events.rs` as a min-time priority queue implemented with `BinaryHeap`. Right now the only scheduled event type is: @@ -122,6 +125,7 @@ Right now the only scheduled event type is: That event carries: +- worker `stage` (`aggregated`, `prefill`, or `decode`) - `worker_idx` - `completed_requests` - `output_signals` @@ -167,7 +171,7 @@ flowchart LR M --> G ``` -### Why KV events are captured only here +### Why KV events are captured only where needed When offline replay uses `kv_router`, workers are created with KV event capture enabled via: @@ -177,6 +181,37 @@ When offline replay uses `kv_router`, workers are created with KV event capture That causes each pass to return router-visible `kv_events`, which the harness applies synchronously to the offline router indexer after the pass completes. In round-robin mode, this capture is skipped because nothing consumes those events. +In offline disagg replay, only the prefill workers capture and publish KV events; the decode workers +run with capture disabled because the decode router is overlap-blind and does not consume router +events. + +## Disaggregated Harness + +The disaggregated runtime in `lib/mocker/src/replay/offline/disagg.rs` models two distinct stages: + +- a prefill router and prefill worker pool +- a decode router and decode worker pool + +It keeps one logical clock and one completion-event heap, but request ownership moves through a +two-stage state machine instead of the aggregated single-pool lifecycle. + +The prefill router is derived from the main router config with `router_track_active_blocks = false`. +The decode router is derived with: + +- overlap disabled +- `assume_kv_reuse = false` +- `track_prefill_tokens = false` + +The prefill stage runs a hidden synthetic one-token bootstrap request. When prefill completes, the +harness: + +1. applies any prefill KV events +2. marks prefill complete in the prefill router +3. frees prefill router state +4. enqueues the original request into decode at the same logical timestamp + +Decode then runs with normal collector visibility. The public replay report remains decode-only, so +TTFT includes prefill queueing and prefill compute. ## Trace vs Concurrency Modes diff --git a/lib/mocker/src/replay/offline/disagg.rs b/lib/mocker/src/replay/offline/disagg.rs new file mode 100644 index 000000000000..e6e955abdba7 --- /dev/null +++ b/lib/mocker/src/replay/offline/disagg.rs @@ -0,0 +1,984 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{BinaryHeap, HashMap, VecDeque}; + +use anyhow::{Result, anyhow, bail}; +use dynamo_kv_router::config::KvRouterConfig; +use dynamo_kv_router::protocols::RouterEvent; +use uuid::Uuid; + +use super::events::{SimulationEvent, SimulationWorkerStage}; +use super::normalize_trace_requests; +use super::runtime_utils::{ + WorkerCompletionPayload, next_timestamp as choose_next_timestamp, pop_next_concurrency_ready, + pop_next_trace_ready, pop_ready_worker_completion, push_worker_completion, +}; +use super::state::OfflineWorkerState; +use crate::common::protocols::{DirectRequest, MockEngineArgs, OutputSignal}; +use crate::loadgen::{ReplayRequestHashes, Trace, WorkloadDriver}; +use crate::replay::router::OfflineReplayRouter; +use crate::replay::{OfflineDisaggReplayConfig, TraceCollector, TraceSimulationReport}; +use crate::scheduler::RouterEventVisibility; + +#[derive(Debug, Clone, Copy)] +enum ReplayMode { + Trace, + Concurrency { max_in_flight: usize }, +} + +enum AdmissionSource { + Requests(VecDeque), + Workload(WorkloadDriver), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum DisaggPhase { + QueuedPrefill, + RunningPrefill, + QueuedDecode, + RunningDecode, + Done, +} + +struct DisaggRequestState { + original: DirectRequest, + #[cfg(test)] + arrival_ms: f64, + phase: DisaggPhase, + prefill_worker_idx: Option, + decode_worker_idx: Option, +} + +#[cfg(test)] +#[derive(Debug, Clone, PartialEq)] +struct DisaggRequestSnapshot { + arrival_ms: f64, + phase: DisaggPhase, + prefill_worker_idx: Option, + decode_worker_idx: Option, +} + +#[cfg(test)] +#[derive(Debug, Default, Clone, PartialEq)] +struct DisaggRuntimeStats { + request_snapshots: HashMap, + prefill_assignments: HashMap, + decode_assignments: HashMap, + handoff_ms: HashMap, + prefill_marked_count: usize, + prefill_freed_count: usize, + decode_freed_count: usize, + max_prefill_router_pending: usize, + max_decode_router_pending: usize, +} + +#[cfg(not(test))] +#[derive(Debug, Default, Clone, PartialEq, Eq)] +struct DisaggRuntimeStats; + +struct DisaggRuntime { + now_ms: f64, + next_event_seq: u64, + admission: AdmissionSource, + prefill_workers: Vec, + decode_workers: Vec, + prefill_router: OfflineReplayRouter, + decode_router: OfflineReplayRouter, + requests: HashMap, + collector: TraceCollector, + events: BinaryHeap, + mode: ReplayMode, + stats: DisaggRuntimeStats, +} + +impl DisaggRuntime { + fn new( + config: &OfflineDisaggReplayConfig, + router_config: Option, + pending: VecDeque, + mode: ReplayMode, + ) -> Result { + Self::new_with_source( + config, + router_config, + AdmissionSource::Requests(pending), + mode, + ) + } + + fn new_workload( + config: &OfflineDisaggReplayConfig, + router_config: Option, + driver: WorkloadDriver, + mode: ReplayMode, + ) -> Result { + Self::new_with_source( + config, + router_config, + AdmissionSource::Workload(driver), + mode, + ) + } + + fn new_with_source( + config: &OfflineDisaggReplayConfig, + router_config: Option, + admission: AdmissionSource, + mode: ReplayMode, + ) -> Result { + let prefill_router_config = + derive_prefill_router_config(&config.prefill_args, router_config.clone()); + let decode_router_config = derive_decode_router_config(&config.decode_args, router_config); + + Ok(Self { + now_ms: 0.0, + next_event_seq: 0, + admission, + prefill_workers: (0..config.num_prefill_workers) + .map(|worker_idx| { + OfflineWorkerState::new(worker_idx, config.prefill_args.clone(), true) + }) + .collect(), + decode_workers: (0..config.num_decode_workers) + .map(|worker_idx| { + OfflineWorkerState::new(worker_idx, config.decode_args.clone(), false) + }) + .collect(), + prefill_router: OfflineReplayRouter::new( + &config.prefill_args, + Some(prefill_router_config), + config.num_prefill_workers, + )?, + decode_router: OfflineReplayRouter::new( + &config.decode_args, + Some(decode_router_config), + config.num_decode_workers, + )?, + requests: HashMap::new(), + collector: TraceCollector::default(), + events: BinaryHeap::new(), + mode, + #[cfg(test)] + stats: DisaggRuntimeStats::default(), + #[cfg(not(test))] + stats: DisaggRuntimeStats, + }) + } + + fn cluster_in_flight(&self) -> usize { + self.prefill_workers + .iter() + .map(OfflineWorkerState::in_flight) + .sum::() + + self + .decode_workers + .iter() + .map(OfflineWorkerState::in_flight) + .sum::() + + self.prefill_router.pending_count() + + self.decode_router.pending_count() + } + + fn record_router_pending(&mut self) { + #[cfg(test)] + { + self.stats.max_prefill_router_pending = self + .stats + .max_prefill_router_pending + .max(self.prefill_router.pending_count()); + self.stats.max_decode_router_pending = self + .stats + .max_decode_router_pending + .max(self.decode_router.pending_count()); + } + } + + fn validate_worker_idx(&self, stage: SimulationWorkerStage, worker_idx: usize) -> Result<()> { + let worker_count = match stage { + SimulationWorkerStage::Prefill => self.prefill_workers.len(), + SimulationWorkerStage::Decode => self.decode_workers.len(), + SimulationWorkerStage::Aggregated => unreachable!("aggregated stage is not used"), + }; + if worker_idx >= worker_count { + bail!("offline disagg replay selected unknown {stage:?} worker index {worker_idx}"); + } + Ok(()) + } + + fn state(&self, uuid: Uuid) -> Result<&DisaggRequestState> { + self.requests + .get(&uuid) + .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}")) + } + + fn state_mut(&mut self, uuid: Uuid) -> Result<&mut DisaggRequestState> { + self.requests + .get_mut(&uuid) + .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}")) + } + + fn prefill_request(state: &DisaggRequestState) -> DirectRequest { + let mut request = state.original.clone(); + request.max_output_tokens = 1; + request + } + + fn decode_request(state: &DisaggRequestState) -> DirectRequest { + state.original.clone() + } + + fn dispatch_prefill(&mut self, uuid: Uuid, worker_idx: usize) -> Result<()> { + self.validate_worker_idx(SimulationWorkerStage::Prefill, worker_idx)?; + let request = Self::prefill_request(self.state(uuid)?); + self.prefill_workers[worker_idx].receive_request(request); + let state = self.state_mut(uuid)?; + state.phase = DisaggPhase::RunningPrefill; + state.prefill_worker_idx = Some(worker_idx); + #[cfg(test)] + { + self.stats.prefill_assignments.insert(uuid, worker_idx); + } + Ok(()) + } + + fn dispatch_decode(&mut self, uuid: Uuid, worker_idx: usize) -> Result<()> { + self.validate_worker_idx(SimulationWorkerStage::Decode, worker_idx)?; + let request = Self::decode_request(self.state(uuid)?); + self.decode_workers[worker_idx].receive_request(request); + let state = self.state_mut(uuid)?; + state.phase = DisaggPhase::RunningDecode; + state.decode_worker_idx = Some(worker_idx); + #[cfg(test)] + { + self.stats.decode_assignments.insert(uuid, worker_idx); + } + Ok(()) + } + + fn dispatch_prefill_admissions(&mut self, admissions: Vec<(Uuid, usize)>) -> Result<()> { + for (uuid, worker_idx) in admissions { + if self.state(uuid)?.phase != DisaggPhase::QueuedPrefill { + bail!("offline disagg replay expected queued prefill request for {uuid}"); + } + self.dispatch_prefill(uuid, worker_idx)?; + } + Ok(()) + } + + fn dispatch_decode_admissions(&mut self, admissions: Vec<(Uuid, usize)>) -> Result<()> { + for (uuid, worker_idx) in admissions { + if self.state(uuid)?.phase != DisaggPhase::QueuedDecode { + bail!("offline disagg replay expected queued decode request for {uuid}"); + } + self.dispatch_decode(uuid, worker_idx)?; + } + Ok(()) + } + + fn enqueue_decode(&mut self, uuid: Uuid) -> Result<()> { + let request = Self::decode_request(self.state(uuid)?); + let maybe_worker_idx = + self.decode_router + .submit_request_with_hashes(&request, None, self.now_ms)?; + self.record_router_pending(); + #[cfg(test)] + { + self.stats.handoff_ms.insert(uuid, self.now_ms); + } + if let Some(worker_idx) = maybe_worker_idx { + self.dispatch_decode(uuid, worker_idx)?; + return Ok(()); + } + + self.state_mut(uuid)?.phase = DisaggPhase::QueuedDecode; + Ok(()) + } + + fn on_external_arrival( + &mut self, + mut request: DirectRequest, + arrival_time_ms: f64, + replay_hashes: Option, + ) -> Result { + let uuid = request.uuid.unwrap_or_else(Uuid::new_v4); + request.uuid = Some(uuid); + request.arrival_timestamp_ms = Some(arrival_time_ms); + + self.collector.on_arrival( + uuid, + arrival_time_ms, + request.tokens.len(), + request.max_output_tokens, + ); + + let maybe_worker_idx = + self.prefill_router + .submit_request_with_hashes(&request, replay_hashes, self.now_ms)?; + let state = DisaggRequestState { + original: request, + #[cfg(test)] + arrival_ms: arrival_time_ms, + phase: DisaggPhase::QueuedPrefill, + prefill_worker_idx: None, + decode_worker_idx: None, + }; + self.requests.insert(uuid, state); + self.record_router_pending(); + if let Some(worker_idx) = maybe_worker_idx { + self.dispatch_prefill(uuid, worker_idx)?; + } + Ok(uuid) + } + + fn is_done(&self) -> bool { + self.events.is_empty() + && self.cluster_in_flight() == 0 + && match &self.admission { + AdmissionSource::Requests(pending) => pending.is_empty(), + AdmissionSource::Workload(driver) => driver.is_drained(), + } + && self + .prefill_workers + .iter() + .all(OfflineWorkerState::is_drained) + && self + .decode_workers + .iter() + .all(OfflineWorkerState::is_drained) + } + + fn next_timestamp(&mut self) -> Option { + let next_event_ms = self.events.peek().map(|event| event.at_ms); + let cluster_in_flight = self.cluster_in_flight(); + let next_arrival_ms = match (&self.mode, &mut self.admission) { + (ReplayMode::Trace, AdmissionSource::Requests(pending)) => pending + .front() + .and_then(|request| request.arrival_timestamp_ms), + (ReplayMode::Trace, AdmissionSource::Workload(driver)) => driver.next_ready_time_ms(), + (ReplayMode::Concurrency { max_in_flight }, AdmissionSource::Workload(driver)) => { + if cluster_in_flight < *max_in_flight { + driver.next_ready_time_ms() + } else { + None + } + } + (ReplayMode::Concurrency { .. }, AdmissionSource::Requests(_)) => None, + }; + + choose_next_timestamp(next_arrival_ms, next_event_ms) + } + + fn apply_prefill_router_events(&mut self, events: Vec) -> Result<()> { + for event in events { + self.prefill_router.apply_event(event)?; + } + Ok(()) + } + + fn process_prefill_signal(&mut self, signal: OutputSignal) -> Result<()> { + if !signal.completed { + return Ok(()); + } + + let prefill_complete_admissions = + self.prefill_router.mark_prefill_completed(signal.uuid)?; + #[cfg(test)] + { + self.stats.prefill_marked_count += 1; + } + self.record_router_pending(); + self.dispatch_prefill_admissions(prefill_complete_admissions)?; + + let admissions = self.prefill_router.free(signal.uuid)?; + #[cfg(test)] + { + self.stats.prefill_freed_count += 1; + } + self.record_router_pending(); + self.dispatch_prefill_admissions(admissions)?; + self.enqueue_decode(signal.uuid)?; + Ok(()) + } + + fn process_decode_signal(&mut self, signal: OutputSignal) -> Result<()> { + if !signal.completed { + return Ok(()); + } + + let admissions = self.decode_router.free(signal.uuid)?; + #[cfg(test)] + { + self.stats.decode_freed_count += 1; + } + self.record_router_pending(); + let state = self.state_mut(signal.uuid)?; + state.phase = DisaggPhase::Done; + if let AdmissionSource::Workload(driver) = &mut self.admission { + driver.on_complete(signal.uuid, self.now_ms)?; + } + self.dispatch_decode_admissions(admissions)?; + Ok(()) + } + + fn process_prefill_pass( + &mut self, + worker_idx: usize, + completed_requests: usize, + output_signals: Vec, + kv_events: Vec, + ) -> Result<()> { + self.prefill_workers[worker_idx].mark_completed(completed_requests); + self.apply_prefill_router_events(kv_events)?; + for signal in output_signals { + self.process_prefill_signal(signal)?; + } + Ok(()) + } + + fn process_decode_pass( + &mut self, + worker_idx: usize, + completed_requests: usize, + output_signals: Vec, + ) -> Result<()> { + self.decode_workers[worker_idx].mark_completed(completed_requests); + for signal in output_signals { + self.process_decode_signal(signal)?; + } + Ok(()) + } + + fn apply_worker_completions(&mut self) -> Result { + let mut changed = false; + while let Some(WorkerCompletionPayload { + stage, + worker_idx, + completed_requests, + output_signals, + kv_events, + }) = pop_ready_worker_completion(&mut self.events, self.now_ms) + { + match stage { + SimulationWorkerStage::Prefill => { + self.prefill_workers[worker_idx].mark_idle(); + self.process_prefill_pass( + worker_idx, + completed_requests, + output_signals, + kv_events, + )?; + } + SimulationWorkerStage::Decode => { + self.decode_workers[worker_idx].mark_idle(); + self.process_decode_pass(worker_idx, completed_requests, output_signals)?; + } + SimulationWorkerStage::Aggregated => { + bail!("offline disagg replay received an aggregated completion event") + } + } + changed = true; + } + Ok(changed) + } + + fn release_trace_arrivals(&mut self) -> Result { + let mut released_any = false; + if matches!(self.admission, AdmissionSource::Requests(_)) { + loop { + let next_ready = match &mut self.admission { + AdmissionSource::Requests(pending) => { + pop_next_trace_ready(pending, self.now_ms) + } + AdmissionSource::Workload(_) => unreachable!(), + }; + let Some((request, arrival_ms)) = next_ready else { + break; + }; + self.on_external_arrival(request, arrival_ms, None)?; + released_any = true; + } + return Ok(released_any); + } + + let ready_requests = match &mut self.admission { + AdmissionSource::Requests(_) => unreachable!(), + AdmissionSource::Workload(driver) => driver.pop_ready(self.now_ms, usize::MAX), + }; + for ready in ready_requests { + self.on_external_arrival( + ready.request, + ready.scheduled_ready_at_ms, + ready.replay_hashes, + )?; + released_any = true; + } + Ok(released_any) + } + + fn top_off_concurrency(&mut self, max_in_flight: usize) -> Result { + let mut released_any = false; + if matches!(self.admission, AdmissionSource::Requests(_)) { + loop { + let cluster_in_flight = self.cluster_in_flight(); + let next_ready = match &mut self.admission { + AdmissionSource::Requests(pending) => pop_next_concurrency_ready( + pending, + self.now_ms, + cluster_in_flight, + max_in_flight, + ), + AdmissionSource::Workload(_) => unreachable!(), + }; + let Some((request, arrival_ms)) = next_ready else { + break; + }; + self.on_external_arrival(request, arrival_ms, None)?; + released_any = true; + } + return Ok(released_any); + } + + let available = max_in_flight.saturating_sub(self.cluster_in_flight()); + if available == 0 { + return Ok(false); + } + + let ready_requests = match &mut self.admission { + AdmissionSource::Requests(_) => unreachable!(), + AdmissionSource::Workload(driver) => driver.pop_ready(self.now_ms, available), + }; + for ready in ready_requests { + self.on_external_arrival(ready.request, self.now_ms, ready.replay_hashes)?; + released_any = true; + } + Ok(released_any) + } + + fn drive_prefill_workers(&mut self) -> Result { + let mut changed = false; + for worker_idx in 0..self.prefill_workers.len() { + loop { + if !self.prefill_workers[worker_idx].is_ready() { + break; + } + + let executed = self.prefill_workers[worker_idx].execute_hidden_pass(self.now_ms); + changed = true; + + let completion_kv_events = + if executed.router_event_visibility == RouterEventVisibility::PassStart { + self.apply_prefill_router_events(executed.kv_events)?; + Vec::new() + } else { + executed.kv_events + }; + + if executed.end_ms == self.now_ms { + self.process_prefill_pass( + worker_idx, + executed.completed_requests, + executed.output_signals, + completion_kv_events, + )?; + continue; + } + + self.prefill_workers[worker_idx].mark_busy(); + push_worker_completion( + &mut self.events, + &mut self.next_event_seq, + executed.end_ms, + WorkerCompletionPayload { + stage: SimulationWorkerStage::Prefill, + worker_idx, + completed_requests: executed.completed_requests, + output_signals: executed.output_signals, + kv_events: completion_kv_events, + }, + ); + break; + } + } + Ok(changed) + } + + fn drive_decode_workers(&mut self) -> Result { + let mut changed = false; + for worker_idx in 0..self.decode_workers.len() { + loop { + if !self.decode_workers[worker_idx].is_ready() { + break; + } + + let executed = { + let (workers, collector) = (&mut self.decode_workers, &mut self.collector); + workers[worker_idx].execute_pass(collector, self.now_ms) + }; + changed = true; + + if executed.end_ms == self.now_ms { + self.process_decode_pass( + worker_idx, + executed.completed_requests, + executed.output_signals, + )?; + continue; + } + + self.decode_workers[worker_idx].mark_busy(); + push_worker_completion( + &mut self.events, + &mut self.next_event_seq, + executed.end_ms, + WorkerCompletionPayload { + stage: SimulationWorkerStage::Decode, + worker_idx, + completed_requests: executed.completed_requests, + output_signals: executed.output_signals, + kv_events: Vec::new(), + }, + ); + break; + } + } + Ok(changed) + } + + fn drain_current_timestamp(&mut self) -> Result<()> { + loop { + let mut changed = self.apply_worker_completions()?; + + changed |= match self.mode { + ReplayMode::Trace => self.release_trace_arrivals()?, + ReplayMode::Concurrency { max_in_flight } => { + self.top_off_concurrency(max_in_flight)? + } + }; + + changed |= self.drive_prefill_workers()?; + changed |= self.drive_decode_workers()?; + + if !changed { + break; + } + } + Ok(()) + } + + fn finish_test_stats(&mut self) { + #[cfg(test)] + { + self.stats.request_snapshots = self + .requests + .iter() + .map(|(uuid, state)| { + ( + *uuid, + DisaggRequestSnapshot { + arrival_ms: state.arrival_ms, + phase: state.phase, + prefill_worker_idx: state.prefill_worker_idx, + decode_worker_idx: state.decode_worker_idx, + }, + ) + }) + .collect(); + } + } + + fn run(mut self) -> Result<(TraceCollector, DisaggRuntimeStats)> { + self.drain_current_timestamp()?; + + while !self.is_done() { + let Some(next_timestamp_ms) = self.next_timestamp() else { + bail!( + "offline disagg replay reached a dead end with {} in-flight requests remaining", + self.cluster_in_flight() + ); + }; + self.now_ms = next_timestamp_ms; + self.drain_current_timestamp()?; + } + + self.finish_test_stats(); + Ok((self.collector, self.stats)) + } +} + +fn base_router_config( + args: &MockEngineArgs, + router_config: Option, +) -> KvRouterConfig { + let mut config = router_config.unwrap_or_default(); + if let Some(policy) = args.router_queue_policy { + config.router_queue_policy = policy; + } + config +} + +fn derive_prefill_router_config( + args: &MockEngineArgs, + router_config: Option, +) -> KvRouterConfig { + let mut config = base_router_config(args, router_config); + config.router_track_active_blocks = false; + config +} + +fn derive_decode_router_config( + args: &MockEngineArgs, + router_config: Option, +) -> KvRouterConfig { + let mut config = base_router_config(args, router_config); + config.overlap_score_weight = 0.0; + config.router_assume_kv_reuse = false; + config.router_track_prefill_tokens = false; + config +} + +pub(crate) fn simulate_trace_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + arrival_speedup_ratio: f64, +) -> Result { + let pending = normalize_trace_requests(requests, arrival_speedup_ratio)?; + let (collector, _) = + DisaggRuntime::new(&config, router_config, pending, ReplayMode::Trace)?.run()?; + Ok(collector.finish()) +} + +pub(crate) fn simulate_concurrency_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + max_in_flight: usize, +) -> Result { + let pending = VecDeque::from(requests); + let (collector, _) = DisaggRuntime::new( + &config, + router_config, + pending, + ReplayMode::Concurrency { max_in_flight }, + )? + .run()?; + Ok(collector.finish()) +} + +pub(crate) fn simulate_trace_workload_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, +) -> Result { + let driver = WorkloadDriver::new_trace(trace)?; + let (collector, _) = + DisaggRuntime::new_workload(&config, router_config, driver, ReplayMode::Trace)?.run()?; + Ok(collector.finish()) +} + +pub(crate) fn simulate_concurrency_workload_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + max_in_flight: usize, +) -> Result { + let driver = WorkloadDriver::new_concurrency(trace)?; + let (collector, _) = DisaggRuntime::new_workload( + &config, + router_config, + driver, + ReplayMode::Concurrency { max_in_flight }, + )? + .run()?; + Ok(collector.finish()) +} + +#[cfg(test)] +fn run_trace_collect( + config: &OfflineDisaggReplayConfig, + requests: Vec, + router_config: Option, + arrival_speedup_ratio: f64, +) -> (TraceCollector, DisaggRuntimeStats) { + let pending = normalize_trace_requests(requests, arrival_speedup_ratio).unwrap(); + DisaggRuntime::new(config, router_config, pending, ReplayMode::Trace) + .unwrap() + .run() + .unwrap() +} + +#[cfg(test)] +fn run_concurrency_collect( + config: &OfflineDisaggReplayConfig, + requests: Vec, + router_config: Option, + max_in_flight: usize, +) -> (TraceCollector, DisaggRuntimeStats) { + DisaggRuntime::new( + config, + router_config, + VecDeque::from(requests), + ReplayMode::Concurrency { max_in_flight }, + ) + .unwrap() + .run() + .unwrap() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::protocols::{MockEngineArgs, WorkerType}; + + fn staged_args(worker_type: WorkerType, speedup_ratio: f64) -> MockEngineArgs { + MockEngineArgs::builder() + .block_size(64) + .num_gpu_blocks(256) + .max_num_batched_tokens(Some(8192)) + .max_num_seqs(Some(8)) + .enable_prefix_caching(true) + .enable_chunked_prefill(true) + .speedup_ratio(speedup_ratio) + .decode_speedup_ratio(speedup_ratio) + .worker_type(worker_type) + .build() + .unwrap() + } + + fn disagg_config() -> OfflineDisaggReplayConfig { + OfflineDisaggReplayConfig { + prefill_args: staged_args(WorkerType::Prefill, 1000.0), + decode_args: staged_args(WorkerType::Decode, 1000.0), + num_prefill_workers: 2, + num_decode_workers: 2, + } + } + + fn router_config() -> KvRouterConfig { + KvRouterConfig { + router_queue_threshold: Some(1.25), + ..KvRouterConfig::default() + } + } + + fn request( + uuid: u128, + prompt_tokens: usize, + output_tokens: usize, + arrival_ms: f64, + ) -> DirectRequest { + DirectRequest { + tokens: vec![1; prompt_tokens], + max_output_tokens: output_tokens, + uuid: Some(Uuid::from_u128(uuid)), + dp_rank: 0, + arrival_timestamp_ms: Some(arrival_ms), + } + } + + #[test] + fn test_derive_stage_router_configs_force_required_overrides() { + let config = KvRouterConfig { + overlap_score_weight: 2.0, + router_track_active_blocks: true, + router_assume_kv_reuse: true, + router_track_prefill_tokens: true, + ..KvRouterConfig::default() + }; + let args = staged_args(WorkerType::Prefill, 1.0); + let prefill = derive_prefill_router_config(&args, Some(config.clone())); + let decode = derive_decode_router_config(&args, Some(config)); + + assert!(!prefill.router_track_active_blocks); + assert_eq!(decode.overlap_score_weight, 0.0); + assert!(!decode.router_assume_kv_reuse); + assert!(!decode.router_track_prefill_tokens); + } + + #[test] + fn test_trace_smoke_reports_decode_only_tokens() { + let config = disagg_config(); + let requests = vec![request(1, 128, 3, 5.0)]; + + let (collector, stats) = run_trace_collect(&config, requests, Some(router_config()), 1.0); + let snapshot = collector.snapshot(Uuid::from_u128(1)).unwrap(); + let report = collector.finish(); + + assert_eq!(snapshot.arrival_time_ms, 0.0); + assert!(snapshot.first_admit_ms.is_some()); + assert!(snapshot.first_token_ms.is_some()); + assert_eq!(snapshot.output_length, 3); + assert_eq!(report.request_counts.completed_requests, 1); + assert_eq!( + stats.request_snapshots[&Uuid::from_u128(1)].phase, + DisaggPhase::Done + ); + } + + #[test] + fn test_prefill_and_decode_use_separate_worker_pools() { + let config = disagg_config(); + let requests = vec![request(1, 128, 2, 0.0), request(2, 128, 2, 10.0)]; + + let (_, stats) = run_trace_collect(&config, requests, Some(router_config()), 1.0); + + for uuid in [Uuid::from_u128(1), Uuid::from_u128(2)] { + assert!(stats.prefill_assignments.contains_key(&uuid)); + assert!(stats.decode_assignments.contains_key(&uuid)); + } + } + + #[test] + fn test_prefill_overlap_prefers_same_worker_after_handoff_delay() { + let config = disagg_config(); + let requests = vec![request(1, 128, 2, 0.0), request(2, 128, 2, 100.0)]; + + let (_, stats) = run_trace_collect(&config, requests, Some(router_config()), 1.0); + + assert_eq!( + stats.prefill_assignments[&Uuid::from_u128(1)], + stats.prefill_assignments[&Uuid::from_u128(2)], + ); + } + + #[test] + fn test_concurrency_backfill_waits_for_decode_completion() { + let config = disagg_config(); + let requests = vec![ + DirectRequest { + tokens: vec![1; 128], + max_output_tokens: 3, + uuid: Some(Uuid::from_u128(1)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + DirectRequest { + tokens: vec![2; 128], + max_output_tokens: 3, + uuid: Some(Uuid::from_u128(2)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + ]; + + let (collector, _) = run_concurrency_collect(&config, requests, Some(router_config()), 1); + let first = collector.snapshot(Uuid::from_u128(1)).unwrap(); + let second = collector.snapshot(Uuid::from_u128(2)).unwrap(); + + assert_eq!(first.arrival_time_ms, 0.0); + assert_eq!(second.arrival_time_ms, first.last_token_ms.unwrap()); + } + + #[test] + fn test_prefill_completion_marks_and_frees_before_decode_handoff() { + let config = disagg_config(); + let requests = vec![request(1, 128, 2, 0.0)]; + + let (_, stats) = run_trace_collect(&config, requests, Some(router_config()), 1.0); + + assert_eq!(stats.prefill_marked_count, 1); + assert_eq!(stats.prefill_freed_count, 1); + assert_eq!(stats.decode_freed_count, 1); + } +} diff --git a/lib/mocker/src/replay/offline/events.rs b/lib/mocker/src/replay/offline/events.rs index d7383a487551..dfed2feaffe2 100644 --- a/lib/mocker/src/replay/offline/events.rs +++ b/lib/mocker/src/replay/offline/events.rs @@ -5,9 +5,17 @@ use std::cmp::Ordering; use crate::common::protocols::OutputSignal; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SimulationWorkerStage { + Aggregated, + Prefill, + Decode, +} + #[derive(Debug)] pub(crate) enum SimulationEventKind { WorkerCompletion { + stage: SimulationWorkerStage, worker_idx: usize, completed_requests: usize, output_signals: Vec, diff --git a/lib/mocker/src/replay/offline/mod.rs b/lib/mocker/src/replay/offline/mod.rs index a84cfd1b0c6f..705bdce21fa5 100644 --- a/lib/mocker/src/replay/offline/mod.rs +++ b/lib/mocker/src/replay/offline/mod.rs @@ -3,11 +3,13 @@ use crate::common::protocols::{DirectRequest, MockEngineArgs}; use crate::loadgen::Trace; +use crate::replay::OfflineDisaggReplayConfig; pub(crate) use crate::replay::normalize_trace_requests; use crate::replay::{ReplayRouterMode, TraceSimulationReport}; use dynamo_kv_router::config::KvRouterConfig; pub(crate) mod core; +pub(crate) mod disagg; pub(crate) mod events; pub(crate) mod multi; pub(crate) mod runtime_utils; @@ -93,3 +95,38 @@ pub(crate) fn simulate_concurrency_workload( ) } } + +pub(crate) fn simulate_trace_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + arrival_speedup_ratio: f64, +) -> anyhow::Result { + disagg::simulate_trace_disagg(config, router_config, requests, arrival_speedup_ratio) +} + +pub(crate) fn simulate_concurrency_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + max_in_flight: usize, +) -> anyhow::Result { + disagg::simulate_concurrency_disagg(config, router_config, requests, max_in_flight) +} + +pub(crate) fn simulate_trace_workload_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, +) -> anyhow::Result { + disagg::simulate_trace_workload_disagg(config, router_config, trace) +} + +pub(crate) fn simulate_concurrency_workload_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + max_in_flight: usize, +) -> anyhow::Result { + disagg::simulate_concurrency_workload_disagg(config, router_config, trace, max_in_flight) +} diff --git a/lib/mocker/src/replay/offline/multi.rs b/lib/mocker/src/replay/offline/multi.rs index 41d939707d82..e9a5315d9f9f 100644 --- a/lib/mocker/src/replay/offline/multi.rs +++ b/lib/mocker/src/replay/offline/multi.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use super::events::SimulationEvent; +use super::events::{SimulationEvent, SimulationWorkerStage}; use super::normalize_trace_requests; use super::runtime_utils::{ WorkerCompletionPayload, next_timestamp as choose_next_timestamp, pop_next_concurrency_ready, @@ -386,12 +386,14 @@ impl OfflineRuntime { fn apply_worker_completions(&mut self) -> anyhow::Result { let mut changed = false; while let Some(WorkerCompletionPayload { + stage, worker_idx, completed_requests, output_signals, kv_events, }) = pop_ready_worker_completion(&mut self.events, self.now_ms) { + debug_assert_eq!(stage, SimulationWorkerStage::Aggregated); self.workers[worker_idx].mark_idle(); self.process_completed_pass(worker_idx, completed_requests, output_signals, kv_events)?; changed = true; @@ -513,6 +515,7 @@ impl OfflineRuntime { &mut self.next_event_seq, executed.end_ms, WorkerCompletionPayload { + stage: SimulationWorkerStage::Aggregated, worker_idx, completed_requests: executed.completed_requests, output_signals: executed.output_signals, diff --git a/lib/mocker/src/replay/offline/runtime_utils.rs b/lib/mocker/src/replay/offline/runtime_utils.rs index 1050c34dab05..67c0f85e6380 100644 --- a/lib/mocker/src/replay/offline/runtime_utils.rs +++ b/lib/mocker/src/replay/offline/runtime_utils.rs @@ -5,11 +5,12 @@ use std::collections::{BinaryHeap, VecDeque}; use dynamo_kv_router::protocols::RouterEvent; -use super::events::{SimulationEvent, SimulationEventKind}; +use super::events::{SimulationEvent, SimulationEventKind, SimulationWorkerStage}; use crate::common::protocols::{DirectRequest, OutputSignal}; #[derive(Debug)] pub(super) struct WorkerCompletionPayload { + pub stage: SimulationWorkerStage, pub worker_idx: usize, pub completed_requests: usize, pub output_signals: Vec, @@ -65,6 +66,7 @@ pub(super) fn push_worker_completion( at_ms, seq_no: *next_event_seq, kind: SimulationEventKind::WorkerCompletion { + stage: payload.stage, worker_idx: payload.worker_idx, completed_requests: payload.completed_requests, output_signals: payload.output_signals, @@ -84,12 +86,14 @@ pub(super) fn pop_ready_worker_completion( } let event = events.pop().expect("event must exist after peek"); let SimulationEventKind::WorkerCompletion { + stage, worker_idx, completed_requests, output_signals, kv_events, } = event.kind; Some(WorkerCompletionPayload { + stage, worker_idx, completed_requests, output_signals, @@ -100,6 +104,7 @@ pub(super) fn pop_ready_worker_completion( #[cfg(test)] mod tests { use super::*; + use crate::replay::offline::events::SimulationWorkerStage; use uuid::Uuid; fn direct_request(uuid: u128, arrival_timestamp_ms: Option) -> DirectRequest { @@ -163,6 +168,7 @@ mod tests { &mut next_event_seq, 10.0, WorkerCompletionPayload { + stage: SimulationWorkerStage::Aggregated, worker_idx: 7, completed_requests: 1, output_signals: vec![OutputSignal { @@ -177,6 +183,7 @@ mod tests { &mut next_event_seq, 10.0, WorkerCompletionPayload { + stage: SimulationWorkerStage::Aggregated, worker_idx: 8, completed_requests: 2, output_signals: vec![OutputSignal { @@ -191,8 +198,10 @@ mod tests { let first = pop_ready_worker_completion(&mut events, 10.0).unwrap(); let second = pop_ready_worker_completion(&mut events, 10.0).unwrap(); + assert_eq!(first.stage, SimulationWorkerStage::Aggregated); assert_eq!(first.worker_idx, 7); assert_eq!(first.completed_requests, 1); + assert_eq!(second.stage, SimulationWorkerStage::Aggregated); assert_eq!(second.worker_idx, 8); assert_eq!(second.completed_requests, 2); assert!(events.is_empty()); diff --git a/lib/mocker/src/replay/offline/state.rs b/lib/mocker/src/replay/offline/state.rs index 0acbacbace53..bd66ddaf56b5 100644 --- a/lib/mocker/src/replay/offline/state.rs +++ b/lib/mocker/src/replay/offline/state.rs @@ -91,6 +91,10 @@ impl OfflineWorkerState { self.core.execute_pass(collector, now_ms) } + pub(crate) fn execute_hidden_pass(&mut self, now_ms: f64) -> EnginePassResult { + self.core.execute_hidden_pass(now_ms) + } + #[cfg(test)] pub(crate) fn debug_snapshot(&self) -> OfflineWorkerSnapshot { OfflineWorkerSnapshot { diff --git a/lib/mocker/src/replay/router/offline.rs b/lib/mocker/src/replay/router/offline.rs index d600a1323e5b..28d2f4c6fed6 100644 --- a/lib/mocker/src/replay/router/offline.rs +++ b/lib/mocker/src/replay/router/offline.rs @@ -103,6 +103,7 @@ struct PendingRequest { token_seq: Option>, isl_tokens: usize, overlaps: OverlapScores, + track_prefill_tokens: bool, expected_output_tokens: Option, } @@ -123,7 +124,7 @@ impl PendingRequest { overlaps: self.overlaps.clone(), decode_blocks, prefill_tokens, - track_prefill_tokens: true, + track_prefill_tokens: self.track_prefill_tokens, router_config_override: None, update_states: true, lora_name: None, @@ -252,7 +253,6 @@ impl OfflineReplayRouter { self.drain_pending() } - #[cfg(test)] pub(crate) fn pending_count(&self) -> usize { self.pending.len() } @@ -363,6 +363,7 @@ impl OfflineReplayRouter { token_seq, isl_tokens: request.tokens.len(), overlaps, + track_prefill_tokens: self.config.router_track_prefill_tokens, expected_output_tokens: Some( u32::try_from(request.max_output_tokens) .context("max_output_tokens does not fit into u32")?, @@ -392,7 +393,7 @@ impl OfflineReplayRouter { token_sequence: request.token_seq, isl: request.isl_tokens, overlap: selection.overlap_blocks, - track_prefill_tokens: true, + track_prefill_tokens: request.track_prefill_tokens, expected_output_tokens: request.expected_output_tokens, worker: selection.worker, lora_name: None, diff --git a/lib/mocker/src/replay/router/shared.rs b/lib/mocker/src/replay/router/shared.rs index 4601f79a3065..56aa85245322 100644 --- a/lib/mocker/src/replay/router/shared.rs +++ b/lib/mocker/src/replay/router/shared.rs @@ -107,7 +107,7 @@ pub(super) fn replay_selector(config: &KvRouterConfig) -> DefaultWorkerSelector DefaultWorkerSelector::new(Some(config.clone()), "replay") } -pub(super) fn replay_router_config( +pub(crate) fn replay_router_config( args: &MockEngineArgs, router_config: Option, ) -> KvRouterConfig { diff --git a/lib/mocker/src/replay/validate.rs b/lib/mocker/src/replay/validate.rs index 824abb5eb9d8..2c2e37c4ace6 100644 --- a/lib/mocker/src/replay/validate.rs +++ b/lib/mocker/src/replay/validate.rs @@ -3,6 +3,7 @@ use anyhow::{Result, bail}; +use super::OfflineDisaggReplayConfig; use super::ReplayRouterMode; use crate::common::protocols::{MockEngineArgs, WorkerType}; @@ -75,3 +76,69 @@ pub(super) fn validate_online_concurrency_args( validate_replay_args(args, num_workers, "online replay") } + +fn validate_disagg_args(config: &OfflineDisaggReplayConfig, mode: &str) -> Result<()> { + if config.num_prefill_workers == 0 { + bail!("{mode} requires num_prefill_workers >= 1"); + } + if config.num_decode_workers == 0 { + bail!("{mode} requires num_decode_workers >= 1"); + } + if config.prefill_args.worker_type != WorkerType::Prefill { + bail!( + "{mode} requires prefill_engine_args.worker_type=prefill, got {:?}", + config.prefill_args.worker_type, + ); + } + if config.decode_args.worker_type != WorkerType::Decode { + bail!( + "{mode} requires decode_engine_args.worker_type=decode, got {:?}", + config.decode_args.worker_type, + ); + } + if config.prefill_args.dp_size != 1 { + bail!( + "{mode} only supports prefill data_parallel_size=1, got {}", + config.prefill_args.dp_size, + ); + } + if config.decode_args.dp_size != 1 { + bail!( + "{mode} only supports decode data_parallel_size=1, got {}", + config.decode_args.dp_size, + ); + } + if config.prefill_args.block_size != config.decode_args.block_size { + bail!( + "{mode} requires matching prefill/decode block_size, got {} and {}", + config.prefill_args.block_size, + config.decode_args.block_size, + ); + } + + Ok(()) +} + +pub(super) fn validate_offline_disagg_replay_args( + config: &OfflineDisaggReplayConfig, + router_mode: ReplayRouterMode, +) -> Result<()> { + if router_mode != ReplayRouterMode::KvRouter { + bail!("offline disagg replay requires router_mode=kv_router"); + } + validate_disagg_args(config, "trace replay") +} + +pub(super) fn validate_offline_disagg_concurrency_args( + config: &OfflineDisaggReplayConfig, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> Result<()> { + if max_in_flight == 0 { + bail!("concurrency replay requires max_in_flight >= 1"); + } + if router_mode != ReplayRouterMode::KvRouter { + bail!("offline disagg replay requires router_mode=kv_router"); + } + validate_disagg_args(config, "concurrency replay") +} diff --git a/lib/mocker/src/scheduler/mod.rs b/lib/mocker/src/scheduler/mod.rs index 79f7c3ef2a31..b86a6e5c09e7 100644 --- a/lib/mocker/src/scheduler/mod.rs +++ b/lib/mocker/src/scheduler/mod.rs @@ -89,6 +89,13 @@ impl EngineCore { Self::Sglang(core) => core.execute_pass(collector, now_ms), } } + + pub(crate) fn execute_hidden_pass(&mut self, now_ms: f64) -> EnginePassResult { + match self { + Self::Vllm(core) => core.execute_hidden_pass(now_ms), + Self::Sglang(core) => core.execute_hidden_pass(now_ms), + } + } } #[derive(Clone)] diff --git a/lib/mocker/src/scheduler/sglang/core.rs b/lib/mocker/src/scheduler/sglang/core.rs index 192d7e3dbc4d..903cb428005b 100644 --- a/lib/mocker/src/scheduler/sglang/core.rs +++ b/lib/mocker/src/scheduler/sglang/core.rs @@ -101,6 +101,10 @@ impl SglangCore { self.execute_pass_internal(Some(collector), now_ms) } + pub(crate) fn execute_hidden_pass(&mut self, now_ms: f64) -> EnginePassResult { + self.execute_pass_internal(None, now_ms) + } + pub(super) fn execute_pass_internal( &mut self, mut collector: Option<&mut TraceCollector>, diff --git a/lib/mocker/src/scheduler/vllm/core.rs b/lib/mocker/src/scheduler/vllm/core.rs index 9057bf4861c3..6986b6a7484c 100644 --- a/lib/mocker/src/scheduler/vllm/core.rs +++ b/lib/mocker/src/scheduler/vllm/core.rs @@ -278,6 +278,10 @@ impl VllmCore { self.execute_pass_internal(Some(collector), now_ms, None) } + pub(crate) fn execute_hidden_pass(&mut self, now_ms: f64) -> EnginePassResult { + self.execute_pass_internal(None, now_ms, None) + } + pub(super) fn execute_pass_internal( &mut self, mut collector: Option<&mut TraceCollector>, From b03b9561dd9a87d7823341244e3ea2e50c1b216a Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 20:37:31 -0700 Subject: [PATCH 04/24] fix: tighten replay guardrails and docs Signed-off-by: PeaBrane --- docs/benchmarks/mocker-trace-replay.md | 8 ++- docs/mocker/mocker.md | 15 ++++- lib/bindings/python/rust/llm/replay.rs | 18 +++++- lib/bindings/python/src/dynamo/replay/api.py | 26 ++++++++ lib/bindings/python/tests/test_replay.py | 17 ++++++ lib/kv-router/src/protocols.rs | 60 ------------------- .../src/replay/offline/runtime_utils.rs | 8 ++- 7 files changed, 87 insertions(+), 65 deletions(-) diff --git a/docs/benchmarks/mocker-trace-replay.md b/docs/benchmarks/mocker-trace-replay.md index 6defdabb57cc..e3c0ec8d0a6c 100644 --- a/docs/benchmarks/mocker-trace-replay.md +++ b/docs/benchmarks/mocker-trace-replay.md @@ -174,7 +174,12 @@ Offline disagg replay uses staged engine args instead of `--extra-engine-args`: - `--decode-engine-args` for the decode worker config - `--num-prefill-workers` and `--num-decode-workers` for pool sizes -The staged JSON accepts `worker_type: "prefill" | "decode"` for CLI convenience. +For offline disagg replay, the staged JSON must set `worker_type` explicitly: + +- `--prefill-engine-args` must use `worker_type: "prefill"` +- `--decode-engine-args` must use `worker_type: "decode"` + +The staged configs must also use the same `block_size`. ### Synthetic Replay @@ -337,6 +342,7 @@ Shared replay constraints: - disagg replay requires both `prefill_engine_args` and `decode_engine_args` - disagg replay requires `router_mode=kv_router` - replay `dp_size` must be `1` +- disagg replay requires matching `block_size` in `prefill_engine_args` and `decode_engine_args` Additional offline constraints: diff --git a/docs/mocker/mocker.md b/docs/mocker/mocker.md index 75b5881bc9b5..12d5486efd39 100644 --- a/docs/mocker/mocker.md +++ b/docs/mocker/mocker.md @@ -261,13 +261,26 @@ python -m dynamo.mocker \ The AIC model automatically uses `--model-path` and `--engine-type` to select the appropriate performance data. Available systems include `h200_sxm`, `h100_sxm`, etc. (see AIC SDK documentation for the full list). -When using `python -m dynamo.replay`, there are no dedicated AIC flags. Pass the equivalent fields directly via `--extra-engine-args`: +When using `python -m dynamo.replay`, there are no dedicated AIC flags. For aggregated replay, +pass the equivalent fields via `--extra-engine-args`: ```bash python -m dynamo.replay /path/to/trace.jsonl \ --extra-engine-args '{"aic_backend":"vllm","aic_system":"h200_sxm","aic_model_path":"nvidia/Llama-3.1-8B-Instruct-FP8","aic_tp_size":1}' ``` +For offline disagg replay, pass the staged engine configs instead: + +```bash +python -m dynamo.replay /path/to/trace.jsonl \ + --replay-mode offline \ + --router-mode kv_router \ + --prefill-engine-args '{"worker_type":"prefill","aic_backend":"vllm","aic_system":"h200_sxm","aic_model_path":"nvidia/Llama-3.1-8B-Instruct-FP8","aic_tp_size":1,"block_size":512}' \ + --decode-engine-args '{"worker_type":"decode","aic_backend":"vllm","aic_system":"h200_sxm","aic_model_path":"nvidia/Llama-3.1-8B-Instruct-FP8","aic_tp_size":1,"block_size":512}' \ + --num-prefill-workers 2 \ + --num-decode-workers 6 +``` + The `aic_backend` field enables the AIC perf model and should match `engine_type` (`"vllm"` or `"sglang"`). The `aic_model_path` field is the equivalent of `--model-path` in `dynamo.mocker`. Example `--reasoning` configuration: diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index 943b23e8919d..c891f0343e78 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -654,8 +654,22 @@ fn load_replay_args_selection( } match (aggregated_args, prefill_args, decode_args) { - (Some(args), None, None) => Ok(ReplayArgsSelection::Aggregated(Box::new(args))), - (None, None, None) => Ok(ReplayArgsSelection::Aggregated(Box::default())), + (Some(args), None, None) => { + if num_prefill_workers != 1 || num_decode_workers != 1 { + return Err(PyException::new_err( + "num_prefill_workers and num_decode_workers are only used for disagg replay; use num_workers for aggregated replay", + )); + } + Ok(ReplayArgsSelection::Aggregated(Box::new(args))) + } + (None, None, None) => { + if num_prefill_workers != 1 || num_decode_workers != 1 { + return Err(PyException::new_err( + "num_prefill_workers and num_decode_workers are only used for disagg replay; use num_workers for aggregated replay", + )); + } + Ok(ReplayArgsSelection::Aggregated(Box::default())) + } (None, Some(prefill_args), Some(decode_args)) => { if num_workers != 1 { return Err(PyException::new_err( diff --git a/lib/bindings/python/src/dynamo/replay/api.py b/lib/bindings/python/src/dynamo/replay/api.py index 049a1afe6f8d..2164078261cc 100644 --- a/lib/bindings/python/src/dynamo/replay/api.py +++ b/lib/bindings/python/src/dynamo/replay/api.py @@ -7,6 +7,20 @@ from dynamo._core import run_mocker_trace_replay as _run_mocker_trace_replay +def _validate_worker_pool_args( + prefill_engine_args, + decode_engine_args, + num_prefill_workers, + num_decode_workers, +): + if prefill_engine_args is not None or decode_engine_args is not None: + return + if num_prefill_workers != 1 or num_decode_workers != 1: + raise ValueError( + "num_prefill_workers and num_decode_workers are only used for disagg replay; use num_workers for aggregated replay" + ) + + def run_trace_replay( trace_file, *, @@ -22,6 +36,12 @@ def run_trace_replay( router_mode="round_robin", arrival_speedup_ratio=1.0, ): + _validate_worker_pool_args( + prefill_engine_args, + decode_engine_args, + num_prefill_workers, + num_decode_workers, + ) return _run_mocker_trace_replay( trace_file, extra_engine_args=extra_engine_args, @@ -60,6 +80,12 @@ def run_synthetic_trace_replay( num_prefix_groups=0, inter_turn_delay_ms=0.0, ): + _validate_worker_pool_args( + prefill_engine_args, + decode_engine_args, + num_prefill_workers, + num_decode_workers, + ) return _run_mocker_synthetic_trace_replay( input_tokens, output_tokens, diff --git a/lib/bindings/python/tests/test_replay.py b/lib/bindings/python/tests/test_replay.py index cb491d63df4f..bd10b21f89f9 100644 --- a/lib/bindings/python/tests/test_replay.py +++ b/lib/bindings/python/tests/test_replay.py @@ -643,6 +643,23 @@ def test_run_trace_replay_rejects_online_disagg(tmp_path): ) +def test_run_trace_replay_rejects_disagg_worker_counts_for_aggregated_mode(tmp_path): + trace_path = _write_trace_and_args(tmp_path) + + with pytest.raises( + Exception, + match="num_prefill_workers and num_decode_workers are only used for disagg replay", + ): + run_trace_replay( + trace_path, + extra_engine_args=MockEngineArgs(block_size=64, speedup_ratio=1000.0), + num_workers=1, + num_prefill_workers=2, + num_decode_workers=2, + replay_mode="offline", + ) + + def test_format_report_table_matches_aiperf_shape(): report = { "mean_ttft_ms": 18.26, diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index 6a3eb838d4af..7a8dcadd76be 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -949,14 +949,6 @@ mod tests { assert_ne!(lora_a[0], lora_b[0]); } - #[test] - fn test_lora_name_none_matches_legacy() { - let tokens: Vec = (0..8).collect(); - let hashes_none = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); - let hashes_none2 = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); - assert_eq!(hashes_none, hashes_none2); - } - #[test] fn test_lora_name_empty_string_normalized_to_none() { let tokens: Vec = (0..4).collect(); @@ -1131,16 +1123,6 @@ mod tests { assert_eq!(deserialized.block_hashes[1].0, 5); } - #[test] - fn test_router_request_mark_free_backwards_compatible_deserialization() { - let request: RouterRequest = serde_json::from_str(r#"{"method":"mark_free"}"#).unwrap(); - - assert!(matches!( - request, - RouterRequest::MarkFree { request_id: None } - )); - } - #[test] fn test_router_request_mark_free_serialization_with_request_id() { let request = RouterRequest::MarkFree { @@ -1161,46 +1143,4 @@ mod tests { } if request_id == "req-123" )); } - - #[test] - fn test_active_sequence_add_request_serialization_preserves_track_prefill_tokens() { - let event = ActiveSequenceEvent { - request_id: "req-123".to_string(), - worker: WorkerWithDpRank::new(7, 0), - data: ActiveSequenceEventData::AddRequest { - token_sequence: Some(vec![11, 22]), - isl: 128, - overlap: 1, - track_prefill_tokens: false, - expected_output_tokens: Some(32), - }, - router_id: 9, - lora_name: None, - }; - - let serialized = serde_json::to_string(&event).unwrap(); - let deserialized: ActiveSequenceEvent = serde_json::from_str(&serialized).unwrap(); - - match deserialized.data { - ActiveSequenceEventData::AddRequest { - track_prefill_tokens, - .. - } => assert!(!track_prefill_tokens), - _ => panic!("expected add request event"), - } - } - - #[test] - fn test_active_sequence_add_request_defaults_track_prefill_tokens_for_legacy_payloads() { - let legacy = r#"{"request_id":"req-123","worker":{"worker_id":7,"dp_rank":0},"data":{"AddRequest":{"token_sequence":[11,22],"isl":128,"overlap":1,"expected_output_tokens":32}},"router_id":9,"lora_name":null}"#; - let deserialized: ActiveSequenceEvent = serde_json::from_str(legacy).unwrap(); - - match deserialized.data { - ActiveSequenceEventData::AddRequest { - track_prefill_tokens, - .. - } => assert!(track_prefill_tokens), - _ => panic!("expected add request event"), - } - } } diff --git a/lib/mocker/src/replay/offline/runtime_utils.rs b/lib/mocker/src/replay/offline/runtime_utils.rs index a67ef71a3d61..24bb37b5831a 100644 --- a/lib/mocker/src/replay/offline/runtime_utils.rs +++ b/lib/mocker/src/replay/offline/runtime_utils.rs @@ -92,7 +92,13 @@ pub(super) fn pop_ready_worker_completion( completed_requests, output_signals, kv_events, - } => (stage, worker_idx, completed_requests, output_signals, kv_events), + } => ( + stage, + worker_idx, + completed_requests, + output_signals, + kv_events, + ), }; Some(WorkerCompletionPayload { stage, From e41b87ef9c61384b0f23dfc45f0158808c03bbbe Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 20:47:22 -0700 Subject: [PATCH 05/24] fix: centralize replay validation before it forks Signed-off-by: PeaBrane --- lib/bindings/c/src/lib.rs | 4 +- lib/bindings/python/rust/llm/replay.rs | 63 +++++++------------- lib/bindings/python/src/dynamo/replay/api.py | 26 -------- lib/mocker/src/replay/mod.rs | 7 +++ lib/mocker/src/replay/validate.rs | 39 +++++++++++- 5 files changed, 68 insertions(+), 71 deletions(-) diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index f1c5b7ca2505..3a2cad51f400 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -877,7 +877,7 @@ pub unsafe extern "C" fn add_request( // Compute overlap_blocks using the public method let overlap_blocks = match decode_router - .get_overlap_blocks(&tokens, None, worker, Some(&router_config_override)) + .get_overlap_blocks(&tokens, None, worker, None) .await { Ok(overlap) => overlap, @@ -896,7 +896,7 @@ pub unsafe extern "C" fn add_request( None, worker, None, // lora_name - Some(router_config_override), + Some(&router_config_override), ) .await; diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index c891f0343e78..85d407b8a32f 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -13,6 +13,7 @@ use dynamo_mocker::common::protocols::{ use dynamo_mocker::loadgen::{ ArrivalSpec, DelaySpec, LengthSpec, SyntheticTraceSpec, Trace as RsTrace, }; +use dynamo_mocker::replay::ReplayArgsMode; use pyo3::{exceptions::PyException, prelude::*}; use pythonize::pythonize; use uuid::Uuid; @@ -647,48 +648,28 @@ fn load_replay_args_selection( let prefill_args = load_optional_replay_mocker_args(py, prefill_engine_args)?; let decode_args = load_optional_replay_mocker_args(py, decode_engine_args)?; - if aggregated_args.is_some() && (prefill_args.is_some() || decode_args.is_some()) { - return Err(PyException::new_err( - "extra_engine_args cannot be combined with prefill_engine_args/decode_engine_args", - )); - } + let replay_args_mode = dynamo_mocker::replay::validate_replay_args_mode( + aggregated_args.as_ref(), + prefill_args.as_ref(), + decode_args.as_ref(), + num_workers, + num_prefill_workers, + num_decode_workers, + ) + .map_err(to_pyerr)?; - match (aggregated_args, prefill_args, decode_args) { - (Some(args), None, None) => { - if num_prefill_workers != 1 || num_decode_workers != 1 { - return Err(PyException::new_err( - "num_prefill_workers and num_decode_workers are only used for disagg replay; use num_workers for aggregated replay", - )); - } - Ok(ReplayArgsSelection::Aggregated(Box::new(args))) - } - (None, None, None) => { - if num_prefill_workers != 1 || num_decode_workers != 1 { - return Err(PyException::new_err( - "num_prefill_workers and num_decode_workers are only used for disagg replay; use num_workers for aggregated replay", - )); - } - Ok(ReplayArgsSelection::Aggregated(Box::default())) - } - (None, Some(prefill_args), Some(decode_args)) => { - if num_workers != 1 { - return Err(PyException::new_err( - "num_workers is only used for aggregated replay; use num_prefill_workers and num_decode_workers for disagg replay", - )); - } - Ok(ReplayArgsSelection::Disagg(Box::new( - dynamo_mocker::replay::OfflineDisaggReplayConfig { - prefill_args, - decode_args, - num_prefill_workers, - num_decode_workers, - }, - ))) - } - (None, Some(_), None) | (None, None, Some(_)) => Err(PyException::new_err( - "prefill_engine_args and decode_engine_args must be provided together", - )), - _ => unreachable!(), + match replay_args_mode { + ReplayArgsMode::Aggregated => Ok(ReplayArgsSelection::Aggregated(Box::new( + aggregated_args.unwrap_or_default(), + ))), + ReplayArgsMode::Disagg => Ok(ReplayArgsSelection::Disagg(Box::new( + dynamo_mocker::replay::OfflineDisaggReplayConfig { + prefill_args: prefill_args.expect("validated disagg prefill args"), + decode_args: decode_args.expect("validated disagg decode args"), + num_prefill_workers, + num_decode_workers, + }, + ))), } } diff --git a/lib/bindings/python/src/dynamo/replay/api.py b/lib/bindings/python/src/dynamo/replay/api.py index 2164078261cc..049a1afe6f8d 100644 --- a/lib/bindings/python/src/dynamo/replay/api.py +++ b/lib/bindings/python/src/dynamo/replay/api.py @@ -7,20 +7,6 @@ from dynamo._core import run_mocker_trace_replay as _run_mocker_trace_replay -def _validate_worker_pool_args( - prefill_engine_args, - decode_engine_args, - num_prefill_workers, - num_decode_workers, -): - if prefill_engine_args is not None or decode_engine_args is not None: - return - if num_prefill_workers != 1 or num_decode_workers != 1: - raise ValueError( - "num_prefill_workers and num_decode_workers are only used for disagg replay; use num_workers for aggregated replay" - ) - - def run_trace_replay( trace_file, *, @@ -36,12 +22,6 @@ def run_trace_replay( router_mode="round_robin", arrival_speedup_ratio=1.0, ): - _validate_worker_pool_args( - prefill_engine_args, - decode_engine_args, - num_prefill_workers, - num_decode_workers, - ) return _run_mocker_trace_replay( trace_file, extra_engine_args=extra_engine_args, @@ -80,12 +60,6 @@ def run_synthetic_trace_replay( num_prefix_groups=0, inter_turn_delay_ms=0.0, ): - _validate_worker_pool_args( - prefill_engine_args, - decode_engine_args, - num_prefill_workers, - num_decode_workers, - ) return _run_mocker_synthetic_trace_replay( input_tokens, output_tokens, diff --git a/lib/mocker/src/replay/mod.rs b/lib/mocker/src/replay/mod.rs index b64a7ad3aecd..9076bedf7b86 100644 --- a/lib/mocker/src/replay/mod.rs +++ b/lib/mocker/src/replay/mod.rs @@ -25,6 +25,12 @@ pub enum ReplayRouterMode { KvRouter, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReplayArgsMode { + Aggregated, + Disagg, +} + #[derive(Clone, Debug)] pub struct OfflineDisaggReplayConfig { pub prefill_args: MockEngineArgs, @@ -62,6 +68,7 @@ pub use entrypoints::{ simulate_trace_requests_with_router_mode, simulate_trace_workload, simulate_trace_workload_disagg_with_router_mode, simulate_trace_workload_with_router_mode, }; +pub use validate::validate_replay_args_mode; pub(crate) fn normalize_trace_requests( mut requests: Vec, diff --git a/lib/mocker/src/replay/validate.rs b/lib/mocker/src/replay/validate.rs index 2c2e37c4ace6..fab2e7d128b0 100644 --- a/lib/mocker/src/replay/validate.rs +++ b/lib/mocker/src/replay/validate.rs @@ -3,10 +3,45 @@ use anyhow::{Result, bail}; -use super::OfflineDisaggReplayConfig; -use super::ReplayRouterMode; +use super::{OfflineDisaggReplayConfig, ReplayArgsMode, ReplayRouterMode}; use crate::common::protocols::{MockEngineArgs, WorkerType}; +pub fn validate_replay_args_mode( + aggregated_args: Option<&MockEngineArgs>, + prefill_args: Option<&MockEngineArgs>, + decode_args: Option<&MockEngineArgs>, + num_workers: usize, + num_prefill_workers: usize, + num_decode_workers: usize, +) -> Result { + if aggregated_args.is_some() && (prefill_args.is_some() || decode_args.is_some()) { + bail!("extra_engine_args cannot be combined with prefill_engine_args/decode_engine_args"); + } + + match (aggregated_args, prefill_args, decode_args) { + (Some(_), None, None) | (None, None, None) => { + if num_prefill_workers != 1 || num_decode_workers != 1 { + bail!( + "num_prefill_workers and num_decode_workers are only used for disagg replay; use num_workers for aggregated replay" + ); + } + Ok(ReplayArgsMode::Aggregated) + } + (None, Some(_), Some(_)) => { + if num_workers != 1 { + bail!( + "num_workers is only used for aggregated replay; use num_prefill_workers and num_decode_workers for disagg replay" + ); + } + Ok(ReplayArgsMode::Disagg) + } + (None, Some(_), None) | (None, None, Some(_)) => { + bail!("prefill_engine_args and decode_engine_args must be provided together") + } + (Some(_), Some(_), _) | (Some(_), _, Some(_)) => unreachable!(), + } +} + fn validate_replay_args(args: &MockEngineArgs, num_workers: usize, mode: &str) -> Result<()> { if num_workers == 0 { bail!("{mode} requires num_workers >= 1"); From e0f4342f0741ced89ac6edcf9585163f6e70ef17 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 21:07:46 -0700 Subject: [PATCH 06/24] feat: let disagg replay stop kv cling Signed-off-by: PeaBrane --- lib/bindings/python/src/dynamo/replay/main.py | 5 +- lib/bindings/python/tests/test_replay.py | 95 ++++-- lib/mocker/src/replay/entrypoints.rs | 19 +- lib/mocker/src/replay/offline/disagg.rs | 288 ++++++++++++++---- lib/mocker/src/replay/offline/mod.rs | 30 +- lib/mocker/src/replay/validate.rs | 10 +- 6 files changed, 335 insertions(+), 112 deletions(-) diff --git a/lib/bindings/python/src/dynamo/replay/main.py b/lib/bindings/python/src/dynamo/replay/main.py index e3422bceb366..d701e8cbbbf5 100644 --- a/lib/bindings/python/src/dynamo/replay/main.py +++ b/lib/bindings/python/src/dynamo/replay/main.py @@ -10,6 +10,7 @@ from collections.abc import Sequence from pathlib import Path from types import SimpleNamespace +from typing import Any os.environ.setdefault("DYNAMO_SKIP_PYTHON_LOG_INIT", "1") @@ -21,7 +22,9 @@ ) except ImportError: - def resolve_planner_profile_data(planner_profile_data: Path): + def resolve_planner_profile_data(planner_profile_data: Path | None) -> Any: + if planner_profile_data is None: + return SimpleNamespace(npz_path=None) return SimpleNamespace( npz_path=planner_profile_data if planner_profile_data.suffix == ".npz" diff --git a/lib/bindings/python/tests/test_replay.py b/lib/bindings/python/tests/test_replay.py index bd10b21f89f9..7964c4907269 100644 --- a/lib/bindings/python/tests/test_replay.py +++ b/lib/bindings/python/tests/test_replay.py @@ -278,18 +278,34 @@ def _assert_replay_cli_outputs(completed, report_path): @pytest.mark.parametrize("engine_type", ["vllm", "sglang"]) @pytest.mark.parametrize("replay_mode", ["offline", "online"]) @pytest.mark.parametrize("router_mode", ["round_robin", "kv_router"]) -def test_run_trace_replay_smoke_matrix(tmp_path, engine_type, replay_mode, router_mode): +@pytest.mark.parametrize("serving_mode", ["agg", "disagg"]) +def test_run_trace_replay_smoke_matrix( + tmp_path, engine_type, replay_mode, router_mode, serving_mode +): trace_path = _write_trace_and_args(tmp_path) - args_path = _vllm_args() if engine_type == "vllm" else _sglang_args() - num_workers = 1 if router_mode == "round_robin" else 2 - - report = run_trace_replay( - trace_path, - extra_engine_args=args_path, - num_workers=num_workers, - replay_mode=replay_mode, - router_mode=router_mode, - ) + if serving_mode == "disagg": + if replay_mode != "offline": + pytest.skip("disagg replay only supports offline mode") + report = run_trace_replay( + trace_path, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode=replay_mode, + router_mode=router_mode, + ) + else: + args_path = _vllm_args() if engine_type == "vllm" else _sglang_args() + num_workers = 1 if router_mode == "round_robin" else 2 + report = run_trace_replay( + trace_path, + extra_engine_args=args_path, + num_workers=num_workers, + replay_mode=replay_mode, + router_mode=router_mode, + ) _assert_basic_report_counts( report, @@ -359,22 +375,39 @@ def test_run_trace_replay_supports_multiturn_sessions(tmp_path, replay_mode): @pytest.mark.parametrize("engine_type", ["vllm", "sglang"]) @pytest.mark.parametrize("replay_mode", ["offline", "online"]) @pytest.mark.parametrize("router_mode", ["round_robin", "kv_router"]) +@pytest.mark.parametrize("serving_mode", ["agg", "disagg"]) def test_run_synthetic_trace_replay_smoke_matrix( - tmp_path, engine_type, replay_mode, router_mode + tmp_path, engine_type, replay_mode, router_mode, serving_mode ): - args_path = _vllm_args() if engine_type == "vllm" else _sglang_args() - num_workers = 1 if router_mode == "round_robin" else 2 - - report = run_synthetic_trace_replay( - 64, - 2, - 2, - extra_engine_args=args_path, - num_workers=num_workers, - replay_mode=replay_mode, - router_mode=router_mode, - arrival_interval_ms=5.0, - ) + if serving_mode == "disagg": + if replay_mode != "offline": + pytest.skip("disagg replay only supports offline mode") + report = run_synthetic_trace_replay( + 64, + 2, + 2, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode=replay_mode, + router_mode=router_mode, + arrival_interval_ms=5.0, + ) + else: + args_path = _vllm_args() if engine_type == "vllm" else _sglang_args() + num_workers = 1 if router_mode == "round_robin" else 2 + report = run_synthetic_trace_replay( + 64, + 2, + 2, + extra_engine_args=args_path, + num_workers=num_workers, + replay_mode=replay_mode, + router_mode=router_mode, + arrival_interval_ms=5.0, + ) _assert_basic_report_counts( report, @@ -567,7 +600,8 @@ def test_run_trace_replay_accepts_partial_extra_engine_args_json(tmp_path, repla ) -def test_run_trace_replay_supports_disagg_offline(tmp_path): +@pytest.mark.parametrize("router_mode", ["round_robin", "kv_router"]) +def test_run_trace_replay_supports_disagg_offline(tmp_path, router_mode): trace_path = _write_trace_and_args(tmp_path) report = run_trace_replay( @@ -578,7 +612,7 @@ def test_run_trace_replay_supports_disagg_offline(tmp_path): num_prefill_workers=2, num_decode_workers=2, replay_mode="offline", - router_mode="kv_router", + router_mode=router_mode, ) _assert_basic_report_counts( @@ -590,7 +624,10 @@ def test_run_trace_replay_supports_disagg_offline(tmp_path): _assert_basic_report_metrics(report) -def test_run_synthetic_trace_replay_disagg_preserves_expected_output_tokens(): +@pytest.mark.parametrize("router_mode", ["round_robin", "kv_router"]) +def test_run_synthetic_trace_replay_disagg_preserves_expected_output_tokens( + router_mode, +): report = run_synthetic_trace_replay( 128, 7, @@ -601,7 +638,7 @@ def test_run_synthetic_trace_replay_disagg_preserves_expected_output_tokens(): num_prefill_workers=2, num_decode_workers=2, replay_mode="offline", - router_mode="kv_router", + router_mode=router_mode, ) _assert_basic_report_counts( diff --git a/lib/mocker/src/replay/entrypoints.rs b/lib/mocker/src/replay/entrypoints.rs index 0032fc2b8794..13f80cd6cdc5 100644 --- a/lib/mocker/src/replay/entrypoints.rs +++ b/lib/mocker/src/replay/entrypoints.rs @@ -70,8 +70,12 @@ pub fn simulate_trace_file_disagg_with_router_mode( .normalize_session_starts()? .speed_up_timing(arrival_speedup_ratio)?; let started_at = Instant::now(); - let report = - crate::replay::offline::simulate_trace_workload_disagg(config, router_config, trace)?; + let report = crate::replay::offline::simulate_trace_workload_disagg( + config, + router_config, + trace, + router_mode, + )?; Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } @@ -168,6 +172,7 @@ pub fn simulate_trace_requests_disagg_with_router_mode( router_config, requests, arrival_speedup_ratio, + router_mode, )?; Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } @@ -407,6 +412,7 @@ pub fn simulate_concurrency_requests_disagg_with_router_mode( router_config, requests, max_in_flight, + router_mode, ) } @@ -453,8 +459,12 @@ pub fn simulate_trace_workload_disagg_with_router_mode( let config = config.normalized()?; validate_offline_disagg_replay_args(&config, router_mode)?; let started_at = Instant::now(); - let report = - crate::replay::offline::simulate_trace_workload_disagg(config, router_config, trace)?; + let report = crate::replay::offline::simulate_trace_workload_disagg( + config, + router_config, + trace, + router_mode, + )?; Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } @@ -534,6 +544,7 @@ pub fn simulate_concurrency_workload_disagg_with_router_mode( router_config, trace, max_in_flight, + router_mode, ) } diff --git a/lib/mocker/src/replay/offline/disagg.rs b/lib/mocker/src/replay/offline/disagg.rs index e6e955abdba7..43ee6c02f4b8 100644 --- a/lib/mocker/src/replay/offline/disagg.rs +++ b/lib/mocker/src/replay/offline/disagg.rs @@ -18,7 +18,9 @@ use super::state::OfflineWorkerState; use crate::common::protocols::{DirectRequest, MockEngineArgs, OutputSignal}; use crate::loadgen::{ReplayRequestHashes, Trace, WorkloadDriver}; use crate::replay::router::OfflineReplayRouter; -use crate::replay::{OfflineDisaggReplayConfig, TraceCollector, TraceSimulationReport}; +use crate::replay::{ + OfflineDisaggReplayConfig, ReplayRouterMode, TraceCollector, TraceSimulationReport, +}; use crate::scheduler::RouterEventVisibility; #[derive(Debug, Clone, Copy)] @@ -42,7 +44,7 @@ enum DisaggPhase { } struct DisaggRequestState { - original: DirectRequest, + original: Option, #[cfg(test)] arrival_ms: f64, phase: DisaggPhase, @@ -79,12 +81,14 @@ struct DisaggRuntimeStats; struct DisaggRuntime { now_ms: f64, + next_prefill_worker_idx: usize, + next_decode_worker_idx: usize, next_event_seq: u64, admission: AdmissionSource, prefill_workers: Vec, decode_workers: Vec, - prefill_router: OfflineReplayRouter, - decode_router: OfflineReplayRouter, + prefill_router: Option, + decode_router: Option, requests: HashMap, collector: TraceCollector, events: BinaryHeap, @@ -98,12 +102,14 @@ impl DisaggRuntime { router_config: Option, pending: VecDeque, mode: ReplayMode, + router_mode: ReplayRouterMode, ) -> Result { Self::new_with_source( config, router_config, AdmissionSource::Requests(pending), mode, + router_mode, ) } @@ -112,12 +118,14 @@ impl DisaggRuntime { router_config: Option, driver: WorkloadDriver, mode: ReplayMode, + router_mode: ReplayRouterMode, ) -> Result { Self::new_with_source( config, router_config, AdmissionSource::Workload(driver), mode, + router_mode, ) } @@ -126,18 +134,43 @@ impl DisaggRuntime { router_config: Option, admission: AdmissionSource, mode: ReplayMode, + router_mode: ReplayRouterMode, ) -> Result { - let prefill_router_config = - derive_prefill_router_config(&config.prefill_args, router_config.clone()); - let decode_router_config = derive_decode_router_config(&config.decode_args, router_config); + let (prefill_router, decode_router) = match router_mode { + ReplayRouterMode::RoundRobin => (None, None), + ReplayRouterMode::KvRouter => { + let prefill_router_config = + derive_prefill_router_config(&config.prefill_args, router_config.clone()); + let decode_router_config = + derive_decode_router_config(&config.decode_args, router_config); + ( + Some(OfflineReplayRouter::new( + &config.prefill_args, + Some(prefill_router_config), + config.num_prefill_workers, + )?), + Some(OfflineReplayRouter::new( + &config.decode_args, + Some(decode_router_config), + config.num_decode_workers, + )?), + ) + } + }; Ok(Self { now_ms: 0.0, + next_prefill_worker_idx: 0, + next_decode_worker_idx: 0, next_event_seq: 0, admission, prefill_workers: (0..config.num_prefill_workers) .map(|worker_idx| { - OfflineWorkerState::new(worker_idx, config.prefill_args.clone(), true) + OfflineWorkerState::new( + worker_idx, + config.prefill_args.clone(), + prefill_router.is_some(), + ) }) .collect(), decode_workers: (0..config.num_decode_workers) @@ -145,16 +178,8 @@ impl DisaggRuntime { OfflineWorkerState::new(worker_idx, config.decode_args.clone(), false) }) .collect(), - prefill_router: OfflineReplayRouter::new( - &config.prefill_args, - Some(prefill_router_config), - config.num_prefill_workers, - )?, - decode_router: OfflineReplayRouter::new( - &config.decode_args, - Some(decode_router_config), - config.num_decode_workers, - )?, + prefill_router, + decode_router, requests: HashMap::new(), collector: TraceCollector::default(), events: BinaryHeap::new(), @@ -173,11 +198,29 @@ impl DisaggRuntime { .sum::() + self .decode_workers - .iter() - .map(OfflineWorkerState::in_flight) - .sum::() - + self.prefill_router.pending_count() - + self.decode_router.pending_count() + .iter() + .map(OfflineWorkerState::in_flight) + .sum::() + + self + .prefill_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count) + + self + .decode_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count) + } + + fn next_prefill_worker(&mut self) -> usize { + let worker_idx = self.next_prefill_worker_idx; + self.next_prefill_worker_idx = (self.next_prefill_worker_idx + 1) % self.prefill_workers.len(); + worker_idx + } + + fn next_decode_worker(&mut self) -> usize { + let worker_idx = self.next_decode_worker_idx; + self.next_decode_worker_idx = (self.next_decode_worker_idx + 1) % self.decode_workers.len(); + worker_idx } fn record_router_pending(&mut self) { @@ -186,11 +229,19 @@ impl DisaggRuntime { self.stats.max_prefill_router_pending = self .stats .max_prefill_router_pending - .max(self.prefill_router.pending_count()); + .max( + self.prefill_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count), + ); self.stats.max_decode_router_pending = self .stats .max_decode_router_pending - .max(self.decode_router.pending_count()); + .max( + self.decode_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count), + ); } } @@ -218,19 +269,26 @@ impl DisaggRuntime { .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}")) } - fn prefill_request(state: &DisaggRequestState) -> DirectRequest { - let mut request = state.original.clone(); + fn original_request(state: &DisaggRequestState) -> Result<&DirectRequest> { + state + .original + .as_ref() + .ok_or_else(|| anyhow!("offline disagg replay request payload was already released")) + } + + fn build_prefill_request(state: &DisaggRequestState) -> Result { + let mut request = Self::original_request(state)?.clone(); request.max_output_tokens = 1; - request + Ok(request) } - fn decode_request(state: &DisaggRequestState) -> DirectRequest { - state.original.clone() + fn build_decode_request(state: &DisaggRequestState) -> Result { + Ok(Self::original_request(state)?.clone()) } fn dispatch_prefill(&mut self, uuid: Uuid, worker_idx: usize) -> Result<()> { self.validate_worker_idx(SimulationWorkerStage::Prefill, worker_idx)?; - let request = Self::prefill_request(self.state(uuid)?); + let request = Self::build_prefill_request(self.state(uuid)?)?; self.prefill_workers[worker_idx].receive_request(request); let state = self.state_mut(uuid)?; state.phase = DisaggPhase::RunningPrefill; @@ -244,7 +302,7 @@ impl DisaggRuntime { fn dispatch_decode(&mut self, uuid: Uuid, worker_idx: usize) -> Result<()> { self.validate_worker_idx(SimulationWorkerStage::Decode, worker_idx)?; - let request = Self::decode_request(self.state(uuid)?); + let request = Self::build_decode_request(self.state(uuid)?)?; self.decode_workers[worker_idx].receive_request(request); let state = self.state_mut(uuid)?; state.phase = DisaggPhase::RunningDecode; @@ -277,10 +335,20 @@ impl DisaggRuntime { } fn enqueue_decode(&mut self, uuid: Uuid) -> Result<()> { - let request = Self::decode_request(self.state(uuid)?); - let maybe_worker_idx = - self.decode_router - .submit_request_with_hashes(&request, None, self.now_ms)?; + let Some(decode_router) = self.decode_router.as_mut() else { + let worker_idx = self.next_decode_worker(); + self.dispatch_decode(uuid, worker_idx)?; + return Ok(()); + }; + let maybe_worker_idx = { + let requests = &self.requests; + let request = Self::original_request( + requests.get(&uuid).ok_or_else(|| { + anyhow!("offline disagg replay missing request state for {uuid}") + })?, + )?; + decode_router.submit_request_with_hashes(request, None, self.now_ms)? + }; self.record_router_pending(); #[cfg(test)] { @@ -312,11 +380,8 @@ impl DisaggRuntime { request.max_output_tokens, ); - let maybe_worker_idx = - self.prefill_router - .submit_request_with_hashes(&request, replay_hashes, self.now_ms)?; let state = DisaggRequestState { - original: request, + original: Some(request), #[cfg(test)] arrival_ms: arrival_time_ms, phase: DisaggPhase::QueuedPrefill, @@ -324,6 +389,20 @@ impl DisaggRuntime { decode_worker_idx: None, }; self.requests.insert(uuid, state); + let Some(prefill_router) = self.prefill_router.as_mut() else { + let worker_idx = self.next_prefill_worker(); + self.dispatch_prefill(uuid, worker_idx)?; + return Ok(uuid); + }; + let maybe_worker_idx = { + let requests = &self.requests; + let request = Self::original_request( + requests.get(&uuid).ok_or_else(|| { + anyhow!("offline disagg replay missing request state for {uuid}") + })?, + )?; + prefill_router.submit_request_with_hashes(request, replay_hashes, self.now_ms)? + }; self.record_router_pending(); if let Some(worker_idx) = maybe_worker_idx { self.dispatch_prefill(uuid, worker_idx)?; @@ -370,8 +449,11 @@ impl DisaggRuntime { } fn apply_prefill_router_events(&mut self, events: Vec) -> Result<()> { + let Some(prefill_router) = self.prefill_router.as_mut() else { + return Ok(()); + }; for event in events { - self.prefill_router.apply_event(event)?; + prefill_router.apply_event(event)?; } Ok(()) } @@ -381,8 +463,15 @@ impl DisaggRuntime { return Ok(()); } - let prefill_complete_admissions = - self.prefill_router.mark_prefill_completed(signal.uuid)?; + let Some(_) = self.prefill_router.as_ref() else { + self.enqueue_decode(signal.uuid)?; + return Ok(()); + }; + let prefill_complete_admissions = self + .prefill_router + .as_mut() + .expect("prefill router must exist after guard") + .mark_prefill_completed(signal.uuid)?; #[cfg(test)] { self.stats.prefill_marked_count += 1; @@ -390,7 +479,11 @@ impl DisaggRuntime { self.record_router_pending(); self.dispatch_prefill_admissions(prefill_complete_admissions)?; - let admissions = self.prefill_router.free(signal.uuid)?; + let admissions = self + .prefill_router + .as_mut() + .expect("prefill router must exist after guard") + .free(signal.uuid)?; #[cfg(test)] { self.stats.prefill_freed_count += 1; @@ -406,17 +499,23 @@ impl DisaggRuntime { return Ok(()); } - let admissions = self.decode_router.free(signal.uuid)?; - #[cfg(test)] - { - self.stats.decode_freed_count += 1; - } + let admissions = if let Some(decode_router) = self.decode_router.as_mut() { + let admissions = decode_router.free(signal.uuid)?; + #[cfg(test)] + { + self.stats.decode_freed_count += 1; + } + admissions + } else { + Vec::new() + }; self.record_router_pending(); - let state = self.state_mut(signal.uuid)?; - state.phase = DisaggPhase::Done; if let AdmissionSource::Workload(driver) = &mut self.admission { driver.on_complete(signal.uuid, self.now_ms)?; } + let state = self.state_mut(signal.uuid)?; + state.phase = DisaggPhase::Done; + state.original = None; self.dispatch_decode_admissions(admissions)?; Ok(()) } @@ -742,10 +841,12 @@ pub(crate) fn simulate_trace_disagg( router_config: Option, requests: Vec, arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, ) -> Result { let pending = normalize_trace_requests(requests, arrival_speedup_ratio)?; let (collector, _) = - DisaggRuntime::new(&config, router_config, pending, ReplayMode::Trace)?.run()?; + DisaggRuntime::new(&config, router_config, pending, ReplayMode::Trace, router_mode)? + .run()?; Ok(collector.finish()) } @@ -754,6 +855,7 @@ pub(crate) fn simulate_concurrency_disagg( router_config: Option, requests: Vec, max_in_flight: usize, + router_mode: ReplayRouterMode, ) -> Result { let pending = VecDeque::from(requests); let (collector, _) = DisaggRuntime::new( @@ -761,6 +863,7 @@ pub(crate) fn simulate_concurrency_disagg( router_config, pending, ReplayMode::Concurrency { max_in_flight }, + router_mode, )? .run()?; Ok(collector.finish()) @@ -770,10 +873,17 @@ pub(crate) fn simulate_trace_workload_disagg( config: OfflineDisaggReplayConfig, router_config: Option, trace: Trace, + router_mode: ReplayRouterMode, ) -> Result { let driver = WorkloadDriver::new_trace(trace)?; - let (collector, _) = - DisaggRuntime::new_workload(&config, router_config, driver, ReplayMode::Trace)?.run()?; + let (collector, _) = DisaggRuntime::new_workload( + &config, + router_config, + driver, + ReplayMode::Trace, + router_mode, + )? + .run()?; Ok(collector.finish()) } @@ -782,6 +892,7 @@ pub(crate) fn simulate_concurrency_workload_disagg( router_config: Option, trace: Trace, max_in_flight: usize, + router_mode: ReplayRouterMode, ) -> Result { let driver = WorkloadDriver::new_concurrency(trace)?; let (collector, _) = DisaggRuntime::new_workload( @@ -789,6 +900,7 @@ pub(crate) fn simulate_concurrency_workload_disagg( router_config, driver, ReplayMode::Concurrency { max_in_flight }, + router_mode, )? .run()?; Ok(collector.finish()) @@ -800,9 +912,10 @@ fn run_trace_collect( requests: Vec, router_config: Option, arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, ) -> (TraceCollector, DisaggRuntimeStats) { let pending = normalize_trace_requests(requests, arrival_speedup_ratio).unwrap(); - DisaggRuntime::new(config, router_config, pending, ReplayMode::Trace) + DisaggRuntime::new(config, router_config, pending, ReplayMode::Trace, router_mode) .unwrap() .run() .unwrap() @@ -814,12 +927,14 @@ fn run_concurrency_collect( requests: Vec, router_config: Option, max_in_flight: usize, + router_mode: ReplayRouterMode, ) -> (TraceCollector, DisaggRuntimeStats) { DisaggRuntime::new( config, router_config, VecDeque::from(requests), ReplayMode::Concurrency { max_in_flight }, + router_mode, ) .unwrap() .run() @@ -896,12 +1011,21 @@ mod tests { assert!(!decode.router_track_prefill_tokens); } - #[test] - fn test_trace_smoke_reports_decode_only_tokens() { + #[rstest::rstest] + #[case(ReplayRouterMode::RoundRobin)] + #[case(ReplayRouterMode::KvRouter)] + fn test_trace_smoke_reports_decode_only_tokens(#[case] router_mode: ReplayRouterMode) { let config = disagg_config(); let requests = vec![request(1, 128, 3, 5.0)]; - let (collector, stats) = run_trace_collect(&config, requests, Some(router_config()), 1.0); + let router_config = (router_mode == ReplayRouterMode::KvRouter).then(router_config); + let (collector, stats) = run_trace_collect( + &config, + requests, + router_config, + 1.0, + router_mode, + ); let snapshot = collector.snapshot(Uuid::from_u128(1)).unwrap(); let report = collector.finish(); @@ -916,12 +1040,21 @@ mod tests { ); } - #[test] - fn test_prefill_and_decode_use_separate_worker_pools() { + #[rstest::rstest] + #[case(ReplayRouterMode::RoundRobin)] + #[case(ReplayRouterMode::KvRouter)] + fn test_prefill_and_decode_use_separate_worker_pools(#[case] router_mode: ReplayRouterMode) { let config = disagg_config(); let requests = vec![request(1, 128, 2, 0.0), request(2, 128, 2, 10.0)]; - let (_, stats) = run_trace_collect(&config, requests, Some(router_config()), 1.0); + let router_config = (router_mode == ReplayRouterMode::KvRouter).then(router_config); + let (_, stats) = run_trace_collect( + &config, + requests, + router_config, + 1.0, + router_mode, + ); for uuid in [Uuid::from_u128(1), Uuid::from_u128(2)] { assert!(stats.prefill_assignments.contains_key(&uuid)); @@ -934,7 +1067,13 @@ mod tests { let config = disagg_config(); let requests = vec![request(1, 128, 2, 0.0), request(2, 128, 2, 100.0)]; - let (_, stats) = run_trace_collect(&config, requests, Some(router_config()), 1.0); + let (_, stats) = run_trace_collect( + &config, + requests, + Some(router_config()), + 1.0, + ReplayRouterMode::KvRouter, + ); assert_eq!( stats.prefill_assignments[&Uuid::from_u128(1)], @@ -942,8 +1081,12 @@ mod tests { ); } - #[test] - fn test_concurrency_backfill_waits_for_decode_completion() { + #[rstest::rstest] + #[case(ReplayRouterMode::RoundRobin)] + #[case(ReplayRouterMode::KvRouter)] + fn test_concurrency_backfill_waits_for_decode_completion( + #[case] router_mode: ReplayRouterMode, + ) { let config = disagg_config(); let requests = vec![ DirectRequest { @@ -962,7 +1105,14 @@ mod tests { }, ]; - let (collector, _) = run_concurrency_collect(&config, requests, Some(router_config()), 1); + let router_config = (router_mode == ReplayRouterMode::KvRouter).then(router_config); + let (collector, _) = run_concurrency_collect( + &config, + requests, + router_config, + 1, + router_mode, + ); let first = collector.snapshot(Uuid::from_u128(1)).unwrap(); let second = collector.snapshot(Uuid::from_u128(2)).unwrap(); @@ -975,7 +1125,13 @@ mod tests { let config = disagg_config(); let requests = vec![request(1, 128, 2, 0.0)]; - let (_, stats) = run_trace_collect(&config, requests, Some(router_config()), 1.0); + let (_, stats) = run_trace_collect( + &config, + requests, + Some(router_config()), + 1.0, + ReplayRouterMode::KvRouter, + ); assert_eq!(stats.prefill_marked_count, 1); assert_eq!(stats.prefill_freed_count, 1); diff --git a/lib/mocker/src/replay/offline/mod.rs b/lib/mocker/src/replay/offline/mod.rs index 705bdce21fa5..64480a579bd1 100644 --- a/lib/mocker/src/replay/offline/mod.rs +++ b/lib/mocker/src/replay/offline/mod.rs @@ -101,8 +101,15 @@ pub(crate) fn simulate_trace_disagg( router_config: Option, requests: Vec, arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, ) -> anyhow::Result { - disagg::simulate_trace_disagg(config, router_config, requests, arrival_speedup_ratio) + disagg::simulate_trace_disagg( + config, + router_config, + requests, + arrival_speedup_ratio, + router_mode, + ) } pub(crate) fn simulate_concurrency_disagg( @@ -110,16 +117,24 @@ pub(crate) fn simulate_concurrency_disagg( router_config: Option, requests: Vec, max_in_flight: usize, + router_mode: ReplayRouterMode, ) -> anyhow::Result { - disagg::simulate_concurrency_disagg(config, router_config, requests, max_in_flight) + disagg::simulate_concurrency_disagg( + config, + router_config, + requests, + max_in_flight, + router_mode, + ) } pub(crate) fn simulate_trace_workload_disagg( config: OfflineDisaggReplayConfig, router_config: Option, trace: Trace, + router_mode: ReplayRouterMode, ) -> anyhow::Result { - disagg::simulate_trace_workload_disagg(config, router_config, trace) + disagg::simulate_trace_workload_disagg(config, router_config, trace, router_mode) } pub(crate) fn simulate_concurrency_workload_disagg( @@ -127,6 +142,13 @@ pub(crate) fn simulate_concurrency_workload_disagg( router_config: Option, trace: Trace, max_in_flight: usize, + router_mode: ReplayRouterMode, ) -> anyhow::Result { - disagg::simulate_concurrency_workload_disagg(config, router_config, trace, max_in_flight) + disagg::simulate_concurrency_workload_disagg( + config, + router_config, + trace, + max_in_flight, + router_mode, + ) } diff --git a/lib/mocker/src/replay/validate.rs b/lib/mocker/src/replay/validate.rs index fab2e7d128b0..3ebafe1056ff 100644 --- a/lib/mocker/src/replay/validate.rs +++ b/lib/mocker/src/replay/validate.rs @@ -156,24 +156,18 @@ fn validate_disagg_args(config: &OfflineDisaggReplayConfig, mode: &str) -> Resul pub(super) fn validate_offline_disagg_replay_args( config: &OfflineDisaggReplayConfig, - router_mode: ReplayRouterMode, + _router_mode: ReplayRouterMode, ) -> Result<()> { - if router_mode != ReplayRouterMode::KvRouter { - bail!("offline disagg replay requires router_mode=kv_router"); - } validate_disagg_args(config, "trace replay") } pub(super) fn validate_offline_disagg_concurrency_args( config: &OfflineDisaggReplayConfig, max_in_flight: usize, - router_mode: ReplayRouterMode, + _router_mode: ReplayRouterMode, ) -> Result<()> { if max_in_flight == 0 { bail!("concurrency replay requires max_in_flight >= 1"); } - if router_mode != ReplayRouterMode::KvRouter { - bail!("offline disagg replay requires router_mode=kv_router"); - } validate_disagg_args(config, "concurrency replay") } From d85a264f6e0965c4389a98368ecfd75ad89b1784 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 21:26:27 -0700 Subject: [PATCH 07/24] cache replay turns and sort stats once Signed-off-by: PeaBrane --- lib/mocker/src/loadgen/driver.rs | 63 ++++++++++----- lib/mocker/src/loadgen/trace.rs | 17 ++-- lib/mocker/src/replay/collector.rs | 51 ++++++------ lib/mocker/src/replay/offline/disagg.rs | 102 ++++++++++-------------- lib/mocker/src/replay/offline/mod.rs | 8 +- 5 files changed, 127 insertions(+), 114 deletions(-) diff --git a/lib/mocker/src/loadgen/driver.rs b/lib/mocker/src/loadgen/driver.rs index ce04d46ff309..6a0c6be8e7db 100644 --- a/lib/mocker/src/loadgen/driver.rs +++ b/lib/mocker/src/loadgen/driver.rs @@ -7,7 +7,8 @@ use std::collections::{BinaryHeap, HashMap}; use anyhow::{Result, anyhow, bail}; use uuid::Uuid; -use super::types::{ReadyTurn, Trace, TurnTrace}; +use super::types::{ReadyTurn, ReplayRequestHashes, Trace}; +use crate::common::protocols::DirectRequest; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum DriverMode { @@ -18,12 +19,20 @@ enum DriverMode { #[derive(Debug)] struct SessionRuntime { session_id: String, - turns: Vec, + turns: Vec, next_turn_index: usize, next_ready_at_ms: Option, in_flight: Option, } +#[derive(Debug)] +struct TurnRuntime { + tokens: Vec, + max_output_tokens: usize, + delay_after_previous_ms: f64, + replay_hashes: ReplayRequestHashes, +} + #[derive(Debug)] struct InFlightTurn { session_index: usize, @@ -66,7 +75,6 @@ impl PartialOrd for ReadySession { #[derive(Debug)] pub struct WorkloadDriver { mode: DriverMode, - block_size: usize, sessions: Vec, in_flight: HashMap, ready_sessions: BinaryHeap, @@ -82,20 +90,36 @@ impl WorkloadDriver { } fn new(trace: Trace, mode: DriverMode) -> Result { + let block_size = trace.block_size; let sessions: Vec = trace .sessions .into_iter() - .map(|session| SessionRuntime { - session_id: session.session_id, - turns: session.turns, - next_turn_index: 0, - next_ready_at_ms: Some(match mode { + .map(|session| -> Result { + let next_ready_at_ms = Some(match mode { DriverMode::Trace => session.first_arrival_timestamp_ms.unwrap_or(0.0), DriverMode::Concurrency => 0.0, - }), - in_flight: None, + }); + let turns = session + .turns + .into_iter() + .map(|turn| -> Result { + Ok(TurnRuntime { + tokens: turn.synthesize_tokens(block_size)?, + max_output_tokens: turn.max_output_tokens, + delay_after_previous_ms: turn.delay_after_previous_ms, + replay_hashes: turn.to_replay_hashes(block_size)?, + }) + }) + .collect::>>()?; + Ok(SessionRuntime { + session_id: session.session_id, + turns, + next_turn_index: 0, + next_ready_at_ms, + in_flight: None, + }) }) - .collect(); + .collect::>>()?; let ready_sessions = sessions .iter() @@ -111,7 +135,6 @@ impl WorkloadDriver { Ok(Self { mode, - block_size: trace.block_size, sessions, in_flight: HashMap::new(), ready_sessions, @@ -146,16 +169,18 @@ impl WorkloadDriver { .next_ready_at_ms .expect("ready session must have a timestamp"); let request_uuid = Uuid::new_v4(); - let replay_hashes = session.turns[turn_index] - .to_replay_hashes(self.block_size) - .expect("validated trace should always synthesize replay hashes"); + let turn = &session.turns[turn_index]; let arrival_timestamp_ms = match self.mode { DriverMode::Trace => Some(scheduled_ready_at_ms), DriverMode::Concurrency => None, }; - let request = session.turns[turn_index] - .to_direct_request(self.block_size, request_uuid, arrival_timestamp_ms) - .expect("validated trace should always synthesize into a direct request"); + let request = DirectRequest { + tokens: turn.tokens.clone(), + max_output_tokens: turn.max_output_tokens, + uuid: Some(request_uuid), + dp_rank: 0, + arrival_timestamp_ms, + }; session.in_flight = Some(request_uuid); session.next_ready_at_ms = None; self.in_flight.insert( @@ -170,7 +195,7 @@ impl WorkloadDriver { session_id: session.session_id.clone(), turn_index, scheduled_ready_at_ms, - replay_hashes: Some(replay_hashes), + replay_hashes: Some(turn.replay_hashes.clone()), request, }); } diff --git a/lib/mocker/src/loadgen/trace.rs b/lib/mocker/src/loadgen/trace.rs index a9931dc18a06..40818daa7c95 100644 --- a/lib/mocker/src/loadgen/trace.rs +++ b/lib/mocker/src/loadgen/trace.rs @@ -59,12 +59,7 @@ impl TurnTrace { Ok(()) } - pub fn to_direct_request( - &self, - block_size: usize, - request_uuid: Uuid, - arrival_timestamp_ms: Option, - ) -> Result { + pub(crate) fn synthesize_tokens(&self, block_size: usize) -> Result> { self.validate_block_size_and_capacity(block_size)?; let mut tokens = Vec::with_capacity(self.input_length); @@ -85,6 +80,16 @@ impl TurnTrace { ); } + Ok(tokens) + } + + pub fn to_direct_request( + &self, + block_size: usize, + request_uuid: Uuid, + arrival_timestamp_ms: Option, + ) -> Result { + let tokens = self.synthesize_tokens(block_size)?; Ok(DirectRequest { tokens, max_output_tokens: self.max_output_tokens, diff --git a/lib/mocker/src/replay/collector.rs b/lib/mocker/src/replay/collector.rs index 39c1ca880a6f..79ed61f71be3 100644 --- a/lib/mocker/src/replay/collector.rs +++ b/lib/mocker/src/replay/collector.rs @@ -309,6 +309,7 @@ impl TraceCollector { } let duration_s = (duration_ms / 1000.0).max(1e-9); + let itl_distribution = build_distribution_stats(&itls); TraceSimulationReport { request_counts: TraceRequestCounts { num_requests: requests.len(), @@ -335,8 +336,8 @@ impl TraceCollector { ttst: build_distribution_stats(&ttsts), tpot: build_distribution_stats(&tpots), itl: TraceInterTokenLatencyStats { - distribution: build_distribution_stats(&itls), - max_ms: max_value(&itls), + max_ms: itl_distribution.max_ms, + distribution: itl_distribution, }, e2e: build_distribution_stats(&e2e_latencies), output_token_throughput_per_user: build_distribution_stats( @@ -386,39 +387,41 @@ fn mean(values: &[f64]) -> f64 { } } -fn max_value(values: &[f64]) -> f64 { - values.iter().copied().reduce(f64::max).unwrap_or(0.0) -} - fn build_distribution_stats(values: &[f64]) -> TraceDistributionStats { + if values.is_empty() { + return TraceDistributionStats { + mean_ms: 0.0, + min_ms: 0.0, + max_ms: 0.0, + median_ms: 0.0, + p75_ms: 0.0, + p90_ms: 0.0, + p95_ms: 0.0, + p99_ms: 0.0, + std_ms: 0.0, + }; + } + + let mut sorted = values.to_vec(); + sorted.sort_by(|left, right| left.total_cmp(right)); TraceDistributionStats { mean_ms: mean(values), - min_ms: min_value(values), - max_ms: max_value(values), - median_ms: percentile(values, 50.0), - p75_ms: percentile(values, 75.0), - p90_ms: percentile(values, 90.0), - p95_ms: percentile(values, 95.0), - p99_ms: percentile(values, 99.0), + min_ms: sorted[0], + max_ms: *sorted.last().expect("sorted values must be non-empty"), + median_ms: percentile_sorted(&sorted, 50.0), + p75_ms: percentile_sorted(&sorted, 75.0), + p90_ms: percentile_sorted(&sorted, 90.0), + p95_ms: percentile_sorted(&sorted, 95.0), + p99_ms: percentile_sorted(&sorted, 99.0), std_ms: std_dev(values), } } -fn percentile(values: &[f64], percentile: f64) -> f64 { - if values.is_empty() { - return 0.0; - } - - let mut sorted = values.to_vec(); - sorted.sort_by(|left, right| left.total_cmp(right)); +fn percentile_sorted(sorted: &[f64], percentile: f64) -> f64 { let rank = ((sorted.len() - 1) as f64 * percentile / 100.0).round() as usize; sorted[rank.min(sorted.len() - 1)] } -fn min_value(values: &[f64]) -> f64 { - values.iter().copied().reduce(f64::min).unwrap_or(0.0) -} - fn std_dev(values: &[f64]) -> f64 { if values.is_empty() { return 0.0; diff --git a/lib/mocker/src/replay/offline/disagg.rs b/lib/mocker/src/replay/offline/disagg.rs index 43ee6c02f4b8..9cc7fc9d3445 100644 --- a/lib/mocker/src/replay/offline/disagg.rs +++ b/lib/mocker/src/replay/offline/disagg.rs @@ -198,9 +198,9 @@ impl DisaggRuntime { .sum::() + self .decode_workers - .iter() - .map(OfflineWorkerState::in_flight) - .sum::() + .iter() + .map(OfflineWorkerState::in_flight) + .sum::() + self .prefill_router .as_ref() @@ -213,7 +213,8 @@ impl DisaggRuntime { fn next_prefill_worker(&mut self) -> usize { let worker_idx = self.next_prefill_worker_idx; - self.next_prefill_worker_idx = (self.next_prefill_worker_idx + 1) % self.prefill_workers.len(); + self.next_prefill_worker_idx = + (self.next_prefill_worker_idx + 1) % self.prefill_workers.len(); worker_idx } @@ -226,22 +227,16 @@ impl DisaggRuntime { fn record_router_pending(&mut self) { #[cfg(test)] { - self.stats.max_prefill_router_pending = self - .stats - .max_prefill_router_pending - .max( - self.prefill_router - .as_ref() - .map_or(0, OfflineReplayRouter::pending_count), - ); - self.stats.max_decode_router_pending = self - .stats - .max_decode_router_pending - .max( - self.decode_router - .as_ref() - .map_or(0, OfflineReplayRouter::pending_count), - ); + self.stats.max_prefill_router_pending = self.stats.max_prefill_router_pending.max( + self.prefill_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count), + ); + self.stats.max_decode_router_pending = self.stats.max_decode_router_pending.max( + self.decode_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count), + ); } } @@ -342,11 +337,9 @@ impl DisaggRuntime { }; let maybe_worker_idx = { let requests = &self.requests; - let request = Self::original_request( - requests.get(&uuid).ok_or_else(|| { - anyhow!("offline disagg replay missing request state for {uuid}") - })?, - )?; + let request = Self::original_request(requests.get(&uuid).ok_or_else(|| { + anyhow!("offline disagg replay missing request state for {uuid}") + })?)?; decode_router.submit_request_with_hashes(request, None, self.now_ms)? }; self.record_router_pending(); @@ -396,11 +389,9 @@ impl DisaggRuntime { }; let maybe_worker_idx = { let requests = &self.requests; - let request = Self::original_request( - requests.get(&uuid).ok_or_else(|| { - anyhow!("offline disagg replay missing request state for {uuid}") - })?, - )?; + let request = Self::original_request(requests.get(&uuid).ok_or_else(|| { + anyhow!("offline disagg replay missing request state for {uuid}") + })?)?; prefill_router.submit_request_with_hashes(request, replay_hashes, self.now_ms)? }; self.record_router_pending(); @@ -844,9 +835,14 @@ pub(crate) fn simulate_trace_disagg( router_mode: ReplayRouterMode, ) -> Result { let pending = normalize_trace_requests(requests, arrival_speedup_ratio)?; - let (collector, _) = - DisaggRuntime::new(&config, router_config, pending, ReplayMode::Trace, router_mode)? - .run()?; + let (collector, _) = DisaggRuntime::new( + &config, + router_config, + pending, + ReplayMode::Trace, + router_mode, + )? + .run()?; Ok(collector.finish()) } @@ -915,10 +911,16 @@ fn run_trace_collect( router_mode: ReplayRouterMode, ) -> (TraceCollector, DisaggRuntimeStats) { let pending = normalize_trace_requests(requests, arrival_speedup_ratio).unwrap(); - DisaggRuntime::new(config, router_config, pending, ReplayMode::Trace, router_mode) - .unwrap() - .run() - .unwrap() + DisaggRuntime::new( + config, + router_config, + pending, + ReplayMode::Trace, + router_mode, + ) + .unwrap() + .run() + .unwrap() } #[cfg(test)] @@ -1019,13 +1021,8 @@ mod tests { let requests = vec![request(1, 128, 3, 5.0)]; let router_config = (router_mode == ReplayRouterMode::KvRouter).then(router_config); - let (collector, stats) = run_trace_collect( - &config, - requests, - router_config, - 1.0, - router_mode, - ); + let (collector, stats) = + run_trace_collect(&config, requests, router_config, 1.0, router_mode); let snapshot = collector.snapshot(Uuid::from_u128(1)).unwrap(); let report = collector.finish(); @@ -1048,13 +1045,7 @@ mod tests { let requests = vec![request(1, 128, 2, 0.0), request(2, 128, 2, 10.0)]; let router_config = (router_mode == ReplayRouterMode::KvRouter).then(router_config); - let (_, stats) = run_trace_collect( - &config, - requests, - router_config, - 1.0, - router_mode, - ); + let (_, stats) = run_trace_collect(&config, requests, router_config, 1.0, router_mode); for uuid in [Uuid::from_u128(1), Uuid::from_u128(2)] { assert!(stats.prefill_assignments.contains_key(&uuid)); @@ -1106,13 +1097,8 @@ mod tests { ]; let router_config = (router_mode == ReplayRouterMode::KvRouter).then(router_config); - let (collector, _) = run_concurrency_collect( - &config, - requests, - router_config, - 1, - router_mode, - ); + let (collector, _) = + run_concurrency_collect(&config, requests, router_config, 1, router_mode); let first = collector.snapshot(Uuid::from_u128(1)).unwrap(); let second = collector.snapshot(Uuid::from_u128(2)).unwrap(); diff --git a/lib/mocker/src/replay/offline/mod.rs b/lib/mocker/src/replay/offline/mod.rs index 64480a579bd1..5bb23f1d6850 100644 --- a/lib/mocker/src/replay/offline/mod.rs +++ b/lib/mocker/src/replay/offline/mod.rs @@ -119,13 +119,7 @@ pub(crate) fn simulate_concurrency_disagg( max_in_flight: usize, router_mode: ReplayRouterMode, ) -> anyhow::Result { - disagg::simulate_concurrency_disagg( - config, - router_config, - requests, - max_in_flight, - router_mode, - ) + disagg::simulate_concurrency_disagg(config, router_config, requests, max_in_flight, router_mode) } pub(crate) fn simulate_trace_workload_disagg( From 00e5b26864f73901823c4516409d9a7fb679d239 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 21:27:52 -0700 Subject: [PATCH 08/24] docs: clarify free subsumes prefill cleanup Signed-off-by: PeaBrane --- lib/kv-router/src/sequences/multi_worker.rs | 4 ++++ lib/kv-router/src/sequences/single.rs | 5 ++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/lib/kv-router/src/sequences/multi_worker.rs b/lib/kv-router/src/sequences/multi_worker.rs index 412f6bc71db7..3ef51298fe9e 100644 --- a/lib/kv-router/src/sequences/multi_worker.rs +++ b/lib/kv-router/src/sequences/multi_worker.rs @@ -533,6 +533,10 @@ impl ActiveSequencesMultiWorker

{ /// /// Note: This operation is idempotent. Calling it multiple times for the same request /// will log a warning but not return an error (double free is allowed). + /// + /// This also performs the underlying prefill-complete cleanup via + /// [`ActiveSequences::free`], so callers do not need to call + /// [`Self::mark_prefill_completed`] before freeing a completed request. pub async fn free(&self, request_id: &RequestId) -> Result<(), SequenceError> { if !self.request_to_worker.contains_key(request_id) { tracing::debug!("Request {request_id} not found, already freed (idempotent)"); diff --git a/lib/kv-router/src/sequences/single.rs b/lib/kv-router/src/sequences/single.rs index f27cadaefa24..f74116a1b192 100644 --- a/lib/kv-router/src/sequences/single.rs +++ b/lib/kv-router/src/sequences/single.rs @@ -271,7 +271,10 @@ impl ActiveSequences { self.new_blocks(token_sequence) + self.active_blocks() } - /// Free all blocks associated with a request + /// Free all blocks associated with a request. + /// + /// This implicitly calls [`Self::mark_prefill_completed`] first, so callers do not need + /// to invoke both when the request is finishing. pub fn free(&mut self, request_id: &RequestId) -> usize { self.mark_prefill_completed(request_id); From cdfb4baa0df09e1667bb7e40ce265298f4dc9948 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 21:46:05 -0700 Subject: [PATCH 09/24] fix replay mypy's identity crisis Signed-off-by: PeaBrane --- lib/bindings/python/src/dynamo/replay/main.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/lib/bindings/python/src/dynamo/replay/main.py b/lib/bindings/python/src/dynamo/replay/main.py index d701e8cbbbf5..60448ccf270a 100644 --- a/lib/bindings/python/src/dynamo/replay/main.py +++ b/lib/bindings/python/src/dynamo/replay/main.py @@ -4,25 +4,32 @@ from __future__ import annotations import argparse +import importlib import json import os import sys from collections.abc import Sequence from pathlib import Path from types import SimpleNamespace -from typing import Any +from typing import Protocol os.environ.setdefault("DYNAMO_SKIP_PYTHON_LOG_INIT", "1") from dynamo.llm import KvRouterConfig, MockEngineArgs +from dynamo.replay import run_synthetic_trace_replay, run_trace_replay +from dynamo.replay.reporting import format_report_table, write_report_json + + +class PlannerProfileDataResult(Protocol): + npz_path: Path | None -try: - from dynamo.mocker.args import ( - resolve_planner_profile_data, # type: ignore[import-not-found] - ) -except ImportError: - def resolve_planner_profile_data(planner_profile_data: Path | None) -> Any: +def resolve_planner_profile_data( + planner_profile_data: Path | None, +) -> PlannerProfileDataResult: + try: + module = importlib.import_module("dynamo.mocker.args") + except ImportError: if planner_profile_data is None: return SimpleNamespace(npz_path=None) return SimpleNamespace( @@ -30,10 +37,7 @@ def resolve_planner_profile_data(planner_profile_data: Path | None) -> Any: if planner_profile_data.suffix == ".npz" else None ) - - -from dynamo.replay import run_synthetic_trace_replay, run_trace_replay -from dynamo.replay.reporting import format_report_table, write_report_json + return module.resolve_planner_profile_data(planner_profile_data) def _load_engine_args(raw_args: str | None): From 253a42766f55bdee08a8a60e49e3994ce36b659b Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 21:52:26 -0700 Subject: [PATCH 10/24] refactor prefill router without moving guardrails Signed-off-by: PeaBrane --- lib/llm/src/kv_router/prefill_router.rs | 773 ------------------ .../kv_router/prefill_router/activation.rs | 188 +++++ .../src/kv_router/prefill_router/execution.rs | 322 ++++++++ lib/llm/src/kv_router/prefill_router/inner.rs | 55 ++ lib/llm/src/kv_router/prefill_router/mod.rs | 259 ++++++ lib/llm/src/kv_router/prefill_router/types.rs | 56 ++ 6 files changed, 880 insertions(+), 773 deletions(-) delete mode 100644 lib/llm/src/kv_router/prefill_router.rs create mode 100644 lib/llm/src/kv_router/prefill_router/activation.rs create mode 100644 lib/llm/src/kv_router/prefill_router/execution.rs create mode 100644 lib/llm/src/kv_router/prefill_router/inner.rs create mode 100644 lib/llm/src/kv_router/prefill_router/mod.rs create mode 100644 lib/llm/src/kv_router/prefill_router/types.rs diff --git a/lib/llm/src/kv_router/prefill_router.rs b/lib/llm/src/kv_router/prefill_router.rs deleted file mode 100644 index b3448c44461e..000000000000 --- a/lib/llm/src/kv_router/prefill_router.rs +++ /dev/null @@ -1,773 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use std::collections::HashSet; -use std::sync::{Arc, OnceLock}; - -use anyhow::Result; -use futures::StreamExt; -use tokio::sync::{OwnedSemaphorePermit, oneshot}; -use tokio_util::sync::CancellationToken; -use tracing::Instrument; - -use dynamo_kv_router::{ - config::{KvRouterConfig, RouterConfigOverride}, - protocols::{BlockExtraInfo, WorkerId}, -}; -use dynamo_runtime::{ - component::Endpoint, - pipeline::{ - AsyncEngine, AsyncEngineContextProvider, Context, ManyOut, Operator, PushRouter, - RouterMode, ServerStreamingEngine, SingleIn, async_trait, - }, - protocols::{EndpointId, annotated::Annotated, maybe_error::MaybeError}, -}; - -use crate::{ - discovery::ModelManager, - kv_router::KvPushRouter, - protocols::common::llm_backend::{LLMEngineOutput, PreprocessedRequest}, - protocols::common::preprocessor::{BootstrapInfo, PrefillResult}, - protocols::common::timing::{RequestPhase, RequestTracker, WORKER_TYPE_PREFILL}, -}; - -/// Errors that can occur during prefill routing -#[derive(Debug, thiserror::Error)] -pub enum PrefillError { - /// Prefill router has not been activated yet - #[error("Prefill router not yet activated")] - NotActivated, - - /// TODO: Separate prefill worker error from prefill router error - /// Error during prefill execution - #[error("Prefill execution failed: {0}")] - PrefillError( - String, - #[source] Option>, - ), - - /// Disaggregated params not found in prefill response - #[error("No disaggregated params in prefill response: {0}")] - NoDisaggregatedParams(String), -} - -/// Result of the prefill phase in `generate()`. -enum PrefillOutcome { - /// Bootstrap optimization: prefill spawned in background, bootstrap info ready - Bootstrap(BootstrapInfo), - /// Synchronous prefill completed with result - Completed(PrefillResult), -} - -fn build_decode_router_override( - existing_override: Option, -) -> RouterConfigOverride { - RouterConfigOverride { - overlap_score_weight: Some(0.0), - assume_kv_reuse: Some(false), - track_prefill_tokens: Some(false), - ..existing_override.unwrap_or_default() - } -} - -/// The inner router used by PrefillRouter -#[derive(Clone)] -enum InnerPrefillRouter { - /// KV-aware routing using KvPushRouter - KvRouter(Arc), - /// Simple routing (RoundRobin, Random, Direct) - /// Note: Per-worker metrics (active_prefill_tokens, active_decode_blocks) are only - /// available in KV routing mode where the router has actual bookkeeping. - SimpleRouter(Arc>>), -} - -impl InnerPrefillRouter { - /// Generate with optional direct routing to specific worker. - /// For KvRouter, target_worker is ignored since prefill_worker_id is already set on the request. - /// For SimpleRouter, target_worker triggers direct routing via router.direct(). - async fn generate_to_worker( - &self, - request: SingleIn, - target_worker: Option, - ) -> Result>> { - match (self, target_worker) { - // KvRouter: prefill_worker_id already set on request, KvPushRouter::select_worker uses it - (InnerPrefillRouter::KvRouter(router), _) => router.generate(request).await, - (InnerPrefillRouter::SimpleRouter(router), Some(worker_id)) => { - router.direct(request, worker_id).await - } - (InnerPrefillRouter::SimpleRouter(router), None) => router.generate(request).await, - } - } - - /// Select next worker (for non-KV modes only) - fn select_next_worker(&self) -> Option { - match self { - InnerPrefillRouter::SimpleRouter(router) => router.select_next_worker(), - InnerPrefillRouter::KvRouter(_) => None, - } - } -} - -/// PrefillRouter is a forward-only operator that sits between Migration and the decode router. -/// It optionally calls a prefill worker before routing to decode, extracting disaggregated_params -/// from the prefill response and injecting them into the decode request. -/// -/// Modes: -/// - Query-only: `query_instance_id` annotation present → returns worker IDs without execution -/// - Pre-routed: `prefill_worker_id`/`decode_worker_id` set → routes to specified workers -/// - Normal: Worker IDs determined by router based on KV cache state -pub struct PrefillRouter { - prefill_router: OnceLock, - model_manager: Arc, - endpoint_id: OnceLock, - cancel_token: CancellationToken, - router_mode: RouterMode, - enforce_disagg: bool, - /// Model name used to look up the worker monitor for prefill client registration - model_name: String, - /// Namespace used to look up the correct WorkerSet's worker monitor - namespace: String, - is_eagle: bool, -} - -impl PrefillRouter { - /// Create a disabled prefill router that will never activate (passthrough only) - pub fn disabled( - model_manager: Arc, - router_mode: RouterMode, - enforce_disagg: bool, - ) -> Arc { - Arc::new(Self { - prefill_router: OnceLock::new(), - model_manager, - endpoint_id: OnceLock::new(), - cancel_token: CancellationToken::new(), - router_mode, - enforce_disagg, - model_name: String::new(), // Not used for disabled router - namespace: String::new(), // Not used for disabled router - is_eagle: false, - }) - } - - #[expect(clippy::too_many_arguments)] - pub fn new( - activation_rx: oneshot::Receiver, - model_manager: Arc, - router_mode: RouterMode, - kv_cache_block_size: u32, - kv_router_config: Option, - enforce_disagg: bool, - model_name: String, - namespace: String, - is_eagle: bool, - ) -> Arc { - let prefill_router = OnceLock::new(); - let cancel_token = CancellationToken::new(); - - let router = Arc::new(Self { - prefill_router, - model_manager: model_manager.clone(), - endpoint_id: OnceLock::new(), - cancel_token: cancel_token.clone(), - router_mode, - enforce_disagg, - model_name, - namespace, - is_eagle, - }); - - // Spawn background task to wait for activation - let router_clone = router.clone(); - tokio::spawn(async move { - tokio::select! { - result = activation_rx => { - let Ok(endpoint) = result else { - tracing::debug!("Prefill router activation channel closed without receiving endpoint"); - return; - }; - - if let Err(e) = router_clone.activate( - endpoint, - model_manager, - kv_cache_block_size, - kv_router_config, - ).await { - tracing::error!(error = %e, "Failed to activate prefill router"); - } - } - _ = cancel_token.cancelled() => { - tracing::debug!("Prefill router activation cancelled"); - } - } - }); - - router - } - - /// Activate the prefill router with the provided endpoint - async fn activate( - &self, - endpoint: Endpoint, - model_manager: Arc, - kv_cache_block_size: u32, - kv_router_config: Option, - ) -> Result<()> { - tracing::info!( - router_mode = ?self.router_mode, - "Activating prefill router" - ); - - // Store endpoint_id for later use in resolve_prefill_worker - let _ = self.endpoint_id.set(endpoint.id()); - - // Start runtime config watcher for this endpoint (needed for get_disaggregated_endpoint) - // This must be done before creating the router so bootstrap info is available - model_manager - .get_or_create_runtime_config_watcher(&endpoint) - .await?; - - let inner_router = if self.router_mode.is_kv_routing() { - // Create KV chooser using the endpoint (this is a prefill router) - let kv_chooser = model_manager - .kv_chooser_for( - &endpoint, - kv_cache_block_size, - kv_router_config, - WORKER_TYPE_PREFILL, - Some(self.model_name.clone()), - self.is_eagle, - ) - .await?; - - // Extract client from kv_chooser to ensure shared state - let client = kv_chooser.client().clone(); - - // Register prefill client with worker monitor for TTFT metric cleanup in disaggregated mode - if let Some(monitor) = - model_manager.get_worker_monitor_for_namespace(&self.model_name, &self.namespace) - { - monitor.set_prefill_client(client.clone()); - } - - // Build the PushRouter for prefill with KV mode using the shared client - let push_router = PushRouter::>::from_client_with_threshold( - client, - RouterMode::KV, - None, // busy_threshold - None, // worker_monitor - ) - .await?; - - // Wrap it in KvPushRouter - InnerPrefillRouter::KvRouter(Arc::new(KvPushRouter::new(push_router, kv_chooser))) - } else { - // Create client for simple router - let client = endpoint.client().await?; - - // Register prefill client with worker monitor for TTFT metric cleanup in disaggregated mode - if let Some(monitor) = - model_manager.get_worker_monitor_for_namespace(&self.model_name, &self.namespace) - { - monitor.set_prefill_client(client.clone()); - } - - // Create simple push router with the frontend's router mode - // Note: Per-worker metrics (active_prefill_tokens, active_decode_blocks) are only - // available in KV routing mode where the router has actual bookkeeping. - let push_router = PushRouter::>::from_client_with_threshold( - client, - self.router_mode, - None, // busy_threshold - None, // worker_monitor - ) - .await?; - - InnerPrefillRouter::SimpleRouter(Arc::new(push_router)) - }; - - // Set the router (ignore error if already set) - let _ = self.prefill_router.set(inner_router); - - tracing::info!( - router_mode = ?self.router_mode, - "Prefill router activated successfully" - ); - - Ok(()) - } - - /// Select a prefill worker and resolve its bootstrap connection info. - /// If preselected_worker is provided (GAIE Stage 2), use it directly. - /// Otherwise, query for the best worker (KV mode) or select next worker (non-KV modes). - async fn resolve_prefill_worker( - &self, - req: &PreprocessedRequest, - preselected_worker: Option, - ) -> Option<(u64, u32, BootstrapInfo)> { - let endpoint_id = self.endpoint_id.get()?; - self.prefill_router.get()?; - - // Worker selection - let (worker_id, dp_rank) = if let Some(id) = preselected_worker { - let dp_rank = req.routing.as_ref().and_then(|r| r.dp_rank).unwrap_or(0); - tracing::debug!( - worker_id = id, - dp_rank = dp_rank, - "Using pre-selected prefill worker for bootstrap" - ); - (id, dp_rank) - } else { - // Use shared worker selection logic (update_states=false for peek behavior) - // Extract LORA name and priority jump from routing hints - let lora_name = req.routing.as_ref().and_then(|r| r.lora_name.clone()); - let priority_jump = req - .routing - .as_ref() - .and_then(|r| r.priority_jump) - .unwrap_or(0.0); - let allowed_worker_ids = req - .routing - .as_ref() - .and_then(|r| r.allowed_worker_ids.clone()); - let (routing_token_ids, block_mm_infos) = req.block_mm_routing_info(); - match self - .query_prefill_worker( - routing_token_ids, - block_mm_infos, - false, - lora_name, - priority_jump, - allowed_worker_ids, - ) - .await - { - Ok((worker_id, dp_rank)) => (worker_id, dp_rank), - Err(_) => return None, - } - }; - - // Get bootstrap info from ModelManager (works for ANY mode) - let endpoint = self - .model_manager - .get_disaggregated_endpoint(endpoint_id, worker_id)?; - let host = endpoint.bootstrap_host?; - let port = endpoint.bootstrap_port?; - - let bootstrap_room: u64 = rand::random_range(0..=i64::MAX.cast_unsigned()); - - tracing::debug!( - worker_id = worker_id, - dp_rank = dp_rank, - bootstrap_host = %host, - bootstrap_port = port, - bootstrap_room = bootstrap_room, - router_mode = ?self.router_mode, - "Built bootstrap_info upfront before prefill" - ); - - Some(( - worker_id, - dp_rank, - BootstrapInfo { - bootstrap_host: host, - bootstrap_port: port, - bootstrap_room, - }, - )) - } - - /// Execute prefill with the given router and extract structured result. - /// - /// Uses direct routing to target_worker when specified (for non-KV modes with bootstrap optimization). - /// - /// If `phase_permit` is provided, it is dropped after the first output is received, - /// allowing subsequent `set_phase` calls to proceed. This is used in the bootstrap - /// optimization path to ensure `record_worker_full` completes before the phase changes. - /// - /// Returns (PrefillResult, Option<(worker_id, dp_rank)>). - async fn execute_prefill( - router: Option, - request: SingleIn, - target_worker: Option, - phase_permit: Option, - ) -> Result<(PrefillResult, Option<(u64, u32)>), PrefillError> { - let router = router.ok_or(PrefillError::NotActivated)?; - let mut prefill_response = router - .generate_to_worker(request, target_worker) - .await - .map_err(|e| { - PrefillError::PrefillError( - "failed to route to prefill worker".to_string(), - Some(e.into()), - ) - })?; - - // Drop phase permit now - routing is complete, record_worker_full was called in select_worker. - // This unblocks set_phase(Decode) in the main task without waiting for prefill output. - drop(phase_permit); - - let Some(first_output) = prefill_response.next().await else { - return Err(PrefillError::PrefillError( - "Prefill router returned no output (stream ended)".to_string(), - None, - )); - }; - - if let Some(err) = first_output.err() { - return Err(PrefillError::PrefillError( - "Prefill router returned error in output".to_string(), - Some(Box::new(err)), - )); - } - - let mut prompt_tokens_details = first_output - .data - .as_ref() - .and_then(|o| o.completion_usage.as_ref()) - .and_then(|u| u.prompt_tokens_details.clone()); - - while let Some(next) = prefill_response.next().await { - if let Some(o) = next.data.as_ref() - && prompt_tokens_details.is_none() - { - prompt_tokens_details = o - .completion_usage - .as_ref() - .and_then(|u| u.prompt_tokens_details.clone()); - } - } - - let Some(output) = &first_output.data else { - return Err(PrefillError::NoDisaggregatedParams( - "Prefill router output has no data field".to_string(), - )); - }; - - let Some(disaggregated_params) = output.disaggregated_params.clone() else { - return Err(PrefillError::NoDisaggregatedParams( - "Prefill router output missing disaggregated_params".to_string(), - )); - }; - - // Extract prefill worker ID and dp_rank from disaggregated_params - let prefill_worker_info = - disaggregated_params - .get("worker_id") - .and_then(|worker_id_json| { - let worker_id = worker_id_json - .get("prefill_worker_id") - .and_then(|v| v.as_u64())?; - let dp_rank = worker_id_json - .get("prefill_dp_rank") - .and_then(|v| v.as_u64()) - .map(|r| r as u32) - .unwrap_or(0); - Some((worker_id, dp_rank)) - }); - Ok(( - PrefillResult { - disaggregated_params, - prompt_tokens_details, - }, - prefill_worker_info, - )) - } - - /// Spawn prefill as a background task. - /// - /// Uses direct routing to target_worker when specified (for non-KV modes with bootstrap optimization). - /// - /// The `phase_permit` is passed to the spawned task and dropped after the first output, - /// allowing the main task's `set_phase(Decode)` to proceed. - fn spawn_prefill_task( - &self, - prefill_request: SingleIn, - target_worker: Option, - phase_permit: OwnedSemaphorePermit, - ) { - let router = self.prefill_router.get().cloned(); - // Capture current span to propagate trace context to the spawned task - let span = tracing::Span::current(); - - tokio::spawn( - async move { - match Self::execute_prefill( - router, - prefill_request, - target_worker, - Some(phase_permit), - ) - .await - { - Ok(_) => { - tracing::debug!("Prefill background task completed"); - } - Err(e) => { - tracing::warn!("Prefill background task error: {e:?}"); - } - } - } - .instrument(span), - ); - } - - /// Query the best prefill worker without executing a request. - /// Returns (worker_id, dp_rank). - /// - /// This is the shared worker selection logic used by both `resolve_prefill_worker` - /// and `query_route`. - /// Register externally-provided workers in the prefill router's slot tracker. - pub fn register_workers(&self, worker_ids: &HashSet) { - if let Some(InnerPrefillRouter::KvRouter(r)) = self.prefill_router.get() { - r.chooser.register_workers(worker_ids); - } - } - - pub async fn query_prefill_worker( - &self, - token_ids: &[u32], - block_mm_infos: Option<&[Option]>, - update_states: bool, - lora_name: Option, - priority_jump: f64, - allowed_worker_ids: Option>, - ) -> Result<(u64, u32)> { - let prefill_router = self - .prefill_router - .get() - .ok_or_else(|| anyhow::anyhow!(PrefillError::NotActivated))?; - - match prefill_router { - InnerPrefillRouter::KvRouter(r) => { - let (worker, _overlap) = r - .chooser - .find_best_match( - None, - token_ids, - block_mm_infos, - None, - update_states, - lora_name, - priority_jump, - None, - allowed_worker_ids, - ) - .await?; - Ok((worker.worker_id, worker.dp_rank)) - } - InnerPrefillRouter::SimpleRouter(r) => { - let worker_id = if update_states { - r.select_next_worker() - } else { - r.peek_next_worker() - } - .ok_or_else(|| anyhow::anyhow!("No workers available for prefill"))?; - Ok((worker_id, 0)) - } - } - } - - /// Check if disaggregated mode is currently active (prefill router activated) - pub fn is_activated(&self) -> bool { - self.prefill_router.get().is_some() - } -} - -impl Drop for PrefillRouter { - fn drop(&mut self) { - tracing::debug!("Dropping PrefillRouter, cancelling background activation task"); - self.cancel_token.cancel(); - } -} - -#[async_trait] -impl - Operator< - SingleIn, - ManyOut>, - SingleIn, - ManyOut>, - > for PrefillRouter -{ - async fn generate( - &self, - request: SingleIn, - next: ServerStreamingEngine>, - ) -> Result>> { - // Extract request data while preserving context - let (mut req, context) = request.into_parts(); - let request_id = context.id().to_string(); - let engine_ctx = context.context(); - - // Save original max_tokens for decode - let original_max_tokens = req.stop_conditions.max_tokens; - - // If prefill router is not activated (no prefill workers discovered), - // this is aggregated mode — route directly to decode. - // With --enforce-disagg, fail instead of falling back. - if self.prefill_router.get().is_none() { - if self.enforce_disagg { - return Err(anyhow::anyhow!(PrefillError::NotActivated)); - } - return next.generate(context.map(|_| req)).await; - } - - // Ensure tracker exists for routing decisions in disaggregated mode. - // Create one if not provided by the upstream DeltaGenerator. - if req.tracker.is_none() { - req.tracker = Some(Arc::new(RequestTracker::new())); - } - let tracker = req.tracker.as_ref().unwrap(); - let prefill_phase_permit = tracker.set_phase(RequestPhase::Prefill).await; - - // Prepare prefill request with max_tokens = 1 (clone after tracker is set) - let mut prefill_req = req.clone(); - prefill_req.stop_conditions.max_tokens = Some(1); - - // Try to resolve prefill worker upfront: if we can get bootstrap info early, - // spawn prefill in background and proceed to decode immediately. - let preselected_worker = prefill_req - .routing - .as_ref() - .and_then(|r| r.prefill_worker_id); - - if self.router_mode.is_direct_routing() && preselected_worker.is_none() { - return Err(anyhow::anyhow!( - "Prefill worker ID required in Direct routing mode but none found in request. \ - Expected prefill_worker_id to be set via x-prefill-instance-id header by external router (e.g., EPP)." - )); - } - - let prefill_result = async { - if let Some((worker_id, dp_rank, bootstrap_info)) = self - .resolve_prefill_worker(&prefill_req, preselected_worker) - .await - { - // Bootstrap optimization path: spawn prefill in background - // We successfully used the peeked worker, so we must now advance the router state - // to ensure the next request gets a different worker. - if !self.router_mode.is_kv_routing() - && let Some(router) = self.prefill_router.get() - { - router.select_next_worker(); - } - - let routing = prefill_req.routing_mut(); - routing.prefill_worker_id = Some(worker_id); - routing.dp_rank = Some(dp_rank); - prefill_req.bootstrap_info = Some(bootstrap_info.clone()); - - let prefill_context = Context::with_id(prefill_req, request_id.clone()); - engine_ctx.link_child(prefill_context.context()); - - // Pass phase permit to spawned task - it drops after first output (record_worker_full complete) - // This allows set_phase(Decode) below to proceed only after prefill routing is done - self.spawn_prefill_task(prefill_context, Some(worker_id), prefill_phase_permit); - - Ok(PrefillOutcome::Bootstrap(bootstrap_info)) - } else { - // Original prefill path: wait for prefill to complete - tracing::debug!("Using original prefill path"); - - // Drop the phase permit - we wait for completion - // so there's no race with set_phase(Decode) below - drop(prefill_phase_permit); - - let prefill_context = Context::with_id(prefill_req, request_id.clone()); - engine_ctx.link_child(prefill_context.context()); - - // In Direct mode, pass preselected_worker so execute_prefill uses - // router.direct() instead of router.generate() (which bails in Direct mode). - let (result, _worker_info) = Self::execute_prefill( - self.prefill_router.get().cloned(), - prefill_context, - preselected_worker, - None, - ) - .await?; - - Ok(PrefillOutcome::Completed(result)) - } - } - .await; - - // Abort if cancelled during prefill - if engine_ctx.is_stopped() || engine_ctx.is_killed() { - tracing::debug!("Abort entering decode after context is stopped or killed"); - return Err(anyhow::anyhow!( - "Context id {} is stopped or killed", - engine_ctx.id() - )); - } - - // Handle prefill result - match prefill_result { - Ok(outcome) => { - tracing::debug!("Prefill completed, proceeding to decode"); - - // Set phase to Decode for the decode request. - // In bootstrap path, this blocks until the spawned prefill task drops its permit - // (after first output / record_worker_full completes), ensuring correct phase for routing. - if let Some(ref tracker) = req.tracker { - let _decode_permit = tracker.set_phase(RequestPhase::Decode).await; - // Permit is dropped immediately - decode proceeds, no need to hold it - } - - let mut decode_req = req; - - match outcome { - PrefillOutcome::Bootstrap(info) => { - decode_req.bootstrap_info = Some(info); - } - PrefillOutcome::Completed(result) => { - decode_req.prefill_result = Some(result); - } - } - - // Restore original max_tokens for decode - decode_req.stop_conditions.max_tokens = original_max_tokens; - - // Set router_config_override for decode: - // - overlap_score_weight = 0 (no KV cache overlap scoring for decode) - // - assume_kv_reuse = false (generate random hashes since decode workers - // may already have blocks cached from prefill transfer) - // - track_prefill_tokens = false (decode router should ignore prompt-side load) - let existing_override = decode_req.router_config_override.take(); - decode_req.router_config_override = - Some(build_decode_router_override(existing_override)); - - // Map the modified request through with preserved context - let decode_request = context.map(|_| decode_req); - next.generate(decode_request).await - } - Err(PrefillError::NotActivated) => { - tracing::error!("Prefill router not activated, failing request"); - Err(anyhow::anyhow!(PrefillError::NotActivated)) - } - Err(e) => { - tracing::error!(error = %e, "Remote prefill failed, failing request"); - Err(anyhow::anyhow!(e)) - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn decode_router_override_disables_overlap_and_prefill_tracking() { - let override_config = build_decode_router_override(Some(RouterConfigOverride { - router_temperature: Some(0.7), - ..Default::default() - })); - - assert_eq!(override_config.overlap_score_weight, Some(0.0)); - assert_eq!(override_config.assume_kv_reuse, Some(false)); - assert_eq!(override_config.track_prefill_tokens, Some(false)); - assert_eq!(override_config.router_temperature, Some(0.7)); - } -} diff --git a/lib/llm/src/kv_router/prefill_router/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs new file mode 100644 index 000000000000..8f2a75c3c6bf --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/activation.rs @@ -0,0 +1,188 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; + +use anyhow::Result; +use tokio::sync::oneshot; + +use dynamo_kv_router::config::KvRouterConfig; +use dynamo_runtime::{ + component::{Client, Endpoint}, + pipeline::{PushRouter, RouterMode}, + protocols::annotated::Annotated, +}; + +use super::{InnerPrefillRouter, PrefillRouter}; +use crate::{ + discovery::ModelManager, + kv_router::KvPushRouter, + protocols::common::{ + llm_backend::{LLMEngineOutput, PreprocessedRequest}, + timing::WORKER_TYPE_PREFILL, + }, +}; + +impl PrefillRouter { + /// Create a disabled prefill router that will never activate (passthrough only) + pub fn disabled( + model_manager: Arc, + router_mode: RouterMode, + enforce_disagg: bool, + ) -> Arc { + Arc::new(Self { + prefill_router: std::sync::OnceLock::new(), + model_manager, + endpoint_id: std::sync::OnceLock::new(), + cancel_token: tokio_util::sync::CancellationToken::new(), + router_mode, + enforce_disagg, + model_name: String::new(), // Not used for disabled router + namespace: String::new(), // Not used for disabled router + is_eagle: false, + }) + } + + #[expect(clippy::too_many_arguments)] + pub fn new( + activation_rx: oneshot::Receiver, + model_manager: Arc, + router_mode: RouterMode, + kv_cache_block_size: u32, + kv_router_config: Option, + enforce_disagg: bool, + model_name: String, + namespace: String, + is_eagle: bool, + ) -> Arc { + let prefill_router = std::sync::OnceLock::new(); + let cancel_token = tokio_util::sync::CancellationToken::new(); + + let router = Arc::new(Self { + prefill_router, + model_manager: model_manager.clone(), + endpoint_id: std::sync::OnceLock::new(), + cancel_token: cancel_token.clone(), + router_mode, + enforce_disagg, + model_name, + namespace, + is_eagle, + }); + + // Spawn background task to wait for activation + let router_clone = router.clone(); + tokio::spawn(async move { + tokio::select! { + result = activation_rx => { + let Ok(endpoint) = result else { + tracing::debug!("Prefill router activation channel closed without receiving endpoint"); + return; + }; + + if let Err(e) = router_clone.activate( + endpoint, + model_manager, + kv_cache_block_size, + kv_router_config, + ).await { + tracing::error!(error = %e, "Failed to activate prefill router"); + } + } + _ = cancel_token.cancelled() => { + tracing::debug!("Prefill router activation cancelled"); + } + } + }); + + router + } + + /// Activate the prefill router with the provided endpoint + async fn activate( + &self, + endpoint: Endpoint, + model_manager: Arc, + kv_cache_block_size: u32, + kv_router_config: Option, + ) -> Result<()> { + tracing::info!( + router_mode = ?self.router_mode, + "Activating prefill router" + ); + + // Store endpoint_id for later use in resolve_prefill_worker + let _ = self.endpoint_id.set(endpoint.id()); + + // Start runtime config watcher for this endpoint (needed for get_disaggregated_endpoint) + // This must be done before creating the router so bootstrap info is available + model_manager + .get_or_create_runtime_config_watcher(&endpoint) + .await?; + + let inner_router = if self.router_mode.is_kv_routing() { + // Create KV chooser using the endpoint (this is a prefill router) + let kv_chooser = model_manager + .kv_chooser_for( + &endpoint, + kv_cache_block_size, + kv_router_config, + WORKER_TYPE_PREFILL, + Some(self.model_name.clone()), + self.is_eagle, + ) + .await?; + + // Extract client from kv_chooser to ensure shared state + let client = kv_chooser.client().clone(); + self.register_prefill_client(model_manager.as_ref(), &client); + + // Build the PushRouter for prefill with KV mode using the shared client + let push_router = PushRouter::>::from_client_with_threshold( + client, + RouterMode::KV, + None, // busy_threshold + None, // worker_monitor + ) + .await?; + + // Wrap it in KvPushRouter + InnerPrefillRouter::KvRouter(Arc::new(KvPushRouter::new(push_router, kv_chooser))) + } else { + // Create client for simple router + let client = endpoint.client().await?; + self.register_prefill_client(model_manager.as_ref(), &client); + + // Create simple push router with the frontend's router mode + // Note: Per-worker metrics (active_prefill_tokens, active_decode_blocks) are only + // available in KV routing mode where the router has actual bookkeeping. + let push_router = PushRouter::>::from_client_with_threshold( + client, + self.router_mode, + None, // busy_threshold + None, // worker_monitor + ) + .await?; + + InnerPrefillRouter::SimpleRouter(Arc::new(push_router)) + }; + + // Set the router (ignore error if already set) + let _ = self.prefill_router.set(inner_router); + + tracing::info!( + router_mode = ?self.router_mode, + "Prefill router activated successfully" + ); + + Ok(()) + } + + fn register_prefill_client(&self, model_manager: &ModelManager, client: &Client) { + if let Some(monitor) = + model_manager.get_worker_monitor_for_namespace(&self.model_name, &self.namespace) + { + monitor.set_prefill_client(client.clone()); + } + } +} diff --git a/lib/llm/src/kv_router/prefill_router/execution.rs b/lib/llm/src/kv_router/prefill_router/execution.rs new file mode 100644 index 000000000000..1b5f52a5aa66 --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/execution.rs @@ -0,0 +1,322 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashSet; +use std::sync::Arc; + +use anyhow::Result; +use futures::StreamExt; +use tokio::sync::OwnedSemaphorePermit; +use tracing::Instrument; + +use dynamo_kv_router::protocols::{BlockExtraInfo, WorkerId}; +use dynamo_runtime::{ + engine::AsyncEngineContext, + pipeline::{AsyncEngineContextProvider, Context, SingleIn}, + protocols::maybe_error::MaybeError, +}; + +use super::{InnerPrefillRouter, PrefillError, PrefillResolveDecision, PrefillRouter}; +use crate::protocols::common::{ + llm_backend::PreprocessedRequest, + preprocessor::{BootstrapInfo, PrefillResult}, +}; + +impl PrefillRouter { + /// Select a prefill worker and resolve its bootstrap connection info. + /// If preselected_worker is provided (GAIE Stage 2), use it directly. + /// Otherwise, query for the best worker (KV mode) or select next worker (non-KV modes). + pub(super) async fn resolve_prefill_worker( + &self, + req: &PreprocessedRequest, + preselected_worker: Option, + ) -> PrefillResolveDecision { + let Some(endpoint_id) = self.endpoint_id.get() else { + return PrefillResolveDecision::NotActivated; + }; + if self.prefill_router.get().is_none() { + return PrefillResolveDecision::NotActivated; + } + + // Worker selection + let (worker_id, dp_rank) = if let Some(id) = preselected_worker { + let dp_rank = req.routing.as_ref().and_then(|r| r.dp_rank).unwrap_or(0); + tracing::debug!( + worker_id = id, + dp_rank = dp_rank, + "Using pre-selected prefill worker for bootstrap" + ); + (id, dp_rank) + } else { + // Use shared worker selection logic (update_states=false for peek behavior) + // Extract LORA name and priority jump from routing hints + let lora_name = req.routing.as_ref().and_then(|r| r.lora_name.clone()); + let priority_jump = req + .routing + .as_ref() + .and_then(|r| r.priority_jump) + .unwrap_or(0.0); + let allowed_worker_ids = req + .routing + .as_ref() + .and_then(|r| r.allowed_worker_ids.clone()); + let (routing_token_ids, block_mm_infos) = req.block_mm_routing_info(); + match self + .query_prefill_worker( + routing_token_ids, + block_mm_infos, + false, + lora_name, + priority_jump, + allowed_worker_ids, + ) + .await + { + Ok((worker_id, dp_rank)) => (worker_id, dp_rank), + Err(_) => return PrefillResolveDecision::Unavailable, + } + }; + + // Get bootstrap info from ModelManager (works for ANY mode) + let Some(endpoint) = self + .model_manager + .get_disaggregated_endpoint(endpoint_id, worker_id) + else { + return PrefillResolveDecision::NoBootstrapEndpoint; + }; + let Some(host) = endpoint.bootstrap_host else { + return PrefillResolveDecision::NoBootstrapEndpoint; + }; + let Some(port) = endpoint.bootstrap_port else { + return PrefillResolveDecision::NoBootstrapEndpoint; + }; + + let bootstrap_room: u64 = rand::random_range(0..=i64::MAX.cast_unsigned()); + + tracing::debug!( + worker_id = worker_id, + dp_rank = dp_rank, + bootstrap_host = %host, + bootstrap_port = port, + bootstrap_room = bootstrap_room, + router_mode = ?self.router_mode, + "Built bootstrap_info upfront before prefill" + ); + + PrefillResolveDecision::Resolved { + worker_id, + dp_rank, + bootstrap_info: BootstrapInfo { + bootstrap_host: host, + bootstrap_port: port, + bootstrap_room, + }, + } + } + + /// Execute prefill with the given router and extract structured result. + /// + /// Uses direct routing to target_worker when specified (for non-KV modes with bootstrap optimization). + /// + /// If `phase_transition_permit` is provided, it is dropped immediately after routing completes, + /// allowing subsequent `set_phase` calls to proceed. This preserves the current synchronization: + /// the prefill route must finish `record_worker_full` before the phase can change to Decode. + /// + /// Returns (PrefillResult, Option<(worker_id, dp_rank)>). + pub(super) async fn execute_prefill( + router: Option, + request: SingleIn, + target_worker: Option, + phase_transition_permit: Option, + ) -> Result<(PrefillResult, Option<(u64, u32)>), PrefillError> { + let router = router.ok_or(PrefillError::NotActivated)?; + let mut prefill_response = router + .generate_to_worker(request, target_worker) + .await + .map_err(|e| { + PrefillError::PrefillError( + "failed to route to prefill worker".to_string(), + Some(e.into()), + ) + })?; + + // Release the phase barrier now that routing completed and record_worker_full already ran. + // Decode may proceed without waiting for prefill output streaming to finish. + drop(phase_transition_permit); + + let Some(first_output) = prefill_response.next().await else { + return Err(PrefillError::PrefillError( + "Prefill router returned no output (stream ended)".to_string(), + None, + )); + }; + + if let Some(err) = first_output.err() { + return Err(PrefillError::PrefillError( + "Prefill router returned error in output".to_string(), + Some(Box::new(err)), + )); + } + + let mut prompt_tokens_details = first_output + .data + .as_ref() + .and_then(|o| o.completion_usage.as_ref()) + .and_then(|u| u.prompt_tokens_details.clone()); + + while let Some(next) = prefill_response.next().await { + if let Some(o) = next.data.as_ref() + && prompt_tokens_details.is_none() + { + prompt_tokens_details = o + .completion_usage + .as_ref() + .and_then(|u| u.prompt_tokens_details.clone()); + } + } + + let Some(output) = &first_output.data else { + return Err(PrefillError::NoDisaggregatedParams( + "Prefill router output has no data field".to_string(), + )); + }; + + let Some(disaggregated_params) = output.disaggregated_params.clone() else { + return Err(PrefillError::NoDisaggregatedParams( + "Prefill router output missing disaggregated_params".to_string(), + )); + }; + + // Extract prefill worker ID and dp_rank from disaggregated_params + let prefill_worker_info = + disaggregated_params + .get("worker_id") + .and_then(|worker_id_json| { + let worker_id = worker_id_json + .get("prefill_worker_id") + .and_then(|v| v.as_u64())?; + let dp_rank = worker_id_json + .get("prefill_dp_rank") + .and_then(|v| v.as_u64()) + .map(|r| r as u32) + .unwrap_or(0); + Some((worker_id, dp_rank)) + }); + Ok(( + PrefillResult { + disaggregated_params, + prompt_tokens_details, + }, + prefill_worker_info, + )) + } + + /// Spawn prefill as a background task. + /// + /// Uses direct routing to target_worker when specified (for non-KV modes with bootstrap optimization). + /// + /// The `phase_transition_permit` is passed to the spawned task and released after routing + /// completes, allowing the main task's `set_phase(Decode)` to proceed. + pub(super) fn spawn_prefill_task( + &self, + prefill_request: SingleIn, + target_worker: Option, + phase_transition_permit: OwnedSemaphorePermit, + ) { + let router = self.prefill_router.get().cloned(); + // Capture current span to propagate trace context to the spawned task + let span = tracing::Span::current(); + + tokio::spawn( + async move { + match Self::execute_prefill( + router, + prefill_request, + target_worker, + Some(phase_transition_permit), + ) + .await + { + Ok(_) => { + tracing::debug!("Prefill background task completed"); + } + Err(e) => { + tracing::warn!("Prefill background task error: {e:?}"); + } + } + } + .instrument(span), + ); + } + + /// Query the best prefill worker without executing a request. + /// Returns (worker_id, dp_rank). + /// + /// This is the shared worker selection logic used by both `resolve_prefill_worker` + /// and `query_route`. + pub async fn query_prefill_worker( + &self, + token_ids: &[u32], + block_mm_infos: Option<&[Option]>, + update_states: bool, + lora_name: Option, + priority_jump: f64, + allowed_worker_ids: Option>, + ) -> Result<(u64, u32)> { + let prefill_router = self + .prefill_router + .get() + .ok_or_else(|| anyhow::anyhow!(PrefillError::NotActivated))?; + + match prefill_router { + InnerPrefillRouter::KvRouter(r) => { + let (worker, _overlap) = r + .chooser + .find_best_match( + None, + token_ids, + block_mm_infos, + None, + update_states, + lora_name, + priority_jump, + None, + allowed_worker_ids, + ) + .await?; + Ok((worker.worker_id, worker.dp_rank)) + } + InnerPrefillRouter::SimpleRouter(r) => { + let worker_id = if update_states { + r.select_next_worker() + } else { + r.peek_next_worker() + } + .ok_or_else(|| anyhow::anyhow!("No workers available for prefill"))?; + Ok((worker_id, 0)) + } + } + } + + /// Register externally-provided workers in the prefill router's slot tracker. + pub fn register_workers(&self, worker_ids: &HashSet) { + if let Some(InnerPrefillRouter::KvRouter(r)) = self.prefill_router.get() { + r.chooser.register_workers(worker_ids); + } + } + + /// Check if disaggregated mode is currently active (prefill router activated) + pub fn is_activated(&self) -> bool { + self.prefill_router.get().is_some() + } +} + +pub(super) fn link_child_context( + engine_ctx: &Arc, + request: T, + request_id: &str, +) -> Context { + let child_context = Context::with_id(request, request_id.to_string()); + engine_ctx.link_child(child_context.context()); + child_context +} diff --git a/lib/llm/src/kv_router/prefill_router/inner.rs b/lib/llm/src/kv_router/prefill_router/inner.rs new file mode 100644 index 000000000000..8b4f3508b67f --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/inner.rs @@ -0,0 +1,55 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; + +use anyhow::Result; + +use dynamo_runtime::{ + pipeline::{AsyncEngine, ManyOut, PushRouter, SingleIn}, + protocols::annotated::Annotated, +}; + +use crate::{ + kv_router::KvPushRouter, + protocols::common::llm_backend::{LLMEngineOutput, PreprocessedRequest}, +}; + +/// The inner router used by PrefillRouter +#[derive(Clone)] +pub(super) enum InnerPrefillRouter { + /// KV-aware routing using KvPushRouter + KvRouter(Arc), + /// Simple routing (RoundRobin, Random, Direct) + /// Note: Per-worker metrics (active_prefill_tokens, active_decode_blocks) are only + /// available in KV routing mode where the router has actual bookkeeping. + SimpleRouter(Arc>>), +} + +impl InnerPrefillRouter { + /// Generate with optional direct routing to specific worker. + /// For KvRouter, target_worker is ignored since prefill_worker_id is already set on the request. + /// For SimpleRouter, target_worker triggers direct routing via router.direct(). + pub(super) async fn generate_to_worker( + &self, + request: SingleIn, + target_worker: Option, + ) -> Result>> { + match (self, target_worker) { + // KvRouter: prefill_worker_id already set on request, KvPushRouter::select_worker uses it + (InnerPrefillRouter::KvRouter(router), _) => router.generate(request).await, + (InnerPrefillRouter::SimpleRouter(router), Some(worker_id)) => { + router.direct(request, worker_id).await + } + (InnerPrefillRouter::SimpleRouter(router), None) => router.generate(request).await, + } + } + + /// Select next worker (for non-KV modes only) + pub(super) fn select_next_worker(&self) -> Option { + match self { + InnerPrefillRouter::SimpleRouter(router) => router.select_next_worker(), + InnerPrefillRouter::KvRouter(_) => None, + } + } +} diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs new file mode 100644 index 000000000000..c2a9aad4e14f --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -0,0 +1,259 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::{Arc, OnceLock}; + +use anyhow::Result; +use tokio_util::sync::CancellationToken; + +use dynamo_runtime::{ + pipeline::{ + AsyncEngineContextProvider, ManyOut, Operator, RouterMode, ServerStreamingEngine, SingleIn, + async_trait, + }, + protocols::{EndpointId, annotated::Annotated}, +}; + +use crate::{ + discovery::ModelManager, + protocols::common::{ + llm_backend::{LLMEngineOutput, PreprocessedRequest}, + timing::{RequestPhase, RequestTracker}, + }, +}; + +mod activation; +mod execution; +mod inner; +mod types; + +use execution::link_child_context; +use inner::InnerPrefillRouter; +pub use types::PrefillError; +use types::{PrefillOutcome, PrefillResolveDecision, build_decode_router_override}; + +/// PrefillRouter is a forward-only operator that sits between Migration and the decode router. +/// It optionally calls a prefill worker before routing to decode, extracting disaggregated_params +/// from the prefill response and injecting them into the decode request. +/// +/// Modes: +/// - Query-only: `query_instance_id` annotation present → returns worker IDs without execution +/// - Pre-routed: `prefill_worker_id`/`decode_worker_id` set → routes to specified workers +/// - Normal: Worker IDs determined by router based on KV cache state +pub struct PrefillRouter { + prefill_router: OnceLock, + model_manager: Arc, + endpoint_id: OnceLock, + cancel_token: CancellationToken, + router_mode: RouterMode, + enforce_disagg: bool, + /// Model name used to look up the worker monitor for prefill client registration + model_name: String, + /// Namespace used to look up the correct WorkerSet's worker monitor + namespace: String, + is_eagle: bool, +} + +impl Drop for PrefillRouter { + fn drop(&mut self) { + tracing::debug!("Dropping PrefillRouter, cancelling background activation task"); + self.cancel_token.cancel(); + } +} + +#[async_trait] +impl + Operator< + SingleIn, + ManyOut>, + SingleIn, + ManyOut>, + > for PrefillRouter +{ + async fn generate( + &self, + request: SingleIn, + next: ServerStreamingEngine>, + ) -> Result>> { + // Extract request data while preserving context + let (mut req, context) = request.into_parts(); + let request_id = context.id().to_string(); + let engine_ctx = context.context(); + + // Save original max_tokens for decode + let original_max_tokens = req.stop_conditions.max_tokens; + + // If prefill router is not activated (no prefill workers discovered), + // this is aggregated mode — route directly to decode. + // With --enforce-disagg, fail instead of falling back. + if self.prefill_router.get().is_none() { + if self.enforce_disagg { + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + } + return next.generate(context.map(|_| req)).await; + } + + // Ensure tracker exists for routing decisions in disaggregated mode. + // Create one if not provided by the upstream DeltaGenerator. + if req.tracker.is_none() { + req.tracker = Some(Arc::new(RequestTracker::new())); + } + let tracker = req.tracker.as_ref().unwrap(); + let prefill_phase_barrier = tracker.set_phase(RequestPhase::Prefill).await; + + // Prepare prefill request with max_tokens = 1 (clone after tracker is set) + let mut prefill_req = req.clone(); + prefill_req.stop_conditions.max_tokens = Some(1); + + // Try to resolve prefill worker upfront: if we can get bootstrap info early, + // spawn prefill in background and proceed to decode immediately. + let preselected_worker = prefill_req + .routing + .as_ref() + .and_then(|r| r.prefill_worker_id); + + if self.router_mode.is_direct_routing() && preselected_worker.is_none() { + return Err(anyhow::anyhow!( + "Prefill worker ID required in Direct routing mode but none found in request. \ + Expected prefill_worker_id to be set via x-prefill-instance-id header by external router (e.g., EPP)." + )); + } + + let prefill_result = match self + .resolve_prefill_worker(&prefill_req, preselected_worker) + .await + { + PrefillResolveDecision::Resolved { + worker_id, + dp_rank, + bootstrap_info, + } => { + // Bootstrap optimization path: spawn prefill in background + // We successfully used the peeked worker, so we must now advance the router state + // to ensure the next request gets a different worker. + if !self.router_mode.is_kv_routing() + && let Some(router) = self.prefill_router.get() + { + router.select_next_worker(); + } + + let routing = prefill_req.routing_mut(); + routing.prefill_worker_id = Some(worker_id); + routing.dp_rank = Some(dp_rank); + prefill_req.bootstrap_info = Some(bootstrap_info.clone()); + + let prefill_context = + link_child_context(&engine_ctx, prefill_req, request_id.as_str()); + + // Pass the phase barrier to the spawned task. It is released after routing + // completes so `record_worker_full` finishes before phase changes to Decode. + self.spawn_prefill_task(prefill_context, Some(worker_id), prefill_phase_barrier); + + Ok(PrefillOutcome::Bootstrap(bootstrap_info)) + } + PrefillResolveDecision::Unavailable + | PrefillResolveDecision::NotActivated + | PrefillResolveDecision::NoBootstrapEndpoint => { + // Original prefill path: wait for prefill to complete + tracing::debug!("Using original prefill path"); + + // Drop the phase barrier because we wait for prefill completion in this task, + // so there is no race with set_phase(Decode) below. + drop(prefill_phase_barrier); + + let prefill_context = + link_child_context(&engine_ctx, prefill_req, request_id.as_str()); + + // In Direct mode, pass preselected_worker so execute_prefill uses + // router.direct() instead of router.generate() (which bails in Direct mode). + let (result, _worker_info) = Self::execute_prefill( + self.prefill_router.get().cloned(), + prefill_context, + preselected_worker, + None, + ) + .await?; + + Ok(PrefillOutcome::Completed(result)) + } + }; + + // Abort if cancelled during prefill + if engine_ctx.is_stopped() || engine_ctx.is_killed() { + tracing::debug!("Abort entering decode after context is stopped or killed"); + return Err(anyhow::anyhow!( + "Context id {} is stopped or killed", + engine_ctx.id() + )); + } + + // Handle prefill result + match prefill_result { + Ok(outcome) => { + tracing::debug!("Prefill completed, proceeding to decode"); + + // Set phase to Decode for the decode request. + // In bootstrap path, this blocks until the spawned prefill task releases its + // phase barrier after routing completes, ensuring correct worker attribution. + if let Some(ref tracker) = req.tracker { + let _decode_permit = tracker.set_phase(RequestPhase::Decode).await; + // Permit is dropped immediately - decode proceeds, no need to hold it + } + + let mut decode_req = req; + + match outcome { + PrefillOutcome::Bootstrap(info) => { + decode_req.bootstrap_info = Some(info); + } + PrefillOutcome::Completed(result) => { + decode_req.prefill_result = Some(result); + } + } + + // Restore original max_tokens for decode + decode_req.stop_conditions.max_tokens = original_max_tokens; + + // Set router_config_override for decode: + // - overlap_score_weight = 0 (no KV cache overlap scoring for decode) + // - assume_kv_reuse = false (generate random hashes since decode workers + // may already have blocks cached from prefill transfer) + // - track_prefill_tokens = false (decode router should ignore prompt-side load) + let existing_override = decode_req.router_config_override.take(); + decode_req.router_config_override = + Some(build_decode_router_override(existing_override)); + + // Map the modified request through with preserved context + let decode_request = context.map(|_| decode_req); + next.generate(decode_request).await + } + Err(PrefillError::NotActivated) => { + tracing::error!("Prefill router not activated, failing request"); + Err(anyhow::anyhow!(PrefillError::NotActivated)) + } + Err(e) => { + tracing::error!(error = %e, "Remote prefill failed, failing request"); + Err(anyhow::anyhow!(e)) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dynamo_kv_router::config::RouterConfigOverride; + + #[test] + fn decode_router_override_disables_overlap_and_prefill_tracking() { + let override_config = build_decode_router_override(Some(RouterConfigOverride { + router_temperature: Some(0.7), + ..Default::default() + })); + + assert_eq!(override_config.overlap_score_weight, Some(0.0)); + assert_eq!(override_config.assume_kv_reuse, Some(false)); + assert_eq!(override_config.track_prefill_tokens, Some(false)); + assert_eq!(override_config.router_temperature, Some(0.7)); + } +} diff --git a/lib/llm/src/kv_router/prefill_router/types.rs b/lib/llm/src/kv_router/prefill_router/types.rs new file mode 100644 index 000000000000..1acf1de53d86 --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/types.rs @@ -0,0 +1,56 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use dynamo_kv_router::config::RouterConfigOverride; + +use crate::protocols::common::preprocessor::{BootstrapInfo, PrefillResult}; + +/// Errors that can occur during prefill routing +#[derive(Debug, thiserror::Error)] +pub enum PrefillError { + /// Prefill router has not been activated yet + #[error("Prefill router not yet activated")] + NotActivated, + + /// TODO: Separate prefill worker error from prefill router error + /// Error during prefill execution + #[error("Prefill execution failed: {0}")] + PrefillError( + String, + #[source] Option>, + ), + + /// Disaggregated params not found in prefill response + #[error("No disaggregated params in prefill response: {0}")] + NoDisaggregatedParams(String), +} + +/// Result of the prefill phase in `generate()`. +pub(super) enum PrefillOutcome { + /// Bootstrap optimization: prefill spawned in background, bootstrap info ready + Bootstrap(BootstrapInfo), + /// Synchronous prefill completed with result + Completed(PrefillResult), +} + +pub(super) enum PrefillResolveDecision { + Resolved { + worker_id: u64, + dp_rank: u32, + bootstrap_info: BootstrapInfo, + }, + Unavailable, + NotActivated, + NoBootstrapEndpoint, +} + +pub(super) fn build_decode_router_override( + existing_override: Option, +) -> RouterConfigOverride { + RouterConfigOverride { + overlap_score_weight: Some(0.0), + assume_kv_reuse: Some(false), + track_prefill_tokens: Some(false), + ..existing_override.unwrap_or_default() + } +} From 3eb2051e42026aa775b2b71d8ed0c0348f5ce601 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 21:59:56 -0700 Subject: [PATCH 11/24] feat: plumb replay planner profile args Signed-off-by: PeaBrane --- components/src/dynamo/mocker/config.py | 87 +++++++++--------- .../src/dynamo/mocker/tests/test_config.py | 88 +++++++++++++++++++ lib/bindings/python/rust/llm/replay.rs | 19 +++- lib/bindings/python/src/dynamo/_core.pyi | 1 + lib/bindings/python/tests/test_replay.py | 78 ++++++++++++++++ 5 files changed, 227 insertions(+), 46 deletions(-) create mode 100644 components/src/dynamo/mocker/tests/test_config.py diff --git a/components/src/dynamo/mocker/config.py b/components/src/dynamo/mocker/config.py index 3d3e4414ba78..fec3a5af6479 100644 --- a/components/src/dynamo/mocker/config.py +++ b/components/src/dynamo/mocker/config.py @@ -41,51 +41,52 @@ def _build_sglang_args(args: argparse.Namespace) -> SglangArgs | None: return SglangArgs(**sglang_args) -def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: - aic_backend = None - aic_system = None - aic_backend_version = None - aic_tp_size = None - aic_model_path = None - if getattr(args, "aic_perf_model", False): - aic_backend = getattr(args, "engine_type", None) or "vllm" - aic_system = getattr(args, "aic_system", None) - aic_backend_version = getattr(args, "aic_backend_version", None) - aic_tp_size = getattr(args, "aic_tp_size", None) - aic_model_path = getattr(args, "model_path", None) - - return MockEngineArgs( - engine_type=getattr(args, "engine_type", None) or "vllm", - num_gpu_blocks=getattr(args, "num_gpu_blocks", _DEFAULT_NUM_GPU_BLOCKS), - block_size=getattr(args, "block_size", 0) or 0, - max_num_seqs=getattr(args, "max_num_seqs", _DEFAULT_MAX_NUM_SEQS), - max_num_batched_tokens=getattr( +def _build_mocker_engine_args_kwargs(args: argparse.Namespace) -> dict[str, object]: + worker_type = ( + "prefill" + if getattr(args, "is_prefill_worker", False) + else "decode" + if getattr(args, "is_decode_worker", False) + else "aggregated" + ) + kwargs: dict[str, object] = { + "engine_type": getattr(args, "engine_type", None) or "vllm", + "num_gpu_blocks": getattr(args, "num_gpu_blocks", _DEFAULT_NUM_GPU_BLOCKS), + "block_size": getattr(args, "block_size", 0) or 0, + "max_num_seqs": getattr(args, "max_num_seqs", _DEFAULT_MAX_NUM_SEQS), + "max_num_batched_tokens": getattr( args, "max_num_batched_tokens", _DEFAULT_MAX_NUM_BATCHED_TOKENS ), - enable_prefix_caching=getattr(args, "enable_prefix_caching", True), - enable_chunked_prefill=getattr(args, "enable_chunked_prefill", True), - preemption_mode=getattr(args, "preemption_mode", "lifo"), - speedup_ratio=getattr(args, "speedup_ratio", 1.0), - decode_speedup_ratio=getattr(args, "decode_speedup_ratio", 1.0), - dp_size=getattr(args, "dp_size", 1), - startup_time=getattr(args, "startup_time", None), - worker_type=( - "prefill" - if getattr(args, "is_prefill_worker", False) - else "decode" - if getattr(args, "is_decode_worker", False) - else "aggregated" - ), - aic_backend=aic_backend, - aic_system=aic_system, - aic_backend_version=aic_backend_version, - aic_tp_size=aic_tp_size, - aic_model_path=aic_model_path, - enable_local_indexer=not getattr(args, "durable_kv_events", False), - kv_transfer_bandwidth=getattr(args, "kv_transfer_bandwidth", None), - reasoning=_parse_reasoning_config(getattr(args, "reasoning", None)), - sglang=_build_sglang_args(args), - ) + "enable_prefix_caching": getattr(args, "enable_prefix_caching", True), + "enable_chunked_prefill": getattr(args, "enable_chunked_prefill", True), + "speedup_ratio": getattr(args, "speedup_ratio", 1.0), + "decode_speedup_ratio": getattr(args, "decode_speedup_ratio", 1.0), + "dp_size": getattr(args, "dp_size", 1), + "startup_time": getattr(args, "startup_time", None), + "worker_type": worker_type, + "planner_profile_data": getattr(args, "planner_profile_data", None), + "aic_backend": None, + "aic_system": None, + "aic_backend_version": None, + "aic_tp_size": None, + "aic_model_path": None, + "enable_local_indexer": not getattr(args, "durable_kv_events", False), + "kv_transfer_bandwidth": getattr(args, "kv_transfer_bandwidth", None), + "reasoning": _parse_reasoning_config(getattr(args, "reasoning", None)), + "sglang": _build_sglang_args(args), + "preemption_mode": getattr(args, "preemption_mode", "lifo"), + } + if getattr(args, "aic_perf_model", False): + kwargs["aic_backend"] = getattr(args, "engine_type", None) or "vllm" + kwargs["aic_system"] = getattr(args, "aic_system", None) + kwargs["aic_backend_version"] = getattr(args, "aic_backend_version", None) + kwargs["aic_tp_size"] = getattr(args, "aic_tp_size", None) + kwargs["aic_model_path"] = getattr(args, "model_path", None) + return kwargs + + +def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: + return MockEngineArgs(**_build_mocker_engine_args_kwargs(args)) def load_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: diff --git a/components/src/dynamo/mocker/tests/test_config.py b/components/src/dynamo/mocker/tests/test_config.py new file mode 100644 index 000000000000..d78424affa72 --- /dev/null +++ b/components/src/dynamo/mocker/tests/test_config.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import json +from pathlib import Path +from types import SimpleNamespace + +from dynamo.mocker import config + + +def test_build_mocker_engine_args_preserves_cli_mapped_fields(monkeypatch): + args = argparse.Namespace( + engine_type="sglang", + num_gpu_blocks=2048, + block_size=128, + max_num_seqs=64, + max_num_batched_tokens=4096, + enable_prefix_caching=False, + enable_chunked_prefill=False, + preemption_mode="fifo", + speedup_ratio=2.0, + decode_speedup_ratio=3.0, + dp_size=4, + startup_time=1.5, + planner_profile_data=Path("/tmp/perf_data.npz"), + is_prefill_worker=True, + is_decode_worker=False, + durable_kv_events=False, + kv_transfer_bandwidth=123.0, + reasoning=json.dumps( + { + "start_thinking_token_id": 11, + "end_thinking_token_id": 12, + "thinking_ratio": 0.25, + } + ), + sglang_schedule_policy="lpm", + sglang_page_size=16, + sglang_max_prefill_tokens=8192, + sglang_chunked_prefill_size=2048, + sglang_clip_max_new_tokens=1024, + sglang_schedule_conservativeness=0.8, + aic_perf_model=True, + aic_system="h200_sxm", + aic_backend_version="0.5.6.post2", + aic_tp_size=8, + model_path="/models/mock", + ) + + captured = {} + + def fake_mock_engine_args(**kwargs): + captured["kwargs"] = kwargs + return SimpleNamespace(**kwargs) + + monkeypatch.setattr(config, "MockEngineArgs", fake_mock_engine_args) + + engine_args = config.build_mocker_engine_args(args) + + assert captured["kwargs"] == { + "engine_type": "sglang", + "num_gpu_blocks": 2048, + "block_size": 128, + "max_num_seqs": 64, + "max_num_batched_tokens": 4096, + "enable_prefix_caching": False, + "enable_chunked_prefill": False, + "preemption_mode": "fifo", + "speedup_ratio": 2.0, + "decode_speedup_ratio": 3.0, + "dp_size": 4, + "startup_time": 1.5, + "worker_type": "prefill", + "planner_profile_data": Path("/tmp/perf_data.npz"), + "enable_local_indexer": True, + "kv_transfer_bandwidth": 123.0, + "reasoning": captured["kwargs"]["reasoning"], + "sglang": captured["kwargs"]["sglang"], + "aic_backend": "sglang", + "aic_system": "h200_sxm", + "aic_backend_version": "0.5.6.post2", + "aic_tp_size": 8, + "aic_model_path": "/models/mock", + } + assert engine_args.planner_profile_data == Path("/tmp/perf_data.npz") + assert engine_args.reasoning is not None + assert engine_args.sglang is not None diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index 85d407b8a32f..df9238f23be1 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -132,7 +132,7 @@ impl MockEngineArgs { #[pymethods] impl MockEngineArgs { #[new] - #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=16384, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, enable_local_indexer=false, bootstrap_port=None, kv_bytes_per_token=None, kv_transfer_bandwidth=None, reasoning=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None))] + #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=16384, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", planner_profile_data=None, aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, enable_local_indexer=false, bootstrap_port=None, kv_bytes_per_token=None, kv_transfer_bandwidth=None, reasoning=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None))] #[allow(clippy::too_many_arguments)] fn new( engine_type: &str, @@ -147,6 +147,7 @@ impl MockEngineArgs { dp_size: u32, startup_time: Option, worker_type: &str, + planner_profile_data: Option, aic_backend: Option, aic_system: Option, aic_backend_version: Option, @@ -174,7 +175,7 @@ impl MockEngineArgs { }) .transpose()?; - let inner = RsMockEngineArgs::builder() + let mut builder = RsMockEngineArgs::builder() .engine_type(engine_type) .num_gpu_blocks(num_gpu_blocks) .block_size(block_size) @@ -201,7 +202,19 @@ impl MockEngineArgs { .zmq_replay_port(zmq_replay_port) .preemption_mode(preemption_mode) .router_queue_policy(router_queue_policy) - .sglang(sglang.map(|config| config.inner())) + .sglang(sglang.map(|config| config.inner())); + + if let Some(npz_path) = planner_profile_data { + let perf_model = PerfModel::from_npz(&npz_path).map_err(|e| { + PyException::new_err(format!( + "Failed to load planner_profile_data from {:?}: {e}", + npz_path + )) + })?; + builder = builder.perf_model(Arc::new(perf_model)); + } + + let inner = builder .build() .map_err(|e| PyException::new_err(format!("Failed to build MockEngineArgs: {e}")))? .normalized() diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index b1adace5a54f..b6939c3b9893 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1216,6 +1216,7 @@ class MockEngineArgs: dp_size: int = 1, startup_time: Optional[float] = None, worker_type: str = "aggregated", + planner_profile_data: Optional[str | os.PathLike[str]] = None, aic_backend: Optional[str] = None, aic_system: Optional[str] = None, aic_backend_version: Optional[str] = None, diff --git a/lib/bindings/python/tests/test_replay.py b/lib/bindings/python/tests/test_replay.py index 7964c4907269..b7bae50b5245 100644 --- a/lib/bindings/python/tests/test_replay.py +++ b/lib/bindings/python/tests/test_replay.py @@ -6,6 +6,7 @@ import subprocess import sys from pathlib import Path +from types import SimpleNamespace import pytest @@ -252,6 +253,13 @@ def _replay_cli_env() -> dict[str, str]: return env +def _planner_profile_data_npz_path() -> Path: + return ( + Path(__file__).resolve().parents[4] + / "benchmarks/results/H200_TP1P_TP1D_perf_data.npz" + ) + + def _run_replay_cli(tmp_path, *args): return subprocess.run( [ @@ -842,6 +850,76 @@ def fake_run(*args, **kwargs): assert captured["kwargs"]["inter_turn_delay_ms"] == 7.0 +@pytest.mark.parametrize( + ("engine_args_flag", "engine_args_kwarg"), + [ + ("--extra-engine-args", "extra_engine_args"), + ("--prefill-engine-args", "prefill_engine_args"), + ("--decode-engine-args", "decode_engine_args"), + ], +) +@pytest.mark.parametrize("workload", ["synthetic", "trace"]) +def test_replay_cli_resolves_planner_profile_data_from_benchmarks( + tmp_path, monkeypatch, engine_args_flag, engine_args_kwarg, workload +): + planner_profile_data = _planner_profile_data_npz_path() + assert planner_profile_data.exists() + + captured = {} + resolved_paths = [] + + def fake_run(*args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return { + "completed_requests": 1, + "request_throughput_rps": 1.0, + "output_throughput_tok_s": 1.0, + } + + def fake_resolve_planner_profile_data(path): + resolved_paths.append(path) + return SimpleNamespace(npz_path=planner_profile_data) + + monkeypatch.setattr("dynamo.replay.main.run_trace_replay", fake_run) + monkeypatch.setattr("dynamo.replay.main.run_synthetic_trace_replay", fake_run) + monkeypatch.setattr( + "dynamo.replay.main.resolve_planner_profile_data", + fake_resolve_planner_profile_data, + ) + + cli_args = [ + engine_args_flag, + json.dumps( + { + "block_size": 64, + "speedup_ratio": 1000.0, + "planner_profile_data": str(planner_profile_data), + } + ), + ] + + if workload == "trace": + trace_path = _write_trace_and_args(tmp_path) + cli_args = [str(trace_path), *cli_args] + else: + cli_args = [ + "--input-tokens", + "16", + "--output-tokens", + "8", + "--request-count", + "1", + *cli_args, + ] + + exit_code = main(cli_args) + + assert exit_code == 0 + assert resolved_paths == [planner_profile_data] + assert captured["kwargs"][engine_args_kwarg] is not None + + @pytest.mark.timeout(30) def test_replay_cli_subprocess_synthetic_smoke(tmp_path): report_path = tmp_path / "synthetic_report.json" From aab6f787889ec01ac58d0c314c12e112422b6dcd Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 22:25:16 -0700 Subject: [PATCH 12/24] keep planner args from vanishing Signed-off-by: PeaBrane --- .../src/dynamo/mocker/tests/test_config.py | 60 ++++--- lib/bindings/python/rust/llm/replay.rs | 56 +++++++ lib/bindings/python/src/dynamo/_core.pyi | 2 + lib/bindings/python/tests/test_replay.py | 154 ------------------ lib/mocker/src/common/protocols.rs | 5 + tests/router/test_router_e2e_with_mockers.py | 51 ++++-- 6 files changed, 140 insertions(+), 188 deletions(-) diff --git a/components/src/dynamo/mocker/tests/test_config.py b/components/src/dynamo/mocker/tests/test_config.py index d78424affa72..e7f2c3aeedf7 100644 --- a/components/src/dynamo/mocker/tests/test_config.py +++ b/components/src/dynamo/mocker/tests/test_config.py @@ -4,12 +4,21 @@ import argparse import json from pathlib import Path -from types import SimpleNamespace from dynamo.mocker import config -def test_build_mocker_engine_args_preserves_cli_mapped_fields(monkeypatch): +def _planner_profile_data_npz_path() -> Path: + return ( + Path(__file__).resolve().parents[5] + / "benchmarks/results/H200_TP1P_TP1D_perf_data.npz" + ) + + +def test_build_mocker_engine_args_preserves_cli_mapped_fields(): + planner_profile_data = _planner_profile_data_npz_path() + assert planner_profile_data.exists() + args = argparse.Namespace( engine_type="sglang", num_gpu_blocks=2048, @@ -23,7 +32,7 @@ def test_build_mocker_engine_args_preserves_cli_mapped_fields(monkeypatch): decode_speedup_ratio=3.0, dp_size=4, startup_time=1.5, - planner_profile_data=Path("/tmp/perf_data.npz"), + planner_profile_data=planner_profile_data, is_prefill_worker=True, is_decode_worker=False, durable_kv_events=False, @@ -36,7 +45,7 @@ def test_build_mocker_engine_args_preserves_cli_mapped_fields(monkeypatch): } ), sglang_schedule_policy="lpm", - sglang_page_size=16, + sglang_page_size=128, sglang_max_prefill_tokens=8192, sglang_chunked_prefill_size=2048, sglang_clip_max_new_tokens=1024, @@ -48,17 +57,10 @@ def test_build_mocker_engine_args_preserves_cli_mapped_fields(monkeypatch): model_path="/models/mock", ) - captured = {} - - def fake_mock_engine_args(**kwargs): - captured["kwargs"] = kwargs - return SimpleNamespace(**kwargs) - - monkeypatch.setattr(config, "MockEngineArgs", fake_mock_engine_args) - engine_args = config.build_mocker_engine_args(args) + payload = json.loads(engine_args.dump_json()) - assert captured["kwargs"] == { + assert payload == { "engine_type": "sglang", "num_gpu_blocks": 2048, "block_size": 128, @@ -66,23 +68,37 @@ def fake_mock_engine_args(**kwargs): "max_num_batched_tokens": 4096, "enable_prefix_caching": False, "enable_chunked_prefill": False, - "preemption_mode": "fifo", "speedup_ratio": 2.0, "decode_speedup_ratio": 3.0, "dp_size": 4, "startup_time": 1.5, "worker_type": "prefill", - "planner_profile_data": Path("/tmp/perf_data.npz"), - "enable_local_indexer": True, - "kv_transfer_bandwidth": 123.0, - "reasoning": captured["kwargs"]["reasoning"], - "sglang": captured["kwargs"]["sglang"], + "planner_profile_data": str(planner_profile_data), "aic_backend": "sglang", "aic_system": "h200_sxm", "aic_backend_version": "0.5.6.post2", "aic_tp_size": 8, "aic_model_path": "/models/mock", + "enable_local_indexer": True, + "bootstrap_port": None, + "kv_bytes_per_token": None, + "kv_transfer_bandwidth": 123.0, + "reasoning": { + "start_thinking_token_id": 11, + "end_thinking_token_id": 12, + "thinking_ratio": 0.25, + }, + "zmq_kv_events_port": None, + "zmq_replay_port": None, + "preemption_mode": "fifo", + "router_queue_policy": None, + "sglang": { + "schedule_policy": "lpm", + "page_size": 128, + "max_prefill_tokens": 8192, + "chunked_prefill_size": 2048, + "clip_max_new_tokens": 1024, + "schedule_conservativeness": 0.8, + }, + "has_perf_model": True, } - assert engine_args.planner_profile_data == Path("/tmp/perf_data.npz") - assert engine_args.reasoning is not None - assert engine_args.sglang is not None diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index df9238f23be1..d3f2cfbbe410 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -16,6 +16,7 @@ use dynamo_mocker::loadgen::{ use dynamo_mocker::replay::ReplayArgsMode; use pyo3::{exceptions::PyException, prelude::*}; use pythonize::pythonize; +use serde_json::json; use uuid::Uuid; use super::aic_callback::create_aic_callback; @@ -188,6 +189,7 @@ impl MockEngineArgs { .dp_size(dp_size) .startup_time(startup_time) .worker_type(worker_type) + .planner_profile_data(planner_profile_data.clone()) .aic_backend(aic_backend) .aic_system(aic_system) .aic_backend_version(aic_backend_version) @@ -232,6 +234,60 @@ impl MockEngineArgs { .map_err(|e| PyException::new_err(format!("Failed to parse MockEngineArgs JSON: {e}"))) } + fn dump_json(&self) -> PyResult { + let worker_type = match self.inner.worker_type { + RsWorkerType::Aggregated => "aggregated", + RsWorkerType::Prefill => "prefill", + RsWorkerType::Decode => "decode", + }; + let engine_type = match self.inner.engine_type { + RsMockerEngineType::Vllm => "vllm", + RsMockerEngineType::Sglang => "sglang", + }; + let preemption_mode = match self.inner.preemption_mode { + RsPreemptionMode::Lifo => "lifo", + RsPreemptionMode::Fifo => "fifo", + }; + let router_queue_policy = self + .inner + .router_queue_policy + .as_ref() + .map(|policy| policy.to_string()); + let payload = json!({ + "engine_type": engine_type, + "num_gpu_blocks": self.inner.num_gpu_blocks, + "block_size": self.inner.block_size, + "max_num_seqs": self.inner.max_num_seqs, + "max_num_batched_tokens": self.inner.max_num_batched_tokens, + "enable_prefix_caching": self.inner.enable_prefix_caching, + "enable_chunked_prefill": self.inner.enable_chunked_prefill, + "speedup_ratio": self.inner.speedup_ratio, + "decode_speedup_ratio": self.inner.decode_speedup_ratio, + "dp_size": self.inner.dp_size, + "startup_time": self.inner.startup_time, + "worker_type": worker_type, + "planner_profile_data": self.inner.planner_profile_data, + "aic_backend": self.inner.aic_backend, + "aic_system": self.inner.aic_system, + "aic_backend_version": self.inner.aic_backend_version, + "aic_tp_size": self.inner.aic_tp_size, + "aic_model_path": self.inner.aic_model_path, + "enable_local_indexer": self.inner.enable_local_indexer, + "bootstrap_port": self.inner.bootstrap_port, + "kv_bytes_per_token": self.inner.kv_bytes_per_token, + "kv_transfer_bandwidth": self.inner.kv_transfer_bandwidth, + "reasoning": self.inner.reasoning, + "zmq_kv_events_port": self.inner.zmq_kv_events_port, + "zmq_replay_port": self.inner.zmq_replay_port, + "preemption_mode": preemption_mode, + "router_queue_policy": router_queue_policy, + "sglang": self.inner.sglang, + "has_perf_model": true, + }); + serde_json::to_string_pretty(&payload) + .map_err(|e| PyException::new_err(format!("Failed to serialize MockEngineArgs: {e}"))) + } + #[getter] fn block_size(&self) -> usize { self.inner.block_size diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index b6939c3b9893..e32ba1a21a0a 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1239,6 +1239,8 @@ class MockEngineArgs: def from_json(config_json: str) -> "MockEngineArgs": ... + def dump_json(self) -> str: ... + @property def block_size(self) -> int: ... diff --git a/lib/bindings/python/tests/test_replay.py b/lib/bindings/python/tests/test_replay.py index b7bae50b5245..5cfa269ac3c6 100644 --- a/lib/bindings/python/tests/test_replay.py +++ b/lib/bindings/python/tests/test_replay.py @@ -6,13 +6,11 @@ import subprocess import sys from pathlib import Path -from types import SimpleNamespace import pytest from dynamo.llm import KvRouterConfig, MockEngineArgs from dynamo.replay import run_synthetic_trace_replay, run_trace_replay -from dynamo.replay.main import main from dynamo.replay.reporting import format_report_table, write_report_json pytestmark = [ @@ -768,158 +766,6 @@ def test_write_report_json_creates_file(tmp_path): ) -def test_replay_cli_prints_table_and_saves_json(tmp_path, monkeypatch, capsys): - report = { - "mean_ttft_ms": 10.0, - "min_ttft_ms": 9.0, - "max_ttft_ms": 12.0, - "p99_ttft_ms": 12.0, - "p90_ttft_ms": 11.0, - "p75_ttft_ms": 10.5, - "std_ttft_ms": 1.0, - "output_throughput_tok_s": 123.0, - "request_throughput_rps": 4.0, - "completed_requests": 3, - } - - def fake_run(*args, **kwargs): - return report - - monkeypatch.setattr("dynamo.replay.main.run_synthetic_trace_replay", fake_run) - report_path = tmp_path / "cli_report.json" - - exit_code = main( - [ - "--input-tokens", - "16", - "--output-tokens", - "8", - "--request-count", - "3", - "--report-json", - str(report_path), - ] - ) - - assert exit_code == 0 - stdout = capsys.readouterr().out - assert "NVIDIA AIPerf | LLM Metrics" in stdout - assert "Saved full report to:" in stdout - assert '"completed_requests"' not in stdout - assert json.loads(report_path.read_text(encoding="utf-8")) == report - - -def test_replay_cli_passes_multiturn_workload_kwargs(monkeypatch): - captured = {} - - def fake_run(*args, **kwargs): - captured["args"] = args - captured["kwargs"] = kwargs - return { - "completed_requests": 4, - "request_throughput_rps": 1.0, - "output_throughput_tok_s": 1.0, - } - - monkeypatch.setattr("dynamo.replay.main.run_synthetic_trace_replay", fake_run) - - exit_code = main( - [ - "--input-tokens", - "16", - "--output-tokens", - "8", - "--request-count", - "2", - "--turns-per-session", - "2", - "--shared-prefix-ratio", - "0.5", - "--num-prefix-groups", - "3", - "--inter-turn-delay-ms", - "7.0", - ] - ) - - assert exit_code == 0 - assert captured["args"] == (16, 8, 2) - assert captured["kwargs"]["turns_per_session"] == 2 - assert captured["kwargs"]["shared_prefix_ratio"] == 0.5 - assert captured["kwargs"]["num_prefix_groups"] == 3 - assert captured["kwargs"]["inter_turn_delay_ms"] == 7.0 - - -@pytest.mark.parametrize( - ("engine_args_flag", "engine_args_kwarg"), - [ - ("--extra-engine-args", "extra_engine_args"), - ("--prefill-engine-args", "prefill_engine_args"), - ("--decode-engine-args", "decode_engine_args"), - ], -) -@pytest.mark.parametrize("workload", ["synthetic", "trace"]) -def test_replay_cli_resolves_planner_profile_data_from_benchmarks( - tmp_path, monkeypatch, engine_args_flag, engine_args_kwarg, workload -): - planner_profile_data = _planner_profile_data_npz_path() - assert planner_profile_data.exists() - - captured = {} - resolved_paths = [] - - def fake_run(*args, **kwargs): - captured["args"] = args - captured["kwargs"] = kwargs - return { - "completed_requests": 1, - "request_throughput_rps": 1.0, - "output_throughput_tok_s": 1.0, - } - - def fake_resolve_planner_profile_data(path): - resolved_paths.append(path) - return SimpleNamespace(npz_path=planner_profile_data) - - monkeypatch.setattr("dynamo.replay.main.run_trace_replay", fake_run) - monkeypatch.setattr("dynamo.replay.main.run_synthetic_trace_replay", fake_run) - monkeypatch.setattr( - "dynamo.replay.main.resolve_planner_profile_data", - fake_resolve_planner_profile_data, - ) - - cli_args = [ - engine_args_flag, - json.dumps( - { - "block_size": 64, - "speedup_ratio": 1000.0, - "planner_profile_data": str(planner_profile_data), - } - ), - ] - - if workload == "trace": - trace_path = _write_trace_and_args(tmp_path) - cli_args = [str(trace_path), *cli_args] - else: - cli_args = [ - "--input-tokens", - "16", - "--output-tokens", - "8", - "--request-count", - "1", - *cli_args, - ] - - exit_code = main(cli_args) - - assert exit_code == 0 - assert resolved_paths == [planner_profile_data] - assert captured["kwargs"][engine_args_kwarg] is not None - - @pytest.mark.timeout(30) def test_replay_cli_subprocess_synthetic_smoke(tmp_path): report_path = tmp_path / "synthetic_report.json" diff --git a/lib/mocker/src/common/protocols.rs b/lib/mocker/src/common/protocols.rs index 08183abbb3f4..e91ede8177bc 100644 --- a/lib/mocker/src/common/protocols.rs +++ b/lib/mocker/src/common/protocols.rs @@ -286,6 +286,10 @@ pub struct MockEngineArgs { #[builder(default = "WorkerType::Aggregated")] pub worker_type: WorkerType, + /// Original planner profile NPZ path used to materialize `perf_model`. + #[builder(default = "None")] + pub planner_profile_data: Option, + /// Performance model for timing predictions (not serialized, loaded from planner_profile_data) #[serde(skip)] #[builder(default = "Arc::new(PerfModel::default())")] @@ -691,6 +695,7 @@ impl MockEngineArgs { && let Some(path_str) = path_str.as_str() { let npz_path = PathBuf::from(path_str); + builder = builder.planner_profile_data(Some(npz_path.clone())); match PerfModel::from_npz(&npz_path) { Ok(model) => { tracing::info!("Successfully loaded performance model from: {:?}", npz_path); diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index e5c3a5da7717..a1619aca63b5 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -8,14 +8,10 @@ # endpoint tables) that races under concurrent xdist workers. Do not add # @pytest.mark.parallel until DRT endpoint registration is confirmed thread-safe. # -# NOTE: TCP request plane is NOT tested here. These tests use --num-workers > 1 which spawns -# multiple workers in a single process sharing one TCP server. The shared TCP server uses -# endpoint_path (e.g., "generate") as the routing key, causing handler collisions when multiple -# workers register the same endpoint. This is a test-only limitation; production deployments -# with separate processes per worker work correctly with TCP. import asyncio import logging import os +from pathlib import Path from typing import Any, Dict, Optional import aiohttp @@ -64,6 +60,9 @@ BASE_PORT_ZMQ = 11100 # Base port for ZMQ KV event publishing NUM_REQUESTS = 100 BLOCK_SIZE = 16 +PLANNER_PROFILE_DATA_DIR = ( + Path(__file__).resolve().parents[1] / "planner/profiling_results/H200_TP1P_TP1D" +) def get_unique_ports( @@ -172,6 +171,20 @@ def _build_mocker_command( command.extend(["--preemption-mode", str(mocker_args["preemption_mode"])]) if "dp_size" in mocker_args: command.extend(["--data-parallel-size", str(mocker_args["dp_size"])]) + if "planner_profile_data" in mocker_args: + command.extend( + ["--planner-profile-data", str(mocker_args["planner_profile_data"])] + ) + if mocker_args.get("aic_perf_model") is True: + command.append("--aic-perf-model") + if "aic_system" in mocker_args: + command.extend(["--aic-system", str(mocker_args["aic_system"])]) + if "aic_backend_version" in mocker_args: + command.extend( + ["--aic-backend-version", str(mocker_args["aic_backend_version"])] + ) + if "aic_tp_size" in mocker_args: + command.extend(["--aic-tp-size", str(mocker_args["aic_tp_size"])]) # Use --durable-kv-events to enable JetStream mode (local indexer disabled) if mocker_args.get("durable_kv_events") is True: command.append("--durable-kv-events") @@ -641,17 +654,29 @@ def __exit__(self, exc_type, exc_val, exc_tb): @pytest.mark.timeout(120) # bumped for xdist contention (was 42s; ~13.80s serial avg) @pytest.mark.parametrize( - "router_mode,durable_kv_events", + "router_mode,durable_kv_events,mocker_args_override", [ - pytest.param("kv", False, id="kv-nondurable"), - pytest.param("kv", True, id="kv-durable"), - pytest.param("round-robin", False, id="roundrobin"), - pytest.param("random", False, id="random"), - pytest.param("power-of-two", False, id="power-of-two"), + pytest.param("kv", False, {}, id="kv-nondurable"), + pytest.param( + "kv", + False, + {"planner_profile_data": PLANNER_PROFILE_DATA_DIR}, + id="kv-planner", + ), + pytest.param( + "kv", + False, + {"aic_perf_model": True, "aic_system": "h200_sxm"}, + id="kv-aic", + ), + pytest.param("kv", True, {}, id="kv-durable"), + pytest.param("round-robin", False, {}, id="roundrobin"), + pytest.param("random", False, {}, id="random"), + pytest.param("power-of-two", False, {}, id="power-of-two"), ], indirect=["durable_kv_events"], ) -@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True) +@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True) def test_mocker_router( request, runtime_services_dynamic_ports, @@ -659,6 +684,7 @@ def test_mocker_router( router_mode, request_plane, durable_kv_events, + mocker_args_override, ): """Test router with multiple mocker engine instances across all router modes. @@ -675,6 +701,7 @@ def test_mocker_router( "block_size": BLOCK_SIZE, "durable_kv_events": durable_kv_events, } + mocker_args.update(mocker_args_override) with MockerProcess( request, From d39ab78de789c65e851f84771d4ad7da40006bd6 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 22:34:28 -0700 Subject: [PATCH 13/24] fail fast when AIC ghosts its DB Signed-off-by: PeaBrane --- components/src/dynamo/mocker/aic_session.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/components/src/dynamo/mocker/aic_session.py b/components/src/dynamo/mocker/aic_session.py index eb715b067bb7..2c719b183f5b 100644 --- a/components/src/dynamo/mocker/aic_session.py +++ b/components/src/dynamo/mocker/aic_session.py @@ -14,7 +14,7 @@ from aiconfigurator.sdk.backends.factory import get_backend from aiconfigurator.sdk.inference_session import InferenceSession from aiconfigurator.sdk.models import get_model -from aiconfigurator.sdk.perf_database import get_database +from aiconfigurator.sdk.perf_database import get_database, get_supported_databases logger = logging.getLogger(__name__) @@ -40,6 +40,14 @@ def __init__( ) database = get_database(system=system, backend=backend_name, version=version) + if database is None: + supported = get_supported_databases().get(system, {}).get(backend_name, []) + supported_versions = ", ".join(supported) if supported else "" + raise RuntimeError( + "AIC perf database not found for " + f"system={system!r}, backend={backend_name!r}, version={version!r}. " + f"Supported versions for this system/backend: {supported_versions}" + ) model_config = config.ModelConfig(tp_size=tp_size) model = get_model( model_path=model_path, From aebd26338ef7f5c774101f5c86af0f47015eb7e1 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 22:40:43 -0700 Subject: [PATCH 14/24] clarify when mocker actually needs AIC Signed-off-by: PeaBrane --- docs/mocker/mocker.md | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/docs/mocker/mocker.md b/docs/mocker/mocker.md index 12d5486efd39..11bd917ae4b2 100644 --- a/docs/mocker/mocker.md +++ b/docs/mocker/mocker.md @@ -96,7 +96,7 @@ python -m dynamo.mocker \ | `--sglang-chunked-prefill-size` | 8192 | SGLang chunked-prefill chunk size | | `--sglang-clip-max-new-tokens` | 4096 | SGLang admission-budget cap for max new tokens | | `--sglang-schedule-conservativeness` | 1.0 | SGLang schedule conservativeness factor | -| `--aic-perf-model` | False | Use AIC SDK for latency prediction instead of interpolated/polynomial models. Requires `aiconfigurator` SDK installed (install with `pip install ai-dynamo[mocker]`) | +| `--aic-perf-model` | False | Use AIC SDK for latency prediction instead of interpolated/polynomial models. Opt-in only: default mocker and replay paths do not use AIC. Requires `aiconfigurator` installed and usable AIC systems/perf data for the requested `system/backend/version` tuple | | `--aic-system` | `h200_sxm` | AIC system name (e.g., `h200_sxm`). Used with `--aic-perf-model` | | `--aic-backend-version` | Auto | AIC backend engine version (e.g., `0.12.0` for vLLM). If not set, uses the default version for the backend | | `--aic-tp-size` | 1 | Tensor parallel size for AIC latency prediction. Only affects AIC performance model lookups, not mocker scheduling | @@ -250,7 +250,7 @@ python -m dynamo.mocker \ To use the AIC SDK for latency prediction: ```bash -pip install ai-dynamo[mocker] +uv pip install '.[mocker]' python -m dynamo.mocker \ --model-path nvidia/Llama-3.1-8B-Instruct-FP8 \ @@ -261,6 +261,13 @@ python -m dynamo.mocker \ The AIC model automatically uses `--model-path` and `--engine-type` to select the appropriate performance data. Available systems include `h200_sxm`, `h100_sxm`, etc. (see AIC SDK documentation for the full list). +Important notes: + +- AIC is opt-in. If you do not pass `--aic-perf-model`, `python -m dynamo.mocker` does not use AIC. +- `python -m dynamo.replay` also does not use AIC unless you explicitly put AIC fields in the engine-args JSON. +- `aiconfigurator` must be able to load the requested performance database for the selected `system/backend/version`. If the SDK is installed but the backing systems data is missing or unreadable, mocker now fails fast at startup with a clear error instead of failing later on first request. +- In development environments, this may require pointing Python at a source checkout of `aiconfigurator` with real Git LFS payloads materialized in its `systems/` directory. + When using `python -m dynamo.replay`, there are no dedicated AIC flags. For aggregated replay, pass the equivalent fields via `--extra-engine-args`: @@ -395,7 +402,7 @@ The mocker supports three timing prediction modes: **Interpolated Model:** Loads actual profiling data from an NPZ file containing measured prefill and decode latencies. The mocker interpolates between data points to predict timing for any input size. This enables high-fidelity simulation matching a specific hardware configuration. -**AIC Model (`--aic-perf-model`):** Uses the NVIDIA AI Configurator (AIC) SDK for latency prediction. AIC provides calibrated performance models for specific GPU/model/engine combinations, predicting prefill and decode latency as a function of batch size, sequence length, and prefix cache hits. The model path is automatically derived from `--model-path`, and the engine type from `--engine-type`. This mode requires the `aiconfigurator` SDK, installable via `pip install ai-dynamo[mocker]`. +**AIC Model (`--aic-perf-model`):** Uses the NVIDIA AI Configurator (AIC) SDK for latency prediction. AIC provides calibrated performance models for specific GPU/model/engine combinations, predicting prefill and decode latency as a function of batch size, sequence length, and prefix cache hits. The model path is automatically derived from `--model-path`, and the engine type from `--engine-type`. This mode is opt-in and requires both the `aiconfigurator` SDK and loadable systems/perf data for the requested tuple. ### Bootstrap Rendezvous (Disaggregated Serving) From 62f50725402c2022dd332ad6d6f3e150f450e75e Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 22:58:58 -0700 Subject: [PATCH 15/24] fix mocker mypy kwargs tantrum Signed-off-by: PeaBrane --- components/src/dynamo/mocker/config.py | 74 +++++++++++++------------- 1 file changed, 37 insertions(+), 37 deletions(-) diff --git a/components/src/dynamo/mocker/config.py b/components/src/dynamo/mocker/config.py index fec3a5af6479..59716325dd82 100644 --- a/components/src/dynamo/mocker/config.py +++ b/components/src/dynamo/mocker/config.py @@ -41,7 +41,7 @@ def _build_sglang_args(args: argparse.Namespace) -> SglangArgs | None: return SglangArgs(**sglang_args) -def _build_mocker_engine_args_kwargs(args: argparse.Namespace) -> dict[str, object]: +def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: worker_type = ( "prefill" if getattr(args, "is_prefill_worker", False) @@ -49,44 +49,44 @@ def _build_mocker_engine_args_kwargs(args: argparse.Namespace) -> dict[str, obje if getattr(args, "is_decode_worker", False) else "aggregated" ) - kwargs: dict[str, object] = { - "engine_type": getattr(args, "engine_type", None) or "vllm", - "num_gpu_blocks": getattr(args, "num_gpu_blocks", _DEFAULT_NUM_GPU_BLOCKS), - "block_size": getattr(args, "block_size", 0) or 0, - "max_num_seqs": getattr(args, "max_num_seqs", _DEFAULT_MAX_NUM_SEQS), - "max_num_batched_tokens": getattr( + aic_backend = None + aic_system = None + aic_backend_version = None + aic_tp_size = None + aic_model_path = None + if getattr(args, "aic_perf_model", False): + aic_backend = getattr(args, "engine_type", None) or "vllm" + aic_system = getattr(args, "aic_system", None) + aic_backend_version = getattr(args, "aic_backend_version", None) + aic_tp_size = getattr(args, "aic_tp_size", None) + aic_model_path = getattr(args, "model_path", None) + return MockEngineArgs( + engine_type=getattr(args, "engine_type", None) or "vllm", + num_gpu_blocks=getattr(args, "num_gpu_blocks", _DEFAULT_NUM_GPU_BLOCKS), + block_size=getattr(args, "block_size", 0) or 0, + max_num_seqs=getattr(args, "max_num_seqs", _DEFAULT_MAX_NUM_SEQS), + max_num_batched_tokens=getattr( args, "max_num_batched_tokens", _DEFAULT_MAX_NUM_BATCHED_TOKENS ), - "enable_prefix_caching": getattr(args, "enable_prefix_caching", True), - "enable_chunked_prefill": getattr(args, "enable_chunked_prefill", True), - "speedup_ratio": getattr(args, "speedup_ratio", 1.0), - "decode_speedup_ratio": getattr(args, "decode_speedup_ratio", 1.0), - "dp_size": getattr(args, "dp_size", 1), - "startup_time": getattr(args, "startup_time", None), - "worker_type": worker_type, - "planner_profile_data": getattr(args, "planner_profile_data", None), - "aic_backend": None, - "aic_system": None, - "aic_backend_version": None, - "aic_tp_size": None, - "aic_model_path": None, - "enable_local_indexer": not getattr(args, "durable_kv_events", False), - "kv_transfer_bandwidth": getattr(args, "kv_transfer_bandwidth", None), - "reasoning": _parse_reasoning_config(getattr(args, "reasoning", None)), - "sglang": _build_sglang_args(args), - "preemption_mode": getattr(args, "preemption_mode", "lifo"), - } - if getattr(args, "aic_perf_model", False): - kwargs["aic_backend"] = getattr(args, "engine_type", None) or "vllm" - kwargs["aic_system"] = getattr(args, "aic_system", None) - kwargs["aic_backend_version"] = getattr(args, "aic_backend_version", None) - kwargs["aic_tp_size"] = getattr(args, "aic_tp_size", None) - kwargs["aic_model_path"] = getattr(args, "model_path", None) - return kwargs - - -def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: - return MockEngineArgs(**_build_mocker_engine_args_kwargs(args)) + enable_prefix_caching=getattr(args, "enable_prefix_caching", True), + enable_chunked_prefill=getattr(args, "enable_chunked_prefill", True), + speedup_ratio=getattr(args, "speedup_ratio", 1.0), + decode_speedup_ratio=getattr(args, "decode_speedup_ratio", 1.0), + dp_size=getattr(args, "dp_size", 1), + startup_time=getattr(args, "startup_time", None), + worker_type=worker_type, + planner_profile_data=getattr(args, "planner_profile_data", None), + aic_backend=aic_backend, + aic_system=aic_system, + aic_backend_version=aic_backend_version, + aic_tp_size=aic_tp_size, + aic_model_path=aic_model_path, + enable_local_indexer=not getattr(args, "durable_kv_events", False), + kv_transfer_bandwidth=getattr(args, "kv_transfer_bandwidth", None), + reasoning=_parse_reasoning_config(getattr(args, "reasoning", None)), + sglang=_build_sglang_args(args), + preemption_mode=getattr(args, "preemption_mode", "lifo"), + ) def load_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: From 7193e9d5572c90b521927dc4d83cfb1e00a95c53 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 23:00:20 -0700 Subject: [PATCH 16/24] split kv_router so it stops hoarding Signed-off-by: PeaBrane --- lib/llm/src/kv_router.rs | 199 +--- lib/llm/src/kv_router/indexer.rs | 233 ++++ .../kv_router/publisher/event_processor.rs | 326 +++++ lib/llm/src/kv_router/publisher/mod.rs | 396 ++++++ .../{publisher.rs => publisher/tests.rs} | 1059 +---------------- .../src/kv_router/publisher/worker_metrics.rs | 116 ++ .../src/kv_router/publisher/zmq_listener.rs | 161 +++ lib/llm/src/kv_router/remote_indexer.rs | 65 - 8 files changed, 1251 insertions(+), 1304 deletions(-) create mode 100644 lib/llm/src/kv_router/indexer.rs create mode 100644 lib/llm/src/kv_router/publisher/event_processor.rs create mode 100644 lib/llm/src/kv_router/publisher/mod.rs rename lib/llm/src/kv_router/{publisher.rs => publisher/tests.rs} (67%) create mode 100644 lib/llm/src/kv_router/publisher/worker_metrics.rs create mode 100644 lib/llm/src/kv_router/publisher/zmq_listener.rs delete mode 100644 lib/llm/src/kv_router/remote_indexer.rs diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index c001a341d3b6..b8ec84b6c632 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -1,20 +1,16 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Instant; use anyhow::Result; use dynamo_kv_router::{ - ConcurrentRadixTree, ThreadPoolIndexer, - approx::PruneConfig, config::{KvRouterConfig, RouterConfigOverride}, - indexer::{GetWorkersRequest, KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError}, + indexer::KvRouterError, protocols::KV_EVENT_SUBJECT, protocols::{ - BlockExtraInfo, BlockHashOptions, DpRank, LocalBlockHash, OverlapScores, RouterEvent, - RouterRequest, RouterResponse, TokensWithHashes, WorkerId, WorkerWithDpRank, - compute_block_hash_for_seq, + BlockExtraInfo, BlockHashOptions, DpRank, RouterEvent, RouterRequest, RouterResponse, + TokensWithHashes, WorkerId, WorkerWithDpRank, compute_block_hash_for_seq, }, }; use dynamo_runtime::{ @@ -29,30 +25,29 @@ use dynamo_runtime::{ traits::DistributedRuntimeProvider, }; use futures::stream; -use tokio::sync::oneshot; use tracing::Instrument; use validator::Validate; pub mod cache_control; +pub mod indexer; mod jetstream; pub mod metrics; pub mod prefill_router; pub mod publisher; pub mod push_router; -pub mod remote_indexer; pub mod scheduler; pub mod sequence; pub mod subscriber; pub mod worker_query; pub use cache_control::{CacheControlClient, spawn_pin_prefix}; +pub use indexer::Indexer; pub use prefill_router::PrefillRouter; pub use push_router::{DirectRoutingRouter, KvPushRouter}; use crate::{ discovery::RuntimeConfigWatch, kv_router::{ - remote_indexer::RemoteIndexer, scheduler::{DefaultWorkerSelector, KvScheduler, PotentialLoad}, sequence::{SequenceError, SequenceRequest}, }, @@ -108,188 +103,6 @@ pub fn router_discovery_query(namespace: String, component: String) -> Discovery } } -#[derive(Clone)] -pub enum Indexer { - /// Single-threaded radix tree with channel-based event processing. - /// Supports TTL-based expiration and size-based pruning. - /// Has the ability to persist and snapshot states. - KvIndexer(KvIndexer), - - /// Concurrent radix tree with a thread pool for event processing. - /// Uses sticky worker routing for per-worker event serialization. - /// Does not support TTL/pruning. - Concurrent(Arc>), - - /// Forwards queries to a standalone KV indexer service via the request plane. - /// The standalone indexer manages its own radix tree and event subscription. - Remote(Arc), - - /// Used when we do not wish to use the indexer at all (e.g., when overlap_score_weight is 0). - /// Note: This will cause KV events to accumulate in JetStream as we do not regularly purge them. - None, -} - -impl Indexer { - pub async fn new( - component: &dynamo_runtime::component::Component, - kv_router_config: &KvRouterConfig, - block_size: u32, - model_name: Option, - ) -> Result { - if kv_router_config.overlap_score_weight == 0.0 { - return Ok(Indexer::None); - } - - // Remote indexer: forward queries to a standalone KV indexer service. - if let Some(ref indexer_component_name) = kv_router_config.remote_indexer_component { - let model_name = model_name.ok_or_else(|| { - anyhow::anyhow!( - "model_name is required when remote_indexer_component is configured" - ) - })?; - tracing::info!( - remote_indexer_component = %indexer_component_name, - model_name, - "Using remote KV indexer" - ); - let remote = RemoteIndexer::new(component, indexer_component_name, model_name).await?; - return Ok(Indexer::Remote(Arc::new(remote))); - } - - // Approximate mode (--no-kv-events): always use single-threaded KvIndexer - // with TTL/pruning regardless of event_threads, since updates come from - // routing decisions only, not live KV events from workers. - if !kv_router_config.use_kv_events { - let kv_indexer_metrics = KvIndexerMetrics::from_component(component); - let cancellation_token = component.drt().primary_token(); - let prune_config = Some(PruneConfig { - ttl: Duration::from_secs_f64(kv_router_config.router_ttl_secs), - max_tree_size: kv_router_config.router_max_tree_size, - prune_target_ratio: kv_router_config.router_prune_target_ratio, - }); - return Ok(Indexer::KvIndexer(KvIndexer::new_with_frequency( - cancellation_token, - None, - block_size, - kv_indexer_metrics, - prune_config, - ))); - } - - if kv_router_config.router_event_threads > 1 { - return Ok(Indexer::Concurrent(Arc::new(ThreadPoolIndexer::new( - ConcurrentRadixTree::new(), - kv_router_config.router_event_threads as usize, - block_size, - )))); - } - - let kv_indexer_metrics = KvIndexerMetrics::from_component(component); - let cancellation_token = component.drt().primary_token(); - - Ok(Indexer::KvIndexer(KvIndexer::new_with_frequency( - cancellation_token, - None, // expiration_duration for frequency tracking - block_size, - kv_indexer_metrics, - None, - ))) - } - - pub(crate) async fn find_matches( - &self, - sequence: Vec, - ) -> Result { - match self { - Indexer::KvIndexer(indexer) => indexer.find_matches(sequence).await, - Indexer::Concurrent(tpi) => tpi.find_matches(sequence).await, - Indexer::Remote(remote) => remote.find_matches(sequence).await.map_err(|e| { - tracing::warn!(error = %e, "Remote indexer query failed"); - KvRouterError::IndexerOffline - }), - Indexer::None => Ok(OverlapScores::new()), - } - } - - pub(crate) async fn dump_events(&self) -> Result, KvRouterError> { - match self { - Indexer::KvIndexer(indexer) => indexer.dump_events().await, - Indexer::Concurrent(tpi) => tpi.dump_events().await, - Indexer::Remote(_) => Ok(Vec::new()), - Indexer::None => { - panic!( - "Cannot dump events: indexer does not exist (is overlap_score_weight set to 0?)" - ); - } - } - } - - pub(crate) async fn process_routing_decision_for_request( - &self, - tokens_with_hashes: &mut TokensWithHashes, - worker: WorkerWithDpRank, - ) -> Result<(), KvRouterError> { - match self { - Indexer::KvIndexer(indexer) => { - indexer - .process_routing_decision_for_request(tokens_with_hashes, worker) - .await - } - Indexer::Concurrent(tpi) => { - tpi.process_routing_decision_for_request(tokens_with_hashes, worker) - .await - } - Indexer::Remote(_) => Ok(()), - Indexer::None => Ok(()), - } - } - - pub(crate) async fn apply_event(&self, event: RouterEvent) { - match self { - Indexer::KvIndexer(indexer) => { - if let Err(e) = indexer.event_sender().send(event).await { - tracing::warn!("Failed to send event to indexer: {e}"); - } - } - Indexer::Concurrent(tpi) => tpi.apply_event(event).await, - Indexer::Remote(_) => {} // standalone indexer gets events directly - Indexer::None => {} - } - } - - pub(crate) async fn remove_worker(&self, worker_id: WorkerId) { - match self { - Indexer::KvIndexer(indexer) => { - if let Err(e) = indexer.remove_worker_sender().send(worker_id).await { - tracing::warn!("Failed to send worker removal for {worker_id}: {e}"); - } - } - Indexer::Concurrent(tpi) => { - KvIndexerInterface::remove_worker(tpi.as_ref(), worker_id).await; - } - Indexer::Remote(_) => {} // standalone indexer manages its own workers - Indexer::None => {} - } - } - - pub(crate) async fn get_workers(&self) -> Vec { - match self { - Indexer::KvIndexer(indexer) => { - let (resp_tx, resp_rx) = oneshot::channel(); - let req = GetWorkersRequest { resp: resp_tx }; - if let Err(e) = indexer.get_workers_sender().send(req).await { - tracing::warn!("Failed to send get_workers request: {e}"); - return Vec::new(); - } - resp_rx.await.unwrap_or_default() - } - Indexer::Concurrent(tpi) => tpi.backend().get_workers(), - Indexer::Remote(_) => Vec::new(), - Indexer::None => Vec::new(), - } - } -} - /// A KvRouter only decides which worker you should use. It doesn't send you there. /// TODO: Rename this to indicate it only selects a worker, it does not route. pub struct KvRouter diff --git a/lib/llm/src/kv_router/indexer.rs b/lib/llm/src/kv_router/indexer.rs new file mode 100644 index 000000000000..84794e2cf8a5 --- /dev/null +++ b/lib/llm/src/kv_router/indexer.rs @@ -0,0 +1,233 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Result; +use futures::StreamExt; + +use dynamo_kv_router::{ + ConcurrentRadixTree, ThreadPoolIndexer, + approx::PruneConfig, + config::KvRouterConfig, + indexer::{ + IndexerQueryRequest, IndexerQueryResponse, KV_INDEXER_QUERY_ENDPOINT, KvIndexer, + KvIndexerInterface, KvIndexerMetrics, KvRouterError, + }, + protocols::{ + LocalBlockHash, OverlapScores, RouterEvent, TokensWithHashes, WorkerId, WorkerWithDpRank, + }, +}; +use dynamo_runtime::{ + component::Component, + pipeline::{ManyOut, RouterMode, SingleIn, network::egress::push_router::PushRouter}, + traits::DistributedRuntimeProvider, +}; +use tokio::sync::oneshot; + +pub struct RemoteIndexer { + router: PushRouter, + model_name: String, + namespace: String, +} + +impl RemoteIndexer { + async fn new( + component: &Component, + indexer_component_name: &str, + model_name: String, + ) -> Result { + let namespace = component.namespace().name(); + let indexer_ns = component.namespace(); + let indexer_component = indexer_ns.component(indexer_component_name)?; + let endpoint = indexer_component.endpoint(KV_INDEXER_QUERY_ENDPOINT); + let client = endpoint.client().await?; + let router = + PushRouter::from_client_no_fault_detection(client, RouterMode::RoundRobin).await?; + Ok(Self { + router, + model_name, + namespace, + }) + } + + async fn find_matches(&self, block_hashes: Vec) -> Result { + let request = IndexerQueryRequest { + model_name: self.model_name.clone(), + namespace: self.namespace.clone(), + block_hashes, + }; + let mut stream: ManyOut = + self.router.round_robin(SingleIn::new(request)).await?; + + match stream.next().await { + Some(IndexerQueryResponse::Scores(scores)) => Ok(scores.into()), + Some(IndexerQueryResponse::Error(msg)) => { + Err(anyhow::anyhow!("Remote indexer error: {}", msg)) + } + None => Err(anyhow::anyhow!("Remote indexer returned empty response")), + } + } +} + +#[derive(Clone)] +pub enum Indexer { + KvIndexer(KvIndexer), + Concurrent(Arc>), + Remote(Arc), + None, +} + +impl Indexer { + pub async fn new( + component: &Component, + kv_router_config: &KvRouterConfig, + block_size: u32, + model_name: Option, + ) -> Result { + if kv_router_config.overlap_score_weight == 0.0 { + return Ok(Self::None); + } + + if let Some(ref indexer_component_name) = kv_router_config.remote_indexer_component { + let model_name = model_name.ok_or_else(|| { + anyhow::anyhow!( + "model_name is required when remote_indexer_component is configured" + ) + })?; + tracing::info!( + remote_indexer_component = %indexer_component_name, + model_name, + "Using remote KV indexer" + ); + let remote = RemoteIndexer::new(component, indexer_component_name, model_name).await?; + return Ok(Self::Remote(Arc::new(remote))); + } + + if !kv_router_config.use_kv_events { + let kv_indexer_metrics = KvIndexerMetrics::from_component(component); + let cancellation_token = component.drt().primary_token(); + let prune_config = Some(PruneConfig { + ttl: Duration::from_secs_f64(kv_router_config.router_ttl_secs), + max_tree_size: kv_router_config.router_max_tree_size, + prune_target_ratio: kv_router_config.router_prune_target_ratio, + }); + return Ok(Self::KvIndexer(KvIndexer::new_with_frequency( + cancellation_token, + None, + block_size, + kv_indexer_metrics, + prune_config, + ))); + } + + if kv_router_config.router_event_threads > 1 { + return Ok(Self::Concurrent(Arc::new(ThreadPoolIndexer::new( + ConcurrentRadixTree::new(), + kv_router_config.router_event_threads as usize, + block_size, + )))); + } + + let kv_indexer_metrics = KvIndexerMetrics::from_component(component); + let cancellation_token = component.drt().primary_token(); + + Ok(Self::KvIndexer(KvIndexer::new_with_frequency( + cancellation_token, + None, + block_size, + kv_indexer_metrics, + None, + ))) + } + + pub(crate) async fn find_matches( + &self, + sequence: Vec, + ) -> Result { + match self { + Self::KvIndexer(indexer) => indexer.find_matches(sequence).await, + Self::Concurrent(tpi) => tpi.find_matches(sequence).await, + Self::Remote(remote) => remote.find_matches(sequence).await.map_err(|e| { + tracing::warn!(error = %e, "Remote indexer query failed"); + KvRouterError::IndexerOffline + }), + Self::None => Ok(OverlapScores::new()), + } + } + + pub(crate) async fn dump_events(&self) -> Result, KvRouterError> { + match self { + Self::KvIndexer(indexer) => indexer.dump_events().await, + Self::Concurrent(tpi) => tpi.dump_events().await, + Self::Remote(_) => Ok(Vec::new()), + Self::None => { + panic!( + "Cannot dump events: indexer does not exist (is overlap_score_weight set to 0?)" + ); + } + } + } + + pub(crate) async fn process_routing_decision_for_request( + &self, + tokens_with_hashes: &mut TokensWithHashes, + worker: WorkerWithDpRank, + ) -> Result<(), KvRouterError> { + match self { + Self::KvIndexer(indexer) => { + indexer + .process_routing_decision_for_request(tokens_with_hashes, worker) + .await + } + Self::Concurrent(tpi) => { + tpi.process_routing_decision_for_request(tokens_with_hashes, worker) + .await + } + Self::Remote(_) | Self::None => Ok(()), + } + } + + pub(crate) async fn apply_event(&self, event: RouterEvent) { + match self { + Self::KvIndexer(indexer) => { + if let Err(e) = indexer.event_sender().send(event).await { + tracing::warn!("Failed to send event to indexer: {e}"); + } + } + Self::Concurrent(tpi) => tpi.apply_event(event).await, + Self::Remote(_) | Self::None => {} + } + } + + pub(crate) async fn remove_worker(&self, worker_id: WorkerId) { + match self { + Self::KvIndexer(indexer) => { + if let Err(e) = indexer.remove_worker_sender().send(worker_id).await { + tracing::warn!("Failed to send worker removal for {worker_id}: {e}"); + } + } + Self::Concurrent(tpi) => { + KvIndexerInterface::remove_worker(tpi.as_ref(), worker_id).await; + } + Self::Remote(_) | Self::None => {} + } + } + + pub(crate) async fn get_workers(&self) -> Vec { + match self { + Self::KvIndexer(indexer) => { + let (resp_tx, resp_rx) = oneshot::channel(); + let req = dynamo_kv_router::indexer::GetWorkersRequest { resp: resp_tx }; + if let Err(e) = indexer.get_workers_sender().send(req).await { + tracing::warn!("Failed to send get_workers request: {e}"); + return Vec::new(); + } + resp_rx.await.unwrap_or_default() + } + Self::Concurrent(tpi) => tpi.backend().get_workers(), + Self::Remote(_) | Self::None => Vec::new(), + } + } +} diff --git a/lib/llm/src/kv_router/publisher/event_processor.rs b/lib/llm/src/kv_router/publisher/event_processor.rs new file mode 100644 index 000000000000..b8e9259f0809 --- /dev/null +++ b/lib/llm/src/kv_router/publisher/event_processor.rs @@ -0,0 +1,326 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::future::Future; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use anyhow::Result; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use dynamo_kv_router::RouterEventSink; +use dynamo_kv_router::indexer::LocalKvIndexer; +use dynamo_kv_router::protocols::*; +use dynamo_runtime::transports::event_plane::EventPublisher; +use dynamo_runtime::transports::nats::NatsQueue; + +use crate::kv_router::KV_EVENT_SUBJECT; + +use super::{DEFAULT_MAX_BATCH_BLOCKS, kv_publisher_metrics}; + +/// Accumulator for in-flight KV cache events that will be merged into a single +/// [`RouterEvent`] before being forwarded to the event sink. +#[derive(Debug)] +pub(super) struct BatchingState { + pub(super) pending_removed: Option, + pub(super) pending_stored: Option, + pub(super) next_publish_id: u64, + pub(super) last_dp_rank: u32, + pub(super) last_flush_time: Instant, +} + +impl BatchingState { + pub(super) fn new() -> Self { + Self { + pending_removed: None, + pending_stored: None, + next_publish_id: 1, + last_dp_rank: 0, + last_flush_time: Instant::now(), + } + } + + pub(super) fn has_pending(&self) -> bool { + self.pending_removed.is_some() || self.pending_stored.is_some() + } + + pub(super) fn pending_block_count(&self) -> usize { + self.pending_removed + .as_ref() + .map(|r| r.block_hashes.len()) + .unwrap_or(0) + + self + .pending_stored + .as_ref() + .map(|s| s.blocks.len()) + .unwrap_or(0) + } + + pub(super) fn record_flush_time(&mut self) { + self.last_flush_time = Instant::now(); + } + + pub(super) fn remaining_timeout(&self, timeout_ms: u64) -> Duration { + let timeout = Duration::from_millis(timeout_ms); + let elapsed = self.last_flush_time.elapsed(); + if elapsed >= timeout { + Duration::ZERO + } else { + timeout - elapsed + } + } + + pub(super) fn is_timeout_elapsed(&self, timeout_ms: u64) -> bool { + self.remaining_timeout(timeout_ms) == Duration::ZERO + } + + async fn flush( + &mut self, + publisher: &P, + local_indexer: &Option>, + worker_id: u64, + ) { + if !self.has_pending() { + return; + } + let id = self.next_publish_id; + let dp_rank = self.last_dp_rank; + if let Some(data) = self.pending_removed.take() { + emit( + publisher, + local_indexer, + worker_id, + KvCacheEvent { + event_id: id, + data: KvCacheEventData::Removed(data), + dp_rank, + }, + ) + .await; + } + if let Some(data) = self.pending_stored.take() { + emit( + publisher, + local_indexer, + worker_id, + KvCacheEvent { + event_id: id, + data: KvCacheEventData::Stored(data), + dp_rank, + }, + ) + .await; + } + self.next_publish_id += 1; + self.record_flush_time(); + } +} + +pub(super) struct EventPlanePublisher(pub(super) EventPublisher); + +impl RouterEventSink for EventPlanePublisher { + fn publish_event(&self, event: &RouterEvent) -> impl Future> + Send { + self.0.publish(event) + } +} + +pub(super) struct JetStreamPublisher(pub(super) NatsQueue); + +impl RouterEventSink for JetStreamPublisher { + fn publish_event(&self, event: &RouterEvent) -> impl Future> + Send { + NatsQueue::publish_event(&self.0, KV_EVENT_SUBJECT, event) + } +} + +async fn emit( + publisher: &P, + local_indexer: &Option>, + worker_id: u64, + event: KvCacheEvent, +) { + let router_event = RouterEvent::new(worker_id, event); + if let Some(indexer) = local_indexer + && let Err(e) = indexer.apply_event_with_buffer(router_event.clone()).await + { + tracing::warn!(worker_id, error = %e, "Failed to apply event to local indexer"); + } + if let Err(e) = publisher.publish_event(&router_event).await { + tracing::error!(worker_id, error = %e, "Failed to publish event"); + } +} + +pub(super) async fn run_event_processor_loop( + publisher: P, + worker_id: u64, + cancellation_token: CancellationToken, + mut rx: mpsc::UnboundedReceiver, + local_indexer: Option>, + timeout_ms: Option, + max_batch_blocks: usize, +) { + let mut batching_state = BatchingState::new(); + let mut last_raw_input_id: Option = None; + + loop { + tokio::select! { + _ = cancellation_token.cancelled() => { + tracing::info!("KV Event source received cancellation signal"); + batching_state.flush(&publisher, &local_indexer, worker_id).await; + break; + } + event = rx.recv() => { + let Some(placement_event) = event else { + tracing::debug!("Event processor channel closed."); + batching_state.flush(&publisher, &local_indexer, worker_id).await; + break; + }; + + let raw_event_id = placement_event.event.event_id; + if let Some(last_id) = last_raw_input_id + && raw_event_id > last_id + 1 + { + let gap = raw_event_id - last_id - 1; + tracing::warn!( + worker_id, + last_raw_input_id = last_id, + raw_event_id, + gap, + "Input event gap detected: raw events dropped before batching" + ); + if let Some(metrics) = kv_publisher_metrics() { + metrics.increment_engines_dropped_events(worker_id, gap); + } else { + tracing::warn!( + worker_id, + gap, + "Failed to record dropped events metric: metrics not initialized" + ); + } + } + last_raw_input_id = Some(raw_event_id); + + if !placement_event.placement.is_local_gpu() { + tracing::trace!( + worker_id, + ?placement_event.placement, + event_id = placement_event.event.event_id, + "Skipping non-local-GPU placement event" + ); + continue; + } + + let event = placement_event.event; + tracing::trace!( + "Event processor for worker_id {} processing event: {:?}", + worker_id, + event.data + ); + + let dp_rank_changed = + batching_state.has_pending() && event.dp_rank != batching_state.last_dp_rank; + + match event.data { + KvCacheEventData::Removed(data) => { + if batching_state.pending_stored.is_some() || dp_rank_changed { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + } + match &mut batching_state.pending_removed { + Some(pending) => pending.block_hashes.extend(data.block_hashes), + None => { + batching_state.pending_removed = Some(data); + } + } + } + KvCacheEventData::Stored(data) => { + let should_flush = dp_rank_changed + || batching_state.pending_removed.is_some() + || batching_state.pending_stored.as_ref().is_some_and(|p| { + data.parent_hash != p.blocks.last().map(|b| b.block_hash) + }); + if should_flush { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + } + match &mut batching_state.pending_stored { + Some(pending) => pending.blocks.extend(data.blocks), + None => { + batching_state.pending_stored = Some(data); + } + } + } + KvCacheEventData::Cleared => { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + emit( + &publisher, + &local_indexer, + worker_id, + KvCacheEvent { + event_id: batching_state.next_publish_id, + data: KvCacheEventData::Cleared, + dp_rank: event.dp_rank, + }, + ) + .await; + batching_state.next_publish_id += 1; + } + } + + batching_state.last_dp_rank = event.dp_rank; + + if batching_state.has_pending() + && (timeout_ms.is_none_or(|ms| batching_state.is_timeout_elapsed(ms)) + || batching_state.pending_block_count() > max_batch_blocks) + { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + } + } + _ = tokio::time::sleep( + timeout_ms + .map(|ms| batching_state.remaining_timeout(ms)) + .unwrap_or(Duration::from_secs(3600)) + ), if timeout_ms.is_some() && batching_state.has_pending() => { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + } + } + } +} + +pub(super) async fn start_event_processor( + publisher: P, + worker_id: u64, + cancellation_token: CancellationToken, + rx: mpsc::UnboundedReceiver, + local_indexer: Option>, + batching_timeout_ms: Option, +) { + run_event_processor_loop( + publisher, + worker_id, + cancellation_token, + rx, + local_indexer, + batching_timeout_ms, + DEFAULT_MAX_BATCH_BLOCKS, + ) + .await +} + +pub(super) async fn start_event_processor_jetstream( + publisher: NatsQueue, + worker_id: u64, + cancellation_token: CancellationToken, + rx: mpsc::UnboundedReceiver, + local_indexer: Option>, + batching_timeout_ms: Option, +) { + run_event_processor_loop( + JetStreamPublisher(publisher), + worker_id, + cancellation_token, + rx, + local_indexer, + batching_timeout_ms, + DEFAULT_MAX_BATCH_BLOCKS, + ) + .await +} diff --git a/lib/llm/src/kv_router/publisher/mod.rs b/lib/llm/src/kv_router/publisher/mod.rs new file mode 100644 index 000000000000..357b7bf13abb --- /dev/null +++ b/lib/llm/src/kv_router/publisher/mod.rs @@ -0,0 +1,396 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::atomic::{AtomicU64, Ordering}; + +use anyhow::Result; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use dynamo_kv_router::indexer::{KvIndexerMetrics, LocalKvIndexer}; +use dynamo_kv_router::protocols::*; +pub use dynamo_kv_router::zmq_wire::create_stored_blocks; +#[cfg(test)] +use dynamo_kv_router::zmq_wire::*; +use dynamo_runtime::config::environment_names::nats as env_nats; +use dynamo_runtime::metrics::MetricsHierarchy; +use dynamo_runtime::metrics::prometheus_names::kv_publisher; +use dynamo_runtime::traits::DistributedRuntimeProvider; +use dynamo_runtime::{ + component::Component, + transports::nats::{NatsQueue, Slug}, +}; + +use crate::kv_router::{ + KV_EVENT_SUBJECT, WORKER_KV_INDEXER_BUFFER_SIZE, worker_query::start_worker_kv_query_endpoint, +}; + +mod event_processor; +#[cfg(test)] +mod tests; +mod worker_metrics; +mod zmq_listener; + +#[cfg(test)] +use event_processor::{BatchingState, run_event_processor_loop}; +use event_processor::{ + EventPlanePublisher, start_event_processor, start_event_processor_jetstream, +}; +pub use worker_metrics::WorkerMetricsPublisher; +use zmq_listener::start_zmq_listener; +#[cfg(test)] +use zmq_listener::{ + INITIAL_BACKOFF_MS, MAX_BACKOFF_EXPONENT, MAX_BACKOFF_MS, MAX_CONSECUTIVE_ERRORS, + calculate_backoff_ms, +}; + +const MAX_BATCHING_TIMEOUT_MS: u64 = 15_000; +pub const DEFAULT_BATCHING_TIMEOUT_MS: Option = None; +const DEFAULT_MAX_BATCH_BLOCKS: usize = 128; + +/// Helper function to create a KV stream name from a component and subject. +/// +/// Generates a slugified stream name in the format: +/// `namespace-{namespace}-component-{component}-{subject}` +fn create_kv_stream_name(component: &Component, subject: &str) -> String { + Slug::slugify(&format!( + "namespace.{}.component.{}.{}", + component.namespace().name(), + component.name(), + subject + )) + .to_string() + .replace("_", "-") +} + +/// Metrics for the KV publisher, created via the MetricsHierarchy API. +/// This provides automatic `dynamo_namespace`, `dynamo_component`, and other +/// hierarchy labels for free. +pub(super) struct KvPublisherMetrics { + /// Total number of raw events dropped by engines before reaching publisher + pub engines_dropped_events_total: prometheus::IntCounterVec, +} + +static KV_PUBLISHER_METRICS: OnceLock> = OnceLock::new(); + +impl KvPublisherMetrics { + /// Create from a Component, memoized in a static OnceLock. + /// Uses the MetricsHierarchy API which auto-prepends `dynamo_component_`, + /// injects hierarchy labels, and registers with the DRT `MetricsRegistry`. + pub fn from_component(component: &Component) -> Arc { + KV_PUBLISHER_METRICS + .get_or_init(|| { + let metrics = component.metrics(); + match metrics.create_intcountervec( + kv_publisher::ENGINES_DROPPED_EVENTS_TOTAL, + "Total number of raw events dropped by engines before reaching publisher (detected via event_id gaps)", + &["worker_id"], + &[], + ) { + Ok(engines_dropped_events_total) => { + Arc::new(Self { engines_dropped_events_total }) + } + Err(e) => { + tracing::warn!("Failed to create kv_publisher metrics from component: {}. Using unregistered metrics as fallback.", e); + Arc::new(Self::new_unregistered()) + } + } + }) + .clone() + } + + /// Creates unregistered metrics for use when the MetricsRegistry is not available. + /// This is used as a fallback when metric creation fails. + pub fn new_unregistered() -> Self { + Self { + engines_dropped_events_total: prometheus::IntCounterVec::new( + prometheus::Opts::new( + kv_publisher::ENGINES_DROPPED_EVENTS_TOTAL, + "Total number of raw events dropped by engines before reaching publisher (detected via event_id gaps)", + ), + &["worker_id"], + ) + .expect("failed to create engines_dropped_events_total counter"), + } + } + + /// Increment the engines dropped events counter by the given amount. + pub fn increment_engines_dropped_events(&self, worker_id: u64, count: u64) { + self.engines_dropped_events_total + .with_label_values(&[&worker_id.to_string()]) + .inc_by(count); + } +} + +fn kv_publisher_metrics() -> Option> { + KV_PUBLISHER_METRICS.get().cloned() +} + +/// Configure the source of KV events. +/// Currently, only ZMQ is supported. +pub enum KvEventSourceConfig { + Zmq { endpoint: String, topic: String }, +} + +enum KvEventSource { + Zmq { + zmq_handle: tokio::task::JoinHandle<()>, + }, +} + +impl KvEventSource { + fn start( + component: Component, + worker_id: WorkerId, + kv_block_size: u32, + source_config: KvEventSourceConfig, + cancellation_token: CancellationToken, + tx: mpsc::UnboundedSender, + next_event_id: Arc, + ) -> Result { + match source_config { + KvEventSourceConfig::Zmq { endpoint, topic } => { + let zmq_handle = component + .drt() + .runtime() + .secondary() + .spawn(start_zmq_listener( + endpoint, + topic, + worker_id, + tx, + cancellation_token.clone(), + kv_block_size, + next_event_id, + )); + + Ok(KvEventSource::Zmq { zmq_handle }) + } + } + } + + fn shutdown(&self) { + match self { + KvEventSource::Zmq { zmq_handle } => { + zmq_handle.abort(); + } + } + } +} + +/// A publisher of KV events. +pub struct KvEventPublisher { + /// The size of the KV block. + kv_block_size: u32, + /// The source of KV events. + /// Can be `None` if all events provided through [`KvEventPublisher::publish`]. + source: Option, + /// The cancellation token. + cancellation_token: CancellationToken, + /// The ID of the local worker emitting placement events. + worker_id: WorkerId, + /// The channel to send events to. + tx: mpsc::UnboundedSender, + /// Internal monotonic event ID counter. Shared with the ZMQ listener if present. + next_event_id: Arc, +} + +impl KvEventPublisher { + pub fn new( + component: Component, + kv_block_size: u32, + source_config: Option, + ) -> Result { + Self::new_with_local_indexer( + component, + kv_block_size, + source_config, + false, + 0, + DEFAULT_BATCHING_TIMEOUT_MS, + ) + } + + pub fn new_with_local_indexer( + component: Component, + kv_block_size: u32, + source_config: Option, + enable_local_indexer: bool, + dp_rank: DpRank, + batching_timeout_ms: Option, + ) -> Result { + let cancellation_token = CancellationToken::new(); + let batching_timeout_ms = batching_timeout_ms + .filter(|&ms| { + if ms > MAX_BATCHING_TIMEOUT_MS { + tracing::warn!( + requested_ms = ms, + max_ms = MAX_BATCHING_TIMEOUT_MS, + "batching_timeout_ms too high, capping to 15s" + ); + } + ms > 0 + }) + .map(|ms| ms.min(MAX_BATCHING_TIMEOUT_MS)); + + let (tx, rx) = mpsc::unbounded_channel::(); + let worker_id = component.drt().connection_id(); + + KvPublisherMetrics::from_component(&component); + + let component_name = component.name(); + tracing::info!( + "Initializing KvEventPublisher for worker {worker_id} in component {component_name}" + ); + + if enable_local_indexer { + tracing::info!( + "LocalKvIndexer enabled for worker {worker_id} in component {component_name}" + ); + } + + let next_event_id = Arc::new(AtomicU64::new(0)); + + let mut source = None; + if let Some(config) = source_config { + source = Some(KvEventSource::start( + component.clone(), + worker_id, + kv_block_size, + config, + cancellation_token.clone(), + tx.clone(), + next_event_id.clone(), + )?); + } + + let local_indexer = if enable_local_indexer { + let metrics = Arc::new(KvIndexerMetrics::new_unregistered()); + Some(Arc::new(LocalKvIndexer::new( + cancellation_token.clone(), + kv_block_size, + metrics, + WORKER_KV_INDEXER_BUFFER_SIZE, + ))) + } else { + None + }; + + let _local_indexer_query_handle = local_indexer.as_ref().map(|local_indexer_ref| { + let component = component.clone(); + let local_indexer = local_indexer_ref.clone(); + + component + .drt() + .runtime() + .secondary() + .spawn(start_worker_kv_query_endpoint( + component, + worker_id, + dp_rank, + local_indexer, + )) + }); + + let cancellation_token_clone = cancellation_token.clone(); + let local_indexer_clone = local_indexer.clone(); + + if enable_local_indexer { + tracing::info!("Using event plane for KV event publishing (local_indexer mode)"); + let component_clone = component.clone(); + component.drt().runtime().secondary().spawn(async move { + let event_publisher = + match dynamo_runtime::transports::event_plane::EventPublisher::for_component( + &component_clone, + KV_EVENT_SUBJECT, + ) + .await + { + Ok(publisher) => publisher, + Err(e) => { + tracing::error!("Failed to create event publisher: {}", e); + return; + } + }; + + start_event_processor( + EventPlanePublisher(event_publisher), + worker_id, + cancellation_token_clone, + rx, + local_indexer_clone, + batching_timeout_ms, + ) + .await + }); + } else { + let stream_name = create_kv_stream_name(&component, KV_EVENT_SUBJECT); + let nats_server = std::env::var(env_nats::NATS_SERVER) + .unwrap_or_else(|_| "nats://localhost:4222".to_string()); + let mut nats_queue = NatsQueue::new_without_consumer( + stream_name, + nats_server, + std::time::Duration::from_secs(60), + ); + + component.drt().runtime().secondary().spawn(async move { + if let Err(e) = nats_queue.connect().await { + tracing::error!("Failed to connect NatsQueue: {e}"); + return; + } + start_event_processor_jetstream( + nats_queue, + worker_id, + cancellation_token_clone, + rx, + local_indexer_clone, + batching_timeout_ms, + ) + .await + }); + } + + Ok(Self { + kv_block_size, + source, + cancellation_token, + worker_id, + tx, + next_event_id, + }) + } + + pub fn publish(&self, event: KvCacheEvent) -> Result<(), mpsc::error::SendError> { + let placement_event = PlacementEvent::local_gpu(self.worker_id, event); + match self.tx.send(placement_event) { + Ok(()) => Ok(()), + Err(err) => Err(mpsc::error::SendError(err.0.event)), + } + } + + pub fn next_event_id(&self) -> u64 { + self.next_event_id.fetch_add(1, Ordering::SeqCst) + } + + pub fn kv_block_size(&self) -> u32 { + self.kv_block_size + } + + pub fn shutdown(&mut self) { + if !self.cancellation_token.is_cancelled() { + self.cancellation_token.cancel(); + } + + if let Some(source) = self.source.take() { + source.shutdown(); + } + } +} + +impl Drop for KvEventPublisher { + fn drop(&mut self) { + self.shutdown(); + } +} diff --git a/lib/llm/src/kv_router/publisher.rs b/lib/llm/src/kv_router/publisher/tests.rs similarity index 67% rename from lib/llm/src/kv_router/publisher.rs rename to lib/llm/src/kv_router/publisher/tests.rs index 5cc7b9015afe..358f7c3532ee 100644 --- a/lib/llm/src/kv_router/publisher.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -1,1049 +1,17 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - +use super::*; +#[allow(unused_imports)] +use bytes::Bytes; +#[allow(unused_imports)] +use dynamo_kv_router::RouterEventSink; +#[allow(unused_imports)] +use rmp_serde as rmps; +#[allow(unused_imports)] use std::future::Future; -use std::sync::Arc; +#[allow(unused_imports)] use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; -use std::time::{Duration, Instant}; - -use anyhow::Result; -use rmp_serde as rmps; -use tokio::sync::mpsc; -use tokio_util::sync::CancellationToken; -use zeromq::{Socket, SocketRecv, SubSocket}; - -use dynamo_runtime::metrics::MetricsHierarchy; -use dynamo_runtime::traits::DistributedRuntimeProvider; -use dynamo_runtime::transports::event_plane::EventPublisher; -use dynamo_runtime::{ - component::{Component, Namespace}, - transports::nats::{NatsQueue, Slug}, -}; - -/// Helper function to create a KV stream name from a component and subject. -/// -/// Generates a slugified stream name in the format: -/// `namespace-{namespace}-component-{component}-{subject}` -fn create_kv_stream_name(component: &Component, subject: &str) -> String { - Slug::slugify(&format!( - "namespace.{}.component.{}.{}", - component.namespace().name(), - component.name(), - subject - )) - .to_string() - .replace("_", "-") -} - -use dynamo_kv_router::indexer::{KvIndexerMetrics, LocalKvIndexer}; -use dynamo_kv_router::protocols::*; -pub use dynamo_kv_router::zmq_wire::create_stored_blocks; -use dynamo_kv_router::zmq_wire::*; - -use crate::kv_router::{ - KV_EVENT_SUBJECT, KV_METRICS_SUBJECT, WORKER_KV_INDEXER_BUFFER_SIZE, - worker_query::start_worker_kv_query_endpoint, -}; -use dynamo_runtime::config::environment_names::nats as env_nats; - -// Error handling configuration for ZMQ operations -const INITIAL_BACKOFF_MS: u64 = 10; -const MAX_BACKOFF_MS: u64 = 5000; -const MAX_CONSECUTIVE_ERRORS: u32 = 10; -const MAX_BACKOFF_EXPONENT: u32 = 8; // Cap at 2^8 = 256x multiplier to prevent overflow - -// Batching configuration -const MAX_BATCHING_TIMEOUT_MS: u64 = 15_000; // 15 seconds, prevents misconfiguration -pub const DEFAULT_BATCHING_TIMEOUT_MS: Option = None; // disabled by default -const DEFAULT_MAX_BATCH_BLOCKS: usize = 128; // Max blocks to batch before flushing - -// --------------------------------------------------------------------------- -// Engines dropped events metric -// --------------------------------------------------------------------------- - -use std::sync::OnceLock; - -use dynamo_runtime::metrics::prometheus_names::kv_publisher; - -/// Metrics for the KV publisher, created via the MetricsHierarchy API. -/// This provides automatic `dynamo_namespace`, `dynamo_component`, and other -/// hierarchy labels for free. -pub struct KvPublisherMetrics { - /// Total number of raw events dropped by engines before reaching publisher - pub engines_dropped_events_total: prometheus::IntCounterVec, -} - -static KV_PUBLISHER_METRICS: OnceLock> = OnceLock::new(); - -impl KvPublisherMetrics { - /// Create from a Component, memoized in a static OnceLock. - /// Uses the MetricsHierarchy API which auto-prepends `dynamo_component_`, - /// injects hierarchy labels, and registers with the DRT `MetricsRegistry`. - pub fn from_component(component: &Component) -> Arc { - KV_PUBLISHER_METRICS - .get_or_init(|| { - let metrics = component.metrics(); - match metrics.create_intcountervec( - kv_publisher::ENGINES_DROPPED_EVENTS_TOTAL, - "Total number of raw events dropped by engines before reaching publisher (detected via event_id gaps)", - &["worker_id"], - &[], - ) { - Ok(engines_dropped_events_total) => { - Arc::new(Self { engines_dropped_events_total }) - } - Err(e) => { - tracing::warn!("Failed to create kv_publisher metrics from component: {}. Using unregistered metrics as fallback.", e); - Arc::new(Self::new_unregistered()) - } - } - }) - .clone() - } - - /// Creates unregistered metrics for use when the MetricsRegistry is not available. - /// This is used as a fallback when metric creation fails. - pub fn new_unregistered() -> Self { - Self { - engines_dropped_events_total: prometheus::IntCounterVec::new( - prometheus::Opts::new( - kv_publisher::ENGINES_DROPPED_EVENTS_TOTAL, - "Total number of raw events dropped by engines before reaching publisher (detected via event_id gaps)", - ), - &["worker_id"], - ) - .expect("failed to create engines_dropped_events_total counter"), - } - } - - /// Increment the engines dropped events counter by the given amount. - pub fn increment_engines_dropped_events(&self, worker_id: u64, count: u64) { - self.engines_dropped_events_total - .with_label_values(&[&worker_id.to_string()]) - .inc_by(count); - } -} - -/// Get the KV publisher metrics if initialized. -fn kv_publisher_metrics() -> Option> { - KV_PUBLISHER_METRICS.get().cloned() -} - -// ------------------------------------------------------------------------- -// Batching State ----------------------------------------------------------- -// ------------------------------------------------------------------------- - -/// Accumulator for in-flight KV cache events that will be merged into a single -/// [`RouterEvent`] before being forwarded to the event sink. -#[derive(Debug)] -struct BatchingState { - /// Block hashes accumulating for the next Removed event. - pending_removed: Option, - /// Blocks accumulating for the next Stored event. - pending_stored: Option, - /// Monotonic published-batch counter. Increments by 1 per flush so downstream - /// consumers always see consecutive event IDs, regardless of how many raw source - /// events were merged into the batch. - next_publish_id: u64, - /// dp_rank of the events in the current pending batch. - /// A change signals that the batch must be flushed before accumulating further. - last_dp_rank: u32, - /// When we last flushed (or initialized). Used to detect stale pending data: - /// if a new event arrives after a long idle period (exceeding timeout), - /// we flush immediately for lower latency on sparse important events. - last_flush_time: Instant, -} - -impl BatchingState { - fn new() -> Self { - Self { - pending_removed: None, - pending_stored: None, - next_publish_id: 1, - last_dp_rank: 0, - last_flush_time: Instant::now(), - } - } - - fn has_pending(&self) -> bool { - self.pending_removed.is_some() || self.pending_stored.is_some() - } - - fn pending_block_count(&self) -> usize { - self.pending_removed - .as_ref() - .map(|r| r.block_hashes.len()) - .unwrap_or(0) - + self - .pending_stored - .as_ref() - .map(|s| s.blocks.len()) - .unwrap_or(0) - } - - /// Records that a flush just happened. Called after every flush to track - /// idle periods for stale-data detection. - fn record_flush_time(&mut self) { - self.last_flush_time = Instant::now(); - } - - /// Returns the time remaining in the current batch window (zero if already elapsed). - fn remaining_timeout(&self, timeout_ms: u64) -> Duration { - let timeout = Duration::from_millis(timeout_ms); - let elapsed = self.last_flush_time.elapsed(); - if elapsed >= timeout { - Duration::ZERO - } else { - timeout - elapsed - } - } - - /// Returns `true` when the batch window has elapsed (or `timeout_ms` is zero). - fn is_timeout_elapsed(&self, timeout_ms: u64) -> bool { - self.remaining_timeout(timeout_ms) == Duration::ZERO - } -} - -// ------------------------------------------------------------------------- -// KV Event Publishers ----------------------------------------------------- -// ------------------------------------------------------------------------- - -/// Configure the source of KV events. -/// Currently, only ZMQ is supported. -pub enum KvEventSourceConfig { - Zmq { endpoint: String, topic: String }, -} - -/// The source of KV events. -enum KvEventSource { - Zmq { - zmq_handle: tokio::task::JoinHandle<()>, - }, -} - -impl KvEventSource { - /// Start the event source from a [`KvEventSourceConfig`]. - fn start( - component: Component, - worker_id: WorkerId, - kv_block_size: u32, - source_config: KvEventSourceConfig, - cancellation_token: CancellationToken, - tx: mpsc::UnboundedSender, - next_event_id: Arc, - ) -> Result { - match source_config { - KvEventSourceConfig::Zmq { endpoint, topic } => { - let zmq_handle = component - .drt() - .runtime() - .secondary() - .spawn(start_zmq_listener( - endpoint, - topic, - worker_id, - tx, - cancellation_token.clone(), - kv_block_size, - next_event_id, - )); - - Ok(KvEventSource::Zmq { zmq_handle }) - } - } - } - - fn shutdown(&self) { - match self { - KvEventSource::Zmq { zmq_handle } => { - zmq_handle.abort(); - } - } - } -} - -/// A publisher of KV events. -pub struct KvEventPublisher { - /// The size of the KV block. - kv_block_size: u32, - /// The source of KV events. - /// Can be `None` if all events provided through [`KvEventPublisher::publish`]. - source: Option, - /// The cancellation token. - cancellation_token: CancellationToken, - /// The ID of the local worker emitting placement events. - worker_id: WorkerId, - /// The channel to send events to. - tx: mpsc::UnboundedSender, - /// Internal monotonic event ID counter - ensures each event gets a unique, incrementing ID. - /// Shared with the ZMQ listener (if any) to maintain consistency. - next_event_id: Arc, -} - -impl KvEventPublisher { - pub fn new( - component: Component, - kv_block_size: u32, - source_config: Option, - ) -> Result { - Self::new_with_local_indexer( - component, - kv_block_size, - source_config, - false, - 0, - DEFAULT_BATCHING_TIMEOUT_MS, - ) - } - - pub fn new_with_local_indexer( - component: Component, - kv_block_size: u32, - source_config: Option, - enable_local_indexer: bool, - dp_rank: DpRank, - batching_timeout_ms: Option, - ) -> Result { - let cancellation_token = CancellationToken::new(); - // None = disabled (flush every event); Some(0) normalised to None; Some(ms) = opt-in. - // Cap at MAX_BATCHING_TIMEOUT_MS to prevent misconfiguration. - let batching_timeout_ms = batching_timeout_ms - .filter(|&ms| { - if ms > MAX_BATCHING_TIMEOUT_MS { - tracing::warn!( - requested_ms = ms, - max_ms = MAX_BATCHING_TIMEOUT_MS, - "batching_timeout_ms too high, capping to 15s" - ); - } - // if ms is 0, treat as disabled (None) - ms > 0 - }) - .map(|ms| ms.min(MAX_BATCHING_TIMEOUT_MS)); - - let (tx, rx) = mpsc::unbounded_channel::(); - - // Infer worker_id from component's connection - let worker_id = component.drt().connection_id(); - - // Initialize the KV publisher metrics via MetricsHierarchy API - // This provides automatic hierarchy labels (dynamo_namespace, dynamo_component, etc.) - KvPublisherMetrics::from_component(&component); - - let component_name = component.name(); - tracing::info!( - "Initializing KvEventPublisher for worker {worker_id} in component {component_name}" - ); - - if enable_local_indexer { - tracing::info!( - "LocalKvIndexer enabled for worker {worker_id} in component {component_name}" - ); - } - - // Internal monotonic event ID counter - shared with ZMQ listener if any - let next_event_id = Arc::new(AtomicU64::new(0)); - - // Create our event source (if any) - let mut source = None; - if let Some(config) = source_config { - source = Some(KvEventSource::start( - component.clone(), - worker_id, - kv_block_size, - config, - cancellation_token.clone(), - tx.clone(), - next_event_id.clone(), - )?); - } - - // Create local indexer if requested - let local_indexer = if enable_local_indexer { - let metrics = Arc::new(KvIndexerMetrics::new_unregistered()); - Some(Arc::new(LocalKvIndexer::new( - cancellation_token.clone(), - kv_block_size, - metrics, - WORKER_KV_INDEXER_BUFFER_SIZE, - ))) - } else { - None - }; - - // Spawn runtime for router->local indexer comm if requested - let _local_indexer_query_handle = local_indexer.as_ref().map(|local_indexer_ref| { - let component = component.clone(); - let local_indexer = local_indexer_ref.clone(); - - component - .drt() - .runtime() - .secondary() - .spawn(start_worker_kv_query_endpoint( - component, - worker_id, - dp_rank, - local_indexer, - )) - }); - - let cancellation_token_clone = cancellation_token.clone(); - let local_indexer_clone = local_indexer.clone(); - - if enable_local_indexer { - // When local indexer is enabled, use the event plane directly. - // EventPublisher handles transport selection (ZMQ or NATS) based on environment. - // Durability is provided by the local indexer's event buffer. - tracing::info!("Using event plane for KV event publishing (local_indexer mode)"); - let component_clone = component.clone(); - component.drt().runtime().secondary().spawn(async move { - let event_publisher = - match EventPublisher::for_component(&component_clone, KV_EVENT_SUBJECT).await { - Ok(publisher) => publisher, - Err(e) => { - tracing::error!("Failed to create event publisher: {}", e); - return; - } - }; - - start_event_processor( - EventPlanePublisher(event_publisher), - worker_id, - cancellation_token_clone, - rx, - local_indexer_clone, - batching_timeout_ms, - ) - .await - }); - } else { - // When local indexer is disabled, use JetStream (NatsQueue) for durability. - let stream_name = create_kv_stream_name(&component, KV_EVENT_SUBJECT); - let nats_server = std::env::var(env_nats::NATS_SERVER) - .unwrap_or_else(|_| "nats://localhost:4222".to_string()); - let mut nats_queue = NatsQueue::new_without_consumer( - stream_name, - nats_server, - std::time::Duration::from_secs(60), // 1 minute timeout - ); - - component.drt().runtime().secondary().spawn(async move { - if let Err(e) = nats_queue.connect().await { - tracing::error!("Failed to connect NatsQueue: {e}"); - return; - } - start_event_processor_jetstream( - JetStreamPublisher(nats_queue), - worker_id, - cancellation_token_clone, - rx, - local_indexer_clone, - batching_timeout_ms, - ) - .await - }); - } - - Ok(Self { - kv_block_size, - source, - cancellation_token, - worker_id, - tx, - next_event_id, - }) - } - - pub fn publish(&self, event: KvCacheEvent) -> Result<(), mpsc::error::SendError> { - let placement_event = PlacementEvent::local_gpu(self.worker_id, event); - match self.tx.send(placement_event) { - Ok(()) => Ok(()), - Err(err) => Err(mpsc::error::SendError(err.0.event)), - } - } - - /// Get and increment the next event ID atomically. - /// Use this to assign monotonically increasing event IDs to events before publishing. - pub fn next_event_id(&self) -> u64 { - self.next_event_id.fetch_add(1, Ordering::SeqCst) - } - - pub fn kv_block_size(&self) -> u32 { - self.kv_block_size - } - - pub fn shutdown(&mut self) { - if !self.cancellation_token.is_cancelled() { - self.cancellation_token.cancel(); - } - - if let Some(source) = self.source.take() { - source.shutdown(); - } - } -} - -impl Drop for KvEventPublisher { - fn drop(&mut self) { - self.shutdown(); - } -} - -use dynamo_kv_router::RouterEventSink; - -struct EventPlanePublisher(EventPublisher); - -impl RouterEventSink for EventPlanePublisher { - fn publish_event(&self, event: &RouterEvent) -> impl Future> + Send { - self.0.publish(event) - } -} - -struct JetStreamPublisher(NatsQueue); - -impl RouterEventSink for JetStreamPublisher { - fn publish_event(&self, event: &RouterEvent) -> impl Future> + Send { - NatsQueue::publish_event(&self.0, KV_EVENT_SUBJECT, event) - } -} - -/// Publishes a single [`KvCacheEvent`] to the event sink and, when present, the local indexer. -/// Errors are logged and swallowed so the caller loop can continue uninterrupted. -async fn emit( - publisher: &P, - local_indexer: &Option>, - worker_id: u64, - event: KvCacheEvent, -) { - let router_event = RouterEvent::new(worker_id, event); - if let Some(indexer) = local_indexer - && let Err(e) = indexer.apply_event_with_buffer(router_event.clone()).await - { - tracing::warn!(worker_id, error = %e, "Failed to apply event to local indexer"); - } - if let Err(e) = publisher.publish_event(&router_event).await { - tracing::error!(worker_id, error = %e, "Failed to publish event"); - } -} - -impl BatchingState { - /// Publishes any pending batch as a single [`RouterEvent`] and advances the monotonic - /// batch ID. No-ops when nothing is pending, so callers may call unconditionally. - async fn flush( - &mut self, - publisher: &P, - local_indexer: &Option>, - worker_id: u64, - ) { - if !self.has_pending() { - return; - } - let id = self.next_publish_id; - let dp_rank = self.last_dp_rank; - if let Some(data) = self.pending_removed.take() { - emit( - publisher, - local_indexer, - worker_id, - KvCacheEvent { - event_id: id, - data: KvCacheEventData::Removed(data), - dp_rank, - }, - ) - .await; - } - if let Some(data) = self.pending_stored.take() { - emit( - publisher, - local_indexer, - worker_id, - KvCacheEvent { - event_id: id, - data: KvCacheEventData::Stored(data), - dp_rank, - }, - ) - .await; - } - // Consecutive batch IDs (1, 2, 3, …) keep downstream gap-detection happy. - self.next_publish_id += 1; - // Record when we flushed for stale-data detection on next event. - self.record_flush_time(); - } -} - -/// Batching loop: accumulates Removed/Stored events and flushes them as a single -/// [`RouterEvent`] when any of the following conditions are met: -/// - Event type switches (Removed ↔ Stored) -/// - `dp_rank` changes between consecutive events -/// - A `Stored` event's `parent_hash` breaks the sequential chain -/// - The batch window expires (`Some(timeout_ms)`; `None` = disabled, flush every event) -/// - Channel is closed or a cancellation signal is received -async fn run_event_processor_loop( - publisher: P, - worker_id: u64, - cancellation_token: CancellationToken, - mut rx: mpsc::UnboundedReceiver, - local_indexer: Option>, - timeout_ms: Option, - max_batch_blocks: usize, -) { - let mut batching_state = BatchingState::new(); - // Track last raw input event_id for gap detection (dropped events before batching). - // The raw event_id is a globally monotonic counter assigned by the ZMQ listener, - // so any gap here means events were silently dropped (e.g. send error on the channel). - let mut last_raw_input_id: Option = None; - - loop { - tokio::select! { - _ = cancellation_token.cancelled() => { - tracing::info!("KV Event source received cancellation signal"); - batching_state.flush(&publisher, &local_indexer, worker_id).await; - break; - } - event = rx.recv() => { - let Some(placement_event) = event else { - tracing::debug!("Event processor channel closed."); - batching_state.flush(&publisher, &local_indexer, worker_id).await; - break; - }; - - // Warn if the raw input event_id is not consecutive — events were dropped - // (e.g. channel send error) before they reached the batching layer. - let raw_event_id = placement_event.event.event_id; - if let Some(last_id) = last_raw_input_id - && raw_event_id > last_id + 1 - { - let gap = raw_event_id - last_id - 1; - tracing::warn!( - worker_id, - last_raw_input_id = last_id, - raw_event_id, - gap, - "Input event gap detected: raw events dropped before batching" - ); - // Increment Prometheus counter for dropped events (if initialized) - if let Some(metrics) = kv_publisher_metrics() { - metrics.increment_engines_dropped_events(worker_id, gap); - } else { - tracing::warn!( - worker_id, - gap, - "Failed to record dropped events metric: metrics not initialized" - ); - } - } - last_raw_input_id = Some(raw_event_id); - - if !placement_event.placement.is_local_gpu() { - tracing::trace!( - worker_id, - ?placement_event.placement, - event_id = placement_event.event.event_id, - "Skipping non-local-GPU placement event" - ); - continue; - } - - let event = placement_event.event; - tracing::trace!("Event processor for worker_id {} processing event: {:?}", worker_id, event.data); - - let dp_rank_changed = batching_state.has_pending() - && event.dp_rank != batching_state.last_dp_rank; - - match event.data { - KvCacheEventData::Removed(data) => { - if batching_state.pending_stored.is_some() || dp_rank_changed { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - } - match &mut batching_state.pending_removed { - Some(pending) => pending.block_hashes.extend(data.block_hashes), - None => { - batching_state.pending_removed = Some(data); - } - } - } - KvCacheEventData::Stored(data) => { - // Flush if: type switch, dp_rank change, or the chain is broken - // (new event's parent_hash doesn't continue from the last stored block). - let should_flush = dp_rank_changed - || batching_state.pending_removed.is_some() - || batching_state.pending_stored.as_ref().is_some_and(|p| { - data.parent_hash != p.blocks.last().map(|b| b.block_hash) - }); - if should_flush { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - } - match &mut batching_state.pending_stored { - // Only extend blocks; parent_hash stays fixed from the first event. - Some(pending) => pending.blocks.extend(data.blocks), - None => { - batching_state.pending_stored = Some(data); - } - } - } - KvCacheEventData::Cleared => { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - emit(&publisher, &local_indexer, worker_id, KvCacheEvent { - event_id: batching_state.next_publish_id, - data: KvCacheEventData::Cleared, - dp_rank: event.dp_rank, - }).await; - batching_state.next_publish_id += 1; - } - } - - // Track dp_rank after the match so in-flight flushes use the old value. - batching_state.last_dp_rank = event.dp_rank; - - // Flush after every event when disabled (None), or when the window has elapsed, - // or when the batch exceeds the max block count. - // The sleep arm only arms when batching is enabled; this covers the disabled path. - if batching_state.has_pending() - && (timeout_ms.is_none_or(|ms| batching_state.is_timeout_elapsed(ms)) - || batching_state.pending_block_count() > max_batch_blocks) - { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - } - } - // if has some pending and has timeout, and no new events come in, then flush when timeout elapsed to prevent stale events - _ = tokio::time::sleep( - timeout_ms.map(|ms| batching_state.remaining_timeout(ms)).unwrap_or(Duration::from_secs(3600)) - ), if timeout_ms.is_some() && batching_state.has_pending() => { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - } - } - } -} - -/// Batched event processor for ephemeral transports (NATS Core / ZMQ). -async fn start_event_processor( - publisher: P, - worker_id: u64, - cancellation_token: CancellationToken, - rx: mpsc::UnboundedReceiver, - local_indexer: Option>, - batching_timeout_ms: Option, -) { - run_event_processor_loop( - publisher, - worker_id, - cancellation_token, - rx, - local_indexer, - batching_timeout_ms, - DEFAULT_MAX_BATCH_BLOCKS, - ) - .await -} - -/// Batched event processor using JetStream (durable). -async fn start_event_processor_jetstream( - publisher: P, - worker_id: u64, - cancellation_token: CancellationToken, - rx: mpsc::UnboundedReceiver, - local_indexer: Option>, - batching_timeout_ms: Option, -) { - run_event_processor_loop( - publisher, - worker_id, - cancellation_token, - rx, - local_indexer, - batching_timeout_ms, - DEFAULT_MAX_BATCH_BLOCKS, - ) - .await -} - -/// Calculate exponential backoff duration based on consecutive error count -fn calculate_backoff_ms(consecutive_errors: u32) -> u64 { - std::cmp::min( - INITIAL_BACKOFF_MS * 2_u64.pow(consecutive_errors.min(MAX_BACKOFF_EXPONENT)), - MAX_BACKOFF_MS, - ) -} - -pub async fn start_zmq_listener( - zmq_endpoint: String, - zmq_topic: String, - worker_id: WorkerId, - tx: mpsc::UnboundedSender, - cancellation_token: CancellationToken, - kv_block_size: u32, - next_event_id: Arc, -) { - tracing::debug!( - "KVEventPublisher connecting to ZMQ endpoint {} (topic '{}')", - zmq_endpoint, - zmq_topic - ); - - let warning_count = Arc::new(AtomicU32::new(0)); - - let mut socket = SubSocket::new(); - - // Subscribe to the requested topic (empty string == all topics) - if let Err(e) = socket.subscribe(&zmq_topic).await { - tracing::error!("Failed to subscribe on ZMQ socket: {}", e); - return; - } - - // Connect to the ZMQ endpoint. SGLang binds locally, Dynamo connects. - // In multi-node setups, each node runs dynamo.sglang alongside local SGLang ranks, - // so ZMQ connections are always local. NATS handles cross-node event distribution. - if let Err(e) = socket.connect(&zmq_endpoint).await { - tracing::error!("Failed to connect ZMQ SUB socket to {zmq_endpoint}: {e}"); - return; - } - - let mut consecutive_errors = 0u32; - #[expect(unused_assignments)] - let mut exit_reason = "unknown"; - let mut messages_processed = 0u64; - - 'main: loop { - tokio::select! { - biased; - - // Check for cancellation - _ = cancellation_token.cancelled() => { - tracing::debug!("ZMQ listener received cancellation signal"); - exit_reason = "cancellation token cancelled"; - break 'main; - } - - // Receive message - msg_result = socket.recv() => { - let Ok(msg) = msg_result else { - let e = msg_result.unwrap_err(); - consecutive_errors += 1; - - if consecutive_errors >= MAX_CONSECUTIVE_ERRORS { - tracing::error!( - error=%e, - consecutive_errors=%consecutive_errors, - "Too many consecutive ZMQ errors, terminating listener" - ); - exit_reason = "too many consecutive errors"; - break 'main; - } - - // Simple exponential backoff with max exponent to prevent overflow - let backoff_ms = calculate_backoff_ms(consecutive_errors); - - tracing::warn!( - error=%e, - consecutive_errors=%consecutive_errors, - backoff_ms=%backoff_ms, - "Error reading from ZMQ socket, applying exponential backoff" - ); - - tokio::time::sleep(Duration::from_millis(backoff_ms)).await; - continue; - }; - // Reset error count on successful message - consecutive_errors = 0; - - // We expect multipart frames: [topic, seq, payload] - let mut frames: Vec> = msg.into_vec().into_iter().map(|frame| frame.to_vec()).collect(); - - if frames.len() != 3 { - tracing::warn!("Received unexpected ZMQ frame count: expected 3, actual {}", frames.len()); - continue; - } - - // Extract the payload and sequence number. - let payload = frames.pop().unwrap(); - let seq_bytes = frames.pop().unwrap(); - - if seq_bytes.len() != 8 { - tracing::warn!("Invalid sequence number byte length: expected 8, actual {}", seq_bytes.len()); - continue; - } - - // Note: We extract the engine's sequence number for logging but use our own - // internal monotonic counter for event_id to ensure per-dp_rank monotonicity - let engine_seq = u64::from_be_bytes(seq_bytes.try_into().unwrap()); - - // Decode our batch of events. - let batch_result = rmps::from_slice::(&payload); - let Ok(batch) = batch_result else { - let e = batch_result.unwrap_err(); - tracing::warn!("Failed to decode KVEventBatch msgpack: {e}"); - continue; - }; - - tracing::trace!( - "ZMQ listener on {} received batch with {} events (engine_seq={}, dp_rank={})", - zmq_endpoint, - batch.events.len(), - engine_seq, - batch.data_parallel_rank.unwrap_or(0) - ); - - let dp_rank = batch.data_parallel_rank.unwrap_or(0).cast_unsigned(); - for raw_event in batch.events.into_iter() { - // Use shared monotonic event_id counter instead of engine's sequence number - let event_id = next_event_id.fetch_add(1, Ordering::SeqCst); - let worker = WorkerWithDpRank::new(worker_id, dp_rank); - let event = convert_event( - raw_event, - event_id, - kv_block_size, - worker, - &warning_count, - ); - if tx.send(event).is_err() { - tracing::warn!("Failed to send message to channel - receiver dropped"); - exit_reason = "channel receiver dropped"; - break 'main; - } - messages_processed += 1; - } - } - } - } - tracing::debug!( - "ZMQ listener exiting, reason: {}, messages processed: {}", - exit_reason, - messages_processed - ); -} - -// ------------------------------------------------------------------------- -// Metrics Publishers ------------------------------------------------------ -// ------------------------------------------------------------------------- - -/// Metrics data passed through the channel for NATS publishing -#[derive(Debug, Clone, Default, PartialEq)] -struct WorkerMetrics { - dp_rank: DpRank, - active_decode_blocks: u64, -} - -pub struct WorkerMetricsPublisher { - tx: tokio::sync::watch::Sender, - rx: tokio::sync::watch::Receiver, -} - -impl WorkerMetricsPublisher { - pub fn new() -> Result { - let (tx, rx) = tokio::sync::watch::channel(WorkerMetrics::default()); - Ok(WorkerMetricsPublisher { tx, rx }) - } - - /// Publish worker metrics for load monitoring. - /// - /// # Arguments - /// * `dp_rank` - Data parallel rank of the worker (None defaults to 0) - /// * `active_decode_blocks` - Number of active KV cache blocks - pub fn publish(&self, dp_rank: Option, active_decode_blocks: u64) -> Result<()> { - let metrics = WorkerMetrics { - dp_rank: dp_rank.unwrap_or(0), - active_decode_blocks, - }; - tracing::trace!( - "Publish metrics: dp_rank={}, active_decode_blocks={}", - metrics.dp_rank, - metrics.active_decode_blocks - ); - self.tx - .send(metrics) - .map_err(|_| anyhow::anyhow!("metrics channel closed")) - } - - pub async fn create_endpoint(&self, component: Component) -> Result<()> { - let worker_id = component.drt().connection_id(); - self.start_nats_metrics_publishing(component.namespace().clone(), worker_id); - Ok(()) - } - - /// Starts a background task to publish metrics over NATS - /// - /// This task monitors metric changes (specifically active_decode_blocks) - /// and publishes stable metrics to NATS after they've been unchanged for 1ms. - fn start_nats_metrics_publishing(&self, namespace: Namespace, worker_id: u64) { - let nats_rx = self.rx.clone(); - - tokio::spawn(async move { - let event_publisher = - match EventPublisher::for_namespace(&namespace, KV_METRICS_SUBJECT).await { - Ok(publisher) => publisher, - Err(e) => { - tracing::error!("Failed to create metrics publisher: {}", e); - return; - } - }; - - let mut rx = nats_rx; - let mut last_metrics: Option = None; - let mut pending_publish: Option = None; - let mut publish_timer = - Box::pin(tokio::time::sleep(tokio::time::Duration::from_secs(0))); - publish_timer.as_mut().reset(tokio::time::Instant::now()); // Complete immediately - - loop { - tokio::select! { - // Handle metrics changes - result = rx.changed() => { - if result.is_err() { - tracing::debug!( - "Metrics publisher sender dropped, stopping NATS background task" - ); - break; - } - - let metrics = rx.borrow_and_update().clone(); - - // Check if metrics have changed - let has_changed = last_metrics.as_ref() != Some(&metrics); - - // If metrics changed, schedule a publish - if has_changed { - pending_publish = Some(metrics.clone()); - last_metrics = Some(metrics); - - // Start the 1ms timer - publish_timer.as_mut().reset( - tokio::time::Instant::now() + tokio::time::Duration::from_millis(1) - ); - } - } - // Timer expired - publish if we have pending metrics - _ = &mut publish_timer => { - if let Some(metrics) = pending_publish.take() { - let active_load = ActiveLoad { - worker_id, - dp_rank: metrics.dp_rank, - active_decode_blocks: Some(metrics.active_decode_blocks), - active_prefill_tokens: None, - }; - - if let Err(e) = event_publisher.publish(&active_load).await { - tracing::warn!("Failed to publish metrics: {}", e); - } - } - - // Reset timer to pending state to avoid tight loop - // It will be reset to 1ms when metrics actually change - publish_timer.as_mut().reset( - tokio::time::Instant::now() + tokio::time::Duration::from_secs(3600) - ); - } - } - } - }); - } -} - -// ------------------------------------------------------------------------- -// Testing ----------------------------------------------------------------- -// ------------------------------------------------------------------------- +use std::time::Duration; +#[allow(unused_imports)] +use zeromq::{PubSocket, Socket, SocketSend, ZmqMessage}; #[cfg(test)] mod test_event_processing { @@ -1459,9 +427,8 @@ mod test_event_processing { #[cfg(test)] mod tests_startup_helpers { use super::*; - use crate::kv_router::KvIndexer; use bytes::Bytes; - use dynamo_kv_router::indexer::{GetWorkersRequest, KvIndexerInterface}; + use dynamo_kv_router::indexer::{GetWorkersRequest, KvIndexer, KvIndexerInterface}; use dynamo_kv_router::protocols::{ExternalSequenceBlockHash, LocalBlockHash}; use std::sync::{Arc, Mutex}; use zeromq::{PubSocket, Socket, SocketSend, ZmqMessage}; diff --git a/lib/llm/src/kv_router/publisher/worker_metrics.rs b/lib/llm/src/kv_router/publisher/worker_metrics.rs new file mode 100644 index 000000000000..d392e8e5cf2e --- /dev/null +++ b/lib/llm/src/kv_router/publisher/worker_metrics.rs @@ -0,0 +1,116 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use anyhow::Result; + +use dynamo_kv_router::protocols::{ActiveLoad, DpRank}; +use dynamo_runtime::component::{Component, Namespace}; +use dynamo_runtime::traits::DistributedRuntimeProvider; +use dynamo_runtime::transports::event_plane::EventPublisher; + +use crate::kv_router::KV_METRICS_SUBJECT; + +#[derive(Debug, Clone, Default, PartialEq)] +struct WorkerMetrics { + dp_rank: DpRank, + active_decode_blocks: u64, +} + +pub struct WorkerMetricsPublisher { + tx: tokio::sync::watch::Sender, + rx: tokio::sync::watch::Receiver, +} + +impl WorkerMetricsPublisher { + pub fn new() -> Result { + let (tx, rx) = tokio::sync::watch::channel(WorkerMetrics::default()); + Ok(Self { tx, rx }) + } + + pub fn publish(&self, dp_rank: Option, active_decode_blocks: u64) -> Result<()> { + let metrics = WorkerMetrics { + dp_rank: dp_rank.unwrap_or(0), + active_decode_blocks, + }; + tracing::trace!( + "Publish metrics: dp_rank={}, active_decode_blocks={}", + metrics.dp_rank, + metrics.active_decode_blocks + ); + self.tx + .send(metrics) + .map_err(|_| anyhow::anyhow!("metrics channel closed")) + } + + pub async fn create_endpoint(&self, component: Component) -> Result<()> { + let worker_id = component.drt().connection_id(); + self.start_nats_metrics_publishing(component.namespace().clone(), worker_id); + Ok(()) + } + + pub(super) fn start_nats_metrics_publishing(&self, namespace: Namespace, worker_id: u64) { + let nats_rx = self.rx.clone(); + + tokio::spawn(async move { + let event_publisher = + match EventPublisher::for_namespace(&namespace, KV_METRICS_SUBJECT).await { + Ok(publisher) => publisher, + Err(e) => { + tracing::error!("Failed to create metrics publisher: {}", e); + return; + } + }; + + let mut rx = nats_rx; + let mut last_metrics: Option = None; + let mut pending_publish: Option = None; + let mut publish_timer = + Box::pin(tokio::time::sleep(tokio::time::Duration::from_secs(0))); + publish_timer.as_mut().reset(tokio::time::Instant::now()); + + loop { + tokio::select! { + result = rx.changed() => { + if result.is_err() { + tracing::debug!( + "Metrics publisher sender dropped, stopping NATS background task" + ); + break; + } + + let metrics = rx.borrow_and_update().clone(); + let has_changed = last_metrics.as_ref() != Some(&metrics); + + if has_changed { + pending_publish = Some(metrics.clone()); + last_metrics = Some(metrics); + publish_timer.as_mut().reset( + tokio::time::Instant::now() + + tokio::time::Duration::from_millis(1) + ); + } + } + _ = &mut publish_timer => { + if let Some(metrics) = pending_publish.take() { + let active_load = ActiveLoad { + worker_id, + dp_rank: metrics.dp_rank, + active_decode_blocks: Some(metrics.active_decode_blocks), + active_prefill_tokens: None, + }; + + if let Err(e) = event_publisher.publish(&active_load).await { + tracing::warn!("Failed to publish metrics: {}", e); + } + } + + publish_timer.as_mut().reset( + tokio::time::Instant::now() + + tokio::time::Duration::from_secs(3600) + ); + } + } + } + }); + } +} diff --git a/lib/llm/src/kv_router/publisher/zmq_listener.rs b/lib/llm/src/kv_router/publisher/zmq_listener.rs new file mode 100644 index 000000000000..1a589604ad80 --- /dev/null +++ b/lib/llm/src/kv_router/publisher/zmq_listener.rs @@ -0,0 +1,161 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; +use std::time::Duration; + +use rmp_serde as rmps; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use zeromq::{Socket, SocketRecv, SubSocket}; + +use dynamo_kv_router::protocols::*; +use dynamo_kv_router::zmq_wire::*; + +pub(super) const INITIAL_BACKOFF_MS: u64 = 10; +pub(super) const MAX_BACKOFF_MS: u64 = 5000; +pub(super) const MAX_CONSECUTIVE_ERRORS: u32 = 10; +pub(super) const MAX_BACKOFF_EXPONENT: u32 = 8; + +pub(super) fn calculate_backoff_ms(consecutive_errors: u32) -> u64 { + std::cmp::min( + INITIAL_BACKOFF_MS * 2_u64.pow(consecutive_errors.min(MAX_BACKOFF_EXPONENT)), + MAX_BACKOFF_MS, + ) +} + +pub(super) async fn start_zmq_listener( + zmq_endpoint: String, + zmq_topic: String, + worker_id: WorkerId, + tx: mpsc::UnboundedSender, + cancellation_token: CancellationToken, + kv_block_size: u32, + next_event_id: Arc, +) { + tracing::debug!( + "KVEventPublisher connecting to ZMQ endpoint {} (topic '{}')", + zmq_endpoint, + zmq_topic + ); + + let warning_count = Arc::new(AtomicU32::new(0)); + let mut socket = SubSocket::new(); + + if let Err(e) = socket.subscribe(&zmq_topic).await { + tracing::error!("Failed to subscribe on ZMQ socket: {}", e); + return; + } + + if let Err(e) = socket.connect(&zmq_endpoint).await { + tracing::error!("Failed to connect ZMQ SUB socket to {zmq_endpoint}: {e}"); + return; + } + + let mut consecutive_errors = 0u32; + #[expect(unused_assignments)] + let mut exit_reason = "unknown"; + let mut messages_processed = 0u64; + + 'main: loop { + tokio::select! { + biased; + + _ = cancellation_token.cancelled() => { + tracing::debug!("ZMQ listener received cancellation signal"); + exit_reason = "cancellation token cancelled"; + break 'main; + } + + msg_result = socket.recv() => { + let Ok(msg) = msg_result else { + let e = msg_result.unwrap_err(); + consecutive_errors += 1; + + if consecutive_errors >= MAX_CONSECUTIVE_ERRORS { + tracing::error!( + error=%e, + consecutive_errors=%consecutive_errors, + "Too many consecutive ZMQ errors, terminating listener" + ); + exit_reason = "too many consecutive errors"; + break 'main; + } + + let backoff_ms = calculate_backoff_ms(consecutive_errors); + + tracing::warn!( + error=%e, + consecutive_errors=%consecutive_errors, + backoff_ms=%backoff_ms, + "Error reading from ZMQ socket, applying exponential backoff" + ); + + tokio::time::sleep(Duration::from_millis(backoff_ms)).await; + continue; + }; + consecutive_errors = 0; + + let mut frames: Vec> = + msg.into_vec().into_iter().map(|frame| frame.to_vec()).collect(); + + if frames.len() != 3 { + tracing::warn!( + "Received unexpected ZMQ frame count: expected 3, actual {}", + frames.len() + ); + continue; + } + + let payload = frames.pop().unwrap(); + let seq_bytes = frames.pop().unwrap(); + + if seq_bytes.len() != 8 { + tracing::warn!( + "Invalid sequence number byte length: expected 8, actual {}", + seq_bytes.len() + ); + continue; + } + + let engine_seq = u64::from_be_bytes(seq_bytes.try_into().unwrap()); + + let batch_result = rmps::from_slice::(&payload); + let Ok(batch) = batch_result else { + let e = batch_result.unwrap_err(); + tracing::warn!("Failed to decode KVEventBatch msgpack: {e}"); + continue; + }; + + tracing::trace!( + "ZMQ listener on {} received batch with {} events (engine_seq={}, dp_rank={})", + zmq_endpoint, + batch.events.len(), + engine_seq, + batch.data_parallel_rank.unwrap_or(0) + ); + + let dp_rank = batch.data_parallel_rank.unwrap_or(0).cast_unsigned(); + for raw_event in batch.events { + let event_id = next_event_id.fetch_add(1, Ordering::SeqCst); + let worker = WorkerWithDpRank::new(worker_id, dp_rank); + let event = + convert_event(raw_event, event_id, kv_block_size, worker, &warning_count); + if tx.send(event).is_err() { + tracing::warn!("Failed to send message to channel - receiver dropped"); + exit_reason = "channel receiver dropped"; + break 'main; + } + messages_processed += 1; + } + } + } + } + + tracing::debug!( + "ZMQ listener exiting, reason: {}, messages processed: {}", + exit_reason, + messages_processed + ); +} diff --git a/lib/llm/src/kv_router/remote_indexer.rs b/lib/llm/src/kv_router/remote_indexer.rs deleted file mode 100644 index 1d0f161f0f73..000000000000 --- a/lib/llm/src/kv_router/remote_indexer.rs +++ /dev/null @@ -1,65 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use anyhow::Result; -use futures::StreamExt; - -use dynamo_runtime::{ - component::Component, - pipeline::{ManyOut, RouterMode, SingleIn, network::egress::push_router::PushRouter}, -}; - -use dynamo_kv_router::{ - indexer::{IndexerQueryRequest, IndexerQueryResponse, KV_INDEXER_QUERY_ENDPOINT}, - protocols::{LocalBlockHash, OverlapScores}, -}; - -/// A remote indexer that queries a standalone KV indexer via the request plane. -/// -/// Used by the frontend when `remote_indexer_component` is configured. Instead of -/// maintaining a local radix tree, this forwards `find_matches` queries to the -/// standalone indexer service over the Dynamo request plane. -pub struct RemoteIndexer { - router: PushRouter, - model_name: String, - namespace: String, -} - -impl RemoteIndexer { - pub async fn new( - component: &Component, - indexer_component_name: &str, - model_name: String, - ) -> Result { - let namespace = component.namespace().name(); - let indexer_ns = component.namespace(); - let indexer_component = indexer_ns.component(indexer_component_name)?; - let endpoint = indexer_component.endpoint(KV_INDEXER_QUERY_ENDPOINT); - let client = endpoint.client().await?; - let router = - PushRouter::from_client_no_fault_detection(client, RouterMode::RoundRobin).await?; - Ok(Self { - router, - model_name, - namespace, - }) - } - - pub async fn find_matches(&self, block_hashes: Vec) -> Result { - let request = IndexerQueryRequest { - model_name: self.model_name.clone(), - namespace: self.namespace.clone(), - block_hashes, - }; - let mut stream: ManyOut = - self.router.round_robin(SingleIn::new(request)).await?; - - match stream.next().await { - Some(IndexerQueryResponse::Scores(scores)) => Ok(scores.into()), - Some(IndexerQueryResponse::Error(msg)) => { - Err(anyhow::anyhow!("Remote indexer error: {}", msg)) - } - None => Err(anyhow::anyhow!("Remote indexer returned empty response")), - } - } -} From a15b8a068d0cf477230b81e6b158a209f87d0007 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 23:04:35 -0700 Subject: [PATCH 17/24] fix tests header before legal notices bite Signed-off-by: PeaBrane --- lib/llm/src/kv_router/publisher/tests.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index 358f7c3532ee..6039f6e127a8 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -1,3 +1,6 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + use super::*; #[allow(unused_imports)] use bytes::Bytes; From 2ae18e17a4201101adbb8c63b6477f14ea7aded2 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Tue, 24 Mar 2026 23:25:02 -0700 Subject: [PATCH 18/24] fix metrics test import after split Signed-off-by: PeaBrane --- lib/llm/src/kv_router/publisher/tests.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/lib/llm/src/kv_router/publisher/tests.rs b/lib/llm/src/kv_router/publisher/tests.rs index 6039f6e127a8..c4c07e67ffbf 100644 --- a/lib/llm/src/kv_router/publisher/tests.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -1172,6 +1172,7 @@ mod test_exponential_backoff { #[cfg(all(test, feature = "integration"))] mod test_integration_publisher { use super::*; + use crate::kv_router::KV_METRICS_SUBJECT; use dynamo_kv_router::protocols::ActiveLoad; use dynamo_runtime::distributed_test_utils::create_test_drt_async; use dynamo_runtime::transports::event_plane::EventSubscriber; From d88451b9a9eeed66e7e9415fcefdde31cab2265f Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 25 Mar 2026 00:16:49 -0700 Subject: [PATCH 19/24] merge mocker config tests before pytest riots Signed-off-by: PeaBrane --- .../src/dynamo/mocker/tests/test_config.py | 104 ------------------ tests/mocker/test_config.py | 98 +++++++++++++++++ 2 files changed, 98 insertions(+), 104 deletions(-) delete mode 100644 components/src/dynamo/mocker/tests/test_config.py diff --git a/components/src/dynamo/mocker/tests/test_config.py b/components/src/dynamo/mocker/tests/test_config.py deleted file mode 100644 index e7f2c3aeedf7..000000000000 --- a/components/src/dynamo/mocker/tests/test_config.py +++ /dev/null @@ -1,104 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import argparse -import json -from pathlib import Path - -from dynamo.mocker import config - - -def _planner_profile_data_npz_path() -> Path: - return ( - Path(__file__).resolve().parents[5] - / "benchmarks/results/H200_TP1P_TP1D_perf_data.npz" - ) - - -def test_build_mocker_engine_args_preserves_cli_mapped_fields(): - planner_profile_data = _planner_profile_data_npz_path() - assert planner_profile_data.exists() - - args = argparse.Namespace( - engine_type="sglang", - num_gpu_blocks=2048, - block_size=128, - max_num_seqs=64, - max_num_batched_tokens=4096, - enable_prefix_caching=False, - enable_chunked_prefill=False, - preemption_mode="fifo", - speedup_ratio=2.0, - decode_speedup_ratio=3.0, - dp_size=4, - startup_time=1.5, - planner_profile_data=planner_profile_data, - is_prefill_worker=True, - is_decode_worker=False, - durable_kv_events=False, - kv_transfer_bandwidth=123.0, - reasoning=json.dumps( - { - "start_thinking_token_id": 11, - "end_thinking_token_id": 12, - "thinking_ratio": 0.25, - } - ), - sglang_schedule_policy="lpm", - sglang_page_size=128, - sglang_max_prefill_tokens=8192, - sglang_chunked_prefill_size=2048, - sglang_clip_max_new_tokens=1024, - sglang_schedule_conservativeness=0.8, - aic_perf_model=True, - aic_system="h200_sxm", - aic_backend_version="0.5.6.post2", - aic_tp_size=8, - model_path="/models/mock", - ) - - engine_args = config.build_mocker_engine_args(args) - payload = json.loads(engine_args.dump_json()) - - assert payload == { - "engine_type": "sglang", - "num_gpu_blocks": 2048, - "block_size": 128, - "max_num_seqs": 64, - "max_num_batched_tokens": 4096, - "enable_prefix_caching": False, - "enable_chunked_prefill": False, - "speedup_ratio": 2.0, - "decode_speedup_ratio": 3.0, - "dp_size": 4, - "startup_time": 1.5, - "worker_type": "prefill", - "planner_profile_data": str(planner_profile_data), - "aic_backend": "sglang", - "aic_system": "h200_sxm", - "aic_backend_version": "0.5.6.post2", - "aic_tp_size": 8, - "aic_model_path": "/models/mock", - "enable_local_indexer": True, - "bootstrap_port": None, - "kv_bytes_per_token": None, - "kv_transfer_bandwidth": 123.0, - "reasoning": { - "start_thinking_token_id": 11, - "end_thinking_token_id": 12, - "thinking_ratio": 0.25, - }, - "zmq_kv_events_port": None, - "zmq_replay_port": None, - "preemption_mode": "fifo", - "router_queue_policy": None, - "sglang": { - "schedule_policy": "lpm", - "page_size": 128, - "max_prefill_tokens": 8192, - "chunked_prefill_size": 2048, - "clip_max_new_tokens": 1024, - "schedule_conservativeness": 0.8, - }, - "has_perf_model": True, - } diff --git a/tests/mocker/test_config.py b/tests/mocker/test_config.py index d8cf92948e5e..405655450118 100644 --- a/tests/mocker/test_config.py +++ b/tests/mocker/test_config.py @@ -1,7 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import argparse import importlib.util +import json from pathlib import Path from types import SimpleNamespace @@ -133,3 +135,99 @@ def test_entrypoint_args_accept_typed_mocker_engine_args(): ) assert entrypoint_args is not None + + +def _planner_profile_data_npz_path() -> Path: + return ( + Path(__file__).resolve().parents[2] + / "benchmarks/results/H200_TP1P_TP1D_perf_data.npz" + ) + + +def test_build_mocker_engine_args_preserves_cli_mapped_fields(): + planner_profile_data = _planner_profile_data_npz_path() + assert planner_profile_data.exists() + + args = argparse.Namespace( + engine_type="sglang", + num_gpu_blocks=2048, + block_size=128, + max_num_seqs=64, + max_num_batched_tokens=4096, + enable_prefix_caching=False, + enable_chunked_prefill=False, + preemption_mode="fifo", + speedup_ratio=2.0, + decode_speedup_ratio=3.0, + dp_size=4, + startup_time=1.5, + planner_profile_data=planner_profile_data, + is_prefill_worker=True, + is_decode_worker=False, + durable_kv_events=False, + kv_transfer_bandwidth=123.0, + reasoning=json.dumps( + { + "start_thinking_token_id": 11, + "end_thinking_token_id": 12, + "thinking_ratio": 0.25, + } + ), + sglang_schedule_policy="lpm", + sglang_page_size=128, + sglang_max_prefill_tokens=8192, + sglang_chunked_prefill_size=2048, + sglang_clip_max_new_tokens=1024, + sglang_schedule_conservativeness=0.8, + aic_perf_model=True, + aic_system="h200_sxm", + aic_backend_version="0.5.6.post2", + aic_tp_size=8, + model_path="/models/mock", + ) + + engine_args = CONFIG.build_mocker_engine_args(args) + payload = json.loads(engine_args.dump_json()) + + assert payload == { + "engine_type": "sglang", + "num_gpu_blocks": 2048, + "block_size": 128, + "max_num_seqs": 64, + "max_num_batched_tokens": 4096, + "enable_prefix_caching": False, + "enable_chunked_prefill": False, + "speedup_ratio": 2.0, + "decode_speedup_ratio": 3.0, + "dp_size": 4, + "startup_time": 1.5, + "worker_type": "prefill", + "planner_profile_data": str(planner_profile_data), + "aic_backend": "sglang", + "aic_system": "h200_sxm", + "aic_backend_version": "0.5.6.post2", + "aic_tp_size": 8, + "aic_model_path": "/models/mock", + "enable_local_indexer": True, + "bootstrap_port": None, + "kv_bytes_per_token": None, + "kv_transfer_bandwidth": 123.0, + "reasoning": { + "start_thinking_token_id": 11, + "end_thinking_token_id": 12, + "thinking_ratio": 0.25, + }, + "zmq_kv_events_port": None, + "zmq_replay_port": None, + "preemption_mode": "fifo", + "router_queue_policy": None, + "sglang": { + "schedule_policy": "lpm", + "page_size": 128, + "max_prefill_tokens": 8192, + "chunked_prefill_size": 2048, + "clip_max_new_tokens": 1024, + "schedule_conservativeness": 0.8, + }, + "has_perf_model": True, + } From 10bfc3f590cae8ac48469390461084f84da5d256 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 25 Mar 2026 01:12:08 -0700 Subject: [PATCH 20/24] test replay planner profile paths Signed-off-by: PeaBrane --- lib/bindings/python/tests/test_replay.py | 67 ++++++++++++++++++++++++ tests/mocker/test_config.py | 19 +++---- 2 files changed, 77 insertions(+), 9 deletions(-) diff --git a/lib/bindings/python/tests/test_replay.py b/lib/bindings/python/tests/test_replay.py index 5cfa269ac3c6..b097ca883db4 100644 --- a/lib/bindings/python/tests/test_replay.py +++ b/lib/bindings/python/tests/test_replay.py @@ -7,6 +7,7 @@ import sys from pathlib import Path +import numpy as np import pytest from dynamo.llm import KvRouterConfig, MockEngineArgs @@ -258,6 +259,26 @@ def _planner_profile_data_npz_path() -> Path: ) +def _planner_profile_data_dir_path() -> Path: + return ( + Path(__file__).resolve().parents[4] + / "tests/planner/profiling_results/H200_TP1P_TP1D" + ) + + +def _write_planner_profile_data_npz(tmp_path: Path) -> Path: + planner_profile_data = tmp_path / "planner_profile_data.npz" + np.savez( + planner_profile_data, + prefill_isl=np.array([128.0, 256.0]), + prefill_ttft_ms=np.array([4.0, 8.0]), + decode_active_kv_tokens=np.array([1024.0, 2048.0]), + decode_context_length=np.array([128.0, 256.0]), + decode_itl=np.array([[1.0, 1.5], [2.0, 2.5]]), + ) + return planner_profile_data + + def _run_replay_cli(tmp_path, *args): return subprocess.run( [ @@ -798,6 +819,52 @@ def test_replay_cli_subprocess_synthetic_smoke(tmp_path): _assert_basic_report_metrics(report) +@pytest.mark.timeout(30) +@pytest.mark.parametrize("planner_profile_data_kind", ["dir", "npz"]) +def test_replay_cli_subprocess_synthetic_smoke_accepts_planner_profile_data( + tmp_path, planner_profile_data_kind +): + report_path = tmp_path / f"synthetic_report_{planner_profile_data_kind}.json" + planner_profile_data = ( + _planner_profile_data_dir_path() + if planner_profile_data_kind == "dir" + else _write_planner_profile_data_npz(tmp_path) + ) + + completed = _run_replay_cli( + tmp_path, + "--input-tokens", + "250", + "--output-tokens", + "25", + "--request-count", + "10", + "--num-workers", + "4", + "--replay-concurrency", + "4", + "--report-json", + str(report_path), + "--extra-engine-args", + json.dumps( + { + "block_size": 64, + "speedup_ratio": 1000.0, + "planner_profile_data": str(planner_profile_data), + } + ), + ) + + report = _assert_replay_cli_outputs(completed, report_path) + _assert_basic_report_counts( + report, + num_requests=10, + input_tokens=250, + output_tokens=25, + ) + _assert_basic_report_metrics(report) + + @pytest.mark.timeout(30) def test_replay_cli_subprocess_synthetic_multiturn_smoke(tmp_path): report_path = tmp_path / "synthetic_multiturn_report.json" diff --git a/tests/mocker/test_config.py b/tests/mocker/test_config.py index 405655450118..06b3fe3be092 100644 --- a/tests/mocker/test_config.py +++ b/tests/mocker/test_config.py @@ -7,6 +7,7 @@ from pathlib import Path from types import SimpleNamespace +import numpy as np import pytest from dynamo.llm import EngineType, EntrypointArgs @@ -137,17 +138,17 @@ def test_entrypoint_args_accept_typed_mocker_engine_args(): assert entrypoint_args is not None -def _planner_profile_data_npz_path() -> Path: - return ( - Path(__file__).resolve().parents[2] - / "benchmarks/results/H200_TP1P_TP1D_perf_data.npz" +def test_build_mocker_engine_args_preserves_cli_mapped_fields(tmp_path): + planner_profile_data = tmp_path / "planner_profile_data.npz" + np.savez( + planner_profile_data, + prefill_isl=np.array([128.0, 256.0]), + prefill_ttft_ms=np.array([4.0, 8.0]), + decode_active_kv_tokens=np.array([1024.0, 2048.0]), + decode_context_length=np.array([128.0, 256.0]), + decode_itl=np.array([[1.0, 1.5], [2.0, 2.5]]), ) - -def test_build_mocker_engine_args_preserves_cli_mapped_fields(): - planner_profile_data = _planner_profile_data_npz_path() - assert planner_profile_data.exists() - args = argparse.Namespace( engine_type="sglang", num_gpu_blocks=2048, From 02e4b1e0a3f8910346822d73acb50023577b5423 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 25 Mar 2026 10:07:25 -0700 Subject: [PATCH 21/24] model decode handoff without prefill fibs Signed-off-by: PeaBrane --- lib/llm/src/mocker.rs | 20 +-- lib/mocker/src/common/protocols.rs | 2 + lib/mocker/src/common/utils.rs | 88 +++++++++++-- lib/mocker/src/replay/offline/disagg.rs | 120 ++++++++++++++---- lib/mocker/src/replay/offline/events.rs | 4 + .../src/replay/offline/runtime_utils.rs | 40 ++++++ lib/mocker/src/scheduler/sglang/config.rs | 4 + lib/mocker/src/scheduler/sglang/decode.rs | 8 ++ lib/mocker/src/scheduler/sglang/tests.rs | 38 +++++- lib/mocker/src/scheduler/vllm/core.rs | 20 ++- lib/mocker/src/scheduler/vllm/tests.rs | 34 +++++ 11 files changed, 322 insertions(+), 56 deletions(-) diff --git a/lib/llm/src/mocker.rs b/lib/llm/src/mocker.rs index 96ed4a718551..c418429463ab 100644 --- a/lib/llm/src/mocker.rs +++ b/lib/llm/src/mocker.rs @@ -23,7 +23,7 @@ use dynamo_mocker::common::protocols::{ DirectRequest, KvCacheEventSink, KvEventPublishers, MockEngineArgs, OutputSignal, RawKvEvent, RawKvEventSink, }; -use dynamo_mocker::common::utils::{compute_kv_transfer_delay, sleep_precise}; +use dynamo_mocker::common::utils::sleep_precise; use dynamo_mocker::engine::create_engine; use dynamo_mocker::scheduler::SchedulerHandle; use dynamo_runtime::DistributedRuntime; @@ -645,14 +645,6 @@ impl AsyncEngine, ManyOut, Error> let bootstrap_server = self.bootstrap_server.clone(); let reasoning = self.engine_args.reasoning.clone(); - // Compute KV transfer delay for prefill workers. - // Simulates the time to transfer KV cache from prefill to decode worker. - let kv_transfer_delay = if is_prefill { - compute_kv_transfer_delay(&self.engine_args, request.token_ids.len()) - } else { - None - }; - // Spawn a task to handle the complex async logic tokio::spawn(async move { let mut token_count = 0; @@ -693,17 +685,15 @@ impl AsyncEngine, ManyOut, Error> if signal.completed { let _ = stream_tx.send(output); - // Simulate KV transfer delay before prefill's first (and only) token. - // This models the time to transfer KV cache to the decode worker. - if token_count == 1 - && let Some(delay) = kv_transfer_delay + // Prefill-to-decode handoff delay is emitted by the shared mocker core. + if is_prefill + && let Some(delay_ms) = signal.handoff_delay_ms { - sleep_precise(delay).await; + sleep_precise(Duration::from_secs_f64(delay_ms / 1000.0)).await; } // Prefill: after first token, mark room complete (unblocks decode) if is_prefill - && token_count == 1 && let (Some(server), Some(room_id)) = (bootstrap_server.get(), bootstrap_room) { server.complete_room(room_id); diff --git a/lib/mocker/src/common/protocols.rs b/lib/mocker/src/common/protocols.rs index e91ede8177bc..3aafa5e55898 100644 --- a/lib/mocker/src/common/protocols.rs +++ b/lib/mocker/src/common/protocols.rs @@ -139,6 +139,8 @@ impl PrefillCost { pub struct OutputSignal { pub uuid: Uuid, pub completed: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub handoff_delay_ms: Option, } /// Preemption policy for evicting decode requests under memory pressure diff --git a/lib/mocker/src/common/utils.rs b/lib/mocker/src/common/utils.rs index 04ecc2c4be45..959d2565d786 100644 --- a/lib/mocker/src/common/utils.rs +++ b/lib/mocker/src/common/utils.rs @@ -3,32 +3,59 @@ use std::time::{Duration, Instant}; -use crate::common::protocols::MockEngineArgs; +use crate::common::protocols::{MockEngineArgs, WorkerType}; -/// Compute the KV transfer delay duration for a given number of input tokens. +/// Compute the modeled handoff delay after a prefill worker emits its terminal token. /// -/// Returns `None` if KV transfer simulation is disabled (bandwidth is 0 or not configured). -pub fn compute_kv_transfer_delay( - args: &MockEngineArgs, +/// NOTE: this intentionally does not model the internal prefill TTFT itself accurately, and the +/// exact prefill/decode boundary is backend dependent. For now we only care about decode-visible +/// TTFT, which is what the client observes, so modeling the delay as prefill-to-decode handoff is +/// good enough. +pub fn compute_prefill_handoff_delay_ms( + worker_type: WorkerType, + completed: bool, num_input_tokens: usize, -) -> Option { - match (args.kv_transfer_bandwidth, args.kv_bytes_per_token) { + kv_transfer_bandwidth: Option, + kv_bytes_per_token: Option, +) -> Option { + if worker_type != WorkerType::Prefill || !completed { + return None; + } + + match (kv_transfer_bandwidth, kv_bytes_per_token) { (Some(bw), Some(bpt)) if bw > 0.0 => { let kv_bytes = num_input_tokens as f64 * bpt as f64; - let delay = Duration::from_secs_f64(kv_bytes / (bw * 1e9)); + let delay_ms = kv_bytes / (bw * 1e9) * 1000.0; tracing::debug!( num_input_tokens, kv_bytes, bandwidth_gb_s = bw, - delay_ms = format!("{:.2}", delay.as_secs_f64() * 1000.0), - "KV transfer delay for prefill" + delay_ms = format!("{delay_ms:.2}"), + "KV handoff delay for prefill completion" ); - Some(delay) + Some(delay_ms) } _ => None, } } +/// Compute the KV transfer delay duration for a given number of input tokens. +/// +/// Returns `None` if KV transfer simulation is disabled (bandwidth is 0 or not configured). +pub fn compute_kv_transfer_delay( + args: &MockEngineArgs, + num_input_tokens: usize, +) -> Option { + compute_prefill_handoff_delay_ms( + args.worker_type, + true, + num_input_tokens, + args.kv_transfer_bandwidth, + args.kv_bytes_per_token, + ) + .map(|delay_ms| Duration::from_secs_f64(delay_ms / 1000.0)) +} + /// Sleep for the specified duration using timerfd on Linux for precision. pub async fn sleep_precise(duration: Duration) { sleep_until_precise(Instant::now() + duration).await; @@ -53,3 +80,42 @@ pub async fn sleep_until_precise(deadline: Instant) { tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await; } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_prefill_handoff_delay_only_applies_to_completed_prefill() { + let delay_ms = compute_prefill_handoff_delay_ms( + WorkerType::Prefill, + true, + 128, + Some(1.0), + Some(1_000_000), + ) + .expect("prefill completion should produce a handoff delay"); + assert!((delay_ms - 128.0).abs() < 1e-9); + + assert!( + compute_prefill_handoff_delay_ms( + WorkerType::Prefill, + false, + 128, + Some(1.0), + Some(1_000_000), + ) + .is_none() + ); + assert!( + compute_prefill_handoff_delay_ms( + WorkerType::Decode, + true, + 128, + Some(1.0), + Some(1_000_000), + ) + .is_none() + ); + } +} diff --git a/lib/mocker/src/replay/offline/disagg.rs b/lib/mocker/src/replay/offline/disagg.rs index 9cc7fc9d3445..5b9dd18cda2e 100644 --- a/lib/mocker/src/replay/offline/disagg.rs +++ b/lib/mocker/src/replay/offline/disagg.rs @@ -12,7 +12,8 @@ use super::events::{SimulationEvent, SimulationWorkerStage}; use super::normalize_trace_requests; use super::runtime_utils::{ WorkerCompletionPayload, next_timestamp as choose_next_timestamp, pop_next_concurrency_ready, - pop_next_trace_ready, pop_ready_worker_completion, push_worker_completion, + pop_next_trace_ready, pop_ready_decode_handoff, pop_ready_worker_completion, + push_decode_handoff, push_worker_completion, }; use super::state::OfflineWorkerState; use crate::common::protocols::{DirectRequest, MockEngineArgs, OutputSignal}; @@ -454,35 +455,31 @@ impl DisaggRuntime { return Ok(()); } - let Some(_) = self.prefill_router.as_ref() else { - self.enqueue_decode(signal.uuid)?; - return Ok(()); - }; - let prefill_complete_admissions = self - .prefill_router - .as_mut() - .expect("prefill router must exist after guard") - .mark_prefill_completed(signal.uuid)?; - #[cfg(test)] - { - self.stats.prefill_marked_count += 1; - } - self.record_router_pending(); - self.dispatch_prefill_admissions(prefill_complete_admissions)?; + if self.prefill_router.is_some() { + let prefill_complete_admissions = { + let prefill_router = self.prefill_router.as_mut().expect("router checked above"); + prefill_router.mark_prefill_completed(signal.uuid)? + }; + #[cfg(test)] + { + self.stats.prefill_marked_count += 1; + } + self.record_router_pending(); + self.dispatch_prefill_admissions(prefill_complete_admissions)?; - let admissions = self - .prefill_router - .as_mut() - .expect("prefill router must exist after guard") - .free(signal.uuid)?; - #[cfg(test)] - { - self.stats.prefill_freed_count += 1; + let admissions = { + let prefill_router = self.prefill_router.as_mut().expect("router checked above"); + prefill_router.free(signal.uuid)? + }; + #[cfg(test)] + { + self.stats.prefill_freed_count += 1; + } + self.record_router_pending(); + self.dispatch_prefill_admissions(admissions)?; } - self.record_router_pending(); - self.dispatch_prefill_admissions(admissions)?; - self.enqueue_decode(signal.uuid)?; - Ok(()) + + self.enqueue_decode_after_handoff(signal.uuid, signal.handoff_delay_ms) } fn process_decode_signal(&mut self, signal: OutputSignal) -> Result<()> { @@ -572,6 +569,35 @@ impl DisaggRuntime { Ok(changed) } + fn apply_decode_handoffs(&mut self) -> Result { + let mut changed = false; + while let Some(uuid) = pop_ready_decode_handoff(&mut self.events, self.now_ms) { + self.enqueue_decode(uuid)?; + changed = true; + } + Ok(changed) + } + + fn enqueue_decode_after_handoff( + &mut self, + uuid: Uuid, + handoff_delay_ms: Option, + ) -> Result<()> { + if let Some(delay_ms) = handoff_delay_ms + && delay_ms > 0.0 + { + push_decode_handoff( + &mut self.events, + &mut self.next_event_seq, + self.now_ms + delay_ms, + uuid, + ); + return Ok(()); + } + + self.enqueue_decode(uuid) + } + fn release_trace_arrivals(&mut self) -> Result { let mut released_any = false; if matches!(self.admission, AdmissionSource::Requests(_)) { @@ -738,6 +764,7 @@ impl DisaggRuntime { fn drain_current_timestamp(&mut self) -> Result<()> { loop { let mut changed = self.apply_worker_completions()?; + changed |= self.apply_decode_handoffs()?; changed |= match self.mode { ReplayMode::Trace => self.release_trace_arrivals()?, @@ -972,6 +999,13 @@ mod tests { } } + fn disagg_config_with_handoff_delay() -> OfflineDisaggReplayConfig { + let mut config = disagg_config(); + config.prefill_args.kv_transfer_bandwidth = Some(1.0); + config.prefill_args.kv_bytes_per_token = Some(1_000_000); + config + } + fn router_config() -> KvRouterConfig { KvRouterConfig { router_queue_threshold: Some(1.25), @@ -1123,4 +1157,34 @@ mod tests { assert_eq!(stats.prefill_freed_count, 1); assert_eq!(stats.decode_freed_count, 1); } + + #[test] + fn test_handoff_delay_increases_decode_visible_ttft() { + let requests = vec![request(1, 128, 2, 0.0)]; + + let (baseline_collector, _) = run_trace_collect( + &disagg_config(), + requests.clone(), + None, + 1.0, + ReplayRouterMode::RoundRobin, + ); + let (delayed_collector, _) = run_trace_collect( + &disagg_config_with_handoff_delay(), + requests, + None, + 1.0, + ReplayRouterMode::RoundRobin, + ); + + let baseline = baseline_collector.snapshot(Uuid::from_u128(1)).unwrap(); + let delayed = delayed_collector.snapshot(Uuid::from_u128(1)).unwrap(); + let baseline_ttft = baseline.first_token_ms.unwrap() - baseline.arrival_time_ms; + let delayed_ttft = delayed.first_token_ms.unwrap() - delayed.arrival_time_ms; + + assert!( + delayed_ttft >= baseline_ttft + 120.0, + "expected delayed TTFT to include roughly 128ms of handoff delay, baseline={baseline_ttft}, delayed={delayed_ttft}" + ); + } } diff --git a/lib/mocker/src/replay/offline/events.rs b/lib/mocker/src/replay/offline/events.rs index dfed2feaffe2..d855f24f4ff1 100644 --- a/lib/mocker/src/replay/offline/events.rs +++ b/lib/mocker/src/replay/offline/events.rs @@ -4,6 +4,7 @@ use std::cmp::Ordering; use crate::common::protocols::OutputSignal; +use uuid::Uuid; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum SimulationWorkerStage { @@ -21,6 +22,9 @@ pub(crate) enum SimulationEventKind { output_signals: Vec, kv_events: Vec, }, + DecodeHandoff { + uuid: Uuid, + }, } #[derive(Debug)] diff --git a/lib/mocker/src/replay/offline/runtime_utils.rs b/lib/mocker/src/replay/offline/runtime_utils.rs index 24bb37b5831a..2f2e29a0165d 100644 --- a/lib/mocker/src/replay/offline/runtime_utils.rs +++ b/lib/mocker/src/replay/offline/runtime_utils.rs @@ -84,6 +84,9 @@ pub(super) fn pop_ready_worker_completion( if event.at_ms != now_ms { return None; } + let SimulationEventKind::WorkerCompletion { .. } = &event.kind else { + return None; + }; let event = events.pop().expect("event must exist after peek"); let (stage, worker_idx, completed_requests, output_signals, kv_events) = match event.kind { SimulationEventKind::WorkerCompletion { @@ -99,6 +102,9 @@ pub(super) fn pop_ready_worker_completion( output_signals, kv_events, ), + SimulationEventKind::DecodeHandoff { .. } => { + unreachable!("peeked worker completion event must match popped event") + } }; Some(WorkerCompletionPayload { stage, @@ -109,6 +115,38 @@ pub(super) fn pop_ready_worker_completion( }) } +pub(super) fn push_decode_handoff( + events: &mut BinaryHeap, + next_event_seq: &mut u64, + at_ms: f64, + uuid: uuid::Uuid, +) { + events.push(SimulationEvent { + at_ms, + seq_no: *next_event_seq, + kind: SimulationEventKind::DecodeHandoff { uuid }, + }); + *next_event_seq += 1; +} + +pub(super) fn pop_ready_decode_handoff( + events: &mut BinaryHeap, + now_ms: f64, +) -> Option { + let event = events.peek()?; + if event.at_ms != now_ms { + return None; + } + let SimulationEventKind::DecodeHandoff { .. } = &event.kind else { + return None; + }; + let event = events.pop().expect("event must exist after peek"); + let SimulationEventKind::DecodeHandoff { uuid } = event.kind else { + unreachable!("peeked decode handoff event must match popped event"); + }; + Some(uuid) +} + #[cfg(test)] mod tests { use super::*; @@ -182,6 +220,7 @@ mod tests { output_signals: vec![OutputSignal { uuid: Uuid::from_u128(7), completed: true, + handoff_delay_ms: None, }], kv_events: Vec::new(), }, @@ -197,6 +236,7 @@ mod tests { output_signals: vec![OutputSignal { uuid: Uuid::from_u128(8), completed: false, + handoff_delay_ms: None, }], kv_events: Vec::new(), }, diff --git a/lib/mocker/src/scheduler/sglang/config.rs b/lib/mocker/src/scheduler/sglang/config.rs index a95e4dc1332e..547a3c764def 100644 --- a/lib/mocker/src/scheduler/sglang/config.rs +++ b/lib/mocker/src/scheduler/sglang/config.rs @@ -36,6 +36,8 @@ pub(super) struct SglangConfig { pub(super) decode_speedup_ratio: f64, pub(super) worker_type: WorkerType, pub(super) block_size: usize, + pub(super) kv_bytes_per_token: Option, + pub(super) kv_transfer_bandwidth: Option, } impl SglangConfig { @@ -81,6 +83,8 @@ impl SglangConfig { decode_speedup_ratio: args.decode_speedup_ratio, worker_type: args.worker_type, block_size: args.block_size, + kv_bytes_per_token: args.kv_bytes_per_token, + kv_transfer_bandwidth: args.kv_transfer_bandwidth, } } } diff --git a/lib/mocker/src/scheduler/sglang/decode.rs b/lib/mocker/src/scheduler/sglang/decode.rs index ec44c0acf599..39e62ea8fefa 100644 --- a/lib/mocker/src/scheduler/sglang/decode.rs +++ b/lib/mocker/src/scheduler/sglang/decode.rs @@ -4,6 +4,7 @@ use std::time::Duration; use crate::common::protocols::OutputSignal; +use crate::common::utils::compute_prefill_handoff_delay_ms; use crate::kv_manager::SglangKvManager; use super::config::{SglangConfig, floor_to_block}; @@ -179,6 +180,13 @@ pub(super) fn simulate_decode_step( output_signals.push(OutputSignal { uuid: req.uuid, completed: is_complete, + handoff_delay_ms: compute_prefill_handoff_delay_ms( + config.worker_type, + is_complete, + req.prompt_len(), + config.kv_transfer_bandwidth, + config.kv_bytes_per_token, + ), }); if is_complete { diff --git a/lib/mocker/src/scheduler/sglang/tests.rs b/lib/mocker/src/scheduler/sglang/tests.rs index f3368b5f51d7..fe4b4cb1af51 100644 --- a/lib/mocker/src/scheduler/sglang/tests.rs +++ b/lib/mocker/src/scheduler/sglang/tests.rs @@ -860,7 +860,6 @@ mod router_events { } } } - assert_eq!(seen, expected); drop(scheduler); drop(sink); @@ -871,4 +870,41 @@ mod router_events { assert!(harness.ok_count(METRIC_EVENT_REMOVED) > 0); harness.shutdown(); } + + #[test] + fn test_prefill_completion_emits_handoff_delay() { + let args = MockEngineArgs::builder() + .engine_type(EngineType::Sglang) + .num_gpu_blocks(64) + .block_size(4) + .worker_type(crate::common::protocols::WorkerType::Prefill) + .kv_transfer_bandwidth(Some(1.0)) + .kv_bytes_per_token(Some(1_000_000)) + .speedup_ratio(0.0) + .sglang(Some(SglangArgs { + page_size: Some(4), + chunked_prefill_size: Some(16), + ..Default::default() + })) + .build() + .unwrap(); + let mut core = SglangCore::new(args); + core.receive(DirectRequest { + tokens: vec![1; 8], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(91)), + dp_rank: 0, + arrival_timestamp_ms: None, + }); + + let mut collector = crate::replay::TraceCollector::default(); + let pass = core.execute_pass(&mut collector, 0.0); + let signal = pass + .output_signals + .first() + .expect("prefill pass should emit one completed signal"); + + assert!(signal.completed); + assert_eq!(signal.handoff_delay_ms, Some(8.0)); + } } diff --git a/lib/mocker/src/scheduler/vllm/core.rs b/lib/mocker/src/scheduler/vllm/core.rs index 6986b6a7484c..a00d83a65a2d 100644 --- a/lib/mocker/src/scheduler/vllm/core.rs +++ b/lib/mocker/src/scheduler/vllm/core.rs @@ -14,6 +14,7 @@ use crate::common::protocols::{ WorkerType, }; use crate::common::sequence::ActiveSequence; +use crate::common::utils::compute_prefill_handoff_delay_ms; use crate::kv_manager::KvManager; use crate::replay::TraceCollector; use crate::scheduler::{ @@ -645,8 +646,25 @@ impl VllmCore { } if let Some(request) = self.state.requests.get(&uuid) { debug_assert_vllm_request_progress(uuid, request); + let handoff_delay_ms = compute_prefill_handoff_delay_ms( + self.args.worker_type, + completed, + request.sequence.num_input_tokens(), + self.args.kv_transfer_bandwidth, + self.args.kv_bytes_per_token, + ); + output_signals.push(OutputSignal { + uuid, + completed, + handoff_delay_ms, + }); + } else { + output_signals.push(OutputSignal { + uuid, + completed, + handoff_delay_ms: None, + }); } - output_signals.push(OutputSignal { uuid, completed }); if completed { self.state.complete(&uuid); } diff --git a/lib/mocker/src/scheduler/vllm/tests.rs b/lib/mocker/src/scheduler/vllm/tests.rs index 87641c53b707..fdfd132aac9d 100644 --- a/lib/mocker/src/scheduler/vllm/tests.rs +++ b/lib/mocker/src/scheduler/vllm/tests.rs @@ -175,6 +175,40 @@ mod core_behavior { ); } + #[test] + fn test_prefill_completion_emits_handoff_delay() { + let args = MockEngineArgs::builder() + .block_size(4) + .num_gpu_blocks(8) + .max_num_batched_tokens(Some(8)) + .max_num_seqs(Some(1)) + .enable_chunked_prefill(true) + .worker_type(crate::common::protocols::WorkerType::Prefill) + .kv_transfer_bandwidth(Some(1.0)) + .kv_bytes_per_token(Some(1_000_000)) + .speedup_ratio(0.0) + .build() + .unwrap(); + let mut core = VllmCore::new(args); + core.receive(DirectRequest { + tokens: vec![1; 8], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(81)), + dp_rank: 0, + arrival_timestamp_ms: None, + }); + + let mut collector = crate::replay::TraceCollector::default(); + let pass = core.execute_pass(&mut collector, 0.0); + let signal = pass + .output_signals + .first() + .expect("prefill pass should emit one completed signal"); + + assert!(signal.completed); + assert_eq!(signal.handoff_delay_ms, Some(8.0)); + } + #[test] fn test_first_token_can_arrive_on_prompt_completion_pass() { let mut core = VllmCore::new(make_args()); From 77ac5e601908145bc8ca62753bdebb60b4322fdd Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 25 Mar 2026 10:29:41 -0700 Subject: [PATCH 22/24] refactor replay state without moving the clock Signed-off-by: PeaBrane --- lib/mocker/src/replay/offline/disagg.rs | 109 ++++------------- lib/mocker/src/replay/offline/multi.rs | 65 +++++++--- lib/mocker/src/replay/offline/state.rs | 152 ++++++++++++++++++++++++ 3 files changed, 225 insertions(+), 101 deletions(-) diff --git a/lib/mocker/src/replay/offline/disagg.rs b/lib/mocker/src/replay/offline/disagg.rs index 5b9dd18cda2e..3875dd5ca4a1 100644 --- a/lib/mocker/src/replay/offline/disagg.rs +++ b/lib/mocker/src/replay/offline/disagg.rs @@ -15,7 +15,11 @@ use super::runtime_utils::{ pop_next_trace_ready, pop_ready_decode_handoff, pop_ready_worker_completion, push_decode_handoff, push_worker_completion, }; -use super::state::OfflineWorkerState; +#[cfg(test)] +use super::state::DisaggPhase; +#[cfg(test)] +use super::state::DisaggRequestSnapshot; +use super::state::{DisaggRequestState, OfflineWorkerState}; use crate::common::protocols::{DirectRequest, MockEngineArgs, OutputSignal}; use crate::loadgen::{ReplayRequestHashes, Trace, WorkloadDriver}; use crate::replay::router::OfflineReplayRouter; @@ -35,33 +39,6 @@ enum AdmissionSource { Workload(WorkloadDriver), } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum DisaggPhase { - QueuedPrefill, - RunningPrefill, - QueuedDecode, - RunningDecode, - Done, -} - -struct DisaggRequestState { - original: Option, - #[cfg(test)] - arrival_ms: f64, - phase: DisaggPhase, - prefill_worker_idx: Option, - decode_worker_idx: Option, -} - -#[cfg(test)] -#[derive(Debug, Clone, PartialEq)] -struct DisaggRequestSnapshot { - arrival_ms: f64, - phase: DisaggPhase, - prefill_worker_idx: Option, - decode_worker_idx: Option, -} - #[cfg(test)] #[derive(Debug, Default, Clone, PartialEq)] struct DisaggRuntimeStats { @@ -265,30 +242,11 @@ impl DisaggRuntime { .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}")) } - fn original_request(state: &DisaggRequestState) -> Result<&DirectRequest> { - state - .original - .as_ref() - .ok_or_else(|| anyhow!("offline disagg replay request payload was already released")) - } - - fn build_prefill_request(state: &DisaggRequestState) -> Result { - let mut request = Self::original_request(state)?.clone(); - request.max_output_tokens = 1; - Ok(request) - } - - fn build_decode_request(state: &DisaggRequestState) -> Result { - Ok(Self::original_request(state)?.clone()) - } - fn dispatch_prefill(&mut self, uuid: Uuid, worker_idx: usize) -> Result<()> { self.validate_worker_idx(SimulationWorkerStage::Prefill, worker_idx)?; - let request = Self::build_prefill_request(self.state(uuid)?)?; + let request = self.state(uuid)?.build_prefill_request()?; self.prefill_workers[worker_idx].receive_request(request); - let state = self.state_mut(uuid)?; - state.phase = DisaggPhase::RunningPrefill; - state.prefill_worker_idx = Some(worker_idx); + self.state_mut(uuid)?.start_prefill(worker_idx); #[cfg(test)] { self.stats.prefill_assignments.insert(uuid, worker_idx); @@ -298,11 +256,9 @@ impl DisaggRuntime { fn dispatch_decode(&mut self, uuid: Uuid, worker_idx: usize) -> Result<()> { self.validate_worker_idx(SimulationWorkerStage::Decode, worker_idx)?; - let request = Self::build_decode_request(self.state(uuid)?)?; + let request = self.state(uuid)?.build_decode_request()?; self.decode_workers[worker_idx].receive_request(request); - let state = self.state_mut(uuid)?; - state.phase = DisaggPhase::RunningDecode; - state.decode_worker_idx = Some(worker_idx); + self.state_mut(uuid)?.start_decode(worker_idx); #[cfg(test)] { self.stats.decode_assignments.insert(uuid, worker_idx); @@ -312,7 +268,7 @@ impl DisaggRuntime { fn dispatch_prefill_admissions(&mut self, admissions: Vec<(Uuid, usize)>) -> Result<()> { for (uuid, worker_idx) in admissions { - if self.state(uuid)?.phase != DisaggPhase::QueuedPrefill { + if !self.state(uuid)?.is_queued_prefill() { bail!("offline disagg replay expected queued prefill request for {uuid}"); } self.dispatch_prefill(uuid, worker_idx)?; @@ -322,7 +278,7 @@ impl DisaggRuntime { fn dispatch_decode_admissions(&mut self, admissions: Vec<(Uuid, usize)>) -> Result<()> { for (uuid, worker_idx) in admissions { - if self.state(uuid)?.phase != DisaggPhase::QueuedDecode { + if !self.state(uuid)?.is_queued_decode() { bail!("offline disagg replay expected queued decode request for {uuid}"); } self.dispatch_decode(uuid, worker_idx)?; @@ -338,9 +294,10 @@ impl DisaggRuntime { }; let maybe_worker_idx = { let requests = &self.requests; - let request = Self::original_request(requests.get(&uuid).ok_or_else(|| { - anyhow!("offline disagg replay missing request state for {uuid}") - })?)?; + let request = requests + .get(&uuid) + .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}"))? + .original_request()?; decode_router.submit_request_with_hashes(request, None, self.now_ms)? }; self.record_router_pending(); @@ -353,7 +310,7 @@ impl DisaggRuntime { return Ok(()); } - self.state_mut(uuid)?.phase = DisaggPhase::QueuedDecode; + self.state_mut(uuid)?.queue_decode(); Ok(()) } @@ -374,15 +331,8 @@ impl DisaggRuntime { request.max_output_tokens, ); - let state = DisaggRequestState { - original: Some(request), - #[cfg(test)] - arrival_ms: arrival_time_ms, - phase: DisaggPhase::QueuedPrefill, - prefill_worker_idx: None, - decode_worker_idx: None, - }; - self.requests.insert(uuid, state); + self.requests + .insert(uuid, DisaggRequestState::new(request, arrival_time_ms)); let Some(prefill_router) = self.prefill_router.as_mut() else { let worker_idx = self.next_prefill_worker(); self.dispatch_prefill(uuid, worker_idx)?; @@ -390,9 +340,10 @@ impl DisaggRuntime { }; let maybe_worker_idx = { let requests = &self.requests; - let request = Self::original_request(requests.get(&uuid).ok_or_else(|| { - anyhow!("offline disagg replay missing request state for {uuid}") - })?)?; + let request = requests + .get(&uuid) + .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}"))? + .original_request()?; prefill_router.submit_request_with_hashes(request, replay_hashes, self.now_ms)? }; self.record_router_pending(); @@ -501,9 +452,7 @@ impl DisaggRuntime { if let AdmissionSource::Workload(driver) = &mut self.admission { driver.on_complete(signal.uuid, self.now_ms)?; } - let state = self.state_mut(signal.uuid)?; - state.phase = DisaggPhase::Done; - state.original = None; + self.state_mut(signal.uuid)?.mark_done(); self.dispatch_decode_admissions(admissions)?; Ok(()) } @@ -789,17 +738,7 @@ impl DisaggRuntime { self.stats.request_snapshots = self .requests .iter() - .map(|(uuid, state)| { - ( - *uuid, - DisaggRequestSnapshot { - arrival_ms: state.arrival_ms, - phase: state.phase, - prefill_worker_idx: state.prefill_worker_idx, - decode_worker_idx: state.decode_worker_idx, - }, - ) - }) + .map(|(uuid, state)| (*uuid, state.debug_snapshot())) .collect(); } } diff --git a/lib/mocker/src/replay/offline/multi.rs b/lib/mocker/src/replay/offline/multi.rs index e9a5315d9f9f..f973f7521d97 100644 --- a/lib/mocker/src/replay/offline/multi.rs +++ b/lib/mocker/src/replay/offline/multi.rs @@ -9,7 +9,7 @@ use super::runtime_utils::{ }; #[cfg(test)] use super::state::OfflineWorkerSnapshot; -use super::state::OfflineWorkerState; +use super::state::{AggRequestState, OfflineWorkerState}; use crate::common::protocols::{DirectRequest, MockEngineArgs, OutputSignal}; use crate::loadgen::{ReplayRequestHashes, Trace, WorkloadDriver}; use crate::replay::router::OfflineReplayRouter; @@ -20,7 +20,7 @@ use crate::scheduler::RouterEventVisibility; use anyhow::bail; use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::protocols::RouterEvent; -use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque}; +use std::collections::{BinaryHeap, HashMap, VecDeque}; use uuid::Uuid; #[derive(Debug, Clone, Copy)] @@ -66,13 +66,13 @@ struct OfflineRuntime { next_worker_idx: usize, next_event_seq: u64, admission: AdmissionSource, - router_pending: HashMap, + requests: HashMap, + queued_requests: usize, workers: Vec, collector: TraceCollector, events: BinaryHeap, mode: ReplayMode, router: Option, - prefill_completed: HashSet, stats: OfflineRuntimeStats, #[cfg(test)] worker_active_requests: Vec>, @@ -139,7 +139,8 @@ impl OfflineRuntime { next_worker_idx: 0, next_event_seq: 0, admission, - router_pending: HashMap::new(), + requests: HashMap::new(), + queued_requests: 0, workers: (0..num_workers) .map(|worker_idx| { OfflineWorkerState::new(worker_idx, args.clone(), capture_kv_events) @@ -149,7 +150,6 @@ impl OfflineRuntime { events: BinaryHeap::new(), mode, router, - prefill_completed: HashSet::new(), #[cfg(test)] stats: OfflineRuntimeStats::default(), #[cfg(not(test))] @@ -166,7 +166,7 @@ impl OfflineRuntime { .iter() .map(OfflineWorkerState::in_flight) .sum::() - + self.router_pending.len() + + self.queued_requests } fn record_in_flight_peak(&mut self) { @@ -224,9 +224,14 @@ impl OfflineRuntime { fn dispatch_router_admissions(&mut self, admissions: Vec<(Uuid, usize)>) -> anyhow::Result<()> { for (uuid, worker_idx) in admissions { - let request = self.router_pending.remove(&uuid).ok_or_else(|| { - anyhow::anyhow!("offline replay missing queued request state for {uuid}") - })?; + let request = self + .requests + .get_mut(&uuid) + .ok_or_else(|| { + anyhow::anyhow!("offline replay missing queued request state for {uuid}") + })? + .take_queued_request(uuid)?; + self.queued_requests = self.queued_requests.saturating_sub(1); self.dispatch_to_worker(request, uuid, worker_idx)?; } Ok(()) @@ -252,6 +257,7 @@ impl OfflineRuntime { ); let Some(router) = self.router.as_mut() else { + self.requests.insert(uuid, AggRequestState::new_running()); let worker_idx = self.next_worker_idx; self.next_worker_idx = (self.next_worker_idx + 1) % self.workers.len(); self.dispatch_to_worker(request, uuid, worker_idx)?; @@ -262,11 +268,14 @@ impl OfflineRuntime { router.submit_request_with_hashes(&request, replay_hashes, self.now_ms)?; self.record_router_pending(); if let Some(worker_idx) = maybe_worker_idx { + self.requests.insert(uuid, AggRequestState::new_running()); self.dispatch_to_worker(request, uuid, worker_idx)?; return Ok(uuid); } - self.router_pending.insert(uuid, request); + self.requests + .insert(uuid, AggRequestState::new_queued(request)); + self.queued_requests += 1; self.record_in_flight_peak(); Ok(uuid) } @@ -329,7 +338,9 @@ impl OfflineRuntime { } self.record_router_pending(); } - self.prefill_completed.remove(&signal.uuid); + self.requests.remove(&signal.uuid).ok_or_else(|| { + anyhow::anyhow!("offline replay missing request state for {}", signal.uuid) + })?; if let AdmissionSource::Workload(driver) = &mut self.admission { driver.on_complete(signal.uuid, self.now_ms)?; } @@ -337,10 +348,23 @@ impl OfflineRuntime { return Ok(()); } - if !self.prefill_completed.insert(signal.uuid) { + let already_marked = self + .requests + .get(&signal.uuid) + .ok_or_else(|| { + anyhow::anyhow!("offline replay missing request state for {}", signal.uuid) + })? + .prefill_completed(); + if already_marked { return Ok(()); } + self.requests + .get_mut(&signal.uuid) + .ok_or_else(|| { + anyhow::anyhow!("offline replay missing request state for {}", signal.uuid) + })? + .mark_prefill_completed(); if let Some(router) = self.router.as_mut() { admissions = router.mark_prefill_completed(signal.uuid)?; #[cfg(test)] @@ -598,10 +622,19 @@ impl OfflineRuntime { #[cfg(test)] fn debug_snapshot(&self) -> OfflineRuntimeSnapshot { - let mut router_pending_request_ids = - self.router_pending.keys().copied().collect::>(); + let mut router_pending_request_ids = self + .requests + .iter() + .filter(|(_, state)| state.is_queued_at_router()) + .map(|(uuid, _)| *uuid) + .collect::>(); router_pending_request_ids.sort_unstable(); - let mut prefill_completed = self.prefill_completed.iter().copied().collect::>(); + let mut prefill_completed = self + .requests + .iter() + .filter(|(_, state)| state.prefill_completed()) + .map(|(uuid, _)| *uuid) + .collect::>(); prefill_completed.sort_unstable(); OfflineRuntimeSnapshot { diff --git a/lib/mocker/src/replay/offline/state.rs b/lib/mocker/src/replay/offline/state.rs index bd66ddaf56b5..47ba06b8b2f9 100644 --- a/lib/mocker/src/replay/offline/state.rs +++ b/lib/mocker/src/replay/offline/state.rs @@ -1,10 +1,162 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use anyhow::{Result, anyhow, bail}; + use crate::common::protocols::DirectRequest; use crate::common::protocols::MockEngineArgs; use crate::replay::TraceCollector; use crate::scheduler::{EngineCore, EnginePassResult}; +use uuid::Uuid; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AggRequestPhase { + QueuedAtRouter, + Running, +} + +pub(crate) struct AggRequestState { + request: Option, + phase: AggRequestPhase, + prefill_completed: bool, +} + +impl AggRequestState { + pub(crate) fn new_queued(request: DirectRequest) -> Self { + Self { + request: Some(request), + phase: AggRequestPhase::QueuedAtRouter, + prefill_completed: false, + } + } + + pub(crate) fn new_running() -> Self { + Self { + request: None, + phase: AggRequestPhase::Running, + prefill_completed: false, + } + } + + pub(crate) fn is_queued_at_router(&self) -> bool { + self.phase == AggRequestPhase::QueuedAtRouter + } + + pub(crate) fn take_queued_request(&mut self, uuid: Uuid) -> Result { + if !self.is_queued_at_router() { + bail!("offline replay expected queued request state for {uuid}"); + } + let request = self + .request + .take() + .ok_or_else(|| anyhow!("offline replay missing queued request payload for {uuid}"))?; + self.phase = AggRequestPhase::Running; + Ok(request) + } + + pub(crate) fn prefill_completed(&self) -> bool { + self.prefill_completed + } + + pub(crate) fn mark_prefill_completed(&mut self) { + self.prefill_completed = true; + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum DisaggPhase { + QueuedPrefill, + RunningPrefill, + QueuedDecode, + RunningDecode, + Done, +} + +pub(crate) struct DisaggRequestState { + original: Option, + #[cfg(test)] + arrival_ms: f64, + phase: DisaggPhase, + prefill_worker_idx: Option, + decode_worker_idx: Option, +} + +#[cfg(test)] +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct DisaggRequestSnapshot { + pub(crate) arrival_ms: f64, + pub(crate) phase: DisaggPhase, + pub(crate) prefill_worker_idx: Option, + pub(crate) decode_worker_idx: Option, +} + +impl DisaggRequestState { + pub(crate) fn new(request: DirectRequest, arrival_ms: f64) -> Self { + #[cfg(not(test))] + let _ = arrival_ms; + Self { + original: Some(request), + #[cfg(test)] + arrival_ms, + phase: DisaggPhase::QueuedPrefill, + prefill_worker_idx: None, + decode_worker_idx: None, + } + } + + pub(crate) fn is_queued_prefill(&self) -> bool { + self.phase == DisaggPhase::QueuedPrefill + } + + pub(crate) fn is_queued_decode(&self) -> bool { + self.phase == DisaggPhase::QueuedDecode + } + + pub(crate) fn original_request(&self) -> Result<&DirectRequest> { + self.original + .as_ref() + .ok_or_else(|| anyhow!("offline disagg replay request payload was already released")) + } + + pub(crate) fn build_prefill_request(&self) -> Result { + let mut request = self.original_request()?.clone(); + request.max_output_tokens = 1; + Ok(request) + } + + pub(crate) fn build_decode_request(&self) -> Result { + Ok(self.original_request()?.clone()) + } + + pub(crate) fn start_prefill(&mut self, worker_idx: usize) { + self.phase = DisaggPhase::RunningPrefill; + self.prefill_worker_idx = Some(worker_idx); + } + + pub(crate) fn queue_decode(&mut self) { + self.phase = DisaggPhase::QueuedDecode; + } + + pub(crate) fn start_decode(&mut self, worker_idx: usize) { + self.phase = DisaggPhase::RunningDecode; + self.decode_worker_idx = Some(worker_idx); + } + + pub(crate) fn mark_done(&mut self) { + self.phase = DisaggPhase::Done; + self.original = None; + } + + #[cfg(test)] + pub(crate) fn debug_snapshot(&self) -> DisaggRequestSnapshot { + DisaggRequestSnapshot { + arrival_ms: self.arrival_ms, + phase: self.phase, + prefill_worker_idx: self.prefill_worker_idx, + decode_worker_idx: self.decode_worker_idx, + } + } +} pub(crate) struct OfflineWorkerState { core: EngineCore, From 6bd84f8e0480e95abbde677abb935e372e8f8870 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 25 Mar 2026 10:58:12 -0700 Subject: [PATCH 23/24] batch online trace arrivals before passes bolt Signed-off-by: PeaBrane --- lib/mocker/src/replay/online/live_runtime.rs | 98 ++++++++++++++++++-- lib/mocker/src/replay/online/task.rs | 66 +++++++++++-- lib/mocker/src/replay/online/tests.rs | 33 ++++++- 3 files changed, 178 insertions(+), 19 deletions(-) diff --git a/lib/mocker/src/replay/online/live_runtime.rs b/lib/mocker/src/replay/online/live_runtime.rs index 2834533dfab7..f065addfd85c 100644 --- a/lib/mocker/src/replay/online/live_runtime.rs +++ b/lib/mocker/src/replay/online/live_runtime.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use anyhow::{Result, anyhow, bail}; use dashmap::DashMap; use dynamo_kv_router::config::KvRouterConfig; -use tokio::sync::{Notify, Semaphore, mpsc}; +use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore, mpsc}; use tokio::task::JoinSet; use tokio::time::Instant; use tokio_util::sync::CancellationToken; @@ -23,7 +23,67 @@ use super::state::{ LiveReplayMode, LiveRuntimeStats, SharedLiveRuntimeStats, WorkloadDispatchState, now_ms, record_arrival, }; -use super::task::{RequestTaskContext, run_request_task, wait_for_workload_progress}; +use super::task::{ + RequestTaskContext, commit_request_dispatch, prepare_request_dispatch, run_request_task, + wait_for_request_completion, wait_for_workload_progress, +}; + +struct PendingDispatch { + request: DirectRequest, + arrival_at_ms: f64, + permit: Option, +} + +pub(super) fn pop_ready_trace_requests( + pending: &mut VecDeque, + now_ms: f64, +) -> Vec { + let mut ready = Vec::new(); + while pending + .front() + .and_then(|request| request.arrival_timestamp_ms) + .is_some_and(|arrival_ms| arrival_ms <= now_ms) + { + ready.push( + pending + .pop_front() + .expect("front request must exist while draining ready arrivals"), + ); + } + ready +} + +async fn dispatch_batch( + arrival_tx: &mpsc::UnboundedSender, + tasks: &mut JoinSet>, + task_ctx: &RequestTaskContext, + batch: Vec, +) -> Result<()> { + if batch.is_empty() { + return Ok(()); + } + + for pending in &batch { + record_arrival(arrival_tx, &pending.request, pending.arrival_at_ms)?; + } + + let mut prepared = Vec::with_capacity(batch.len()); + for pending in batch { + let prepared_request = prepare_request_dispatch(task_ctx, pending.request).await?; + prepared.push((prepared_request, pending.permit)); + } + + for (prepared_request, permit) in prepared { + let dispatched = commit_request_dispatch(task_ctx, prepared_request)?; + tasks.spawn(wait_for_request_completion( + task_ctx.clone(), + dispatched, + permit, + )); + } + + Ok(()) +} struct LiveRuntime { pending: VecDeque, @@ -122,13 +182,26 @@ impl LiveRuntime { match self.mode { LiveReplayMode::Trace => { - while let Some(request) = self.pending.pop_front() { - let arrival_ms = request.arrival_timestamp_ms.unwrap_or(0.0); + while !self.pending.is_empty() { + let arrival_ms = self + .pending + .front() + .and_then(|request| request.arrival_timestamp_ms) + .ok_or_else(|| { + anyhow!("online trace replay request missing arrival timestamp") + })?; let deadline = start + tokio::time::Duration::from_secs_f64(arrival_ms / 1000.0); tokio::time::sleep_until(deadline).await; - record_arrival(&arrival_tx, &request, arrival_ms)?; - tasks.spawn(run_request_task(task_ctx.clone(), request, None)); + let batch = pop_ready_trace_requests(&mut self.pending, now_ms(start)) + .into_iter() + .map(|request| PendingDispatch { + arrival_at_ms: request.arrival_timestamp_ms.unwrap_or(0.0), + request, + permit: None, + }) + .collect(); + dispatch_batch(&arrival_tx, &mut tasks, &task_ctx, batch).await?; } } LiveReplayMode::Concurrency { max_in_flight } => { @@ -222,6 +295,19 @@ impl LiveRuntime { .unwrap() .pop_ready(now, dispatch_limit); if !ready_turns.is_empty() { + if matches!(self.mode, LiveReplayMode::Trace) { + let batch = ready_turns + .into_iter() + .map(|ready_turn| PendingDispatch { + arrival_at_ms: ready_turn.scheduled_ready_at_ms, + request: ready_turn.request, + permit: None, + }) + .collect(); + dispatch_batch(&arrival_tx, &mut tasks, &task_ctx, batch).await?; + continue; + } + for ready_turn in ready_turns { let permit = match &semaphore { Some(semaphore) => { diff --git a/lib/mocker/src/replay/online/task.rs b/lib/mocker/src/replay/online/task.rs index ea8952a9734a..ac1291165a4f 100644 --- a/lib/mocker/src/replay/online/task.rs +++ b/lib/mocker/src/replay/online/task.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use anyhow::{Result, anyhow, bail}; use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use tokio::time::Instant; +use uuid::Uuid; use crate::common::protocols::DirectRequest; use crate::replay::router::ReplayRouter; @@ -26,6 +27,18 @@ pub(super) struct RequestTaskContext { pub(super) workload: Option>, } +pub(super) struct PreparedRequestDispatch { + request: DirectRequest, + uuid: Uuid, + worker_idx: usize, + state: Arc, +} + +pub(super) struct DispatchedRequest { + uuid: Uuid, + state: Arc, +} + pub(super) async fn wait_for_workload_progress( mode: LiveReplayMode, semaphore: Option<&Semaphore>, @@ -71,8 +84,16 @@ pub(super) async fn run_request_task( request: DirectRequest, permit: Option, ) -> Result<()> { - let uuid = request_uuid(&request)?; + let prepared = prepare_request_dispatch(&ctx, request).await?; + let dispatched = commit_request_dispatch(&ctx, prepared)?; + wait_for_request_completion(ctx, dispatched, permit).await +} +pub(super) async fn prepare_request_dispatch( + ctx: &RequestTaskContext, + request: DirectRequest, +) -> Result { + let uuid = request_uuid(&request)?; let worker_idx = ctx .router .select_worker(&request, ctx.senders.len()) @@ -81,26 +102,51 @@ pub(super) async fn run_request_task( bail!("online replay selected unknown worker index {worker_idx}"); } - let state = Arc::new(RequestState::default()); - ctx.requests.insert(uuid, Arc::clone(&state)); - if let Err(error) = ctx.senders[worker_idx].send(request) { - ctx.requests.remove(&uuid); + Ok(PreparedRequestDispatch { + request, + uuid, + worker_idx, + state: Arc::new(RequestState::default()), + }) +} + +pub(super) fn commit_request_dispatch( + ctx: &RequestTaskContext, + prepared: PreparedRequestDispatch, +) -> Result { + ctx.requests + .insert(prepared.uuid, Arc::clone(&prepared.state)); + if let Err(error) = ctx.senders[prepared.worker_idx].send(prepared.request) { + ctx.requests.remove(&prepared.uuid); return Err(anyhow!( - "online replay failed to dispatch request to worker {worker_idx}: {error}" + "online replay failed to dispatch request to worker {}: {}", + prepared.worker_idx, + error )); } - ctx.stats.record_dispatch(worker_idx); - state.wait_for_completion().await; + ctx.stats.record_dispatch(prepared.worker_idx); + Ok(DispatchedRequest { + uuid: prepared.uuid, + state: prepared.state, + }) +} + +pub(super) async fn wait_for_request_completion( + ctx: RequestTaskContext, + dispatched: DispatchedRequest, + permit: Option, +) -> Result<()> { + dispatched.state.wait_for_completion().await; ctx.stats.record_completion(); - ctx.requests.remove(&uuid); + ctx.requests.remove(&dispatched.uuid); if let Some(workload) = ctx.workload.as_ref() { let completion_ms = now_ms(workload.start); workload .driver .lock() .unwrap() - .on_complete(uuid, completion_ms)?; + .on_complete(dispatched.uuid, completion_ms)?; workload.wakeup.notify_waiters(); } drop(permit); diff --git a/lib/mocker/src/replay/online/tests.rs b/lib/mocker/src/replay/online/tests.rs index e9f21ac31ba8..5417e3cabe5b 100644 --- a/lib/mocker/src/replay/online/tests.rs +++ b/lib/mocker/src/replay/online/tests.rs @@ -17,9 +17,9 @@ use crate::replay::ReplayRouterMode; use crate::replay::router::ReplayRouter; use super::live_runtime::{ - simulate_concurrency_requests_with_stats, simulate_concurrency_workload_with_stats, - simulate_trace_requests, simulate_trace_requests_with_stats, - simulate_trace_workload_with_stats, + pop_ready_trace_requests, simulate_concurrency_requests_with_stats, + simulate_concurrency_workload_with_stats, simulate_trace_requests, + simulate_trace_requests_with_stats, simulate_trace_workload_with_stats, }; use super::state::{LiveReplayMode, SharedLiveRuntimeStats, WorkloadDispatchState, record_arrival}; use super::task::{RequestTaskContext, run_request_task, wait_for_workload_progress}; @@ -151,6 +151,33 @@ async fn test_record_arrival_uses_caller_arrival_timestamp() { assert_eq!(arrival.at_ms, arrival_at_ms); } +#[test] +fn test_pop_ready_trace_requests_drains_all_ready_arrivals() { + let mut pending = VecDeque::from(vec![ + request(1, 11, Some(0.0)), + request(2, 22, Some(0.0)), + request(3, 33, Some(0.1)), + request(4, 44, Some(0.2)), + ]); + + let ready = pop_ready_trace_requests(&mut pending, 0.1); + + assert_eq!( + ready + .iter() + .map(|request| request.uuid.unwrap()) + .collect::>(), + vec![Uuid::from_u128(1), Uuid::from_u128(2), Uuid::from_u128(3),] + ); + assert_eq!( + pending + .iter() + .map(|request| request.uuid.unwrap()) + .collect::>(), + vec![Uuid::from_u128(4)] + ); +} + #[tokio::test] async fn test_trace_arrivals_are_not_blocked_by_queued_router_selection() { let args = MockEngineArgs::builder() From d09a1fc39fb02a73f72d8216bad33c58a9144247 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 25 Mar 2026 11:02:13 -0700 Subject: [PATCH 24/24] Revert "batch online trace arrivals before passes bolt" This reverts commit 6bd84f8e0480e95abbde677abb935e372e8f8870. --- lib/mocker/src/replay/online/live_runtime.rs | 98 ++------------------ lib/mocker/src/replay/online/task.rs | 66 ++----------- lib/mocker/src/replay/online/tests.rs | 33 +------ 3 files changed, 19 insertions(+), 178 deletions(-) diff --git a/lib/mocker/src/replay/online/live_runtime.rs b/lib/mocker/src/replay/online/live_runtime.rs index f065addfd85c..2834533dfab7 100644 --- a/lib/mocker/src/replay/online/live_runtime.rs +++ b/lib/mocker/src/replay/online/live_runtime.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use anyhow::{Result, anyhow, bail}; use dashmap::DashMap; use dynamo_kv_router::config::KvRouterConfig; -use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore, mpsc}; +use tokio::sync::{Notify, Semaphore, mpsc}; use tokio::task::JoinSet; use tokio::time::Instant; use tokio_util::sync::CancellationToken; @@ -23,67 +23,7 @@ use super::state::{ LiveReplayMode, LiveRuntimeStats, SharedLiveRuntimeStats, WorkloadDispatchState, now_ms, record_arrival, }; -use super::task::{ - RequestTaskContext, commit_request_dispatch, prepare_request_dispatch, run_request_task, - wait_for_request_completion, wait_for_workload_progress, -}; - -struct PendingDispatch { - request: DirectRequest, - arrival_at_ms: f64, - permit: Option, -} - -pub(super) fn pop_ready_trace_requests( - pending: &mut VecDeque, - now_ms: f64, -) -> Vec { - let mut ready = Vec::new(); - while pending - .front() - .and_then(|request| request.arrival_timestamp_ms) - .is_some_and(|arrival_ms| arrival_ms <= now_ms) - { - ready.push( - pending - .pop_front() - .expect("front request must exist while draining ready arrivals"), - ); - } - ready -} - -async fn dispatch_batch( - arrival_tx: &mpsc::UnboundedSender, - tasks: &mut JoinSet>, - task_ctx: &RequestTaskContext, - batch: Vec, -) -> Result<()> { - if batch.is_empty() { - return Ok(()); - } - - for pending in &batch { - record_arrival(arrival_tx, &pending.request, pending.arrival_at_ms)?; - } - - let mut prepared = Vec::with_capacity(batch.len()); - for pending in batch { - let prepared_request = prepare_request_dispatch(task_ctx, pending.request).await?; - prepared.push((prepared_request, pending.permit)); - } - - for (prepared_request, permit) in prepared { - let dispatched = commit_request_dispatch(task_ctx, prepared_request)?; - tasks.spawn(wait_for_request_completion( - task_ctx.clone(), - dispatched, - permit, - )); - } - - Ok(()) -} +use super::task::{RequestTaskContext, run_request_task, wait_for_workload_progress}; struct LiveRuntime { pending: VecDeque, @@ -182,26 +122,13 @@ impl LiveRuntime { match self.mode { LiveReplayMode::Trace => { - while !self.pending.is_empty() { - let arrival_ms = self - .pending - .front() - .and_then(|request| request.arrival_timestamp_ms) - .ok_or_else(|| { - anyhow!("online trace replay request missing arrival timestamp") - })?; + while let Some(request) = self.pending.pop_front() { + let arrival_ms = request.arrival_timestamp_ms.unwrap_or(0.0); let deadline = start + tokio::time::Duration::from_secs_f64(arrival_ms / 1000.0); tokio::time::sleep_until(deadline).await; - let batch = pop_ready_trace_requests(&mut self.pending, now_ms(start)) - .into_iter() - .map(|request| PendingDispatch { - arrival_at_ms: request.arrival_timestamp_ms.unwrap_or(0.0), - request, - permit: None, - }) - .collect(); - dispatch_batch(&arrival_tx, &mut tasks, &task_ctx, batch).await?; + record_arrival(&arrival_tx, &request, arrival_ms)?; + tasks.spawn(run_request_task(task_ctx.clone(), request, None)); } } LiveReplayMode::Concurrency { max_in_flight } => { @@ -295,19 +222,6 @@ impl LiveRuntime { .unwrap() .pop_ready(now, dispatch_limit); if !ready_turns.is_empty() { - if matches!(self.mode, LiveReplayMode::Trace) { - let batch = ready_turns - .into_iter() - .map(|ready_turn| PendingDispatch { - arrival_at_ms: ready_turn.scheduled_ready_at_ms, - request: ready_turn.request, - permit: None, - }) - .collect(); - dispatch_batch(&arrival_tx, &mut tasks, &task_ctx, batch).await?; - continue; - } - for ready_turn in ready_turns { let permit = match &semaphore { Some(semaphore) => { diff --git a/lib/mocker/src/replay/online/task.rs b/lib/mocker/src/replay/online/task.rs index ac1291165a4f..ea8952a9734a 100644 --- a/lib/mocker/src/replay/online/task.rs +++ b/lib/mocker/src/replay/online/task.rs @@ -8,7 +8,6 @@ use std::sync::Arc; use anyhow::{Result, anyhow, bail}; use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use tokio::time::Instant; -use uuid::Uuid; use crate::common::protocols::DirectRequest; use crate::replay::router::ReplayRouter; @@ -27,18 +26,6 @@ pub(super) struct RequestTaskContext { pub(super) workload: Option>, } -pub(super) struct PreparedRequestDispatch { - request: DirectRequest, - uuid: Uuid, - worker_idx: usize, - state: Arc, -} - -pub(super) struct DispatchedRequest { - uuid: Uuid, - state: Arc, -} - pub(super) async fn wait_for_workload_progress( mode: LiveReplayMode, semaphore: Option<&Semaphore>, @@ -84,16 +71,8 @@ pub(super) async fn run_request_task( request: DirectRequest, permit: Option, ) -> Result<()> { - let prepared = prepare_request_dispatch(&ctx, request).await?; - let dispatched = commit_request_dispatch(&ctx, prepared)?; - wait_for_request_completion(ctx, dispatched, permit).await -} - -pub(super) async fn prepare_request_dispatch( - ctx: &RequestTaskContext, - request: DirectRequest, -) -> Result { let uuid = request_uuid(&request)?; + let worker_idx = ctx .router .select_worker(&request, ctx.senders.len()) @@ -102,51 +81,26 @@ pub(super) async fn prepare_request_dispatch( bail!("online replay selected unknown worker index {worker_idx}"); } - Ok(PreparedRequestDispatch { - request, - uuid, - worker_idx, - state: Arc::new(RequestState::default()), - }) -} - -pub(super) fn commit_request_dispatch( - ctx: &RequestTaskContext, - prepared: PreparedRequestDispatch, -) -> Result { - ctx.requests - .insert(prepared.uuid, Arc::clone(&prepared.state)); - if let Err(error) = ctx.senders[prepared.worker_idx].send(prepared.request) { - ctx.requests.remove(&prepared.uuid); + let state = Arc::new(RequestState::default()); + ctx.requests.insert(uuid, Arc::clone(&state)); + if let Err(error) = ctx.senders[worker_idx].send(request) { + ctx.requests.remove(&uuid); return Err(anyhow!( - "online replay failed to dispatch request to worker {}: {}", - prepared.worker_idx, - error + "online replay failed to dispatch request to worker {worker_idx}: {error}" )); } - ctx.stats.record_dispatch(prepared.worker_idx); - Ok(DispatchedRequest { - uuid: prepared.uuid, - state: prepared.state, - }) -} - -pub(super) async fn wait_for_request_completion( - ctx: RequestTaskContext, - dispatched: DispatchedRequest, - permit: Option, -) -> Result<()> { - dispatched.state.wait_for_completion().await; + ctx.stats.record_dispatch(worker_idx); + state.wait_for_completion().await; ctx.stats.record_completion(); - ctx.requests.remove(&dispatched.uuid); + ctx.requests.remove(&uuid); if let Some(workload) = ctx.workload.as_ref() { let completion_ms = now_ms(workload.start); workload .driver .lock() .unwrap() - .on_complete(dispatched.uuid, completion_ms)?; + .on_complete(uuid, completion_ms)?; workload.wakeup.notify_waiters(); } drop(permit); diff --git a/lib/mocker/src/replay/online/tests.rs b/lib/mocker/src/replay/online/tests.rs index 5417e3cabe5b..e9f21ac31ba8 100644 --- a/lib/mocker/src/replay/online/tests.rs +++ b/lib/mocker/src/replay/online/tests.rs @@ -17,9 +17,9 @@ use crate::replay::ReplayRouterMode; use crate::replay::router::ReplayRouter; use super::live_runtime::{ - pop_ready_trace_requests, simulate_concurrency_requests_with_stats, - simulate_concurrency_workload_with_stats, simulate_trace_requests, - simulate_trace_requests_with_stats, simulate_trace_workload_with_stats, + simulate_concurrency_requests_with_stats, simulate_concurrency_workload_with_stats, + simulate_trace_requests, simulate_trace_requests_with_stats, + simulate_trace_workload_with_stats, }; use super::state::{LiveReplayMode, SharedLiveRuntimeStats, WorkloadDispatchState, record_arrival}; use super::task::{RequestTaskContext, run_request_task, wait_for_workload_progress}; @@ -151,33 +151,6 @@ async fn test_record_arrival_uses_caller_arrival_timestamp() { assert_eq!(arrival.at_ms, arrival_at_ms); } -#[test] -fn test_pop_ready_trace_requests_drains_all_ready_arrivals() { - let mut pending = VecDeque::from(vec![ - request(1, 11, Some(0.0)), - request(2, 22, Some(0.0)), - request(3, 33, Some(0.1)), - request(4, 44, Some(0.2)), - ]); - - let ready = pop_ready_trace_requests(&mut pending, 0.1); - - assert_eq!( - ready - .iter() - .map(|request| request.uuid.unwrap()) - .collect::>(), - vec![Uuid::from_u128(1), Uuid::from_u128(2), Uuid::from_u128(3),] - ); - assert_eq!( - pending - .iter() - .map(|request| request.uuid.unwrap()) - .collect::>(), - vec![Uuid::from_u128(4)] - ); -} - #[tokio::test] async fn test_trace_arrivals_are_not_blocked_by_queued_router_selection() { let args = MockEngineArgs::builder()