From 9beb4811f46d1a60b6130b11da0b989ddce694cc Mon Sep 17 00:00:00 2001 From: Harry Lee Date: Thu, 22 Oct 2026 15:41:18 -0700 Subject: [PATCH] feat: run AgentX with native Dynamo routing interfaces Signed-off-by: Harry Lee --- .github/codeowners/areas.yaml | 5 + CODEOWNERS | 3 + THIRD_PARTY_NOTICES.md | 35 + crates/core/src/perfmodel/fpm/config.rs | 7 + crates/core/src/python.rs | 143 ++- crates/core/src/python/external_policy.rs | 821 ++++++++++++++++++ crates/core/src/replay/agg.rs | 120 ++- crates/core/src/replay/agg_tests.rs | 412 +++++++++ crates/core/src/replay/components/engine.rs | 24 + crates/core/src/replay/core/mod.rs | 43 +- crates/core/src/replay/disagg.rs | 106 ++- crates/core/src/replay/disagg_tests.rs | 351 +++++++- crates/core/src/replay/loadgen/driver.rs | 5 +- crates/core/src/replay/loadgen/lineage.rs | 63 ++ crates/core/src/replay/loadgen/mod.rs | 1 + crates/core/src/replay/loadgen/snapshot.rs | 11 + .../core/src/replay/loadgen/snapshot_tests.rs | 91 ++ crates/core/src/replay/loadgen/trace.rs | 2 + crates/core/src/replay/loadgen/types.rs | 1 + crates/core/src/replay/mod.rs | 9 +- crates/core/src/replay/protocol.rs | 31 +- crates/core/src/replay/replayer.rs | 20 + crates/core/src/replay/runtime_utils.rs | 279 ++++++ crates/core/tests/agentx_qualification.rs | 84 +- docs/adapters/README.md | 1 + docs/adapters/native-composition.md | 32 + docs/replay/agentic/conversation-lineage.md | 39 + docs/replay/agentic/quickstart.md | 64 +- docs/replay/dynamo.md | 39 +- python/aisimulate/THIRD_PARTY_NOTICES.md | 35 + .../src/aisimulate/config/common.py | 1 + .../src/aisimulate/config_adapter.py | 21 +- python/aisimulate/src/aisimulate/main.py | 5 +- python/aisimulate/src/aisimulate/resources.py | 60 +- python/aisimulate/src/aisimulate/router.py | 468 ++++++++++ python/aisimulate/src/aisimulate/runner.py | 26 +- .../aisimulate/src/aisimulate/supervision.py | 5 +- tests/sweeper/test_packaging.py | 36 +- tests/test_conversation_routing_cli.py | 478 ++++++++++ tests/test_host_resources.py | 166 +++- tests/test_runner.py | 19 +- 41 files changed, 3988 insertions(+), 174 deletions(-) create mode 100644 crates/core/src/python/external_policy.rs create mode 100644 crates/core/src/replay/loadgen/lineage.rs create mode 100644 docs/replay/agentic/conversation-lineage.md create mode 100644 python/aisimulate/src/aisimulate/router.py create mode 100644 tests/test_conversation_routing_cli.py diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 50339c25d..1c195b898 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -75,6 +75,7 @@ areas: - tests/test_agentx_g2.py - tests/test_agentx_snapshot_export.py - tests/test_agentx_live_warmup.py + - tests/test_conversation_routing_cli.py - tests/test_runner.py - tests/test_traffic.py - tests/e2e/ @@ -289,6 +290,8 @@ shared: owners: [replay, maintainers] - glob: tests/test_agentx_live_warmup.py owners: [replay, maintainers] + - glob: tests/test_conversation_routing_cli.py + owners: [replay, maintainers] - glob: tests/test_runner.py owners: [replay, maintainers] - glob: tests/test_traffic.py @@ -313,6 +316,8 @@ shared: owners: [forward-pass-engine, maintainers] - glob: crates/core/src/python.rs owners: [mocker, replay, forward-pass-engine, maintainers] + - glob: crates/core/src/python/** + owners: [mocker, replay, forward-pass-engine, maintainers] - glob: Cargo.lock owners: [infra, maintainers] - glob: Cargo.toml diff --git a/CODEOWNERS b/CODEOWNERS index 396409d4f..48f89b3fb 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -84,6 +84,7 @@ /tests/test_output_adapter_context.py @ai-dynamo/aisimulate-replay-codeowners /tests/fixtures/power-contract-v1.json @ai-dynamo/aisimulate-replay-codeowners /tests/test_unified_traffic_runtime.py @ai-dynamo/aisimulate-replay-codeowners +/tests/test_conversation_routing_cli.py @ai-dynamo/aisimulate-replay-codeowners /tests/test_output_adapter_discovery.py @ai-dynamo/aisimulate-replay-codeowners # === mocker (@ai-dynamo/aisimulate-mocker-codeowners) === @@ -181,6 +182,7 @@ /tests/test_cli_config.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_pages_site.py @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /crates/core/src/python.rs @ai-dynamo/aisimulate-mocker-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain +/crates/core/src/python/** @ai-dynamo/aisimulate-mocker-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/docker/ @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/uv.lock @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_unified_cli.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain @@ -231,6 +233,7 @@ /tests/fixtures/power-contract-v1.json @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_spline_estimator_config.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_unified_traffic_runtime.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain +/tests/test_conversation_routing_cli.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_output_adapter_discovery.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/src/aisimulate/replay/ @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain /.claude/skills/aisimulate-fpe-performance @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index d1af3fb73..132aa81ea 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -1432,3 +1432,38 @@ https://github.com/alibaba/tair-kvcache/blob/85402f1a012fb4279be5bb4b57613bb0a06 The AISimulate files are modified re-implementations, not verbatim copies. Both derived files carry the source URLs and revisions in their headers. + +## NVIDIA Dynamo offline routing adapter and acceptance + +The repository-root file `python/aisimulate/src/aisimulate/router.py` is adapted +and modified from Dynamo's +`components/src/dynamo/router/simulation/config.py` and `provider.py` at commit +`6d6d36d7bfa63d80a0302df637c08f85523aa5b6`: +https://github.com/ai-dynamo/dynamo/tree/6d6d36d7bfa63d80a0302df637c08f85523aa5b6/components/src/dynamo/router/simulation +The adaptation retains the configuration vocabulary and field mapping, and +transports policy calls to the separately installed Dynamo runtime. It does not +copy the Router, KV index, or session-affinity algorithms into AISimulate. +The AISim lifecycle and conversation-metadata translation was also adapted and +modified from the same Apache-2.0 NVIDIA project at +`2c4be58bdd6a2dc78c409bdfd67ad54d1b6fedf8`, original files +`lib/bindings/python/rust/llm/offline_router/{affinity,mod,wire}.rs`: +https://github.com/ai-dynamo/dynamo/tree/2c4be58bdd6a2dc78c409bdfd67ad54d1b6fedf8/lib/bindings/python/rust/llm/offline_router + +The repository-root file `tests/test_conversation_routing_cli.py` is adapted and +modified from Dynamo's +`components/src/dynamo/replay/tests/e2e/test_unified_cli_dynamo.py` and +`components/src/dynamo/replay/tests/e2e/configs/unified_cli/predict/dynamo/13-trace-dynamo-agentic.yaml` +at commit `35ea1e7f689198333b1af0a2db8f94f545817792`: +https://github.com/ai-dynamo/dynamo/tree/35ea1e7f689198333b1af0a2db8f94f545817792/components/src/dynamo/replay/tests/e2e +The adaptation runs the conversation fixtures through AISimulate's engine CLI, +uses the standalone policy API, and rejects dependence on Dynamo's whole-run +plugin. The synthetic request tree contains no external dataset material. + +Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +These source revisions are licensed under Apache-2.0. The full license is +reproduced in `LICENSE`; the applicable upstream licenses are: +https://github.com/ai-dynamo/dynamo/blob/6d6d36d7bfa63d80a0302df637c08f85523aa5b6/LICENSE +https://github.com/ai-dynamo/dynamo/blob/35ea1e7f689198333b1af0a2db8f94f545817792/LICENSE +https://github.com/ai-dynamo/dynamo/blob/2c4be58bdd6a2dc78c409bdfd67ad54d1b6fedf8/LICENSE +The source files' copyright/license headers are preserved. These source +revisions have no root NOTICE file applicable to these files. diff --git a/crates/core/src/perfmodel/fpm/config.rs b/crates/core/src/perfmodel/fpm/config.rs index 9c5f27da7..6bf59477e 100644 --- a/crates/core/src/perfmodel/fpm/config.rs +++ b/crates/core/src/perfmodel/fpm/config.rs @@ -69,6 +69,13 @@ pub struct ForwardPassSpeculationConfig { } impl ForwardPassSpeculationConfig { + /// Resolve executor depth through the SDK scheme registry, using the cost + /// identity returned by `ForwardPassPerfModel::best_available` provenance. + #[cfg(feature = "python")] + pub fn replay_depth(&self) -> Result { + crate::py::speculation_replay_depth(self) + } + /// Preserve the existing prompt-lookup contract for canonical and replay callers. pub(crate) fn ngram_depth(&self) -> Result { if self.kind != "ngram" { diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index eca57f8c9..a7b41ef63 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -16,7 +16,8 @@ use crate::replay::{ EncoderSpec, POWER_DATA_COVERAGE_THRESHOLD, ReplayArtifactKvEventVisibility, ReplayArtifacts, ReplayCaptureOptions, ReplayDeterminism, ReplayEngineConfig, ReplayEngineFactory, ReplayOperationPowerDiagnostics, ReplayPhasePowerDiagnostics, ReplayPowerDiagnostics, - ReplayRoleConfig, ReplayRuntimeInput, ReplaySpec, ReplayTopology, Replayer, TracePowerStats, + ReplayReport, ReplayRoleConfig, ReplayRuntimeInput, ReplaySpec, ReplayTopology, Replayer, + TracePowerStats, loadgen::{ AgenticProfileOptions, AgenticSnapshotOptions, AgenticWarmupPolicy, ArrivalSpec, DelaySpec, DynamoRequestTrace, LengthSpec, SyntheticImages, SyntheticTraceSpec, Trace, @@ -35,6 +36,8 @@ use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict, PyModule}; use serde::Deserialize; +mod external_policy; + #[derive(Debug, Deserialize)] #[serde(untagged)] enum ExecutionPayload { @@ -455,7 +458,7 @@ impl AicTimingConfig { speculation.kind ), } - return Ok(crate::py::speculation_replay_depth(speculation)?); + return Ok(speculation.replay_depth()?); } Ok(self.nextn) } @@ -2616,6 +2619,68 @@ fn replay_fpm_coverage( } fn execute_json(payload: &str, capture_artifacts: bool) -> Result { + execute_replay_with_dispatch(payload, capture_artifacts, run_with_input)?.into_json() +} + +/// Shared report serialization for built-in and external placement policies. +#[derive(Debug)] +struct ReplayExecutionResult { + /// Includes per-request records and execution-owned runtime evidence. + report: ReplayReport, + /// Single-worker parity artifacts, when requested through the built-in runtime. + artifacts: Option, + /// Additional execution diagnostics and input provenance that are not + /// part of the topology runtime's typed report. + report_fields: serde_json::Map, + include_per_request: bool, +} + +impl ReplayExecutionResult { + fn into_json(self) -> Result { + let mut report_json = serde_json::to_value(&self.report) + .context("serializing AISimulate replay report summary")?; + let object = report_json + .as_object_mut() + .context("AISimulate replay report did not serialize as an object")?; + object.extend(self.report_fields); + if self.include_per_request || !self.report.per_request.is_empty() { + // Keep records typed until serialization, avoiding a second full + // in-memory JSON tree for large replay reports. + object.insert("per_request".to_string(), serde_json::Value::Null); + } + let report_json = ReportWithRecords { + report: &report_json, + records: &self.report.per_request, + }; + if let Some(artifacts) = self.artifacts { + let output = serde_json::json!({"report": report_json, "artifacts": artifacts}); + return serde_json::to_string(&output).context("serializing AISimulate replay output"); + } + serde_json::to_string(&report_json).context("serializing AISimulate replay output") + } +} + +/// Resolve and finalize a replay while the host selects its execution composition. +/// The callback receives the same prepared inputs as the built-in executor, +/// after canonical parsing, timing construction, and workload preparation. +fn execute_replay_with_dispatch( + payload: &str, + capture_artifacts: bool, + run: F, +) -> Result +where + F: FnOnce( + ReplaySpec, + ReplayEngineFactory, + Option, + bool, + ReplayDeterminism, + Option, + ) -> crate::replay::ReplayResult<( + crate::replay::ReplayReport, + Option, + )>, +{ let (mut spec, mut traffic, capture_performance_diagnostics, determinism, tokenization) = match serde_json::from_str(payload).context("invalid AISimulate execution ReplaySpec")? { ExecutionPayload::Configured { @@ -2764,7 +2829,7 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { if let Some(timing) = &encoder_timing { factory = factory.with_encoder_timing(Arc::clone(timing)); } - run_with_input(spec, factory, input, capture_artifacts, determinism, tokenization) + run(spec, factory, input, capture_artifacts, determinism, tokenization) .map(|(report, artifacts)| (report, artifacts, resolved_basis)) } ReplayTopology::Disaggregated { .. } => { @@ -2884,7 +2949,7 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { if let Some(timing) = &encoder_timing { factory = factory.with_encoder_timing(Arc::clone(timing)); } - run_with_input(spec, factory, input, capture_artifacts, determinism, tokenization) + run(spec, factory, input, capture_artifacts, determinism, tokenization) .map(|(report, artifacts)| (report, artifacts, resolved_basis)) } } @@ -2921,19 +2986,23 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { timing_evidence.as_ref(), unavailable_reason, )?)); - let mut report_json = - serde_json::to_value(&report).context("serializing AISimulate replay report summary")?; - report_json["prediction_provenance"] = prediction_provenance.into(); + let mut report_fields = serde_json::Map::new(); + report_fields.insert( + "prediction_provenance".to_string(), + prediction_provenance.into(), + ); if let Some(coverage) = replay_fpm_coverage(&coverage_sources, Some(&report))? { - report_json["fpm_query_coverage"] = coverage; + report_fields.insert("fpm_query_coverage".to_string(), coverage); } if capture_performance_diagnostics { let performance_evidence = replay_timing_evidence_with_fpm(&power_sources, true)?; - report_json["performance_diagnostics"] = - replay_performance_diagnostics(performance_evidence.as_ref()); + report_fields.insert( + "performance_diagnostics".to_string(), + replay_performance_diagnostics(performance_evidence.as_ref()), + ); } if let Some(evidence) = replay_fpm_query_evidence(&power_sources)? { - report_json["fpm_query_evidence"] = evidence; + report_fields.insert("fpm_query_evidence".to_string(), evidence); } if report.agentic_graph.is_some() && let Some((input_format, agentic_lanes, execution_model)) = agentic_input @@ -2949,29 +3018,26 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { .map(str::trim) .filter(|model| !model.is_empty()) .context("agentic execution did not declare its configured target model")?; - let object = report_json - .as_object_mut() - .context("AISimulate replay report did not serialize as an object")?; - object.insert( + report_fields.insert( "agentic_qualification".to_string(), serde_json::Value::String("functional_only".to_string()), ); - object.insert( + report_fields.insert( "agentic_input_format".to_string(), serde_json::Value::String(input_format), ); - object.insert( + report_fields.insert( "agentic_lanes".to_string(), serde_json::to_value(agentic_lanes) .context("serializing configured agentic lane count")?, ); if let Some(resolved_basis) = resolved_weka_timestamp_basis { - object.insert( + report_fields.insert( "weka_nested_timestamp_basis".to_string(), serde_json::Value::String(resolved_basis.as_str().to_string()), ); } - object.insert( + report_fields.insert( "agentic_model_projection".to_string(), serde_json::json!({ "policy": AGENTIC_MODEL_PROJECTION_POLICY, @@ -2980,26 +3046,12 @@ fn execute_json(payload: &str, capture_artifacts: bool) -> Result { }), ); } - if record_per_request || !report.per_request.is_empty() { - // Holds the key's place in the map; the records are written there during - // serialization, one `Value` at a time. - report_json - .as_object_mut() - .context("AISimulate replay report did not serialize as an object")? - .insert("per_request".to_string(), serde_json::Value::Null); - } - let report_json = ReportWithRecords { - report: &report_json, - records: &report.per_request, - }; - if let Some(artifacts) = artifacts { - let output = serde_json::json!({ - "report": report_json, - "artifacts": artifacts, - }); - return serde_json::to_string(&output).context("serializing AISimulate replay output"); - } - serde_json::to_string(&report_json).context("serializing AISimulate replay output") + Ok(ReplayExecutionResult { + report, + artifacts, + report_fields, + include_per_request: record_per_request, + }) } /// The report object with its `per_request` placeholder serialized as the @@ -3078,6 +3130,18 @@ fn run_replay_json(py: Python<'_>, payload: &str) -> PyResult { .map_err(replay_python_error) } +/// Execute with an independently installed placement policy while retaining +/// AISimulate's workload, engine, virtual clock and report ownership. +#[pyfunction] +fn run_replay_with_policy_json( + py: Python<'_>, + payload: &str, + factory: Py, +) -> PyResult { + py.allow_threads(|| external_policy::execute(payload, factory)) + .map_err(replay_python_error) +} + /// Execute one fixed aggregated ReplaySpec and return report plus parity artifacts. #[pyfunction] fn run_replay_with_artifacts_json(py: Python<'_>, payload: &str) -> PyResult { @@ -3144,6 +3208,7 @@ fn _runtime(module: &Bound<'_, PyModule>) -> PyResult<()> { )?)?; module.add_function(wrap_pyfunction!(canonical_timing_config_json, module)?)?; module.add_function(wrap_pyfunction!(run_replay_json, module)?)?; + module.add_function(wrap_pyfunction!(run_replay_with_policy_json, module)?)?; module.add_function(wrap_pyfunction!(run_replay_with_artifacts_json, module)?)?; crate::perfmodel::register_python(module)?; Ok(()) diff --git a/crates/core/src/python/external_policy.rs b/crates/core/src/python/external_policy.rs new file mode 100644 index 000000000..821639ac6 --- /dev/null +++ b/crates/core/src/python/external_policy.rs @@ -0,0 +1,821 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Transport for independently installed placement policies. The replay owns +//! execution and event visibility; the external object owns placement state. + +use std::cell::RefCell; +use std::ops::Deref; +use std::sync::{Arc, Mutex}; + +use anyhow::{Context, Result, bail, ensure}; +use pyo3::prelude::*; +use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use serde_json::{Map, Value, json}; +use uuid::Uuid; + +use crate::engine::{KvEvent, KvEventData}; +use crate::replay::loadgen::{ReplayRequestHashes, ReplayRequestPayload}; +use crate::replay::{ + EngineEventBatch, Placement, PlacementDecision, PlacementEffects, PlacementPolicy, + ReplayAdmissionMetadata, ReplayCaptureOptions, ReplayComposition, ReplayDeterminism, + ReplayEngineConfig, ReplayEngineObservation, ReplayError, ReplayResult, ReplaySpec, Replayer, + WorkerStage, WorkerTopology, +}; + +/// Keep every owned reference's destruction on the synchronous replay thread. +/// Dropping a plain `Py` without the GIL queues its decref in PyO3's global +/// reference pool, which another thread can flush. External native policies may +/// be unsendable, so even error cleanup must acquire the GIL before decref. +struct PolicyHandle(Option>); + +impl PolicyHandle { + fn new(object: Py) -> Self { + Self(Some(object)) + } + + fn clone_owned(&self) -> Self { + Python::with_gil(|py| Self::new(self.clone_ref(py))) + } +} + +impl Deref for PolicyHandle { + type Target = Py; + + fn deref(&self) -> &Self::Target { + self.0.as_ref().expect("live policy handle") + } +} + +impl Drop for PolicyHandle { + fn drop(&mut self) { + Python::with_gil(|_| drop(self.0.take())); + } +} + +#[derive(Default)] +struct SharedState { + error: Option, + evidence: Map, + policies: Vec<(&'static str, PolicyHandle)>, +} + +type Shared = Arc>; + +fn check_state(shared: &Shared) -> Result<()> { + let state = shared + .lock() + .map_err(|_| anyhow::anyhow!("policy state lock poisoned"))?; + if let Some(error) = &state.error { + bail!("external placement policy failed: {error}"); + } + Ok(()) +} + +/// Execute through the same preparation, event loop and report finalization as +/// the built-in composition. No engine handles or Rust traits cross modules. +pub(super) fn execute(payload: &str, factory: Py) -> Result { + let shared = Shared::default(); + let composition = ExternalComposition { + factory: PolicyHandle::new(factory), + spec: RefCell::new(None), + determinism: ReplayDeterminism::default(), + shared: shared.clone(), + }; + let mut output = super::execute_replay_with_dispatch( + payload, + false, + move |spec, factory, input, _, determinism, tokenization| { + let mut replayer = Replayer::with_composition(spec, factory, composition)? + .with_capture_options(ReplayCaptureOptions { + determinism, + ..Default::default() + }); + if let Some(input) = input { + replayer = replayer.with_runtime_input(input); + } + if let Some(tokenization) = tokenization { + replayer = replayer.with_tokenization( + tokenization.workers, + tokenization.base_ms, + tokenization.per_input_token_ms, + )?; + } + replayer.run().map(|report| (report, None)) + }, + )?; + // cancel_pending has an infallible legacy trait signature. An exception in + // that callback must still invalidate an otherwise successful final report. + check_state(&shared)?; + // Request the potentially large decision evidence once, after execution. + // Returning it with every lifecycle callback would make capture quadratic. + let policies = std::mem::take(&mut shared.lock().unwrap().policies); + for (role, object) in policies { + ExternalPlacement { + object, + role, + pending_count: 0, + next_wakeup_ms: None, + shared: shared.clone(), + } + .call::<()>(json!({"op": "snapshot"}))?; + } + output.report_fields.insert( + "external_policy".into(), + Value::Object(std::mem::take(&mut shared.lock().unwrap().evidence)), + ); + output.into_json() +} + +struct ExternalComposition { + factory: PolicyHandle, + spec: RefCell>, + determinism: ReplayDeterminism, + shared: Shared, +} + +impl ExternalComposition { + fn placement( + &self, + stage: WorkerStage, + dp_size: u32, + topology: Vec, + ) -> Result { + check_state(&self.shared)?; + let spec = self.spec.borrow(); + let (engine, placement, capture) = + spec.as_ref().context("policy spec was not validated")?; + let role = stage_name(stage); + let context = json!({ + "schema_version": 1, + "role": role, + "dp_size": dp_size, + "topology": topology, + "engine": engine.role(stage), + "placement": placement, + "selector_seed": self.determinism.selector_seed(), + "capture_decisions": capture, + }); + let object = Python::with_gil(|py| self.factory.call1(py, (context.to_string(),))) + .with_context(|| format!("constructing external {role} placement policy"))?; + let mut policy = ExternalPlacement { + object: PolicyHandle::new(object), + role, + pending_count: 0, + next_wakeup_ms: None, + shared: self.shared.clone(), + }; + policy.call::<()>(json!({"op": "status"}))?; + self.shared + .lock() + .unwrap() + .policies + .push((role, policy.object.clone_owned())); + Ok(policy) + } +} + +impl ReplayComposition for ExternalComposition { + type Metadata = ExternalMetadata; + type Observation = ExternalObservation; + type AggregatedPlacement = ExternalPlacement; + type DisaggregatedPlacement = ExternalPlacement; + + fn validate_spec(&self, spec: &ReplaySpec) -> ReplayResult<()> { + if spec.adapters.scaling.provider != "none" || !spec.adapters.scaling.config.is_null() { + return Err(ReplayError::InvalidSpec( + "external placement composition does not provide scaling".into(), + )); + } + *self.spec.borrow_mut() = Some(( + ReplayEngineConfig::parse(&spec.engine)?, + serde_json::to_value(&spec.adapters.placement) + .map_err(|error| ReplayError::InvalidSpec(error.to_string()))?, + spec.record_per_request, + )); + Ok(()) + } + + fn create_aggregated_placement( + &mut self, + dp_size: u32, + topology: Vec, + ) -> Result { + self.placement(WorkerStage::Aggregated, dp_size, topology) + } + + fn create_disaggregated_placements( + &mut self, + prefill_dp_size: u32, + prefill_topology: Vec, + decode_dp_size: u32, + decode_topology: Vec, + ) -> Result<(ExternalPlacement, ExternalPlacement)> { + Ok(( + self.placement(WorkerStage::Prefill, prefill_dp_size, prefill_topology)?, + self.placement(WorkerStage::Decode, decode_dp_size, decode_topology)?, + )) + } + + fn set_determinism(&mut self, determinism: ReplayDeterminism) -> ReplayResult<()> { + self.determinism = determinism; + Ok(()) + } +} + +struct ExternalMetadata { + hashes: Option, + output_limit: Option, +} + +impl ReplayAdmissionMetadata for ExternalMetadata { + fn from_hashes(hashes: Option) -> Self { + Self { + hashes, + output_limit: None, + } + } + + fn for_prefill(mut self) -> Self { + self.output_limit = Some(1); + self + } + + fn max_output_tokens_override(&self) -> Option { + self.output_limit + } + + fn into_hashes(self) -> Option { + self.hashes + } +} + +#[derive(Default, Serialize)] +struct ObservationBatch(Vec); + +#[derive(Serialize)] +struct Observation { + stage: &'static str, + worker_id: usize, + dp_rank: u32, + events: Vec, +} + +impl EngineEventBatch for ObservationBatch { + fn is_empty(&self) -> bool { + self.0.is_empty() + } + + fn append(&mut self, mut other: Self) { + self.0.append(&mut other.0); + } +} + +struct ExternalObservation; + +impl ReplayEngineObservation for ExternalObservation { + type Batch = ObservationBatch; + const CAPTURE_ENGINE_KV_EVENTS: bool = true; + + fn observe_engine_events( + stage: WorkerStage, + worker_id: usize, + dp_rank: u32, + events: Vec, + ) -> ObservationBatch { + if events.is_empty() { + ObservationBatch::default() + } else { + ObservationBatch(vec![Observation { + stage: stage_name(stage), + worker_id, + dp_rank, + events, + }]) + } + } + + fn stored_hashes(batch: &ObservationBatch) -> Vec { + batch + .0 + .iter() + .flat_map(|observation| &observation.events) + .flat_map(|event| match &event.data { + KvEventData::Stored(stored) => stored.blocks.as_slice(), + KvEventData::Removed { .. } => &[], + }) + .map(|block| block.tokens_hash) + .collect() + } +} + +#[derive(Deserialize)] +struct Envelope { + result: Value, + pending_count: usize, + next_wakeup_ms: Option, + evidence: Option>, +} + +#[derive(Deserialize)] +struct WireEffects { + placement: Option, + released: Vec, +} + +struct ExternalPlacement { + object: PolicyHandle, + role: &'static str, + pending_count: usize, + next_wakeup_ms: Option, + shared: Shared, +} + +impl ExternalPlacement { + fn call(&mut self, command: Value) -> Result { + let invoke = || -> Result<(T, Envelope)> { + // Cleanup callbacks remain callable after a partial policy commit + // failed. The original poison still prevents a successful report. + if !matches!( + command["op"].as_str(), + Some("dispatch_aborted" | "cancel_pending") + ) { + check_state(&self.shared)?; + } + let encoded = Python::with_gil(|py| -> PyResult { + self.object + .call_method1(py, "call_json", (command.to_string(),))? + .extract(py) + })?; + let value: Value = + serde_json::from_str(&encoded).context("invalid external policy response JSON")?; + ensure!( + value.get("next_wakeup_ms").is_some(), + "external policy response omitted next_wakeup_ms" + ); + let envelope: Envelope = serde_json::from_value(value)?; + if command["op"] == "snapshot" { + ensure!( + envelope.evidence.is_some(), + "external policy snapshot omitted evidence" + ); + } + if let Some(wakeup) = envelope.next_wakeup_ms { + ensure!( + wakeup.is_finite() && wakeup >= 0.0, + "external policy wakeup must be finite and nonnegative" + ); + } + let result = serde_json::from_value(envelope.result.clone())?; + Ok((result, envelope)) + }; + match invoke().with_context(|| format!("external {} policy {}", self.role, command["op"])) { + Ok((result, envelope)) => { + self.pending_count = envelope.pending_count; + self.next_wakeup_ms = envelope.next_wakeup_ms; + if let Some(evidence) = envelope.evidence { + self.shared + .lock() + .unwrap() + .evidence + .insert(self.role.into(), Value::Object(evidence)); + } + Ok(result) + } + Err(error) => { + self.shared + .lock() + .unwrap() + .error + .get_or_insert_with(|| format!("{error:#}")); + Err(error) + } + } + } +} + +impl PlacementPolicy for ExternalPlacement { + type Metadata = ExternalMetadata; + type Observation = ObservationBatch; + + fn place( + &mut self, + request: &ReplayRequestPayload, + metadata: ExternalMetadata, + session_id: Option, + now_ms: f64, + ) -> Result { + let effects: WireEffects = self.call(json!({ + "op": "place", "now_ms": now_ms, "request": request.metadata(), + "input_length": request.input_length(), "hashes": metadata.hashes, + "max_output_tokens_override": metadata.output_limit, "session_id": session_id, + }))?; + Ok(PlacementEffects { + decision: effects + .placement + .map_or(PlacementDecision::Queued, PlacementDecision::Immediate), + released: effects.released, + }) + } + + fn observe(&mut self, observation: ObservationBatch, now_ms: f64) -> Result> { + self.call(json!({"op": "observe", "now_ms": now_ms, "events": observation})) + } + + fn dispatch_committed(&mut self, request_id: Uuid, now_ms: f64) -> Result<()> { + self.call(json!({"op": "dispatch_committed", "request_id": request_id, "now_ms": now_ms})) + } + + fn dispatch_aborted(&mut self, request_id: Uuid, now_ms: f64) -> Result<()> { + self.call(json!({"op": "dispatch_aborted", "request_id": request_id, "now_ms": now_ms})) + } + + fn advance_clock(&mut self, now_ms: f64) -> Result> { + self.call(json!({"op": "advance_clock", "now_ms": now_ms})) + } + + fn next_wakeup_ms(&self) -> Option { + self.next_wakeup_ms + } + + fn pending_count(&self) -> usize { + self.pending_count + } + + fn cancel_pending(&mut self, request_id: Uuid) -> bool { + self.call(json!({"op": "cancel_pending", "request_id": request_id})) + .unwrap_or(false) + } + + fn request_terminal(&mut self, request_id: Uuid, now_ms: f64) -> Result> { + self.call(json!({"op": "request_terminal", "request_id": request_id, "now_ms": now_ms})) + } + + fn prefill_completed(&mut self, request_id: Uuid, now_ms: f64) -> Result> { + self.call(json!({"op": "prefill_completed", "request_id": request_id, "now_ms": now_ms})) + } + + fn worker_ready(&mut self, worker: WorkerTopology, now_ms: f64) -> Result> { + self.call(json!({"op": "worker_ready", "now_ms": now_ms, "worker": worker})) + } + + fn worker_draining(&mut self, worker: WorkerTopology, now_ms: f64) -> Result> { + self.call(json!({"op": "worker_draining", "now_ms": now_ms, "worker": worker})) + } + + fn worker_removed(&mut self, worker: WorkerTopology, now_ms: f64) -> Result> { + self.call(json!({"op": "worker_removed", "now_ms": now_ms, "worker": worker})) + } + + fn topology_settled(&mut self, now_ms: f64) -> Result> { + self.call(json!({"op": "topology_settled", "now_ms": now_ms})) + } +} + +fn stage_name(stage: WorkerStage) -> &'static str { + match stage { + WorkerStage::Aggregated => "aggregated", + WorkerStage::Prefill => "prefill", + WorkerStage::Decode => "decode", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::engine::{KvBlock, KvEventTier, StoredBlocks}; + use crate::replay::{DirectRequest, ReplayTopology, WorkerPoolSpec}; + use pyo3::exceptions::PyValueError; + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[derive(Default)] + struct DropCounts { + factories: AtomicUsize, + policies: AtomicUsize, + } + + #[pyclass(unsendable)] + struct RecordingFactory { + fail_op: Option, + drops: Option>, + } + + impl Drop for RecordingFactory { + fn drop(&mut self) { + if let Some(drops) = &self.drops { + drops.factories.fetch_add(1, Ordering::SeqCst); + } + } + } + + #[pymethods] + impl RecordingFactory { + fn __call__(&self, py: Python<'_>, context: &str) -> PyResult> { + if self.fail_op.as_deref() == Some("construct") { + return Err(PyValueError::new_err("injected construction failure")); + } + Py::new( + py, + RecordingPolicy { + context: serde_json::from_str(context).unwrap(), + calls: Vec::new(), + fail_op: self.fail_op.clone(), + drops: self.drops.clone(), + }, + ) + } + } + + #[pyclass(unsendable)] + struct RecordingPolicy { + context: Value, + calls: Vec, + fail_op: Option, + drops: Option>, + } + + impl Drop for RecordingPolicy { + fn drop(&mut self) { + if let Some(drops) = &self.drops { + drops.policies.fetch_add(1, Ordering::SeqCst); + } + } + } + + #[pymethods] + impl RecordingPolicy { + fn call_json(&mut self, command: &str) -> PyResult { + let command: Value = serde_json::from_str(command).unwrap(); + let op = command["op"].as_str().unwrap(); + if self.fail_op.as_deref() == Some(op) { + return Err(PyValueError::new_err(format!("injected {op} failure"))); + } + let result = match op { + "place" => json!({"placement": { + "request_id": command["request"]["uuid"], + "scheduler_id": self.context["topology"][0]["scheduler_ids"][0], + "reported_overlap_tokens": 0, + }, "released": []}), + "status" | "snapshot" | "dispatch_committed" | "dispatch_aborted" => Value::Null, + "cancel_pending" => json!(true), + _ => json!([]), + }; + let snapshot = op == "snapshot"; + self.calls.push(command); + let mut envelope = + json!({"result": result, "pending_count": 0, "next_wakeup_ms": null}); + if snapshot { + envelope["evidence"] = json!({"construction": self.context, "calls": self.calls}); + } + Ok(envelope.to_string()) + } + } + + fn factory(fail_op: Option<&str>) -> Py { + pyo3::prepare_freethreaded_python(); + Python::with_gil(|py| { + Py::new( + py, + RecordingFactory { + fail_op: fail_op.map(str::to_owned), + drops: None, + }, + ) + .unwrap() + .into_any() + }) + } + + fn policy(fail_op: Option<&str>) -> ExternalPlacement { + let composition = ExternalComposition { + factory: PolicyHandle::new(factory(fail_op)), + spec: RefCell::new(Some(( + ReplayEngineConfig::default(), + json!({"provider": "test_policy", "config": {}}), + true, + ))), + determinism: ReplayDeterminism::default(), + shared: Shared::default(), + }; + composition + .placement( + WorkerStage::Aggregated, + 1, + vec![WorkerTopology { + worker_id: 0, + scheduler_ids: vec![0], + router_hint: false, + }], + ) + .unwrap() + } + + fn snapshot(policy: &mut ExternalPlacement) -> Value { + policy.call::<()>(json!({"op": "snapshot"})).unwrap(); + policy.shared.lock().unwrap().evidence[policy.role].clone() + } + + #[test] + fn external_policy_uses_shared_execution_for_both_topologies() { + for topology in [ + ReplayTopology::aggregated(1), + ReplayTopology::Disaggregated { + prefill: WorkerPoolSpec::default(), + decode: WorkerPoolSpec::default(), + handoff_latency_ms: 0.0, + }, + ] { + let payload = json!({"spec": { + "topology": topology, + "adapters": {"placement": {"provider": "test_policy", "config": {"opaque": 42}}}, + "engine": {"rank": {"num_gpu_blocks": 16, "block_size": 4, "timing_model": {"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0}}}, + "requests": [{"id": "turn", "arrival_time_ms": 0.0, "input_tokens": 8, "output_tokens": 2}], + }, "determinism": "canonical_v1", "tokenization": {"workers": 1, "base_ms": 2.0, "per_input_token_ms": 0.0}}).to_string(); + let output: Value = + serde_json::from_str(&execute(&payload, factory(None)).unwrap()).unwrap(); + assert_eq!(output["per_request"].as_array().unwrap().len(), 1); + for (role, evidence) in output["external_policy"].as_object().unwrap() { + assert_eq!(evidence["construction"]["role"], *role); + assert_eq!( + evidence["construction"]["selector_seed"], + crate::replay::CANONICAL_SELECTOR_SEED + ); + assert_eq!( + evidence["construction"]["placement"]["config"]["opaque"], + 42 + ); + assert_eq!( + evidence["construction"]["engine"]["rank"]["num_gpu_blocks"], + 16 + ); + let calls = evidence["calls"].as_array().unwrap(); + let commit = calls + .iter() + .position(|call| call["op"] == "dispatch_committed") + .unwrap(); + let terminal = calls + .iter() + .position(|call| call["op"] == "request_terminal") + .unwrap(); + assert!(commit < terminal); + assert_eq!(calls.last().unwrap()["op"], "snapshot"); + assert_eq!( + calls.iter().filter(|call| call["op"] == "snapshot").count(), + 1 + ); + if role != "decode" { + assert!(calls.iter().any(|call| call["op"] == "observe")); + } + } + } + } + + #[test] + fn external_policy_preserves_deferred_hashes_and_prefill_limit() { + let mut policy = policy(None); + let request = ReplayRequestPayload::deferred( + DirectRequest { + uuid: Some(Uuid::from_u128(7)), + max_output_tokens: 12, + ..Default::default() + }, + 1_000_000, + vec![9], + 1_000_000, + ); + let hashes = ReplayRequestHashes { + local_block_hashes: vec![81], + sequence_hashes: vec![82], + }; + let metadata = ExternalMetadata::from_hashes(Some(hashes)).for_prefill(); + policy + .place(&request, metadata, Some("conversation".into()), 7.5) + .unwrap(); + let evidence = snapshot(&mut policy); + let command = &evidence["calls"][1]; + assert_eq!(command["request"]["tokens"], json!([])); + assert_eq!(command["input_length"], 1_000_000); + assert_eq!(command["hashes"]["sequence_hashes"], json!([82])); + assert_eq!(command["max_output_tokens_override"], 1); + assert_eq!(command["session_id"], "conversation"); + assert_eq!(command["now_ms"], 7.5); + } + + #[test] + fn external_policy_preserves_physical_kv_event_identity_and_tiers() { + let mut policy = policy(None); + let events = vec![ + KvEvent { + event_id: 6, + dp_rank: 2, + tier: KvEventTier::HostPinned, + data: KvEventData::Stored(StoredBlocks { + parent_hash: Some(5), + start_position: None, + session_id: Some("session".into()), + blocks: vec![KvBlock { + block_hash: 7, + tokens_hash: 8, + token_ids: Some(vec![9]), + }], + }), + }, + KvEvent { + event_id: 7, + dp_rank: 2, + tier: KvEventTier::Device, + data: KvEventData::Removed { + block_hashes: vec![10], + }, + }, + ]; + let expected = serde_json::to_value(&events).unwrap(); + let observation = + ExternalObservation::observe_engine_events(WorkerStage::Prefill, 3, 2, events); + assert_eq!(ExternalObservation::stored_hashes(&observation), vec![8]); + policy.observe(observation, 4.5).unwrap(); + let evidence = snapshot(&mut policy); + let observation = &evidence["calls"][1]["events"][0]; + assert_eq!(observation["events"], expected); + assert_eq!(observation["stage"], "prefill"); + assert_eq!(observation["worker_id"], 3); + assert_eq!(observation["dp_rank"], 2); + } + + #[test] + fn external_policy_infallible_cancel_failure_poison_is_not_hidden() { + let mut policy = policy(Some("cancel_pending")); + assert!(!policy.cancel_pending(Uuid::from_u128(1))); + assert!( + check_state(&policy.shared) + .unwrap_err() + .to_string() + .contains("injected cancel_pending failure") + ); + assert!(policy.advance_clock(1.0).is_err()); + } + + #[test] + fn external_policy_partial_commit_still_allows_abort_cleanup() { + let mut policy = policy(Some("dispatch_committed")); + assert!(policy.dispatch_committed(Uuid::from_u128(1), 0.0).is_err()); + policy.dispatch_aborted(Uuid::from_u128(1), 0.0).unwrap(); + assert!(check_state(&policy.shared).is_err()); + Python::with_gil(|py| { + let recorder = policy + .object + .bind(py) + .extract::>() + .unwrap(); + assert_eq!(recorder.calls.last().unwrap()["op"], "dispatch_aborted"); + }); + } + + #[test] + fn external_policy_drops_unsendable_handles_before_leaving_execution_thread() { + pyo3::prepare_freethreaded_python(); + let payload = json!({ + "topology": ReplayTopology::Disaggregated { + prefill: WorkerPoolSpec::default(), decode: WorkerPoolSpec::default(), + handoff_latency_ms: 0.0, + }, + "adapters": {"placement": {"provider": "test_policy", "config": {}}}, + "engine": {"rank": {"num_gpu_blocks": 16, "block_size": 4, + "timing_model": {"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0}}}, + "requests": [{"id": "turn", "arrival_time_ms": 0.0, "input_tokens": 8, "output_tokens": 2}], + }).to_string(); + for (fail_op, expected_policies) in [ + (None, 2), + (Some("construct"), 0), + (Some("status"), 1), + (Some("place"), 2), + (Some("dispatch_committed"), 2), + (Some("snapshot"), 2), + ] { + let drops = Arc::new(DropCounts::default()); + Python::with_gil(|py| { + let factory = Py::new( + py, + RecordingFactory { + fail_op: fail_op.map(str::to_owned), + drops: Some(drops.clone()), + }, + ) + .unwrap() + .into_any(); + py.allow_threads(|| { + let output = execute(&payload, factory); + assert_eq!(output.is_ok(), fail_op.is_none(), "{fail_op:?}"); + // Assert before reacquiring the caller's GIL: deferred PyO3 + // decrefs must not await another thread's reference-pool flush. + assert_eq!(drops.factories.load(Ordering::SeqCst), 1, "{fail_op:?}"); + assert_eq!( + drops.policies.load(Ordering::SeqCst), + expected_policies, + "{fail_op:?}" + ); + }); + }); + } + } +} diff --git a/crates/core/src/replay/agg.rs b/crates/core/src/replay/agg.rs index 718cb0f0b..4850698cb 100644 --- a/crates/core/src/replay/agg.rs +++ b/crates/core/src/replay/agg.rs @@ -14,10 +14,10 @@ use super::evidence::{ }; use super::progress::ReplayProgress; use super::runtime_utils::{ - ReplayStepOutcome, next_non_telemetry_event_ms, next_timestamp as choose_next_timestamp, - pop_ready_scaling_tick, pop_ready_telemetry_tick, pop_ready_worker_completions, - pop_ready_worker_ready, push_scaling_tick, push_telemetry_tick, push_worker_completions, - push_worker_ready, settle_internal_work, + DispatchFailure, ReplayStepOutcome, next_non_telemetry_event_ms, + next_timestamp as choose_next_timestamp, pop_ready_scaling_tick, pop_ready_telemetry_tick, + pop_ready_worker_completions, pop_ready_worker_ready, push_scaling_tick, push_telemetry_tick, + push_worker_completions, push_worker_ready, settle_internal_work, validate_policy_wakeup, }; use super::scaling::{LatestFpmBuffer, ReplayScalingPolicy, ReplayScalingSnapshot}; use super::telemetry::{ @@ -65,6 +65,7 @@ where PlacementPolicyImpl: PlacementPolicy, { now_ms: f64, + dispatch_failure: DispatchFailure, dp_size: u32, next_event_seq: u64, next_scaling_tick_ordinal: u64, @@ -156,6 +157,7 @@ where Ok(Self { now_ms: 0.0, + dispatch_failure: DispatchFailure::default(), dp_size, next_event_seq: 0, next_scaling_tick_ordinal: 0, @@ -326,8 +328,7 @@ where (request, uuid, worker_idx, kv_hint), )?; } else { - self.engine - .dispatch(worker_idx, request, kv_hint, self.now_ms)?; + self.submit_to_engine(request, uuid, worker_idx, kv_hint)?; } // Aggregated replay uses a single pool. Treat the assignment as the // decode_worker_idx so per-request records consistently carry the @@ -337,6 +338,41 @@ where Ok(()) } + /// Commit policy state only after the engine accepts the tokenized request. + fn submit_to_engine( + &mut self, + request: DirectRequest, + uuid: Uuid, + worker_idx: usize, + kv_hint: Option, + ) -> anyhow::Result<()> { + if let Err(error) = self + .engine + .dispatch(worker_idx, request, kv_hint, self.now_ms) + { + // A rejected duplicate still belongs to the original submission; + // only an accepted dispatch may be rolled back here. + let abort = self.placement.dispatch_aborted(uuid, self.now_ms); + self.requests.remove(&uuid); + return Err(self.dispatch_failure.record(error, Ok(()), abort)); + } + if let Err(error) = self.placement.dispatch_committed(uuid, self.now_ms) { + let rollback = self.engine.rollback_dispatch( + worker_idx, + uuid, + Command::CancelRequest { + request_id: uuid, + discard_pending_output: true, + }, + self.now_ms, + ); + let abort = self.placement.dispatch_aborted(uuid, self.now_ms); + self.requests.remove(&uuid); + return Err(self.dispatch_failure.record(error, rollback, abort)); + } + Ok(()) + } + fn record_placement(&mut self, placement: Placement) { if let Some(sample) = placement.cache_sample { self.traffic @@ -548,9 +584,12 @@ where }; let next_internal_deadline_ms = choose_next_timestamp( self.engine.next_internal_deadline_ms(), - self.encoder - .as_ref() - .and_then(EncoderPool::next_deadline_ms), + choose_next_timestamp( + self.encoder + .as_ref() + .and_then(EncoderPool::next_deadline_ms), + self.placement.next_wakeup_ms(), + ), ); ( choose_next_timestamp( @@ -886,8 +925,7 @@ where let changed = !released.is_empty(); for (tokenized_ms, (request, uuid, worker_idx, kv_hint)) in released { let kv_hint = self.engine.live_kv_hint(kv_hint); - self.engine - .dispatch(worker_idx, request, kv_hint, self.now_ms)?; + self.submit_to_engine(request, uuid, worker_idx, kv_hint)?; self.collector.on_tokenized(uuid, tokenized_ms); } Ok(changed) @@ -1032,6 +1070,7 @@ where /// Repeatedly process all work that becomes possible without advancing logical time. fn drain_current_timestamp(&mut self) -> anyhow::Result<()> { + self.dispatch_failure.ensure_healthy()?; #[cfg(test)] { self.stats.semantic_drain_count += 1; @@ -1039,9 +1078,18 @@ where let mut consecutive_internal_steps = 0usize; loop { let mut changed = false; - // Settle idle deadlines first: a completion below may release a - // queued placement and submit it to that same idle worker. + validate_policy_wakeup( + self.placement.next_wakeup_ms(), + self.now_ms, + "aggregated", + false, + )?; + // Settle idle deadlines before policy wakeups or completions can + // release a queued placement to that same idle worker. changed |= self.settle_internal_work(&mut consecutive_internal_steps)?; + let placements = self.placement.advance_clock(self.now_ms)?; + changed |= !placements.is_empty(); + self.dispatch_placements(placements)?; let completed = self.apply_worker_completions()?; changed |= completed; if completed { @@ -1125,6 +1173,12 @@ where changed |= self.apply_scaling_ticks()?; } + validate_policy_wakeup( + self.placement.next_wakeup_ms(), + self.now_ms, + "aggregated", + !changed, + )?; if !changed { break; } @@ -1652,6 +1706,7 @@ where } fn ensure_drive_started(&mut self) -> anyhow::Result { + self.dispatch_failure.ensure_healthy()?; if self.drive_started { return Ok(false); } @@ -2761,6 +2816,7 @@ mod host_offload_tests { struct ReleaseToIdleWorker { pending: Option, + release_at_ms: Option, } impl PlacementPolicy for ReleaseToIdleWorker { @@ -2817,12 +2873,24 @@ mod host_offload_tests { } } + fn next_wakeup_ms(&self) -> Option { + self.pending.and(self.release_at_ms) + } + + fn advance_clock(&mut self, now_ms: f64) -> anyhow::Result> { + if self.release_at_ms.is_some_and(|at_ms| at_ms <= now_ms) { + self.release_at_ms = None; + return self.request_terminal(COMPLETING_A, now_ms); + } + Ok(Vec::new()) + } + fn request_terminal( &mut self, request_id: Uuid, _now_ms: f64, ) -> anyhow::Result> { - if request_id != COMPLETING_A { + if request_id != COMPLETING_A || self.release_at_ms.is_some() { return Ok(Vec::new()); } Ok(self @@ -2891,8 +2959,13 @@ mod host_offload_tests { } } - #[test] - fn idle_internal_deadline_settles_before_completion_releases_placement() { + #[rstest::rstest] + #[case::completion(None)] + #[case::coincident_wakeup(Some(5.0))] + #[case::later_wakeup(Some(5.25))] + fn idle_internal_deadline_settles_before_releasing_placement( + #[case] release_at_ms: Option, + ) { let config = ReplayEngineConfig { rank: EngineConfig { num_gpu_blocks: 1, @@ -2938,7 +3011,10 @@ mod host_offload_tests { .collect::>(), vec![&[0][..], &[1][..]] ); - Ok(ReleaseToIdleWorker { pending: None }) + Ok(ReleaseToIdleWorker { + pending: None, + release_at_ms, + }) }, ) .unwrap() @@ -2962,7 +3038,7 @@ mod host_offload_tests { .find(|request| request.uuid == QUEUED_FOR_B.to_string()) .expect("released request record"); assert_eq!(released.decode_worker_idx, Some(1)); - assert_eq!(released.first_admit_ms, Some(5.0)); + assert_eq!(released.first_admit_ms, Some(release_at_ms.unwrap_or(5.0))); assert_eq!(released.terminal_status, ReplayTerminalStatus::Completed); } } @@ -3013,8 +3089,11 @@ where } /// Admit `request` at the current simulated time. The returned id - /// correlates the request with later measurements. + /// correlates the request with later measurements. A dispatch failure + /// poisons this runtime: its policy/collector may have changed even after + /// engine rollback. Construct a new runtime instead of retrying submission. pub(crate) fn submit_dynamic(&mut self, request: DirectRequest) -> anyhow::Result { + self.dispatch_failure.ensure_healthy()?; let arrival_time_ms = self.now_ms; let uuid = self.assign_request( ReplayRequestPayload::materialized(request), @@ -3035,6 +3114,7 @@ where &mut self, uuid: Uuid, ) -> anyhow::Result> { + self.dispatch_failure.ensure_healthy()?; let Some((phase, scheduler_id)) = self .requests .get(&uuid) @@ -3148,6 +3228,7 @@ where /// bypasses the terminal would break the invariant silently instead of /// failing a test. pub(crate) fn step_dynamic_until(&mut self, until_ms: f64) -> anyhow::Result { + self.dispatch_failure.ensure_healthy()?; if until_ms.is_nan() || until_ms < self.now_ms { bail!( "aggregated step deadline {until_ms}ms precedes runtime time {}ms", @@ -3196,6 +3277,7 @@ where &mut self, wall_ms: f64, ) -> anyhow::Result { + self.dispatch_failure.ensure_healthy()?; anyhow::ensure!(wall_ms.is_finite(), "replay report wall_ms must be finite"); anyhow::ensure!( self.is_workload_done(), diff --git a/crates/core/src/replay/agg_tests.rs b/crates/core/src/replay/agg_tests.rs index d5b9fbad1..ee03e5a27 100644 --- a/crates/core/src/replay/agg_tests.rs +++ b/crates/core/src/replay/agg_tests.rs @@ -184,6 +184,230 @@ fn mismatched_placement_does_not_retain_arrival_or_offered_traffic() { ); } +#[derive(Debug, Clone, PartialEq)] +enum DispatchHookRecord { + Selected(Uuid, f64), + Clock(f64), + Released(Uuid, f64), + Committed(Uuid, f64), + Terminal(Uuid, f64), +} + +type DispatchHookRecords = std::sync::Arc>>; + +struct RecordingDispatchPlacement { + records: DispatchHookRecords, + release_at_ms: Option, + pending: Option, +} + +impl RecordingDispatchPlacement { + fn placement(request_id: Uuid) -> Placement { + Placement { + request_id, + scheduler_id: 0, + reported_overlap_tokens: 0, + cache_sample: None, + kv_hints: None, + placement_replica_id: None, + } + } + + fn record(&self, event: DispatchHookRecord) { + self.records.lock().unwrap().push(event); + } +} + +impl PlacementPolicy for RecordingDispatchPlacement { + type Metadata = NoReplayMetadata; + type Observation = (); + + fn place( + &mut self, + request: &ReplayRequestPayload, + _: (), + _: Option, + now_ms: f64, + ) -> anyhow::Result { + let uuid = request.metadata().uuid.unwrap(); + self.record(DispatchHookRecord::Selected(uuid, now_ms)); + let decision = if self.release_at_ms.is_some() { + self.pending = Some(uuid); + PlacementDecision::Queued + } else { + PlacementDecision::Immediate(Self::placement(uuid)) + }; + Ok(crate::replay::PlacementEffects { + decision, + released: Vec::new(), + }) + } + fn observe(&mut self, _: (), _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn dispatch_committed(&mut self, uuid: Uuid, now_ms: f64) -> anyhow::Result<()> { + self.record(DispatchHookRecord::Committed(uuid, now_ms)); + Ok(()) + } + fn advance_clock(&mut self, now_ms: f64) -> anyhow::Result> { + self.record(DispatchHookRecord::Clock(now_ms)); + if self.release_at_ms.is_some_and(|at| at <= now_ms) + && let Some(uuid) = self.pending.take() + { + self.release_at_ms = None; + self.record(DispatchHookRecord::Released(uuid, now_ms)); + return Ok(vec![Self::placement(uuid)]); + } + Ok(Vec::new()) + } + fn next_wakeup_ms(&self) -> Option { + self.pending.and(self.release_at_ms) + } + fn cancel_pending(&mut self, uuid: Uuid) -> bool { + if self.pending == Some(uuid) { + self.pending = None; + true + } else { + false + } + } + fn request_terminal(&mut self, uuid: Uuid, now_ms: f64) -> anyhow::Result> { + self.record(DispatchHookRecord::Terminal(uuid, now_ms)); + Ok(Vec::new()) + } + fn prefill_completed(&mut self, _: Uuid, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn pending_count(&self) -> usize { + usize::from(self.pending.is_some()) + } + fn worker_ready(&mut self, _: WorkerTopology, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn worker_draining(&mut self, _: WorkerTopology, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn worker_removed(&mut self, _: WorkerTopology, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn topology_settled(&mut self, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } +} + +fn dispatch_hook_runtime( + backend: crate::engine::Backend, + release_at_ms: Option, + records: DispatchHookRecords, +) -> AggRuntimeImpl { + let config = ReplayEngineConfig { + rank: EngineConfig { + timing_model: TimingModelConfig::Fixed { + prefill_ms: 1.0, + decode_ms: 1.0, + }, + ..EngineConfig::for_backend(backend) + }, + ..ReplayEngineConfig::default() + }; + let factory = ReplayEngineFactory::new() + .role_factory(&config, WorkerStage::Aggregated, false) + .unwrap(); + AggRuntimeImpl::new_composed( + factory, + AdmissionQueue::new_requests(VecDeque::from([request(1, 0.0)]), ReplayMode::Trace), + 1, + None, + |_, _| { + Ok(RecordingDispatchPlacement { + records, + release_at_ms, + pending: None, + }) + }, + ) + .unwrap() + .with_per_request_records(true) +} + +#[rstest::rstest] +#[case(crate::engine::Backend::Vllm)] +#[case(crate::engine::Backend::Sglang)] +fn placement_dispatch_hooks_commit_after_clock_releases_a_queued_request( + #[case] backend: crate::engine::Backend, + #[values(0.0, 3.0)] tokenization_ms: f64, + #[values(0.0, 7.0)] wakeup_ms: f64, +) { + let records = DispatchHookRecords::default(); + let report = dispatch_hook_runtime(backend, Some(wakeup_ms), records.clone()) + .with_tokenization(TokenizationLatency::new(1, tokenization_ms, 0.0)) + .run() + .unwrap() + .0 + .finish(); + assert_eq!(report.request_counts.completed_requests, 1); + assert_eq!(report.per_request[0].arrival_time_ms, 0.0); + assert_eq!( + report.per_request[0].first_admit_ms, + Some(wakeup_ms + tokenization_ms) + ); + let terminal_ms = report.per_request[0].terminal_time_ms; + assert!(terminal_ms > wakeup_ms); + let records = records.lock().unwrap(); + let semantic = records + .iter() + .filter(|event| !matches!(event, DispatchHookRecord::Clock(_))) + .cloned() + .collect::>(); + let uuid = Uuid::from_u128(1); + assert_eq!( + semantic, + vec![ + DispatchHookRecord::Selected(uuid, 0.0), + DispatchHookRecord::Released(uuid, wakeup_ms), + DispatchHookRecord::Committed(uuid, wakeup_ms + tokenization_ms), + DispatchHookRecord::Terminal(uuid, terminal_ms) + ] + ); + assert!(records.contains(&DispatchHookRecord::Clock(wakeup_ms))); +} + +#[rstest::rstest] +fn canceling_tokenization_releases_policy_without_committing_dispatch( + #[values(crate::engine::Backend::Vllm, crate::engine::Backend::Sglang)] + backend: crate::engine::Backend, +) { + let records = DispatchHookRecords::default(); + let mut runtime = dispatch_hook_runtime(backend, None, records.clone()) + .with_tokenization(TokenizationLatency::new(1, 3.0, 0.0)); + let uuid = Uuid::from_u128(1); + runtime.step().unwrap(); + assert_eq!(runtime.now_ms, 0.0); + assert_eq!(runtime.engine.in_flight(), 0); + assert_eq!(runtime.cluster_in_flight(), 1); + assert_eq!( + runtime.cancel_dynamic(uuid).unwrap(), + Some(ReplayTerminalStatus::Canceled) + ); + assert_eq!(runtime.cluster_in_flight(), 0); + runtime.now_ms = 3.0; + assert!(!runtime.release_tokenized().unwrap()); + let semantic = records + .lock() + .unwrap() + .iter() + .filter(|event| !matches!(event, DispatchHookRecord::Clock(_))) + .cloned() + .collect::>(); + assert_eq!( + semantic, + [ + DispatchHookRecord::Selected(uuid, 0.0), + DispatchHookRecord::Terminal(uuid, 0.0), + ] + ); +} + // G3 fixtures independently adapt vLLM behavioral contracts (Apache-2.0). // Upstream: https://github.com/vllm-project/vllm/tree/6e448d0ea9bf3d88d898b65449ca6dc2aec170ac // Paths: vllm/v1/core/sched/{scheduler,request_queue}.py and @@ -867,3 +1091,191 @@ fn agg_first_step_settles_initial_scaling_and_zero_delay_worker_startup() { assert!(stepped.events.iter().all(|event| event.at_ms > 0.0)); assert!(stepped.next_timestamps().1.unwrap() > 0.0); } + +#[rstest::rstest] +#[case(f64::NAN)] +#[case(f64::INFINITY)] +#[case(f64::NEG_INFINITY)] +#[case(-1.0)] +#[case(0.0)] +fn agg_rejects_invalid_policy_wakeup_before_advancing_clock(#[case] wakeup_ms: f64) { + use crate::replay::runtime_utils::wakeup_test_policy::WakeupPlacement; + let config = ReplayEngineConfig::default(); + let factory = ReplayEngineFactory::new() + .role_factory(&config, WorkerStage::Aggregated, false) + .unwrap(); + let mut runtime = AggRuntimeImpl::<_, NoEngineEvents, NoReplayMetadata>::new_composed( + factory, + AdmissionQueue::new_requests(VecDeque::from([request(1, 0.0)]), ReplayMode::Trace), + 1, + None, + |_, _| Ok(WakeupPlacement::new(Some(wakeup_ms), false)), + ) + .unwrap(); + let error = runtime + .step() + .expect_err("invalid or unconsumed wakeup must fail while settling"); + assert!( + error + .to_string() + .contains("aggregated placement policy wakeup"), + "{error:#}" + ); + assert_eq!( + runtime.now_ms, 0.0, + "invalid policy must not corrupt replay time" + ); +} + +fn failing_dispatch_runtime( + backend: crate::engine::Backend, + fail_commit: bool, + fail_abort: bool, +) -> AggRuntimeImpl< + crate::replay::runtime_utils::dispatch_failure_test_policy::FailingDispatchPlacement, + NoEngineEvents, + NoReplayMetadata, +> { + use crate::replay::runtime_utils::dispatch_failure_test_policy::FailingDispatchPlacement; + let config = ReplayEngineConfig { + rank: EngineConfig { + timing_model: TimingModelConfig::Fixed { + prefill_ms: 1.0, + decode_ms: 1.0, + }, + ..EngineConfig::for_backend(backend) + }, + ..ReplayEngineConfig::default() + }; + let factory = ReplayEngineFactory::new() + .role_factory(&config, WorkerStage::Aggregated, false) + .unwrap(); + AggRuntimeImpl::new_composed( + factory, + AdmissionQueue::new_requests(VecDeque::new(), ReplayMode::Trace), + 1, + None, + |_, _| { + Ok(FailingDispatchPlacement { + fail_commit, + fail_abort, + ..Default::default() + }) + }, + ) + .unwrap() + .into_steppable() +} + +#[rstest::rstest] +fn dispatch_commit_failure_rolls_back_and_prevents_runtime_reuse( + #[values(crate::engine::Backend::Vllm, crate::engine::Backend::Sglang)] + backend: crate::engine::Backend, + #[values(false, true)] fail_abort: bool, + #[values(0.0, 3.0)] tokenization_ms: f64, +) { + let mut runtime = failing_dispatch_runtime(backend, true, fail_abort) + .with_tokenization(TokenizationLatency::new(1, tokenization_ms, 0.0)); + let error = if tokenization_ms == 0.0 { + runtime.submit_dynamic(request(1, 0.0)).unwrap_err() + } else { + runtime.submit_dynamic(request(1, 0.0)).unwrap(); + assert_eq!(runtime.engine.in_flight(), 0); + assert!(runtime.placement.calls.lock().unwrap().is_empty()); + runtime.step_dynamic_until(10.0).unwrap_err() + }; + assert_eq!( + runtime.engine.in_flight(), + 0, + "failed submission retained engine ownership" + ); + assert!(!runtime.engine.has_runnable_worker()); + assert!(runtime.requests.is_empty()); + assert_eq!( + *runtime.placement.calls.lock().unwrap(), + ["commit", "abort"] + ); + assert!( + error + .chain() + .any(|cause| cause.to_string() == "injected commit failure") + ); + assert!( + error.to_string().contains("injected commit failure"), + "{error:#}" + ); + assert_eq!( + format!("{error:#}").contains("injected abort failure"), + fail_abort + ); + runtime.placement.fail_commit = false; + runtime.placement.fail_abort = false; + for retry_id in [1, 2] { + assert!( + runtime + .submit_dynamic(request(retry_id, 0.0)) + .unwrap_err() + .to_string() + .contains("poisoned") + ); + } + assert!( + runtime + .step_dynamic_until(10.0) + .unwrap_err() + .to_string() + .contains("poisoned") + ); + assert!(runtime.step().unwrap_err().to_string().contains("poisoned")); + assert!( + runtime + .take_report_dynamic(0.0) + .unwrap_err() + .to_string() + .contains("poisoned") + ); + assert!( + runtime + .run() + .err() + .unwrap() + .to_string() + .contains("poisoned") + ); +} + +#[rstest::rstest] +fn dispatch_abort_preserves_native_error_and_existing_ownership( + #[values(crate::engine::Backend::Vllm, crate::engine::Backend::Sglang)] + backend: crate::engine::Backend, + #[values(false, true)] fail_abort: bool, +) { + let mut runtime = failing_dispatch_runtime(backend, false, fail_abort); + runtime + .engine + .dispatch(0, request(1, 0.0), None, 0.0) + .unwrap(); + let error = runtime.submit_dynamic(request(1, 0.0)).unwrap_err(); + assert!( + error + .chain() + .any(|cause| cause.to_string().contains("already active")), + "{error:#}" + ); + assert!(error.to_string().contains("already active"), "{error:#}"); + assert_eq!( + format!("{error:#}").contains("injected abort failure"), + fail_abort + ); + // A rejected duplicate must never cancel the engine's original request. + assert_eq!(runtime.engine.in_flight(), 1); + assert!(runtime.requests.is_empty()); + assert_eq!(*runtime.placement.calls.lock().unwrap(), ["abort"]); + assert!( + runtime + .step_dynamic_until(10.0) + .unwrap_err() + .to_string() + .contains("poisoned") + ); +} diff --git a/crates/core/src/replay/components/engine.rs b/crates/core/src/replay/components/engine.rs index a848be18b..549ac445a 100644 --- a/crates/core/src/replay/components/engine.rs +++ b/crates/core/src/replay/components/engine.rs @@ -773,6 +773,30 @@ where Ok(()) } + /// Undo an accepted dispatch before abandoning a failed runtime. Effects + /// are deliberately not fed back into a policy whose callback already + /// failed; the caller must poison replay and forbid subsequent reporting. + pub(crate) fn rollback_dispatch( + &mut self, + scheduler_id: usize, + request_id: Uuid, + cancel: Command, + now_ms: f64, + ) -> Result<()> { + let effects = self.apply_command(scheduler_id, cancel, now_ms)?; + anyhow::ensure!( + matches!(effects.result, CommandResult::Applied | CommandResult::Noop), + "native engine returned an unexpected dispatch rollback result" + ); + let owner = self.scheduler_owner(scheduler_id)?; + anyhow::ensure!( + !self.required_worker(owner.worker_id)?.in_flight_by_rank[owner.dp_rank as usize] + .contains(&request_id), + "native engine retained request {request_id} after dispatch rollback" + ); + Ok(()) + } + pub(crate) fn apply_command( &mut self, scheduler_id: usize, diff --git a/crates/core/src/replay/core/mod.rs b/crates/core/src/replay/core/mod.rs index 7d6b7ccf4..25c0640b0 100644 --- a/crates/core/src/replay/core/mod.rs +++ b/crates/core/src/replay/core/mod.rs @@ -7,6 +7,7 @@ //! can later become a standalone crate without deployment-specific APIs. use anyhow::Result; +use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::replay::ReplayTerminalStatus; @@ -64,7 +65,7 @@ pub trait AdmissionSource { fn total_requests(&self) -> usize; } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] pub struct PlacementCacheSample { /// Prefix blocks available on the selected worker. pub overlap_blocks: u32, @@ -75,7 +76,7 @@ pub struct PlacementCacheSample { /// Request-scoped KV hints a router attaches to a placement (vLLM `kv_hints`, /// Dynamo `kv_hint`). Only `kv.fetch` is modeled. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] pub struct KvHints { pub fetch: KvFetchAction, } @@ -86,12 +87,12 @@ pub struct KvHints { /// identifies one worker and attention-DP rank. It uses the scheduler-id space /// of the stage that placed the request. The hint's block list is not carried /// because it does not affect behavior with opportunistic query. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] pub struct KvFetchAction { pub source_scheduler_id: usize, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] pub struct Placement { pub request_id: Uuid, pub scheduler_id: usize, @@ -120,7 +121,7 @@ pub struct PlacementEffects { pub released: Vec, } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct WorkerTopology { pub worker_id: usize, pub scheduler_ids: Vec, @@ -141,6 +142,38 @@ pub trait PlacementPolicy { now_ms: f64, ) -> Result; fn observe(&mut self, observation: Self::Observation, now_ms: f64) -> Result>; + /// The selected engine accepted ownership (destination reservation for P/D + /// decode). Policies may commit a tentative binding only at this boundary. + /// If this fails, replay attempts to cancel the accepted engine request and + /// calls `dispatch_aborted`, then permanently fails the runtime. Callbacks + /// need not be transactional; replay cannot safely resume or report after a + /// potentially partial policy commit, even when engine cleanup succeeds. + fn dispatch_committed(&mut self, _request_id: Uuid, _now_ms: f64) -> Result<()> { + Ok(()) + } + /// The selected engine rejected ownership, or rollback was attempted after + /// `dispatch_committed` failed. Discard tentative policy + /// state; tolerate partial commits. An abort error is reported together with + /// the original dispatch error and permanently fails the runtime. + fn dispatch_aborted(&mut self, _request_id: Uuid, _now_ms: f64) -> Result<()> { + Ok(()) + } + /// Advance policy time using the replay clock and release any newly ready + /// placements. Wall time must not drive offline policy state. + fn advance_clock(&mut self, _now_ms: f64) -> Result> { + Ok(Vec::new()) + } + /// A concrete wakeup for policy-owned work, finite and no earlier than the + /// current replay time. A same-time wakeup is allowed while immediate work + /// is draining (for example, after `dispatch_committed`); `advance_clock` + /// must consume it or release work so the drain makes progress. Once the + /// timestamp is settled, any remaining wakeup must be strictly in the future. + /// Invalid or unconsumed same-time wakeups fail replay rather than advancing + /// its clock. Idle housekeeping that cannot release work should be performed + /// lazily in `advance_clock`. + fn next_wakeup_ms(&self) -> Option { + None + } fn cancel_pending(&mut self, request_id: Uuid) -> bool; fn request_terminal(&mut self, request_id: Uuid, now_ms: f64) -> Result>; fn prefill_completed(&mut self, request_id: Uuid, now_ms: f64) -> Result>; diff --git a/crates/core/src/replay/disagg.rs b/crates/core/src/replay/disagg.rs index 9236a8c93..513162876 100644 --- a/crates/core/src/replay/disagg.rs +++ b/crates/core/src/replay/disagg.rs @@ -38,10 +38,11 @@ use super::evidence::{ }; use super::progress::ReplayProgress; use super::runtime_utils::{ - ReplayStepOutcome, next_non_telemetry_event_ms, next_timestamp as choose_next_timestamp, - pop_ready_scaling_tick, pop_ready_telemetry_tick, pop_ready_transfer_complete, - pop_ready_worker_completions, pop_ready_worker_ready, push_scaling_tick, push_telemetry_tick, - push_transfer_complete, push_worker_completions, push_worker_ready, settle_internal_work, + DispatchFailure, ReplayStepOutcome, next_non_telemetry_event_ms, + next_timestamp as choose_next_timestamp, pop_ready_scaling_tick, pop_ready_telemetry_tick, + pop_ready_transfer_complete, pop_ready_worker_completions, pop_ready_worker_ready, + push_scaling_tick, push_telemetry_tick, push_transfer_complete, push_worker_completions, + push_worker_ready, settle_internal_work, validate_policy_wakeup, }; use super::scaling::{LatestFpmBuffer, ReplayScalingPolicy, ReplayScalingSnapshot}; #[cfg(test)] @@ -936,6 +937,7 @@ where PlacementPolicyImpl: PlacementPolicy, { now_ms: f64, + dispatch_failure: DispatchFailure, next_event_seq: u64, next_scaling_tick_ordinal: u64, admission: AdmissionQueue, @@ -1087,6 +1089,7 @@ where Ok(Self { now_ms: 0.0, + dispatch_failure: DispatchFailure::default(), next_event_seq: 0, next_scaling_tick_ordinal: 0, admission, @@ -1394,6 +1397,9 @@ where ) { Ok(effects) => effects, Err(error) => { + if let Err(abort) = self.prefill_placement.dispatch_aborted(uuid, self.now_ms) { + return Err(self.dispatch_failure.record(error, Ok(()), Err(abort))); + } self.acknowledge_action( uuid, action, @@ -1403,7 +1409,26 @@ where } }; if !matches!(effects.result, CommandResult::Submitted(id) if id == uuid) { - bail!("offline disagg replay prefill submission returned an unexpected result"); + let error = + anyhow!("offline disagg replay prefill submission returned an unexpected result"); + let rollback = self.prefill_engine.rollback_dispatch( + worker_idx, + uuid, + Command::CancelSource { handoff_id }, + self.now_ms, + ); + let abort = self.prefill_placement.dispatch_aborted(uuid, self.now_ms); + return Err(self.dispatch_failure.record(error, rollback, abort)); + } + if let Err(error) = self.prefill_placement.dispatch_committed(uuid, self.now_ms) { + let rollback = self.prefill_engine.rollback_dispatch( + worker_idx, + uuid, + Command::CancelSource { handoff_id }, + self.now_ms, + ); + let abort = self.prefill_placement.dispatch_aborted(uuid, self.now_ms); + return Err(self.dispatch_failure.record(error, rollback, abort)); } self.flow.finish_prefill_submission( uuid, @@ -1434,6 +1459,9 @@ where ) { Ok(effects) => effects, Err(error) => { + if let Err(abort) = self.decode_placement.dispatch_aborted(uuid, self.now_ms) { + return Err(self.dispatch_failure.record(error, Ok(()), Err(abort))); + } self.acknowledge_action( uuid, action, @@ -1446,7 +1474,27 @@ where effects.result, CommandResult::DestinationAccepted { request_id } if request_id == uuid ) { - bail!("offline disagg replay destination acceptance returned an unexpected result"); + let error = anyhow!( + "offline disagg replay destination acceptance returned an unexpected result" + ); + let rollback = self.decode_engine.rollback_dispatch( + worker_idx, + uuid, + Command::CancelDestination { handoff_id }, + self.now_ms, + ); + let abort = self.decode_placement.dispatch_aborted(uuid, self.now_ms); + return Err(self.dispatch_failure.record(error, rollback, abort)); + } + if let Err(error) = self.decode_placement.dispatch_committed(uuid, self.now_ms) { + let rollback = self.decode_engine.rollback_dispatch( + worker_idx, + uuid, + Command::CancelDestination { handoff_id }, + self.now_ms, + ); + let abort = self.decode_placement.dispatch_aborted(uuid, self.now_ms); + return Err(self.dispatch_failure.record(error, rollback, abort)); } let stored_hashes = self .flow @@ -2089,16 +2137,21 @@ where } else { next_event_ms }; - // Host transfers and frontend deliveries are engine-internal deadlines - // for either role; the encoder pool delivers embeddings on its own clock. + // Engine, encoder and placement deadlines all advance replay semantics. let next_internal_deadline_ms = choose_next_timestamp( choose_next_timestamp( self.prefill_engine.next_internal_deadline_ms(), self.decode_engine.next_internal_deadline_ms(), ), - self.encoder - .as_ref() - .and_then(EncoderPool::next_deadline_ms), + choose_next_timestamp( + self.encoder + .as_ref() + .and_then(EncoderPool::next_deadline_ms), + choose_next_timestamp( + self.prefill_placement.next_wakeup_ms(), + self.decode_placement.next_wakeup_ms(), + ), + ), ); ( choose_next_timestamp( @@ -2732,12 +2785,25 @@ where /// Repeatedly process all work that becomes possible without advancing logical time. fn drain_current_timestamp(&mut self) -> Result<()> { + self.dispatch_failure.ensure_healthy()?; #[cfg(test)] { self.stats.semantic_drain_count += 1; } let mut consecutive_internal_steps = 0usize; loop { + validate_policy_wakeup( + self.prefill_placement.next_wakeup_ms(), + self.now_ms, + "prefill", + false, + )?; + validate_policy_wakeup( + self.decode_placement.next_wakeup_ms(), + self.now_ms, + "decode", + false, + )?; let now_ms = self.now_ms; let mut settle = |runtime: &mut Self| { settle_internal_work(now_ms, &mut consecutive_internal_steps, || { @@ -2748,6 +2814,11 @@ where // request to an idle prefill worker at this same timestamp. let mut changed = settle(self)?; changed |= self.prune_stale_transfer_events(); + let prefill = self.prefill_placement.advance_clock(self.now_ms)?; + let decode = self.decode_placement.advance_clock(self.now_ms)?; + changed |= !prefill.is_empty() || !decode.is_empty(); + self.dispatch_prefill_placements(prefill)?; + self.dispatch_decode_placements(decode)?; let completed = self.apply_worker_completions()?; changed |= completed; if completed { @@ -2873,6 +2944,18 @@ where changed |= self.apply_scaling_ticks()?; } + validate_policy_wakeup( + self.prefill_placement.next_wakeup_ms(), + self.now_ms, + "prefill", + !changed, + )?; + validate_policy_wakeup( + self.decode_placement.next_wakeup_ms(), + self.now_ms, + "decode", + !changed, + )?; if !changed { break; } @@ -3656,6 +3739,7 @@ where } fn ensure_drive_started(&mut self) -> Result { + self.dispatch_failure.ensure_healthy()?; if self.drive_started { return Ok(false); } diff --git a/crates/core/src/replay/disagg_tests.rs b/crates/core/src/replay/disagg_tests.rs index b4f5f7a27..bc4a3f3a1 100644 --- a/crates/core/src/replay/disagg_tests.rs +++ b/crates/core/src/replay/disagg_tests.rs @@ -2492,8 +2492,7 @@ fn test_source_first_handoff_waits_for_decode_scale_up() { assert_eq!(stats.request_snapshots[&uuid].phase, DisaggPhase::Done); } -#[test] -fn shared_g2_outlives_a_prefill_role_scaled_to_zero_and_back() { +fn shared_g2_disagg_config() -> TestDisaggConfig { use crate::engine::{NativeHostOffloadConfig, TimingModelConfig}; let args = |worker_type| MockEngineArgs { block_size: 4, @@ -2516,14 +2515,19 @@ fn shared_g2_outlives_a_prefill_role_scaled_to_zero_and_back() { .with_bandwidths(1.0, 1.0) .cluster_shared("tp1"), ); - let config = TestDisaggConfig { + TestDisaggConfig { prefill_args, decode_args: args(WorkerType::Decode), prefill_dp_size: 1, decode_dp_size: 1, num_prefill_workers: 1, num_decode_workers: 1, - }; + } +} + +#[test] +fn shared_g2_outlives_a_prefill_role_scaled_to_zero_and_back() { + let config = shared_g2_disagg_config(); let prompt = |id, arrival_ms| DirectRequest { tokens: (1..=9).collect(), max_output_tokens: 1, @@ -3849,6 +3853,345 @@ mod agentic_pd_qualification { } } +#[rstest::rstest] +#[case(f64::NAN)] +#[case(f64::INFINITY)] +#[case(f64::NEG_INFINITY)] +#[case(-1.0)] +#[case(0.0)] +fn disagg_rejects_invalid_policy_wakeup_before_advancing_clock( + #[case] wakeup_ms: f64, + #[values(true, false)] invalid_prefill: bool, +) { + use crate::replay::runtime_utils::wakeup_test_policy::WakeupPlacement; + let config = disagg_config().runtime_config(false).unwrap(); + let mut runtime = DisaggRuntimeImpl::<_, NoEngineEvents, NoReplayMetadata>::new_composed( + &config, + AdmissionQueue::new_requests(VecDeque::from([request(1, 64, 2, 0.0)]), ReplayMode::Trace), + false, + |_, _, _, _| { + Ok(( + WakeupPlacement::new(invalid_prefill.then_some(wakeup_ms), false), + WakeupPlacement::new((!invalid_prefill).then_some(wakeup_ms), false), + )) + }, + ) + .unwrap(); + // Decode selection happens after prefill. Bound the number of semantic + // steps so a broken same-time wakeup fails the test without hanging it. + let mut failure = None; + let mut previous_ms = 0.0; + for _ in 0..8 { + match runtime.step() { + Err(error) => { + failure = Some(error); + break; + } + Ok(ReplayStepOutcome::Settled { now_ms }) => { + assert!(now_ms.is_finite() && now_ms >= previous_ms); + previous_ms = now_ms; + } + Ok(_) => break, + } + } + let error = failure.expect("invalid or unconsumed wakeup must fail while settling"); + let role = if invalid_prefill { "prefill" } else { "decode" }; + assert!( + error + .to_string() + .contains(&format!("{role} placement policy wakeup")), + "{error:#}" + ); + assert!( + runtime.now_ms.is_finite() && runtime.now_ms >= previous_ms, + "invalid policy must not corrupt replay time" + ); +} + +#[rstest::rstest] +#[case::coincident(0.5)] +#[case::later(0.75)] +fn disagg_internal_deadline_settles_before_policy_wakeup(#[case] delay_ms: f64) { + use crate::replay::runtime_utils::wakeup_test_policy::WakeupPlacement; + let config = shared_g2_disagg_config().runtime_config(false).unwrap(); + let mut runtime = DisaggRuntimeImpl::<_, NoEngineEvents, NoReplayMetadata>::new_composed( + &config, + AdmissionQueue::new_requests( + VecDeque::from([request(1, 9, 1, 0.0), request(2, 9, 1, 2.5)]), + ReplayMode::Trace, + ), + false, + |_, _, _, _| { + Ok(( + WakeupPlacement::new(None, true), + WakeupPlacement::new(None, true), + )) + }, + ) + .unwrap() + .with_per_request_records(true); + // Seed two host blocks: the 1 ms prefill starts a 2 ms D2H transfer. + // Request 1 finishes at 2 ms, while the prefill worker is busy until 3 ms. + for expected_ms in [0.0, 1.0, 2.0] { + assert_eq!( + runtime.step().unwrap(), + ReplayStepOutcome::Settled { + now_ms: expected_ms + } + ); + } + runtime.prefill_placement = WakeupPlacement::new(Some(delay_ms), true); + let report = runtime.run().unwrap().0.finish(); + assert_eq!(report.request_counts.completed_requests, 2); + let released = report + .per_request + .iter() + .find(|request| request.uuid == Uuid::from_u128(2).to_string()) + .unwrap(); + assert_eq!(released.prefill_admit_ms, Some(2.5 + delay_ms)); + assert!( + report + .per_request + .iter() + .all(|request| { request.terminal_status == ReplayTerminalStatus::Completed }) + ); +} + +#[rstest::rstest] +#[case(0.0)] +#[case(7.0)] +fn disagg_consumes_same_now_and_future_policy_wakeups( + #[case] wakeup_ms: f64, + #[values(true, false)] queued_prefill: bool, +) { + use crate::replay::runtime_utils::wakeup_test_policy::WakeupPlacement; + let config = disagg_config().runtime_config(false).unwrap(); + let (collector, _) = DisaggRuntimeImpl::<_, NoEngineEvents, NoReplayMetadata>::new_composed( + &config, + AdmissionQueue::new_requests(VecDeque::from([request(1, 64, 2, 0.0)]), ReplayMode::Trace), + false, + |_, _, _, _| { + Ok(( + WakeupPlacement::new(queued_prefill.then_some(wakeup_ms), true), + WakeupPlacement::new((!queued_prefill).then_some(wakeup_ms), true), + )) + }, + ) + .unwrap() + .with_per_request_records(true) + .run() + .unwrap(); + let report = collector.finish(); + assert_eq!(report.request_counts.completed_requests, 1); + assert!(report.per_request[0].terminal_time_ms >= wakeup_ms); +} + +#[rstest::rstest] +fn disagg_canceling_tokenization_releases_policy_without_committing_prefill( + #[values(EngineType::Vllm, EngineType::Sglang)] backend: EngineType, +) { + use crate::replay::runtime_utils::dispatch_failure_test_policy::FailingDispatchPlacement; + let config = match backend { + EngineType::Vllm => disagg_config(), + EngineType::Sglang => sglang_disagg_config(), + _ => unreachable!(), + } + .runtime_config(false) + .unwrap(); + let mut runtime = DisaggRuntimeImpl::<_, NoEngineEvents, NoReplayMetadata>::new_composed( + &config, + AdmissionQueue::new_requests(VecDeque::from([request(1, 64, 2, 0.0)]), ReplayMode::Trace), + false, + |_, _, _, _| { + Ok(( + FailingDispatchPlacement::default(), + FailingDispatchPlacement::default(), + )) + }, + ) + .unwrap() + .with_tokenization(TokenizationLatency::new(1, 3.0, 0.0)); + let uuid = Uuid::from_u128(1); + runtime.step().unwrap(); + assert_eq!(runtime.now_ms, 0.0); + assert_eq!(runtime.prefill_engine.in_flight(), 0); + assert!(runtime.prefill_placement.calls.lock().unwrap().is_empty()); + assert_eq!(runtime.tokenization.as_ref().unwrap().len(), 1); + let handoff_id = runtime.state(uuid).unwrap().handoff_id; + runtime + .apply_handoff_fact(uuid, HandoffFact::Canceled { handoff_id }) + .unwrap(); + runtime.drain_current_timestamp().unwrap(); + assert_eq!(runtime.cluster_in_flight(), 0); + assert_eq!(runtime.tokenization.as_ref().unwrap().len(), 0); + runtime.now_ms = 3.0; + assert!(!runtime.release_tokenized().unwrap()); + assert_eq!( + *runtime.prefill_placement.calls.lock().unwrap(), + ["terminal"] + ); +} + +#[rstest::rstest] +fn disagg_dispatch_commit_failure_rolls_back_and_prevents_runtime_reuse( + #[values(EngineType::Vllm, EngineType::Sglang)] backend: EngineType, + #[values(false, true)] fail_prefill: bool, + #[values(false, true)] fail_abort: bool, + #[values(0.0, 3.0)] tokenization_ms: f64, +) { + use crate::replay::runtime_utils::dispatch_failure_test_policy::FailingDispatchPlacement; + let config = match backend { + EngineType::Vllm => disagg_config(), + EngineType::Sglang => sglang_disagg_config(), + _ => unreachable!(), + } + .runtime_config(false) + .unwrap(); + let mut runtime = DisaggRuntimeImpl::<_, NoEngineEvents, NoReplayMetadata>::new_composed( + &config, + AdmissionQueue::new_requests(VecDeque::from([request(1, 64, 2, 0.0)]), ReplayMode::Trace), + false, + |_, _, _, _| { + Ok(( + FailingDispatchPlacement { + fail_commit: fail_prefill, + fail_abort: fail_prefill && fail_abort, + ..Default::default() + }, + FailingDispatchPlacement { + fail_commit: !fail_prefill, + fail_abort: !fail_prefill && fail_abort, + ..Default::default() + }, + )) + }, + ) + .unwrap() + .with_tokenization(TokenizationLatency::new(1, tokenization_ms, 0.0)); + let error = loop { + match runtime.step() { + Err(error) => break error, + Ok(ReplayStepOutcome::Settled { .. }) => {} + Ok(other) => panic!("failure was not surfaced: {other:?}"), + } + }; + let (engine, policy) = if fail_prefill { + (&runtime.prefill_engine, &runtime.prefill_placement) + } else { + (&runtime.decode_engine, &runtime.decode_placement) + }; + assert_eq!( + engine.in_flight(), + 0, + "failed dispatch retained engine ownership" + ); + assert!(!engine.has_runnable_worker()); + assert_eq!(*policy.calls.lock().unwrap(), ["commit", "abort"]); + assert!( + error + .chain() + .any(|cause| cause.to_string() == "injected commit failure"), + "{error:#}" + ); + assert!( + error.to_string().contains("injected commit failure"), + "{error:#}" + ); + assert_eq!( + format!("{error:#}").contains("injected abort failure"), + fail_abort + ); + assert!(runtime.step().unwrap_err().to_string().contains("poisoned")); + assert!( + runtime + .run() + .err() + .unwrap() + .to_string() + .contains("poisoned") + ); +} + +#[rstest::rstest] +fn disagg_dispatch_abort_failure_preserves_native_rejection( + #[values(EngineType::Vllm, EngineType::Sglang)] backend: EngineType, + #[values(false, true)] fail_prefill: bool, +) { + use crate::replay::runtime_utils::dispatch_failure_test_policy::FailingDispatchPlacement; + let config = match backend { + EngineType::Vllm => disagg_config(), + EngineType::Sglang => sglang_disagg_config(), + _ => unreachable!(), + } + .runtime_config(false) + .unwrap(); + let mut runtime = DisaggRuntimeImpl::<_, NoEngineEvents, NoReplayMetadata>::new_composed( + &config, + AdmissionQueue::new_requests(VecDeque::from([request(1, 64, 2, 0.0)]), ReplayMode::Trace), + false, + |_, _, _, _| { + Ok(( + FailingDispatchPlacement { + fail_abort: fail_prefill, + ..Default::default() + }, + FailingDispatchPlacement { + fail_abort: !fail_prefill, + ..Default::default() + }, + )) + }, + ) + .unwrap(); + let engine = if fail_prefill { + &mut runtime.prefill_engine + } else { + &mut runtime.decode_engine + }; + if fail_prefill { + engine + .dispatch(0, request(1, 64, 2, 0.0), None, 0.0) + .unwrap(); + } else { + // Keep the duplicate destination parked: an ordinary decode submit + // would execute before the coordinator attempts its reservation. + engine + .apply_command( + 0, + Command::ReserveDestination { + handoff_id: HandoffId::new(Uuid::from_u128(999)), + request: direct_to_native(request(1, 64, 2, 0.0), None).unwrap(), + }, + 0.0, + ) + .unwrap(); + } + let error = loop { + match runtime.step() { + Err(error) => break error, + Ok(ReplayStepOutcome::Settled { .. }) => {} + Ok(other) => panic!("failure was not surfaced: {other:?}"), + } + }; + assert!( + error + .chain() + .any(|cause| cause.to_string().contains("already active")), + "{error:#}" + ); + assert!(error.to_string().contains("already active"), "{error:#}"); + assert!(format!("{error:#}").contains("injected abort failure")); + assert!(runtime.step().unwrap_err().to_string().contains("poisoned")); + assert!( + runtime + .run() + .err() + .unwrap() + .to_string() + .contains("poisoned") + ); +} + /// Admission metadata that carries replay hashes through unchanged, so each /// placement call sees exactly what the runtime handed it. struct HashFlag(Option); diff --git a/crates/core/src/replay/loadgen/driver.rs b/crates/core/src/replay/loadgen/driver.rs index 247514736..16fdb366e 100644 --- a/crates/core/src/replay/loadgen/driver.rs +++ b/crates/core/src/replay/loadgen/driver.rs @@ -2236,7 +2236,8 @@ impl WorkloadDriver { let mut identities = trace .nodes .iter() - .map(|node| { + .enumerate() + .map(|(node_index, node)| { let parent_id = node .dependencies .iter() @@ -2251,6 +2252,7 @@ impl WorkloadDriver { root_id: root_id_by_play.get(&node.play_id).cloned().flatten(), parent_id, cache_id: None, + lineage: trace.conversation_lineage[node_index].clone(), } }) .collect::>(); @@ -3629,6 +3631,7 @@ mod tests { digest: "fixture".into(), }, graph_digest: "fixture".into(), + conversation_lineage: super::super::lineage::conversation_lineage(&nodes), nodes, plays, hash_values: (0..hash_id_count).collect(), diff --git a/crates/core/src/replay/loadgen/lineage.rs b/crates/core/src/replay/loadgen/lineage.rs new file mode 100644 index 000000000..4d4b0dcaa --- /dev/null +++ b/crates/core/src/replay/loadgen/lineage.rs @@ -0,0 +1,63 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Derive conversation ancestry once, before snapshot filtering. General request +//! DAGs remain valid even when they do not describe a conversation tree; their +//! ambiguous lineage is unavailable rather than guessed from arrival order. + +use std::collections::{BTreeMap, BTreeSet}; + +use super::{AgenticDependencyRelation, AgenticNode}; +use crate::replay::{AGENTIC_CONVERSATION_LINEAGE_SCHEMA_V1, AgenticConversationLineage}; + +pub(super) fn conversation_lineage( + nodes: &[AgenticNode], +) -> Vec> { + let by_request = nodes + .iter() + .map(|node| (node.request_id.as_str(), node)) + .collect::>(); + // Session names can be reused in different authored plays. + let mut parents = BTreeMap::<(&str, &str), BTreeSet<&str>>::new(); + for node in nodes { + let candidates = parents + .entry((&node.play_id, &node.session_id)) + .or_default(); + for edge in &node.dependencies { + if edge.relation == AgenticDependencyRelation::Spawn { + candidates.insert(by_request[edge.request_id.as_str()].session_id.as_str()); + } + } + } + let mut resolved = BTreeMap::new(); + for &(play, conversation) in parents.keys() { + let mut ancestry = BTreeSet::new(); + let mut cursor = conversation; + let mut root = None; + while ancestry.insert(cursor) { + let candidates = &parents[&(play, cursor)]; + match candidates.len() { + 0 => { + root = Some(cursor); + break; + } + 1 => cursor = candidates.first().copied().expect("one parent"), + _ => break, + } + } + resolved.insert( + (play, conversation), + root.map(|root| AgenticConversationLineage { + schema: AGENTIC_CONVERSATION_LINEAGE_SCHEMA_V1.into(), + root_conversation_id: root.into(), + parent_conversation_id: parents[&(play, conversation)] + .first() + .map(|parent| (*parent).into()), + }), + ); + } + nodes + .iter() + .map(|node| resolved[&(node.play_id.as_str(), node.session_id.as_str())].clone()) + .collect() +} diff --git a/crates/core/src/replay/loadgen/mod.rs b/crates/core/src/replay/loadgen/mod.rs index beb51c718..718316bfd 100644 --- a/crates/core/src/replay/loadgen/mod.rs +++ b/crates/core/src/replay/loadgen/mod.rs @@ -5,6 +5,7 @@ mod arrival; mod driver; mod dynamo; mod generated; +mod lineage; mod phase; mod profile; mod snapshot; diff --git a/crates/core/src/replay/loadgen/snapshot.rs b/crates/core/src/replay/loadgen/snapshot.rs index d5a31085c..03de5fd33 100644 --- a/crates/core/src/replay/loadgen/snapshot.rs +++ b/crates/core/src/replay/loadgen/snapshot.rs @@ -597,6 +597,17 @@ impl AgenticPlaySnapshot { .then(|| self.instance_request_id(&graph.nodes[play.root_nodes[0]].request_id)), parent_id: parent.map(|e| self.instance_request_id(&e.request_id)), cache_id: Some(self.evidence.cache_id.clone()), + lineage: graph.conversation_lineage[index].as_ref().map(|lineage| { + crate::replay::AgenticConversationLineage { + schema: lineage.schema.clone(), + root_conversation_id: self + .instance_conversation_id(&lineage.root_conversation_id), + parent_conversation_id: lineage + .parent_conversation_id + .as_ref() + .map(|parent| self.instance_conversation_id(parent)), + } + }), } } pub(super) fn token_ids(&self, index: usize) -> Vec { diff --git a/crates/core/src/replay/loadgen/snapshot_tests.rs b/crates/core/src/replay/loadgen/snapshot_tests.rs index 3808209b5..31cc3b17c 100644 --- a/crates/core/src/replay/loadgen/snapshot_tests.rs +++ b/crates/core/src/replay/loadgen/snapshot_tests.rs @@ -72,6 +72,97 @@ fn boundary_graph() -> ValidatedAgenticGraph { graph(vec![root, child, background, boundary, child_next, future]) } +#[test] +fn conversation_lineage_survives_history_and_later_turns_with_fresh_play_namespaces() { + use AgenticDependencyRelation::{Sequence, Spawn}; + use AgenticDependencyTrigger::{Completion, Dispatch}; + let root = row("root", "main", 0.0, Some(10.0)); + let mut child = row("child", "worker", 10.0, Some(10.0)); + child.dependencies = vec![edge("root", Spawn, Dispatch, 10.0)]; + let mut later = row("later", "worker", 30.0, Some(10.0)); + later.dependencies = vec![edge("child", Sequence, Completion, 10.0)]; + let mut nested = row("nested", "grandchild", 40.0, Some(10.0)); + nested.dependencies = vec![edge("later", Spawn, Dispatch, 10.0)]; + let graph = graph(vec![root, child, later, nested]); + let snapshot = prepare(&graph, 25.0); + assert!(request(&snapshot, "root").historical); + assert!(request(&snapshot, "child").historical); + let root = snapshot.identity("root").unwrap(); + let child = snapshot.identity("child").unwrap(); + let later = snapshot.identity("later").unwrap(); + let nested = snapshot.identity("nested").unwrap(); + assert_eq!(later.parent_id, None); // Legacy request-edge semantics stay intact. + assert_eq!(later.lineage, child.lineage); + assert_eq!(child.parent_id.as_deref(), Some(root.request_id.as_str())); + let lineage = nested.lineage.as_ref().unwrap(); + assert_eq!( + lineage.parent_conversation_id.as_deref(), + Some(child.conversation_id.as_str()) + ); + assert_eq!(lineage.root_conversation_id, root.conversation_id); + assert_ne!(lineage.root_conversation_id, root.request_id); + let encoded = serde_json::to_value(&nested).unwrap(); + assert_eq!( + serde_json::from_value::(encoded).unwrap(), + nested + ); + let next = snapshot.context().prepare_play_from_start(0, 1).unwrap(); + assert_ne!(next.identity("nested").unwrap().lineage, nested.lineage); + assert_eq!( + next.identity("nested") + .unwrap() + .lineage + .unwrap() + .parent_conversation_id, + Some(next.identity("child").unwrap().conversation_id) + ); + let turns = drain(WorkloadDriver::new_agentic_trace(graph, 64).unwrap()); + let static_lineage = turns["nested"] + .request + .replay_context + .as_ref() + .unwrap() + .agentic + .as_ref() + .unwrap() + .lineage + .as_ref() + .unwrap(); + assert_eq!( + static_lineage.parent_conversation_id.as_deref(), + Some("worker") + ); + assert_eq!(static_lineage.root_conversation_id, "main"); +} + +#[test] +fn ambiguous_conversation_parent_does_not_invent_a_routing_tree() { + use AgenticDependencyRelation::Spawn; + use AgenticDependencyTrigger::Dispatch; + let a = row("a", "parent-a", 0.0, Some(10.0)); + let b = row("b", "parent-b", 0.0, Some(10.0)); + let mut child = row("child", "child", 20.0, Some(10.0)); + child.dependencies = vec![ + edge("a", Spawn, Dispatch, 20.0), + edge("b", Spawn, Dispatch, 20.0), + ]; + let mut nested = row("nested", "nested", 30.0, Some(10.0)); + nested.dependencies = vec![edge("child", Spawn, Dispatch, 10.0)]; + let graph = graph(vec![a, b, child, nested]); + let snapshot = prepare(&graph, 0.0); + assert!(snapshot.identity("child").unwrap().lineage.is_none()); + assert!(snapshot.identity("nested").unwrap().lineage.is_none()); + assert_eq!( + snapshot + .identity("a") + .unwrap() + .lineage + .unwrap() + .parent_conversation_id, + None + ); +} + fn request<'a>(play: &'a AgenticPlaySnapshot, source: &str) -> &'a AgenticSnapshotRequest { play.evidence() .requests diff --git a/crates/core/src/replay/loadgen/trace.rs b/crates/core/src/replay/loadgen/trace.rs index 3b0745a48..74f8b28dd 100644 --- a/crates/core/src/replay/loadgen/trace.rs +++ b/crates/core/src/replay/loadgen/trace.rs @@ -1718,6 +1718,7 @@ impl AgenticTraceBuilder { let hash_values = renumber_hash_ids(&mut self.nodes, self.hash_ids); let graph_digest = canonical_agentic_graph_digest(self.header.block_size, &mut self.nodes, &hash_values)?; + let conversation_lineage = super::lineage::conversation_lineage(&self.nodes); Ok(AgenticTrace { block_size: self.header.block_size, @@ -1725,6 +1726,7 @@ impl AgenticTraceBuilder { graph_digest, nodes: self.nodes, plays, + conversation_lineage, hash_values, }) } diff --git a/crates/core/src/replay/loadgen/types.rs b/crates/core/src/replay/loadgen/types.rs index 2c8d08448..e9718a598 100644 --- a/crates/core/src/replay/loadgen/types.rs +++ b/crates/core/src/replay/loadgen/types.rs @@ -48,6 +48,7 @@ pub struct ValidatedAgenticGraph { pub(super) graph_digest: String, pub(super) nodes: Vec, pub(super) plays: Vec, + pub(super) conversation_lineage: Vec>, /// Source hash for each graph-local hash ID in `AgenticNode::hash_ids`. pub(super) hash_values: Arc<[u64]>, } diff --git a/crates/core/src/replay/mod.rs b/crates/core/src/replay/mod.rs index 6e18fbdf6..d3542e001 100644 --- a/crates/core/src/replay/mod.rs +++ b/crates/core/src/replay/mod.rs @@ -1,13 +1,13 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Dynamo-neutral, deterministic offline replay. +//! Engine-neutral, deterministic offline replay. //! //! [`Replayer`] owns virtual time, the event queue, logical-worker lifecycle, //! placement/scaling policy composition, and report collection. Engine //! scheduling and timing are supplied by [`crate::engine`] through its -//! generalized-engine contract. Dynamo integrations depend on `aisimulate-core` and -//! provide placement/scaling policies; replay never depends back on Dynamo. +//! generalized-engine contract. Downstream integrations supply routing adapters +//! through neutral contracts; AISimulate owns the execution loop. mod agg; mod artifact; @@ -153,7 +153,8 @@ pub use handoff::{ }; #[doc(hidden)] pub use protocol::{ - AgenticRuntimeIdentity, DirectRequest, ReplayPromptTokenSource, ReplayRequestContext, + AGENTIC_CONVERSATION_LINEAGE_SCHEMA_V1, AgenticConversationLineage, AgenticRuntimeIdentity, + DirectRequest, ReplayPromptTokenSource, ReplayRequestContext, }; pub use protocol::{AuthoredKvHints, ForwardPassSnapshot}; #[doc(hidden)] diff --git a/crates/core/src/replay/protocol.rs b/crates/core/src/replay/protocol.rs index de29a73ee..a27af38c0 100644 --- a/crates/core/src/replay/protocol.rs +++ b/crates/core/src/replay/protocol.rs @@ -10,6 +10,19 @@ use uuid::Uuid; use crate::engine::ImageSpec; +pub const AGENTIC_CONVERSATION_LINEAGE_SCHEMA_V1: &str = + "aisimulate.agentic.conversation-lineage.v1"; + +/// Conversation ancestry derived from the full source graph, including history +/// before a replay snapshot. These are conversation IDs, never request IDs. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AgenticConversationLineage { + pub schema: String, + pub root_conversation_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parent_conversation_id: Option, +} + /// Stable Agentic correlation data carried across replay-owned boundaries. /// /// The static graph already owns request, play, and conversation identity. @@ -28,6 +41,11 @@ pub struct AgenticRuntimeIdentity { pub parent_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_id: Option, + /// Absent for legacy inputs or graphs with ambiguous conversation ancestry. + /// Consumers requiring parent affinity must reject missing/unknown schemas. + /// Existing `root_id` and `parent_id` continue to identify requests. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub lineage: Option, } /// Where the prompt token identities carried by a lowered request came from. @@ -194,7 +212,18 @@ impl DirectRequest { #[cfg(test)] mod tests { - use super::DirectRequest; + use super::{AgenticRuntimeIdentity, DirectRequest}; + + #[test] + fn legacy_agentic_identity_preserves_request_parent_semantics() { + let old = serde_json::json!({ + "request_id": "request-2", "play_id": "play", "conversation_id": "child", + "root_id": "request-0", "parent_id": "request-1" + }); + let identity: AgenticRuntimeIdentity = serde_json::from_value(old.clone()).unwrap(); + assert!(identity.lineage.is_none()); + assert_eq!(serde_json::to_value(identity).unwrap(), old); + } #[test] fn output_plan_is_authoritative_when_limited() { diff --git a/crates/core/src/replay/replayer.rs b/crates/core/src/replay/replayer.rs index 160dfb8f0..44a2a2d67 100644 --- a/crates/core/src/replay/replayer.rs +++ b/crates/core/src/replay/replayer.rs @@ -122,6 +122,26 @@ where self.0.cancel_pending(request_id) } + fn dispatch_committed(&mut self, request_id: Uuid, now_ms: f64) -> AnyResult<()> { + self.0 + .dispatch_committed(request_id, now_ms) + .map_err(placement_boundary) + } + + fn dispatch_aborted(&mut self, request_id: Uuid, now_ms: f64) -> AnyResult<()> { + self.0 + .dispatch_aborted(request_id, now_ms) + .map_err(placement_boundary) + } + + fn advance_clock(&mut self, now_ms: f64) -> AnyResult> { + self.0.advance_clock(now_ms).map_err(placement_boundary) + } + + fn next_wakeup_ms(&self) -> Option { + self.0.next_wakeup_ms() + } + fn request_terminal( &mut self, request_id: Uuid, diff --git a/crates/core/src/replay/runtime_utils.rs b/crates/core/src/replay/runtime_utils.rs index a9ee28c1c..36e9aaf24 100644 --- a/crates/core/src/replay/runtime_utils.rs +++ b/crates/core/src/replay/runtime_utils.rs @@ -14,6 +14,43 @@ use crate::engine::HandoffId; #[cfg(test)] use crate::replay::protocol::DirectRequest; +/// A failed dispatch may leave arbitrary policy and reporting side effects. +/// Cleanup is best-effort; no subsequent execution or report is trustworthy. +#[derive(Default)] +pub(super) struct DispatchFailure { + reason: Option, +} + +impl DispatchFailure { + pub(super) fn ensure_healthy(&self) -> anyhow::Result<()> { + if let Some(reason) = &self.reason { + anyhow::bail!( + "replay is poisoned after a failed dispatch: {reason}; construct a new runtime" + ); + } + Ok(()) + } + + pub(super) fn record( + &mut self, + mut error: anyhow::Error, + rollback: anyhow::Result<()>, + abort: anyhow::Result<()>, + ) -> anyhow::Error { + if let Err(cleanup) = rollback { + error = error.context(format!("engine dispatch rollback failed: {cleanup:#}")); + } + if let Err(cleanup) = abort { + error = error.context(format!("placement dispatch_aborted failed: {cleanup:#}")); + } + let reason = format!("{error:#}"); + self.reason.get_or_insert_with(|| reason.clone()); + error.context(format!( + "replay is poisoned after a failed dispatch: {reason}; construct a new runtime" + )) + } +} + /// Result of advancing a replay runtime to its next settled semantic boundary. #[derive(Debug, Clone, Copy, PartialEq)] pub(crate) enum ReplayStepOutcome { @@ -56,6 +93,28 @@ pub(super) fn next_timestamp( } } +/// A same-time wakeup is allowed while draining immediate policy work, but +/// must be consumed before the timestamp is settled. Check each policy before +/// merging deadlines: `f64::min` can otherwise hide a NaN behind a valid event. +pub(super) fn validate_policy_wakeup( + wakeup_ms: Option, + now_ms: f64, + role: &str, + settled: bool, +) -> anyhow::Result<()> { + if let Some(wakeup_ms) = wakeup_ms { + anyhow::ensure!( + wakeup_ms.is_finite() && wakeup_ms >= now_ms, + "{role} placement policy wakeup must be finite and not precede replay time {now_ms}ms; got {wakeup_ms}ms" + ); + anyhow::ensure!( + !settled || wakeup_ms > now_ms, + "{role} placement policy wakeup at {wakeup_ms}ms made no progress at replay time {now_ms}ms; a settled wakeup must be strictly in the future" + ); + } + Ok(()) +} + /// Return the earliest scheduled event that can advance replay semantics. /// /// At most one telemetry heartbeat is armed at a time. Temporarily removing @@ -477,3 +536,223 @@ mod tests { ); } } + +// Minimal queued placement used to exercise scheduler wakeups through both engines. +#[cfg(test)] +pub(super) mod wakeup_test_policy { + use super::super::core::{ + Placement, PlacementDecision, PlacementEffects, PlacementPolicy, WorkerTopology, + }; + use super::super::loadgen::ReplayRequestPayload; + use uuid::Uuid; + + pub(crate) struct WakeupPlacement { + delay_ms: Option, + wakeup_ms: Option, + consume_wakeup: bool, + pending: Vec, + } + + impl WakeupPlacement { + pub(crate) fn new(wakeup_ms: Option, consume_wakeup: bool) -> Self { + Self { + delay_ms: wakeup_ms, + wakeup_ms: None, + consume_wakeup, + pending: Vec::new(), + } + } + fn placement(request_id: Uuid) -> Placement { + Placement { + request_id, + scheduler_id: 0, + reported_overlap_tokens: 0, + cache_sample: None, + kv_hints: None, + placement_replica_id: None, + } + } + } + + impl PlacementPolicy for WakeupPlacement { + type Metadata = (); + type Observation = (); + fn place( + &mut self, + request: &ReplayRequestPayload, + _: (), + _: Option, + now_ms: f64, + ) -> anyhow::Result { + let id = request.metadata().uuid.unwrap(); + let decision = if let Some(delay_ms) = self.delay_ms { + self.wakeup_ms = Some(now_ms + delay_ms); + self.pending.push(id); + PlacementDecision::Queued + } else { + PlacementDecision::Immediate(Self::placement(id)) + }; + Ok(PlacementEffects { + decision, + released: Vec::new(), + }) + } + fn next_wakeup_ms(&self) -> Option { + if self.pending.is_empty() { + None + } else { + self.wakeup_ms + } + } + fn advance_clock(&mut self, now_ms: f64) -> anyhow::Result> { + if self.consume_wakeup + && self.wakeup_ms.is_some_and(|wake| wake <= now_ms) + && !self.pending.is_empty() + { + self.wakeup_ms = None; + self.delay_ms = None; + return Ok(self.pending.drain(..).map(Self::placement).collect()); + } + Ok(Vec::new()) + } + fn observe(&mut self, _: (), _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn cancel_pending(&mut self, id: Uuid) -> bool { + let before = self.pending.len(); + self.pending.retain(|pending| *pending != id); + self.pending.len() != before + } + fn request_terminal(&mut self, _: Uuid, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn prefill_completed(&mut self, _: Uuid, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn pending_count(&self) -> usize { + self.pending.len() + } + fn worker_ready(&mut self, _: WorkerTopology, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn worker_draining(&mut self, _: WorkerTopology, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn worker_removed(&mut self, _: WorkerTopology, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn topology_settled(&mut self, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + } +} + +#[cfg(test)] +pub(super) mod dispatch_failure_test_policy { + use super::super::core::{ + Placement, PlacementDecision, PlacementEffects, PlacementPolicy, WorkerTopology, + }; + use crate::replay::loadgen::ReplayRequestPayload; + use std::sync::{Arc, Mutex}; + use uuid::Uuid; + + #[derive(Default)] + pub struct FailingDispatchPlacement { + pub fail_commit: bool, + pub fail_abort: bool, + pub calls: Arc>>, + } + + impl PlacementPolicy for FailingDispatchPlacement { + type Metadata = (); + type Observation = (); + + fn place( + &mut self, + request: &ReplayRequestPayload, + _: (), + _: Option, + _: f64, + ) -> anyhow::Result { + Ok(PlacementEffects { + decision: PlacementDecision::Immediate(Placement { + request_id: request.metadata().uuid.unwrap(), + scheduler_id: 0, + reported_overlap_tokens: 0, + cache_sample: None, + kv_hints: None, + placement_replica_id: None, + }), + released: Vec::new(), + }) + } + fn dispatch_committed(&mut self, _: Uuid, _: f64) -> anyhow::Result<()> { + self.calls.lock().unwrap().push("commit"); + anyhow::ensure!(!self.fail_commit, "injected commit failure"); + Ok(()) + } + fn dispatch_aborted(&mut self, _: Uuid, _: f64) -> anyhow::Result<()> { + self.calls.lock().unwrap().push("abort"); + anyhow::ensure!(!self.fail_abort, "injected abort failure"); + Ok(()) + } + fn observe(&mut self, _: (), _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn cancel_pending(&mut self, _: Uuid) -> bool { + false + } + fn request_terminal(&mut self, _: Uuid, _: f64) -> anyhow::Result> { + self.calls.lock().unwrap().push("terminal"); + Ok(Vec::new()) + } + fn prefill_completed(&mut self, _: Uuid, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn pending_count(&self) -> usize { + 0 + } + fn worker_ready(&mut self, _: WorkerTopology, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn worker_draining(&mut self, _: WorkerTopology, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn worker_removed(&mut self, _: WorkerTopology, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + fn topology_settled(&mut self, _: f64) -> anyhow::Result> { + Ok(Vec::new()) + } + } +} + +#[cfg(test)] +mod dispatch_failure_tests { + use super::DispatchFailure; + + #[test] + fn dispatch_failure_keeps_primary_and_both_cleanup_diagnostics() { + let mut failure = DispatchFailure::default(); + let error = failure.record( + anyhow::anyhow!("original dispatch failure"), + Err(anyhow::anyhow!("rollback failure")), + Err(anyhow::anyhow!("abort failure")), + ); + assert_eq!(error.root_cause().to_string(), "original dispatch failure"); + for message in [ + "original dispatch failure", + "rollback failure", + "abort failure", + ] { + assert!(error.to_string().contains(message)); + assert!( + failure + .ensure_healthy() + .unwrap_err() + .to_string() + .contains(message) + ); + } + } +} diff --git a/crates/core/tests/agentx_qualification.rs b/crates/core/tests/agentx_qualification.rs index eadaf5738..6a747e5b0 100644 --- a/crates/core/tests/agentx_qualification.rs +++ b/crates/core/tests/agentx_qualification.rs @@ -614,6 +614,53 @@ fn snapshot_pin_report( run_g2(spec, driver) } +fn assert_snapshot_lineage(name: &str, report: &ReplayReport) { + let identities = report + .agentic_snapshots + .as_ref() + .unwrap() + .iter() + .flat_map(|snapshot| snapshot.requests.iter().map(|request| &request.identity)) + .chain( + report + .agentic_phases + .iter() + .flat_map(|phase| phase.requests.iter().map(|request| &request.identity)), + ) + .chain( + report + .per_request + .iter() + .map(|request| request.agentic.as_ref().expect("agentic request identity")), + ); + for identity in identities { + let prefix = format!("{}:conversation:", identity.play_id); + let conversation = identity.conversation_id.strip_prefix(&prefix).unwrap(); + // Derive the expected ancestry from the authored fixtures, independently + // of the graph/snapshot lineage implementation, including historical roots. + let (root, child) = match name { + "g2" => { + assert!(matches!(conversation, "parent" | "child")); + ("parent".to_string(), conversation == "child") + } + "two-plays" => { + if let Some(play) = conversation.strip_suffix(":session:subagent:1:worker") { + (format!("{play}:session:root"), true) + } else { + assert!(conversation.ends_with(":session:root")); + (conversation.to_string(), false) + } + } + _ => unreachable!("unknown fixture"), + }; + let root = format!("{prefix}{root}"); + let lineage = identity.lineage.as_ref().expect("conversation lineage"); + assert_eq!(lineage.schema, "aisimulate.agentic.conversation-lineage.v1"); + assert_eq!(lineage.root_conversation_id, root); + assert_eq!(lineage.parent_conversation_id, child.then_some(root)); + } +} + /// Snapshot, warmup and recycling-profile replays must stay byte-identical /// across changes to how snapshot plays are built and retained. #[test] @@ -630,6 +677,7 @@ fn snapshot_warmup_and_profile_reports_are_pinned() { for mode in ["snapshot", "warmup", "profile", "profile-cold"] { for disaggregated in [false, true] { let report = snapshot_pin_report(graph, mode, disaggregated); + assert_snapshot_lineage(name, &report); if mode.starts_with("profile") { let profile = report.agentic_profile.as_ref().unwrap(); assert!(profile.plays_started >= 6, "{name}: {profile:?}"); @@ -654,26 +702,30 @@ fn snapshot_warmup_and_profile_reports_are_pinned() { } } } + // These pins include conversation lineage. Comparing all 16 scenarios with + // main 88c7d426 and pre-merge PR b8692dc8 established that lineage was the + // only added JSON field; lifecycle bytes and every previous field were + // unchanged, and the merged output matched the pre-merge PR byte-for-byte. // To re-pin, run this test and copy the `left:` list from the failure. A // mismatch the change did not intend must first be diffed semantically // against the previous report, per-request and lifecycle output. let expected = [ - "g2 snapshot agg: 5 5e02e9f17a77658f", - "g2 snapshot disagg: 5 a6573f2379f39c15", - "g2 warmup agg: 5 21a10bcd0843bc96", - "g2 warmup disagg: 5 a0834d087069f92e", - "g2 profile agg: 40 a69d736cdea338aa", - "g2 profile disagg: 40 fa167c22cb9d4b34", - "g2 profile-cold agg: 40 f63f6e7cf21a4de3", - "g2 profile-cold disagg: 40 29b72dc9224ae354", - "two-plays snapshot agg: 3 19563e804401e8a2", - "two-plays snapshot disagg: 3 88c3919df99f7ce1", - "two-plays warmup agg: 3 2bb148e6548bb88a", - "two-plays warmup disagg: 3 f784df450822621a", - "two-plays profile agg: 821 35351cf48f3958c3", - "two-plays profile disagg: 667 6b252248d748175a", - "two-plays profile-cold agg: 820 20ee4611dbaa1e83", - "two-plays profile-cold disagg: 601 e2d4f4e93360d22e", + "g2 snapshot agg: 5 20ab2cfe1f97cb15", + "g2 snapshot disagg: 5 8c6c217a93820def", + "g2 warmup agg: 5 e5a8dda39159fa01", + "g2 warmup disagg: 5 2dea629178c141df", + "g2 profile agg: 40 68e2b208e87b3cad", + "g2 profile disagg: 40 131164177078405b", + "g2 profile-cold agg: 40 dd3790d0738a7d50", + "g2 profile-cold disagg: 40 879186a642482fab", + "two-plays snapshot agg: 3 c1d26ce77a455ecf", + "two-plays snapshot disagg: 3 1dc21fce69c7690a", + "two-plays warmup agg: 3 3d0cc958afd8216d", + "two-plays warmup disagg: 3 263766310612fe41", + "two-plays profile agg: 821 ef8276602d9df84b", + "two-plays profile disagg: 667 ae91137630178cc3", + "two-plays profile-cold agg: 820 bd322ba76617e24a", + "two-plays profile-cold disagg: 601 f8d9ba58ae29ec49", ]; assert_eq!(actual, expected); } diff --git a/docs/adapters/README.md b/docs/adapters/README.md index 77adb7369..7a61b439b 100644 --- a/docs/adapters/README.md +++ b/docs/adapters/README.md @@ -16,6 +16,7 @@ the existing Dynamo integration is documented [with Replay](../replay/dynamo.md) | [Simulation configuration](configuration-abi.md) | Compile prediction, search plan and candidate hooks | `aisimulate.config_adapters`; config API 3 | | [SDK sweep provider](configuration-abi.md#sweepconfigprovider-v1) | SDK search-space preparation and materialization | `aisimulate.sweep_config_providers`; provider API 1 | | [Recommendation output](output-abi.md) | Optional live callbacks and additional final artifacts | `aisimulate.output_adapters`; output API 1 | +| [External policy](native-composition.md#external-policy-transport) | Independently installed placement policy called by the AISimulate engine | Versioned JSON callbacks; AISimulate owns execution | | [Native composition](native-composition.md) | Placement and optional scaling in Rust Replayer | Rust traits in `aisimulate_core::replay`; crate-version compatibility | Here ABI means the versioned cross-package data/interface contract. These diff --git a/docs/adapters/native-composition.md b/docs/adapters/native-composition.md index 621abe1de..1dbe0364b 100644 --- a/docs/adapters/native-composition.md +++ b/docs/adapters/native-composition.md @@ -11,6 +11,35 @@ The built-in `RoundRobinComposition` has fixed workers. Dynamo constructs its own Router/Planner policies against the same neutral contracts. These are Rust crate interfaces, not a stable dynamically loaded binary ABI. +## External policy transport + +The engine stack can call an independently built policy through +`aisimulate._runtime.run_replay_with_policy_json(payload_json, factory)` while +retaining AISimulate's executor. The factory receives JSON describing one +role, worker topology, resolved engine configuration, placement descriptor, +selector seed and capture settings. It returns an object implementing +`call_json(command_json) -> result_json`. + +Commands carry schema-versioned neutral data for placement, KV observations, +dispatch commit/abort, terminal requests, virtual-time advancement, cancellation, +and worker lifecycle events. A result supplies the operation result, pending +count and next wakeup. A final snapshot adds provider-owned evidence without +copying a growing decision log on every callback. Policy errors fail simulation; +they cannot silently select another policy. Evidence is returned by role in +`external_policy` and may be presented by the Python adapter as `routing_policy`. + +The built-in `engine.router` adapter implements this transport in AISimulate and +calls Dynamo's native `OfflineRouter`, `SessionAffinity`, and existing `RadixTree` +Python interfaces. Dynamo does not interpret AISimulate requests, conversation +lineage, lifecycle messages or reports; the adapter translates those host fields. The `ai-dynamo` whole-run plugin and its AISimulate version +update are optional for this path. The AISimulate CLI and its acceptance suite +exercise the complete AgentX simulation. See the +[AgentX quickstart](../replay/agentic/quickstart.md). + +This Python/JSON boundary does not link Dynamo into the AISimulate wheel. The +Rust composition interface below remains available to consumers that choose to +compile against AISimulate's core. + ## Composition construction The trait is defined in @@ -43,6 +72,9 @@ into stable user APIs. Public input/report contracts are in the - `observe(observation, now_ms)`, `request_terminal(request_id, now_ms)` and `prefill_completed(request_id, now_ms)`, returning newly releasable placements. - `cancel_pending(request_id) -> bool` and `pending_count()` for queued ownership. +- Optional `dispatch_committed`, `dispatch_aborted`, `advance_clock` and + `next_wakeup_ms` hooks for tentative bindings and policy wakeups; see the + [conversation lineage contract](../replay/agentic/conversation-lineage.md). - `worker_ready`, `worker_draining`, `worker_removed`, and `topology_settled` lifecycle callbacks, also returning releasable placements. diff --git a/docs/replay/agentic/conversation-lineage.md b/docs/replay/agentic/conversation-lineage.md new file mode 100644 index 000000000..13852b36b --- /dev/null +++ b/docs/replay/agentic/conversation-lineage.md @@ -0,0 +1,39 @@ +# Agentic conversation lineage + +The 0.13.0 replay identity adds an optional `lineage` object with schema +`aisimulate.agentic.conversation-lineage.v1`, `root_conversation_id`, and optional +`parent_conversation_id`. The existing `root_id` and `parent_id` remain request +IDs. A later conversation turn can have no request-level spawn dependency while +still retaining its conversation parent. + +Lineage is compiled from all spawn edges in the validated source graph, before +snapshot history is removed. All turns in a conversation share its lineage. +Snapshot instances namespace every lineage reference with the same play instance +as the request, including references to historical conversations. Recycling a +lane therefore does not reuse the prior play's conversation or affinity keys. + +A conversation without a spawn parent is a root and names itself as +`root_conversation_id`. Multiple distinct spawn parents, self-parenting, or a +cycle between conversations leave lineage absent for that conversation and its +descendants. Such request DAGs remain valid replay inputs; an adapter requiring +conversation-tree routing must reject unavailable or unsupported lineage instead +of choosing a parent by request order. + +Old JSON without `lineage` still deserializes and round-trips without the field. +New Rust struct literals must set `lineage` explicitly. Consumers must check the +schema before relying on ancestry. No router policy is enabled by carrying these +identities: a versioned consumer must invoke its native policy and independently +qualify actual placement, cache behavior, and supported topologies. + +Placement policies receive optional default no-op `dispatch_committed` and +`dispatch_aborted` callbacks. A successful dispatch means the engine accepted +ownership; for P/D decode this is destination reservation, before transfer or +token generation. A policy can hold a tentative binding during selection and +commit it only after acceptance. `advance_clock(now_ms)` runs during semantic +settlement and can release policy waiters after a commit or abort. +`next_wakeup_ms()` must identify concrete future policy work. TTL housekeeping +that cannot release a request should run lazily and must not extend the run. + +This contract uses AISimulate's existing graph edges and snapshot namespaces. It +does not alter graph digests, sampling, request dependencies, cache allocation, +warmup behavior, or the methodology qualification status of replay. diff --git a/docs/replay/agentic/quickstart.md b/docs/replay/agentic/quickstart.md index 843b9ffd9..79fe6c0e3 100644 --- a/docs/replay/agentic/quickstart.md +++ b/docs/replay/agentic/quickstart.md @@ -12,6 +12,11 @@ remaining requests and recycles completed plays to keep two client lanes occupie for a 3,600-second simulated admission window. The simulator runs offline on your CPU; you do not need to allocate those GPUs or download model weights. +AISimulate runs the simulation and calls the separately installed Dynamo native +policy for worker selection, KV indexing, and conversation affinity. No Dynamo +serving process is required, and its routing implementation is not packaged in +the AISimulate wheel. + Run the commands below from the repository root. This path is experimental and qualified for functional behavior, not hardware performance accuracy. @@ -33,6 +38,24 @@ This installs the Python environment and builds the native simulator. The commands below use that environment explicitly. Internet access is needed for initial dependency installation, the trace download, and model metadata. +Build the `ai-dynamo-runtime` native bindings from +[Dynamo #16854](https://github.com/ai-dynamo/dynamo/pull/16854), revision +[`c2e854fe829445dfb22128f082a053f9f37a425c`](https://github.com/ai-dynamo/dynamo/tree/c2e854fe829445dfb22128f082a053f9f37a425c/lib/bindings/python#readme), +following the linked build instructions and using the AISimulate environment at +`python/aisimulate/.venv`. Dynamo must provide this API first; AISimulate #306 +then consumes and accepts it. The `ai-dynamo` runner package and its version +update in Dynamo #15240 are optional. Check the installed interfaces: + +```bash +python/aisimulate/.venv/bin/python - <<'PYTHON' +from aisimulate import _runtime +from dynamo._core import OfflineRouter, SessionAffinity +assert callable(_runtime.run_replay_with_policy_json) +assert callable(OfflineRouter.submit) +assert callable(SessionAffinity.try_acquire) +PYTHON +``` + ## 2. Download a reproducible Weka workload The input is SemiAnalysis's @@ -93,6 +116,11 @@ traffic: agentic_warmup: true agentic_profile: duration_seconds: 3600 +router: + policy: kv_router + affinity: + mode: sibling_group # or session + ttl_seconds: 3600 engine: mode: disaggregated model: Qwen/Qwen3-4B-Instruct-2507 @@ -125,12 +153,17 @@ execution: memory_limit_gb: 4 ``` +The default `engine` stack uses native hard affinity: `sibling_group` binds +children of the same parent conversation together, while `session` binds each +conversation separately. Omit `router` for round-robin replay without Dynamo. +See [router fields and supported configurations](../dynamo.md#router-policy-on-the-aisimulate-engine-stack) +for worker/rank binding, virtual-time TTL, and capability restrictions. + The trace's embedded hash blocks contain 64 tokens. The explicit `traffic.source.block_size: 64` keeps host resource inspection aligned with those blocks; it is separate from the worker KV-cache block setting. -Use a host with at least 5 GB of available RAM: the example allows a 4 GB -execution-process budget and keeps the default 1 GB host reserve. Initial trace -materialization is estimated at about 2.81 GB. The full profile has no qualified +Allow headroom for the 4 GB execution-process budget, the default 1 GB host +reserve, and the CLI coordinator. The full profile has no qualified static peak-memory bound because recycled plays retain lifecycle evidence. The CLI runs it under live resource supervision and can stop with `resource_limited` if that budget is exhausted. Increasing the duration or lane @@ -221,6 +254,15 @@ qualification boundaries. | `prediction.json` | Aggregate predictions, snapshot information, preparation/barrier audit, profile duration/cutoff accounting, and play outcomes | | `requests.jsonl` | Measured profile requests, including per-request timing and identity | +`routing_policy.roles` records each pool's native policy, Dynamo version, +placement decisions and physical KV-event count. Match `request_id`, `group_key`, +`worker_id`, and `dp_rank` to per-request `routing_history`; decision counts also +include preparation requests. Decode selection is load-based and its index event +count is zero. Measure destination-cache reuse through the difference between +`destination_reserved_ms` and `destination_activated_ms`, comparing +`destination_missing` and `full_prompt` transfer timing with a cache-disabled +control. A routing overlap score alone does not prove physical cache reuse. + Check `agentic_phases` for preparation success and barrier state, and `agentic_play_outcomes` for completed or incomplete plays. Check `agentic_profile` for the resolved duration and grace settings, admission cutoff, recycled play @@ -251,3 +293,19 @@ To run the initial snapshot suffixes once without recycling, remove `agentic_profile` from the YAML. To replay those finite plays from turn zero, also remove `agentic_snapshot` and `agentic_warmup`. Continuous profiles require seeded snapshots for the initial lanes. + +## Acceptance from AISimulate + +Run the CPU-only suite with this AISimulate build and the native Dynamo runtime, +without installing the `ai-dynamo` runner package: + +```bash +AISIM_REQUIRE_DYNAMO_ROUTER=1 python/aisimulate/.venv/bin/python -m pytest \ + tests/test_conversation_routing_cli.py +``` + +Required mode fails if the native API is absent. The suite blocks Dynamo runner +imports and checks vLLM/SGLang aggregated/P-D routing, session/sibling affinity, +physical cache reuse, Weka lineage, tokenization, and snapshot/warmup profiles. +These checks establish functional simulation, not hardware accuracy or bitwise +parity with the live AgentX harness. diff --git a/docs/replay/dynamo.md b/docs/replay/dynamo.md index b1e234fe8..fce5f9fee 100644 --- a/docs/replay/dynamo.md +++ b/docs/replay/dynamo.md @@ -9,7 +9,12 @@ SPDX-License-Identifier: Apache-2.0 # Dynamo integration -Select `--stack dynamo` to use Dynamo-owned Router/Planner behavior. Install the +For conversation-aware AgentX simulation, use AISimulate's default engine stack +with the standalone Dynamo policy API; see the [AgentX quickstart](agentic/quickstart.md). +That path needs neither the Dynamo runner plugin nor an update to its AISimulate +pin. AISimulate owns its complete simulation and acceptance suite. + +Select `--stack dynamo` to use the optional Dynamo-owned Router/Planner integration. Install the integration in the same environment as AISimulate and select the same stack when replaying a saved recommendation. The YAML does not embed stack selection. The engine stack remains sufficient for reading native Dynamo trace files. @@ -104,7 +109,37 @@ run `python3 -m pip check`, and repeat this prediction. A successful dependency check, basic prediction without Planner, or `predict --help` alone does not verify that Planner loads. -## Router configuration +## Router policy on the AISimulate engine stack + +The built-in `engine.router` adapter accepts a fixed `router` block with the +default `--stack engine`. `kv_router` requires a separately installed Dynamo +runtime exposing `OfflineRouter`; its production selector, cache index and +session affinity remain in Dynamo. A missing or incompatible API fails instead +of switching to another policy. + +| Knob | Default | Rules | +| --- | --- | --- | +| `router.policy` | `round_robin` | `round_robin` or `kv_router`; fixed in prediction and recommendation. | +| `router.prefill_load_model.type` | `none` | Only `none`; custom prefill estimators belong to the optional Dynamo integration. | +| `router.overlap_score_credit` | Dynamo policy default | Finite value in `[0, 1]`, KV-router-only. | +| `router.prefill_load_scale` | Dynamo policy default | Finite, nonnegative, KV-router-only. | +| `router.temperature` | Dynamo policy default | Finite, nonnegative, KV-router-only. | +| `router.affinity.mode` | Unset | `session` or `sibling_group`; requires `kv_router`. | +| `router.affinity.ttl_seconds` | `3600` | Finite value from `1` through `31536000`, including fractions. Uses simulation time. | + +Affinity uses vLLM/SGLang with static aggregated or P/D workers, device-only KV +cache, no grouped caches or speculative decoding, and default scoring. It binds +the selected worker and DP rank separately for each pool; dispatch acceptance +commits the binding. Session keys distinguish conversations within each play; +sibling keys group conversations with the same parent. TTL starts after the +last active request releases its lease. Seeded snapshots, saved-frontier warmup, +continuous profiles, and worker tokenization remain owned by AISimulate. + +The following sections describe the optional Dynamo-stack adapter. Its search +and Planner capabilities do not establish support in the engine-stack policy +adapter. + +## Router configuration on the Dynamo stack Router is not part of the AISimulate core schema. The `dynamo.router` config adapter owns this section's concrete model, defaults, recommendation domains, validation, and runtime lowering. The diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index d1af3fb73..132aa81ea 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -1432,3 +1432,38 @@ https://github.com/alibaba/tair-kvcache/blob/85402f1a012fb4279be5bb4b57613bb0a06 The AISimulate files are modified re-implementations, not verbatim copies. Both derived files carry the source URLs and revisions in their headers. + +## NVIDIA Dynamo offline routing adapter and acceptance + +The repository-root file `python/aisimulate/src/aisimulate/router.py` is adapted +and modified from Dynamo's +`components/src/dynamo/router/simulation/config.py` and `provider.py` at commit +`6d6d36d7bfa63d80a0302df637c08f85523aa5b6`: +https://github.com/ai-dynamo/dynamo/tree/6d6d36d7bfa63d80a0302df637c08f85523aa5b6/components/src/dynamo/router/simulation +The adaptation retains the configuration vocabulary and field mapping, and +transports policy calls to the separately installed Dynamo runtime. It does not +copy the Router, KV index, or session-affinity algorithms into AISimulate. +The AISim lifecycle and conversation-metadata translation was also adapted and +modified from the same Apache-2.0 NVIDIA project at +`2c4be58bdd6a2dc78c409bdfd67ad54d1b6fedf8`, original files +`lib/bindings/python/rust/llm/offline_router/{affinity,mod,wire}.rs`: +https://github.com/ai-dynamo/dynamo/tree/2c4be58bdd6a2dc78c409bdfd67ad54d1b6fedf8/lib/bindings/python/rust/llm/offline_router + +The repository-root file `tests/test_conversation_routing_cli.py` is adapted and +modified from Dynamo's +`components/src/dynamo/replay/tests/e2e/test_unified_cli_dynamo.py` and +`components/src/dynamo/replay/tests/e2e/configs/unified_cli/predict/dynamo/13-trace-dynamo-agentic.yaml` +at commit `35ea1e7f689198333b1af0a2db8f94f545817792`: +https://github.com/ai-dynamo/dynamo/tree/35ea1e7f689198333b1af0a2db8f94f545817792/components/src/dynamo/replay/tests/e2e +The adaptation runs the conversation fixtures through AISimulate's engine CLI, +uses the standalone policy API, and rejects dependence on Dynamo's whole-run +plugin. The synthetic request tree contains no external dataset material. + +Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +These source revisions are licensed under Apache-2.0. The full license is +reproduced in `LICENSE`; the applicable upstream licenses are: +https://github.com/ai-dynamo/dynamo/blob/6d6d36d7bfa63d80a0302df637c08f85523aa5b6/LICENSE +https://github.com/ai-dynamo/dynamo/blob/35ea1e7f689198333b1af0a2db8f94f545817792/LICENSE +https://github.com/ai-dynamo/dynamo/blob/2c4be58bdd6a2dc78c409bdfd67ad54d1b6fedf8/LICENSE +The source files' copyright/license headers are preserved. These source +revisions have no root NOTICE file applicable to these files. diff --git a/python/aisimulate/src/aisimulate/config/common.py b/python/aisimulate/src/aisimulate/config/common.py index c3022b25f..96172cf4a 100644 --- a/python/aisimulate/src/aisimulate/config/common.py +++ b/python/aisimulate/src/aisimulate/config/common.py @@ -222,6 +222,7 @@ def load_yaml(path: str | Path) -> dict[str, Any]: CONFIG_ADAPTER_ENTRY_POINT_GROUP = "aisimulate.config_adapters" +BUILTIN_CONFIG_ADAPTER_NAMES = frozenset({"engine.router"}) PREDICTION_CORE_SECTIONS = frozenset({"traffic", "engine", "evaluation", "execution"}) RECOMMENDATION_CORE_SECTIONS = frozenset({*PREDICTION_CORE_SECTIONS, "optimization", "optimizer"}) diff --git a/python/aisimulate/src/aisimulate/config_adapter.py b/python/aisimulate/src/aisimulate/config_adapter.py index 6e3033c4d..c6a3c89aa 100644 --- a/python/aisimulate/src/aisimulate/config_adapter.py +++ b/python/aisimulate/src/aisimulate/config_adapter.py @@ -11,7 +11,7 @@ from dataclasses import dataclass from typing import Any, Protocol, runtime_checkable -from .config.common import CONFIG_ADAPTER_ENTRY_POINT_GROUP +from .config.common import BUILTIN_CONFIG_ADAPTER_NAMES, CONFIG_ADAPTER_ENTRY_POINT_GROUP from .sweeper.provider import API_VERSION as SWEEP_PROVIDER_API_VERSION from .sweeper.provider import ( AdapterReplaySpec, @@ -153,19 +153,26 @@ def resolve_config_adapters( requested = list(dict.fromkeys(names)) injected = injected or {} - installed = ( - list(entry_points) - if entry_points is not None - else list(importlib.metadata.entry_points().select(group=CONFIG_ADAPTER_ENTRY_POINT_GROUP)) - ) + installed = None resolved: dict[str, SimulationConfigAdapter] = {} for name in requested: if name in injected: resolved[name] = validate_config_adapter(injected[name], requested_name=name) continue + if name == "engine.router": + from .router import EngineRouterConfigAdapter + + resolved[name] = validate_config_adapter(EngineRouterConfigAdapter(), requested_name=name) + continue + if installed is None: + installed = ( + list(entry_points) + if entry_points is not None + else list(importlib.metadata.entry_points().select(group=CONFIG_ADAPTER_ENTRY_POINT_GROUP)) + ) matches = [entry for entry in installed if entry.name == name] if not matches: - available = sorted(set(injected) | {entry.name for entry in installed}) + available = sorted(set(injected) | BUILTIN_CONFIG_ADAPTER_NAMES | {entry.name for entry in installed}) raise ConfigAdapterResolutionError( f"config adapter {name!r} is unavailable; installed adapters: " f"{', '.join(available) if available else ''}" diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py index f2ce37559..4eaf841c6 100644 --- a/python/aisimulate/src/aisimulate/main.py +++ b/python/aisimulate/src/aisimulate/main.py @@ -167,7 +167,7 @@ def _predict(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: raise ValueError("ngram speculation requires offline --stack engine without adapters") if config.engine.speculation is not None and args.online: raise ValueError("speculation requires offline execution") - plan = _resource_plan(args, config, factory) + plan = _resource_plan(args, config, factory, adapter_raw) require_plan(plan) from .supervision import checkpoint, mark_execution_ready, mark_shutdown @@ -389,13 +389,14 @@ def _recommend(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: return 0 -def _resource_plan(args, config, factory) -> dict[str, Any]: +def _resource_plan(args, config, factory, adapter_configs=None) -> dict[str, Any]: return build_plan( workload_bounds(config), stack=args.stack, policy=config.execution.resources, requested_parallelism=config.optimizer.parallelism if isinstance(config, CoreRecommendationConfig) else 1, factory=factory, + native_router=args.stack == "engine" and (adapter_configs or {}).get("router", {}).get("policy") == "kv_router", ) diff --git a/python/aisimulate/src/aisimulate/resources.py b/python/aisimulate/src/aisimulate/resources.py index ea9742cd6..83d39f92a 100644 --- a/python/aisimulate/src/aisimulate/resources.py +++ b/python/aisimulate/src/aisimulate/resources.py @@ -21,7 +21,7 @@ from collections.abc import Mapping from contextlib import contextmanager from contextvars import ContextVar -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass, field, replace from pathlib import Path from typing import Any @@ -881,13 +881,13 @@ def _estimate_trace( or workload.get("agentic_warmup") ): return _estimate_trace_legacy( - workload, format_name=format_name, inspection_budget_bytes=inspection_budget_bytes + workload, stack=stack, format_name=format_name, inspection_budget_bytes=inspection_budget_bytes ) return _estimate_finite_trace(workload, format_name=format_name, inspection_budget_bytes=inspection_budget_bytes) def _estimate_trace_legacy( - workload: Mapping[str, Any], *, format_name: str, inspection_budget_bytes: int | None + workload: Mapping[str, Any], *, stack: str, format_name: str, inspection_budget_bytes: int | None ) -> ResourceEstimate: """Keep the conservative storage estimate for snapshot, warmup and profile replays.""" unqualified = lambda reason: ResourceEstimate("trace-unqualified-v1", None, 0, 0, None, reason) @@ -971,28 +971,34 @@ def _estimate_trace_legacy( cumulative = count if format_name in {"mooncake-delta", "applied_compute_agentic"} else 1 lanes = int(workload.get("agentic_lanes") or 1) peak = WORKER_BASELINE_BYTES + 128 * total_bytes + lanes * (32 * tokens * cumulative + 65536 * count) - if workload.get("agentic_profile") is not None: - # Retired plays release their large payloads, but retain identities and - # lifecycle rows. Their count depends on simulated completion times, so - # neither the finite corpus nor duration alone bounds the full profile. + profile = workload.get("agentic_profile") is not None + if profile or stack == "dynamo": + # Native routing retains additional state outside the engine estimate. + # Profiles also retain retired identities and lifecycle rows, whose + # count depends on simulated completion times rather than corpus size. # Keep the initial-materialization refusal before admitting unknown peaks. + model = "agentic-profile" if profile else "dynamo" if inspection_budget_bytes is not None and peak > inspection_budget_bytes: return ResourceEstimate( - "agentic-profile-materialization-v1", + f"{model}-materialization-v1", None, 0, 0, peak, - "initial profile trace materialization estimate exceeds live headroom", + "initial trace materialization estimate exceeds live headroom", ) return ResourceEstimate( - "agentic-profile-unqualified-v1", + f"{model}-unqualified-v1", None, 0, 0, None, - "agentic profile retains evidence beyond the initial corpus; total memory has no qualified static bound " - "and requires supervised serial execution", + ( + "agentic profile retains evidence beyond the initial corpus" + if profile + else "Dynamo policy retains native router state beyond the engine trace estimate" + ) + + "; total memory has no qualified static bound and requires supervised serial execution", ) return ResourceEstimate( "trace-json-metadata-v1", @@ -1059,6 +1065,7 @@ def build_plan( host: HostResources | None = None, factory: Any = None, concurrency: int | None = None, + native_router: bool = False, ) -> dict[str, Any]: policy = policy or _POLICY.get() or ResourceConfig() host = host or discover_host() @@ -1085,6 +1092,17 @@ def build_plan( raise ResourceLimitError("runner returned an invalid host resource estimate") if peak is not None and (type(peak) is not int or peak <= 0 or peak < estimate.lower_bound_bytes): raise ResourceLimitError("runner returned an invalid host resource estimate") + if native_router and ((peak is not None and peak <= free) or estimate.allocation_model == OBSERVED_TRACE_MODEL): + # Preserve initial-materialization refusals and lower bounds, but neither + # an engine-only estimate nor an observed trace peak qualifies router state. + estimate = replace( + estimate, + allocation_model="dynamo-unqualified-v1", + estimated_peak_bytes=None, + reason="Dynamo policy retains native router state beyond the engine estimate; " + "total memory requires supervised serial execution", + ) + peak = None workers = min(requested_parallelism, budget["cpu_limit"], free // peak) if peak else 0 reason = "" if workers else "candidate exceeds the host memory budget" if peak is None: @@ -1136,8 +1154,23 @@ def resource_policy(policy: ResourceConfig): _POLICY.reset(token) +def _uses_native_router(spec: Any) -> bool: + return any( + hook.provider == "engine.router" + and hook.kind == "placement_policy" + and hook.config.get("provider") == "dynamo_kv_router" + for hook in getattr(spec, "runtime_hooks", ()) + ) + + def guard_replay(spec: Any, *, stack: str, factory: Any = None) -> dict[str, Any]: - plan = build_plan(spec.workload, stack=stack, concurrency=spec.concurrency, factory=factory) + plan = build_plan( + spec.workload, + stack=stack, + concurrency=spec.concurrency, + factory=factory, + native_router=_uses_native_router(spec), + ) require_plan(plan) return plan @@ -1235,6 +1268,7 @@ def admit_wave(self, specs: list[Any]) -> dict[str, Any]: host=host, factory=self.factory, concurrency=spec.concurrency, + native_router=_uses_native_router(spec), ) plans.append(plan) estimate = plan["estimate"] diff --git a/python/aisimulate/src/aisimulate/router.py b/python/aisimulate/src/aisimulate/router.py new file mode 100644 index 000000000..e6f91a106 --- /dev/null +++ b/python/aisimulate/src/aisimulate/router.py @@ -0,0 +1,468 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Configuration schema and field mapping adapted from ai-dynamo/dynamo, +# components/src/dynamo/router/simulation/{config.py,provider.py}, revision +# 6d6d36d7bfa63d80a0302df637c08f85523aa5b6 (Apache-2.0): +# https://github.com/ai-dynamo/dynamo/tree/6d6d36d7bfa63d80a0302df637c08f85523aa5b6/components/src/dynamo/router/simulation +# Modified for AISimulate execution using an externally installed Dynamo policy API. + +"""Configuration and transport for externally installed Dynamo routing policy.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from typing import Annotated, Any, Literal + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + +from .config_adapter import PredictionAdapterContext, RecommendationAdapterContext +from .sweeper.provider import AdapterReplaySpec, AdapterSearchPlan, CandidateContext, JSONValue, RuntimeHookSpec +from .sweeper.replay import ReplaySpec + +_NATIVE_HOOK = RuntimeHookSpec(provider="engine.router", kind="placement_policy", api_version=1) +_FiniteNonnegative = Annotated[float, Field(strict=True, ge=0, allow_inf_nan=False)] + + +class ConversationAffinityConfig(BaseModel): + """Native conversation affinity whose lifetime follows simulation time.""" + + model_config = ConfigDict(extra="forbid") + + mode: Literal["session", "sibling_group"] + ttl_seconds: float = Field(default=3600, strict=True, ge=1, le=31_536_000, allow_inf_nan=False) + + +class _PrefillLoadModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + # Custom AIS load estimators belong to the optional integration stack. + type: Literal["none"] = "none" + + +class RouterConfig(BaseModel): + """Concrete native policy; parameter search and custom load models are unsupported.""" + + model_config = ConfigDict(extra="forbid") + + policy: Literal["round_robin", "kv_router"] = "round_robin" + prefill_load_model: _PrefillLoadModel = Field(default_factory=_PrefillLoadModel) + overlap_score_credit: Annotated[float, Field(strict=True, ge=0, le=1, allow_inf_nan=False)] | None = None + prefill_load_scale: _FiniteNonnegative | None = None + temperature: _FiniteNonnegative | None = None + affinity: ConversationAffinityConfig | None = None + + @field_validator("overlap_score_credit", "prefill_load_scale", "temperature", mode="before") + @classmethod + def _reject_explicit_null(cls, value: Any) -> Any: + if value is None: + raise ValueError("Router numeric knobs do not accept null") + return value + + @model_validator(mode="after") + def _validate_policy(self) -> RouterConfig: + if self.policy == "round_robin": + if self.affinity is not None: + raise ValueError("router.affinity requires policy='kv_router'") + knobs = {"overlap_score_credit", "prefill_load_scale", "temperature"} & self.model_fields_set + if knobs: + raise ValueError(f"round_robin rejects KV-router knobs {sorted(knobs)}") + elif self.affinity is not None: + for name, default in (("overlap_score_credit", 1.0), ("prefill_load_scale", 1.0), ("temperature", 0.0)): + if getattr(self, name) not in (None, default): + raise ValueError("router.affinity requires default KV-router scoring parameters") + return self + + +def _compile(config: Mapping[str, JSONValue]) -> AdapterReplaySpec: + public = RouterConfig.model_validate(config) + concrete = public.model_dump(mode="json", exclude_none=True) + if public.policy == "round_robin": + return AdapterReplaySpec(config=concrete) + # Omit unspecified fields so the installed production policy owns its defaults. + router_config = { + native: getattr(public, name) + for name, native in ( + ("overlap_score_credit", "overlap_score_credit"), + ("prefill_load_scale", "prefill_load_scale"), + ("temperature", "router_temperature"), + ) + if getattr(public, name) is not None + } + native: dict[str, JSONValue] = {"router_config": router_config} + if public.affinity is not None: + native["affinity"] = public.affinity.model_dump(mode="json") + hook = RuntimeHookSpec( + provider=_NATIVE_HOOK.provider, + kind=_NATIVE_HOOK.kind, + api_version=_NATIVE_HOOK.api_version, + config={"provider": "dynamo_kv_router", "config": native}, + ) + return AdapterReplaySpec(config=concrete, runtime_hooks=(hook,)) + + +class EngineRouterConfigAdapter: + """Compile a fixed routing policy without importing any Dynamo Python package.""" + + name = "engine.router" + section = "router" + config_adapter_api_version = 3 + + def compile_prediction( + self, config: Mapping[str, JSONValue], context: PredictionAdapterContext + ) -> AdapterReplaySpec: + del context + return _compile(config) + + def compile_recommendation( + self, config: Mapping[str, JSONValue], context: RecommendationAdapterContext + ) -> AdapterSearchPlan: + del context + compiled = _compile(config) + return AdapterSearchPlan(state=compiled.config, potential_runtime_hooks=compiled.runtime_hooks) + + def materialize_candidate( + self, plan: AdapterSearchPlan, selection: Mapping[str, JSONValue], context: CandidateContext + ) -> AdapterReplaySpec: + del context + if selection: + raise ValueError("engine.router supports a fixed policy, not router parameter search") + if not isinstance(plan.state, Mapping): + raise ValueError("engine.router search plan requires a concrete configuration") + return _compile(plan.state) + + +def native_placement(spec: ReplaySpec) -> dict[str, JSONValue]: + """Validate the built-in adapter and lower it to the core placement descriptor.""" + + unknown = set(spec.adapters) - {"engine.router"} + if unknown: + raise ValueError(f"EngineReplayRunner does not support component adapters {sorted(unknown)}") + adapter = spec.adapters.get("engine.router") + if adapter is None: + return {"provider": "round_robin", "config": None} + expected = _compile(adapter.config) + if adapter.runtime_hooks != expected.runtime_hooks: + raise ValueError("engine.router runtime hooks must match the configured routing policy") + deployment = spec.backend_deployment + if deployment.deployment_mode not in {"agg", "disagg"} or deployment.encoder is not None: + raise ValueError("engine.router requires aggregated or disaggregated text replay") + if not expected.runtime_hooks: + return {"provider": "round_robin", "config": None} + if spec.execution_mode != "offline": + raise ValueError("engine.router requires offline execution") + if adapter.config.get("affinity") is not None: + _require_affinity_scope(spec) + return expected.runtime_hooks[0].config + + +def _require_affinity_scope(spec: ReplaySpec) -> None: + deployment = spec.backend_deployment + if deployment.backend not in {"vllm", "sglang"}: + raise ValueError("router.affinity requires backend=vllm or sglang") + if any(hook.kind == "scaling_policy" for hook in spec.runtime_hooks): + raise ValueError("router.affinity requires static worker pools") + for args in (deployment.agg_engine_args, deployment.prefill_engine_args, deployment.decode_engine_args): + if not args: + continue + rank = args.get("rank", args) + if not isinstance(rank, Mapping): + continue # The engine descriptor validator reports malformed ranks. + if any(rank.get(key) is not None for key in ("native_host_offload", "g3_offload")): + raise ValueError("router.affinity requires device-only KV cache without host or G3 offload") + if rank.get("kv_cache_groups"): + raise ValueError("router.affinity does not support grouped KV caches") + timing = rank.get("timing_model") + identity = timing.get("config", {}) if isinstance(timing, Mapping) else {} + if not isinstance(identity, Mapping): + identity = {} + if any( + source.get(key) not in (None, 0) + for source in (args, rank, identity) + for key in ("aic_nextn", "nextn", "speculation") + ): + raise ValueError("router.affinity requires speculative decoding disabled") + if any(source.get("startup_time") not in (None, 0) for source in (args, rank)): + raise ValueError("router.affinity requires static worker pools without startup delay") + + +class _DynamoPolicy: + """Translate AISim lifecycle events into existing Dynamo policy operations.""" + + def __init__(self, context_json: str): + from dynamo import _core as native + + context = json.loads(context_json) + if context["schema_version"] != 1: + raise ValueError("unsupported external placement protocol version") + self.role = context["role"] + self.dp_size = context["dp_size"] + self.workers = set(range(len(context["topology"]))) + for worker_id, worker in enumerate(context["topology"]): + if worker["worker_id"] != worker_id or worker["scheduler_ids"] != list( + range(worker_id * self.dp_size, (worker_id + 1) * self.dp_size) + ): + raise ValueError("Dynamo policy requires contiguous worker and rank IDs") + rank = context["engine"]["rank"] + self.block_size = rank["block_size"] + config = context["placement"]["config"] + router_config = dict(config.get("router_config", {})) + # Each pool reports its own load: prefill has no decode occupancy, and + # decode allocates independent blocks without prefill-cache credit. + if self.role == "prefill": + router_config["router_track_active_blocks"] = False + elif self.role == "decode": + router_config.update( + overlap_score_credit=0.0, + router_assume_kv_reuse=False, + router_track_prefill_tokens=False, + router_prefill_load_model="none", + ) + self.native = native + self.router = native.OfflineRouter( + self.block_size, + len(self.workers), + self.dp_size, + rank["num_gpu_blocks"], + rank["max_num_batched_tokens"], + json.dumps(router_config), + context["selector_seed"], + ) + self.tree = native.RadixTree() + self.affinity_config = config.get("affinity") + self.affinity = native.SessionAffinity(self.affinity_config["ttl_seconds"]) if self.affinity_config else None + self.waiting = {} + self.pending = {} + self.staged = {} + self.active = {} + self.now_ms = 0.0 + self.ready = False + self.evidence = dict( + native_policy="dynamo.DefaultWorkerSelector", + dynamo_version=native.__version__, + physical_kv_events=0, + decision_count=0, + decisions_captured=context["capture_decisions"], + decisions=[], + post_dispatch_checks=0, + dispatch_aborts=0, + ) + + def _group(self, command): + if self.affinity is None: + return None + context = command["request"].get("replay_context") or {} + agentic = context.get("agentic") or {} + session = agentic.get("conversation_id") or command.get("session_id") or context.get("session_id") + if not session: + raise ValueError("conversation affinity requires nonempty session identity") + play = agentic.get("play_id") + identity = ["session", play, session] + if self.affinity_config["mode"] == "sibling_group": + lineage = agentic.get("lineage") or {} + if lineage.get("schema") != "aisimulate.agentic.conversation-lineage.v1" or not lineage.get( + "root_conversation_id" + ): + raise ValueError("sibling_group affinity requires unambiguous conversation lineage") + parent = lineage.get("parent_conversation_id") + identity = ["siblings" if parent else "root", play, lineage["root_conversation_id"], parent or session] + return self.native.subagent_group_affinity_id(json.dumps(identity)) + + def _submit(self, command): + request = command["request"] + request_id = request["uuid"] + group = self._group(command) + hold = self.affinity.try_acquire(group) if self.affinity else None + if self.affinity and hold is None: + self.waiting[request_id] = command + return + length = command["input_length"] + hashes = command.get("hashes") + tokens = request.get("tokens", []) + if hashes is None: + source = (request.get("replay_context") or {}).get("prompt_token_source") + if source not in (None, "materialized") or len(tokens) != length: + raise ValueError("Dynamo routing requires authored prompt tokens or complete block hashes") + local = self.native.compute_block_hash_for_seq(tokens, self.block_size) + sequence = None + else: + local, sequence = hashes["local_block_hashes"], hashes["sequence_hashes"] + if len(local) != length // self.block_size or len(sequence) != len(local): + raise ValueError("request hashes must cover exactly the complete prompt blocks") + limit = ( + len(request["output_token_ids"]) + if request.get("output_token_ids") is not None + else request["max_output_tokens"] + ) + if command.get("max_output_tokens_override") is not None: + limit = min(limit, command["max_output_tokens_override"]) + target = hold.target if hold else None + self.router.submit( + request_id, + length, + limit, + self.tree.find_matches(local), + local, + sequence, + tokens if len(tokens) == length else None, + target, + max(0, request.get("priority", 0)), + request.get("strict_priority", 0), + request.get("policy_class"), + command.get("session_id"), + ) + self.pending[request_id] = (length, group, hold) + + def _drain(self): + for request_id in list(self.waiting): + self._submit(self.waiting.pop(request_id)) + result = [] + for request_id, worker, rank, overlap, best in self.router.poll(): + length, group, hold = self.pending.pop(request_id) + self.staged[request_id] = (hold, (worker, rank)) + self.evidence["decision_count"] += 1 + if self.evidence["decisions_captured"]: + self.evidence["decisions"].append( + dict( + request_id=request_id, + role=self.role, + native_policy=self.evidence["native_policy"], + worker_id=worker, + dp_rank=rank, + group_key=group, + binding_reused=hold is not None and hold.target is not None, + overlap_blocks=overlap, + best_available_overlap_blocks=best, + ) + ) + result.append( + dict( + request_id=request_id, + scheduler_id=worker * self.dp_size + rank, + reported_overlap_tokens=overlap * self.block_size, + cache_sample=dict( + overlap_blocks=overlap, + best_available_overlap_blocks=best, + isl_blocks=(length + self.block_size - 1) // self.block_size, + ), + ) + ) + self.ready = False + return result + + def _cancel(self, request_id): + removed = self.waiting.pop(request_id, None) is not None + pending = self.pending.pop(request_id, None) + if pending is not None: + self.router.cancel(request_id) + if pending[2] is not None: + pending[2].release() + removed = True + self.ready |= removed + return removed + + def call_json(self, command_json: str) -> str: + command = json.loads(command_json) + op, request_id = command["op"], command.get("request_id") + if "now_ms" in command: + self.router.advance_clock(command["now_ms"]) + if self.affinity: + self.affinity.advance_clock(command["now_ms"]) + self.now_ms = command["now_ms"] + result = None + if op == "place": + request = command["request"] + request_id = request["uuid"] + if any(request_id in state for state in (self.waiting, self.pending, self.staged, self.active)): + raise ValueError("duplicate routing request") + if request.get("preferred_dp_rank") is not None or request.get("preferred_prefill_dp_rank") is not None: + raise ValueError("Dynamo policy does not support authored DP-rank pins") + if request.get("kv_hints") is not None: + raise ValueError("Dynamo policy does not support authored KV-fetch hints") + self._submit(command) + ready = self._drain() + result = dict( + placement=next((row for row in ready if row["request_id"] == request_id), None), + released=[row for row in ready if row["request_id"] != request_id], + ) + elif op == "dispatch_committed": + hold, target = self.staged.pop(request_id) + self.active[request_id] = self.affinity.commit(hold, *target) if hold else None + self.evidence["post_dispatch_checks"] += 1 + self.ready = True + elif op in ("request_terminal", "dispatch_aborted"): + self._cancel(request_id) + hold, _ = self.staged.pop(request_id, (None, None)) + lease = self.active.pop(request_id, None) + for owned in (hold, lease): + if owned is not None: + owned.release() + self.router.free(request_id) + self.ready = True + if op == "dispatch_aborted": + self.evidence["dispatch_aborts"] += 1 + else: + result = self._drain() + elif op == "cancel_pending": + result = self._cancel(request_id) + elif op == "prefill_completed": + self.router.prefill_complete(request_id) + result = self._drain() + elif op == "observe": + for batch in command["events"]: + if ( + batch["stage"] != self.role + or batch["worker_id"] not in self.workers + or not 0 <= batch["dp_rank"] < self.dp_size + ): + raise ValueError("cache observation has an invalid role, worker or rank") + for event in batch["events"]: + if event.get("tier", "device") != "device" or event["dp_rank"] != batch["dp_rank"]: + raise ValueError("routing requires device KV events with matching rank") + if self.role != "decode": + self.tree.apply_event(batch["worker_id"], json.dumps(event).encode()) + self.evidence["physical_kv_events"] += 1 + result = [] + elif op in ("advance_clock", "topology_settled"): + result = self._drain() + elif op in ("worker_ready", "worker_draining", "worker_removed"): + worker = command["worker"]["worker_id"] + if op == "worker_ready": + self.router.add_worker(worker) + self.workers.add(worker) + elif op == "worker_draining": + self.router.remove_worker(worker) + else: + self.router.unregister_worker(worker) + self.tree.remove_worker(worker) + self.workers.remove(worker) + result = [] + elif op not in ("status", "snapshot"): + raise ValueError(f"unknown policy operation {op!r}") + pending = len(self.waiting) + len(self.pending) + response = dict( + result=result, pending_count=pending, next_wakeup_ms=self.now_ms if self.ready and pending else None + ) + if op == "snapshot": + response["evidence"] = self.evidence + return json.dumps(response, separators=(",", ":")) + + +def run_dynamo_policy(runtime: Any, payload: str) -> str: + """Run AISimulate's executor using independently installed native policy APIs.""" + try: + from dynamo._core import OfflineRouter, SessionAffinity + except ImportError as error: + raise ValueError( + "router.policy=kv_router requires ai-dynamo-runtime with OfflineRouter and SessionAffinity" + ) from error + if not hasattr(OfflineRouter, "submit") or not hasattr(SessionAffinity, "try_acquire"): + raise ValueError("installed Dynamo runtime predates the standalone policy bindings") + report = json.loads(runtime.run_replay_with_policy_json(payload, _DynamoPolicy)) + roles = report.pop("external_policy") + first = next(iter(roles.values())) + report["routing_policy"] = dict( + api_version=1, native_policy=first["native_policy"], dynamo_version=first["dynamo_version"], roles=roles + ) + return json.dumps(report, separators=(",", ":")) diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index ae1d1bc9a..32e717a24 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Dynamo-free engine replay implementation of the Sweeper Runner contract.""" +"""Native replay implementation hosted by AISimulate.""" from __future__ import annotations @@ -24,6 +24,7 @@ from .config.common import ENGINE_MODEL_CONTROL_FIELDS, is_active_engine_model_control from .config.engine import StateCacheConfig, _validate_replay_estimation_mode from .power import normalize_power_summary, power_metadata +from .router import native_placement from .sweeper.afd_engine import AFDForegroundEngine from .sweeper.afd_parallel import AFDPhase, AFDTopology from .sweeper.afd_perfmodel import AFDLayerTimes @@ -31,6 +32,7 @@ from .sweeper.replay import ( BackendDeploymentSpec, EncoderPoolSpec, + HookCapability, ReplayOutputRequirements, ReplayReport, ReplaySpec, @@ -359,7 +361,7 @@ def measure(self, spec: ReplaySpec) -> AFDCompanionTiming: @dataclass(frozen=True) class EngineReplayRunnerFactory: - """Create reusable engine-only Runners for Sweeper candidates. + """Create reusable AISimulate runners for replay and Sweeper candidates. The optional runtime is a test seam. Production factories leave it unset so the dataclass remains process-serializable and each worker loads its local @@ -382,6 +384,7 @@ def capabilities(self) -> RunnerCapabilities: return RunnerCapabilities( replay_spec_api_version=1, supported_backend_topologies=_SUPPORTED_BACKEND_TOPOLOGIES, + supported_hooks=(HookCapability("engine.router", "placement_policy", 1),), supports_disaggregated_attention_dp=True, supports_analytical_epd=True, supports_native_epd=True, @@ -479,6 +482,9 @@ def run( if spec.workload.get("source_type") is not None and "length_sampler" in spec.workload: raise ValueError("length_sampler requires materialized direct synthetic replay without source_type") self.capabilities.require_compatible(spec) + if set(spec.adapters) - {"engine.router"}: + raise InvalidRunnerError("EngineReplayRunner does not support component adapters other than engine.router") + placement = native_placement(spec) encoder = spec.backend_deployment.encoder if ( encoder is None @@ -593,7 +599,13 @@ def run( separators=(",", ":"), ) try: - report_json = self._resolve_runtime().run_replay_json(execution_spec_json) + runtime = self._resolve_runtime() + if placement["provider"] == "dynamo_kv_router": + from .router import run_dynamo_policy + + report_json = run_dynamo_policy(runtime, execution_spec_json) + else: + report_json = runtime.run_replay_json(execution_spec_json) except MemoryError as error: # Resource-aware installations classify a report storage failure as host # exhaustion, not a failed candidate. Older SDKs retain MemoryError. @@ -1130,10 +1142,7 @@ def _materialize_engine_execution_spec( "topology": topology, "engine": engine, "adapters": { - "placement": { - "provider": "round_robin", - "config": None, - }, + "placement": native_placement(spec), "scaling": { "provider": "none", "config": None, @@ -1272,7 +1281,8 @@ def _execution_target_model( model = config.get("model", config.get("model_path")) if isinstance(model, str) and model.strip(): return model.strip() - timing = raw_engine_args.get("timing_model") + rank = raw_engine_args.get("rank", raw_engine_args) + timing = rank.get("timing_model") if isinstance(rank, Mapping) else None if isinstance(timing, Mapping) and timing.get("provider") == "aic": config = timing.get("config") if isinstance(config, Mapping): diff --git a/python/aisimulate/src/aisimulate/supervision.py b/python/aisimulate/src/aisimulate/supervision.py index d286074c5..3d914b5b7 100644 --- a/python/aisimulate/src/aisimulate/supervision.py +++ b/python/aisimulate/src/aisimulate/supervision.py @@ -24,6 +24,7 @@ import psutil from .config.common import ( + BUILTIN_CONFIG_ADAPTER_NAMES, CONFIG_ADAPTER_ENTRY_POINT_GROUP, PREDICTION_CORE_SECTIONS, RECOMMENDATION_CORE_SECTIONS, @@ -366,7 +367,7 @@ def _require_output_adapters(names: Iterable[str]) -> None: def _require_known_sections(raw: dict[str, Any], *, command: str, stack: str) -> None: """Reject top-level sections that no installed config adapter owns.""" core = PREDICTION_CORE_SECTIONS if command == "predict" else RECOMMENDATION_CORE_SECTIONS - installed = _installed_names(CONFIG_ADAPTER_ENTRY_POINT_GROUP) + installed = [*_installed_names(CONFIG_ADAPTER_ENTRY_POINT_GROUP), *BUILTIN_CONFIG_ADAPTER_NAMES] for section in raw: # split_config_sections reports malformed keys and predict-only restrictions. if not isinstance(section, str) or not section or "." in section: @@ -374,6 +375,8 @@ def _require_known_sections(raw: dict[str, Any], *, command: str, stack: str) -> if section in RECOMMENDATION_CORE_SECTIONS: continue name = f"{stack}.{section}" + if name in BUILTIN_CONFIG_ADAPTER_NAMES: + continue if installed.count(name) > 1: raise ValueError(f"config adapter {name!r} has multiple installed providers") if name not in installed: diff --git a/tests/sweeper/test_packaging.py b/tests/sweeper/test_packaging.py index dcb0ecd5a..ed025d0f0 100644 --- a/tests/sweeper/test_packaging.py +++ b/tests/sweeper/test_packaging.py @@ -12,8 +12,6 @@ import pytest from packaging.requirements import Requirement -from aisimulate.sweeper.replay import BackendDeploymentSpec, ReplaySpec - try: import tomllib except ModuleNotFoundError: # Python 3.10 @@ -101,13 +99,14 @@ def test_aisimulate_has_no_dynamo_or_component_adapter_dependencies(): assert "prophet" not in names -def test_importing_sweeper_does_not_import_dynamo(): +@pytest.mark.parametrize("module", ["aisimulate.sweeper", "aisimulate.router"]) +def test_importing_simulation_adapters_does_not_import_dynamo(module): subprocess.run( [ sys.executable, "-c", ( - "import sys; import aisimulate.sweeper; " + f"import sys; import {module}; " "assert not any(name == 'dynamo' or name.startswith('dynamo.') " "for name in sys.modules)" ), @@ -133,32 +132,6 @@ def test_ai_dynamo_registers_optional_sweeper_providers(): } -@pytest.mark.parametrize(("field", "bound"), [("ttft_ms", 800.0), ("itl_ms", 30.0)]) -def test_ai_dynamo_runner_preserves_independent_sla_bounds(field: str, bound: float) -> None: - _ai_dynamo_distribution_or_skip() - from dynamo.replay.simulation import DynamoReplayRunner - - spec = ReplaySpec( - backend_deployment=BackendDeploymentSpec( - deployment_mode="agg", - backend="vllm", - backend_version="test", - agg_engine_args={}, - num_workers=1, - ), - workload={}, - goal={"target": "throughput", "strict_sla": False, "sla": {field: bound}}, - ) - expected = { - "sla_ttft_ms": None, - "sla_itl_ms": None, - "sla_e2e_ms": None, - } - expected[f"sla_{field}"] = bound - - assert DynamoReplayRunner._goodput_sla_kwargs(spec) == expected - - def test_aisimulate_source_versions_are_synchronized(): root, repo_root = _source_checkout_roots() project = tomllib.loads((root / "pyproject.toml").read_text()) @@ -170,6 +143,9 @@ def test_aisimulate_source_versions_are_synchronized(): assert core["package"]["version"] == workspace_version assert workspace["workspace"]["members"] == ["crates/core"] assert project["tool"]["maturin"]["manifest-path"] == "../../crates/core/Cargo.toml" + lock = tomllib.loads((repo_root / "Cargo.lock").read_text()) + assert not any(package["name"].startswith("dynamo-") for package in lock["package"]) + assert not (repo_root / "crates/python").exists() def test_profiler_does_not_publish_or_reexport_sweeper(): diff --git a/tests/test_conversation_routing_cli.py b/tests/test_conversation_routing_cli.py new file mode 100644 index 000000000..8f85f6ad8 --- /dev/null +++ b/tests/test_conversation_routing_cli.py @@ -0,0 +1,478 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Modified from ai-dynamo/dynamo at 35ea1e7f689198333b1af0a2db8f94f545817792: +# components/src/dynamo/replay/tests/e2e/test_unified_cli_dynamo.py and its +# configs/unified_cli/predict/dynamo/13-trace-dynamo-agentic.yaml fixture. +# https://github.com/ai-dynamo/dynamo/tree/35ea1e7f689198333b1af0a2db8f94f545817792/components/src/dynamo/replay/tests/e2e +# Adapted for AISim-owned CLI execution with only the standalone policy API. + +"""Conversation routing acceptance: AISim runs AgentX and calls Dynamo policy. + +Set AISIM_REQUIRE_DYNAMO_ROUTER=1 for the required cross-repository acceptance +job. The job installs ai-dynamo-runtime with OfflineRouter and no ai-dynamo +runner package. Ordinary unit environments may omit this optional runtime. +""" + +from __future__ import annotations + +import importlib +import importlib.metadata +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import yaml + +pytestmark = [pytest.mark.integration, pytest.mark.gpu_0, pytest.mark.timeout(240)] +_REPO_ROOT = Path(__file__).resolve().parents[1] + + +@pytest.fixture(scope="module") +def dynamo_policy(): + required = os.environ.get("AISIM_REQUIRE_DYNAMO_ROUTER") == "1" + try: + core = importlib.import_module("dynamo._core") + except ImportError as error: + if required: + pytest.fail(f"Dynamo policy acceptance requires ai-dynamo-runtime with OfflineRouter: {error}") + pytest.skip("optional Dynamo policy runtime is not installed") + if not hasattr(getattr(core, "OfflineRouter", None), "submit") or not hasattr(core, "SessionAffinity"): + if required: + pytest.fail( + "installed ai-dynamo-runtime lacks native scheduler/affinity bindings; build the policy API revision" + ) + pytest.skip("installed Dynamo runtime predates the standalone OfflineRouter API") + if required: + try: + importlib.metadata.distribution("ai-dynamo") + except importlib.metadata.PackageNotFoundError: + pass + else: + pytest.fail("acceptance must run without the ai-dynamo whole-run plugin package") + return core.OfflineRouter + + +def _run_cli(tmp_path: Path, *args: str) -> subprocess.CompletedProcess[str]: + # Sitecustomize is inherited by supervised child interpreters as well as the + # coordinator. Neither may delegate execution to Dynamo's replay plugin. + guard = tmp_path / "import-guard" + guard.mkdir(exist_ok=True) + (guard / "sitecustomize.py").write_text( + "import sys\n" + "class NoDynamoRunner:\n" + " def find_spec(self, fullname, path=None, target=None):\n" + " if any(fullname == name or fullname.startswith(name + '.') " + "for name in ('dynamo.replay', 'dynamo.router', 'dynamo.planner')):\n" + " raise AssertionError('AISim must own simulation: ' + fullname)\n" + "sys.meta_path.insert(0, NoDynamoRunner())\n" + ) + env = dict(os.environ) + env["PYTHONPATH"] = str(guard) + os.pathsep + env.get("PYTHONPATH", "") + result = subprocess.run( + [sys.executable, "-m", "aisimulate", *args], + cwd=_REPO_ROOT, + env=env, + text=True, + capture_output=True, + timeout=180, + check=False, + ) + assert result.returncode == 0, f"AISim CLI failed:\n{result.stdout}\n{result.stderr}" + return result + + +def _conversation_prediction(tmp_path: Path, backend: str, topology: str, mode: str): + """Self-authored two-play tree; no weights, profiles, or network access.""" + config = { + "engine": { + "mode": "aggregated", + "model": str(_REPO_ROOT / "tests/e2e/configs/unified_cli/fixtures/tiny-model"), + "hardware": "h200_sxm", + "backend": "vllm", + "context_length": 2048, + "workers": { + "aggregated": { + "parallelism": {"replicas": 1, "tensor": 1, "attention_data": 1}, + "scheduler": {"max_batched_tokens": 8192, "max_sequences": 256}, + "kv_cache": { + "block_size": 64, + "prefix_caching": True, + "capacity": {"type": "fixed", "blocks": 256}, + }, + "timing": {"type": "fixed", "prefill_ms": 1, "decode_ms": 1}, + } + }, + } + } + rows = [ + { + "schema": "dynamo.agentic_mooncake", + "version": 2, + "block_size": 64, + "hash_id_scope": "local", + "source": { + "format": "aisimulate-policy-api-test", + "digest": "two-branches-two-plays-v1", + }, + } + ] + for play in ("first", "second"): + for index, session, predecessor, relation in ( + (0, "owner", None, None), + (1, "branch-a", 0, "spawn"), + (2, "branch-b", 0, "spawn"), + (3, "branch-a", 1, "sequence"), + (4, "branch-b", 2, "sequence"), + ): + rows.append( + { + "request_id": f"{play}-{index}", + "play_id": play, + "session_id": session, + "model": config["engine"]["model"], + "input_length": 128, + "output_length": 2, + "hash_ids": [101, 202], + "not_before_ms": index * 20, + "recorded_api_time_ms": 2, + "dependencies": [] + if predecessor is None + else [ + { + "request_id": f"{play}-{predecessor}", + "relation": relation, + "trigger": "dispatch" if relation == "spawn" else "completion", + "delay_ms": 20, + } + ], + } + ) + trace = tmp_path / "conversation-tree.jsonl" + trace.write_text("".join(json.dumps(row) + "\n" for row in rows)) + config["traffic"] = { + "source": { + "type": "trace", + "format": "agentic_mooncake", + "paths": [str(trace)], + "block_size": 64, + }, + "load": {"type": "trace_timestamps"}, + } + worker = config["engine"]["workers"]["aggregated"] + worker["parallelism"].update(replicas=2, attention_data=2) + config["engine"].update(mode=topology, backend=backend) + config["engine"]["workers"] = { + role: json.loads(json.dumps(worker)) + for role in (["aggregated"] if topology == "aggregated" else ["prefill", "decode"]) + } + config["router"] = { + "policy": "kv_router", + "affinity": {"mode": mode, "ttl_seconds": 3600}, + } + return config + + +def _predict_conversations(tmp_path: Path, config: dict, name: str): + path = tmp_path / f"{name}.yaml" + path.write_text(yaml.safe_dump(config)) + output = tmp_path / name + args = [ + "predict", + "--config", + str(path), + "--output-dir", + str(output), + "--capture-per-request", + "--format", + "json", + ] + result = _run_cli(tmp_path, *args) + report = json.loads((output / "prediction.json").read_text()) + assert json.loads(result.stdout)["completed_requests"] == report["completed_requests"] + assert report["per_request"] == [json.loads(line) for line in (output / "requests.jsonl").read_text().splitlines()] + return report + + +def _assert_conversation_bindings(report: dict, topology: str, mode: str, *, block_size=64): + policies = report["routing_policy"]["roles"] + roles = {"aggregated"} if topology == "aggregated" else {"prefill", "decode"} + assert set(policies) == roles + decisions = {} + expected_requests = {row["uuid"] for row in report["per_request"]} | { + row["uuid"] for row in report.get("agentic_phases", {}).get("requests", []) + } + for role, policy in policies.items(): + assert policy["native_policy"] == "dynamo.DefaultWorkerSelector" + if role == "decode": + # The existing decode adapter uses native load selection. Its engine + # cache reuse is checked independently through transfer timing below. + assert policy["physical_kv_events"] == 0 + else: + assert policy["physical_kv_events"] > 0 + assert policy["decisions_captured"] is True + assert policy["decision_count"] == len(policy["decisions"]) + assert policy["post_dispatch_checks"] == policy["decision_count"] + assert policy["dispatch_aborts"] == 0 + assert {row["request_id"] for row in policy["decisions"]} == expected_requests + for row in policy["decisions"]: + assert row["native_policy"] == "dynamo.DefaultWorkerSelector" + assert row["role"] == role + key = (row["request_id"], role) + assert key not in decisions + decisions[key] = row + bindings: dict[tuple, set] = {} + group_keys: dict[tuple, set] = {} + observations = 0 + for record in report["per_request"]: + assert record["terminal_status"] == "completed" + identity = record["agentic"] + lineage = identity["lineage"] + conversation = identity["conversation_id"] + group = ( + (identity["play_id"], conversation) + if mode == "session" + else ( + identity["play_id"], + lineage["root_conversation_id"], + lineage.get("parent_conversation_id") or conversation, + bool(lineage.get("parent_conversation_id")), + ) + ) + observed_roles = set() + for route in record["routing_history"]: + role = "aggregated" if route["pool"] == "agg" else route["pool"] + observed_roles.add(role) + decision = decisions[(record["uuid"], role)] + pair = decision["worker_id"], decision["dp_rank"] + assert pair == (route["logical_worker_id"], route["dp_rank"]) + assert all(value in (0, 1) for value in pair) + assert route["outcome"] == "immediate" + assert route["selected_overlap_blocks"] == decision["overlap_blocks"] + assert route["best_available_overlap_blocks"] == decision["best_available_overlap_blocks"] + bindings.setdefault((group, role), set()).add(pair) + group_keys.setdefault(group, set()).add(decision["group_key"]) + observations += 1 + assert observed_roles == roles + assert observations > len(bindings) + assert all(len(pairs) == 1 for pairs in bindings.values()) + assert all(len(keys) == 1 for keys in group_keys.values()) + assert len({next(iter(keys)) for keys in group_keys.values()}) == len(group_keys) + assert any(row["binding_reused"] for row in decisions.values()) + assert report["first_admission_prefix_cache_reused_ratio"] > 0 + records = report["per_request"] + actual_reused_tokens = sum(row["reused_input_tokens"] for row in records) + assert actual_reused_tokens > 0 + # Per-worker tokenization runs after placement. Other requests can grow the + # physical cache before admission, beyond the router observation at placement. + # Without that delay, check the existing fixture's overlap bound. + assert all( + row["reused_input_tokens"] <= route["selected_overlap_blocks"] * block_size + for row in records + for route in row["routing_history"] + if route["pool"] in ("agg", "prefill") and "tokenized_ms" not in row + ) + assert report["first_admission_prefix_cache_reused_ratio"] == pytest.approx( + actual_reused_tokens / sum(row["input_length"] for row in records) + ) + + +@pytest.mark.parametrize("backend", ["vllm", "sglang"]) +@pytest.mark.parametrize("topology", ["aggregated", "disaggregated"]) +@pytest.mark.parametrize("mode", ["session", "sibling_group"]) +def test_conversation_cli_calls_dynamo_policy(tmp_path, backend, topology, mode, dynamo_policy): + config = _conversation_prediction(tmp_path, backend, topology, mode) + report = _predict_conversations(tmp_path, config, "affinity") + assert report["completed_requests"] == 10 + _assert_conversation_bindings(report, topology, mode) + for worker in config["engine"]["workers"].values(): + worker["kv_cache"]["prefix_caching"] = False + cold = _predict_conversations(tmp_path, config, "cache-disabled") + assert cold["completed_requests"] == report["completed_requests"] + assert cold["first_admission_prefix_cache_reused_ratio"] == 0 + assert all(row["reused_input_tokens"] == 0 for row in cold["per_request"]) + + +@pytest.mark.parametrize("topology", ["aggregated", "disaggregated"]) +def test_conversation_tokenization_precedes_engine_admission(tmp_path, topology, dynamo_policy): + config = _conversation_prediction(tmp_path, "vllm", topology, "session") + config["engine"]["tokenization"] = { + "workers": 1, + "base_ms": 3.0, + "per_input_token_us": 500.0, + } + report = _predict_conversations(tmp_path, config, "tokenization") + assert report["completed_requests"] == 10 + _assert_conversation_bindings(report, topology, "session") + for record in report["per_request"]: + service_ms = 3.0 + 0.5 * record["input_length"] + assert record["tokenized_ms"] - record["arrival_time_ms"] >= service_ms + assert record["first_admit_ms"] >= record["tokenized_ms"] + + +@pytest.mark.parametrize("backend", ["vllm", "sglang"]) +def test_conversation_cli_reuses_actual_decode_destination_cache(tmp_path, backend, dynamo_policy): + config = _conversation_prediction(tmp_path, backend, "disaggregated", "session") + trace = Path(config["traffic"]["source"]["paths"][0]) + rows = [json.loads(line) for line in trace.read_text().splitlines()] + first = rows[1] + repeated = { + **first, + "request_id": "first-repeat", + "not_before_ms": 1000, + "dependencies": [ + { + "request_id": first["request_id"], + "relation": "sequence", + "trigger": "completion", + "delay_ms": 1000, + } + ], + } + trace.write_text("".join(json.dumps(row) + "\n" for row in [rows[0], first, repeated])) + config["engine"]["kv_transfer"] = { + "bandwidth_gb_per_second": 1, + "timing_mode": "destination_missing", + } + missing = _predict_conversations(tmp_path, config, "destination-missing") + _assert_conversation_bindings(missing, "disaggregated", "session") + config["engine"]["kv_transfer"]["timing_mode"] = "full_prompt" + full = _predict_conversations(tmp_path, config, "full-prompt") + config["engine"]["kv_transfer"]["timing_mode"] = "destination_missing" + for worker in config["engine"]["workers"].values(): + worker["kv_cache"]["prefix_caching"] = False + cold = _predict_conversations(tmp_path, config, "uncached-destination") + + def transfer_span(report): + assert report["completed_requests"] == 2 + record = max(report["per_request"], key=lambda row: row["arrival_time_ms"]) + return record["destination_activated_ms"] - record["destination_reserved_ms"] + + # This reads actual engine handoff timestamps. A policy overlap estimate or + # the presence of an affinity configuration cannot satisfy these assertions. + assert transfer_span(missing) < transfer_span(full) - 1e-6 + assert transfer_span(cold) == pytest.approx(transfer_span(full), abs=1e-6) + assert all(row["reused_input_tokens"] == 0 for row in cold["per_request"]) + + +@pytest.mark.parametrize("backend", ["vllm", "sglang"]) +@pytest.mark.parametrize("topology", ["aggregated", "disaggregated"]) +@pytest.mark.parametrize("mode", ["session", "sibling_group"]) +def test_conversation_duration_reuses_snapshot_warmup_pipeline(tmp_path, backend, topology, mode, dynamo_policy): + config = _conversation_prediction(tmp_path, backend, topology, mode) + config["traffic"]["load"].update( + agentic_lanes=2, + agentic_snapshot={"seed": 42}, + agentic_warmup=True, + agentic_profile={"duration_seconds": 0.4}, + ) + report = _predict_conversations(tmp_path, config, "duration") + _assert_conversation_bindings(report, topology, mode) + profile = report["agentic_profile"] + phases = report["agentic_phases"] + assert profile["admission_closed"] + assert profile["profile_start_ms"] == phases["profile_start_ms"] + assert profile["admission_cutoff_ms"] - profile["profile_start_ms"] == pytest.approx(400) + assert profile["plays_started"] > 2 + assert profile["unsettled_server_requests"] == 0 + assert not profile["cancel_drain_timed_out"] + # Warmup requests traverse the same native policy before measured traffic. + measured = {row["uuid"] for row in report["per_request"]} + prepared = {row["uuid"] for row in phases["requests"]} + assert prepared - measured + for policy in report["routing_policy"]["roles"].values(): + decisions = {row["request_id"] for row in policy["decisions"]} + assert measured | prepared == decisions + + +@pytest.mark.parametrize("topology", ["aggregated", "disaggregated"]) +@pytest.mark.parametrize("mode", ["session", "sibling_group"]) +def test_weka_agentx_cli_keeps_parent_and_child_lineage(tmp_path, topology, mode, dynamo_policy): + config = _conversation_prediction(tmp_path, "vllm", topology, mode) + config["traffic"] = { + "source": { + "type": "trace", + "format": "weka", + "block_size": 4, + "nested_timestamp_basis": "absolute", + "paths": [str(_REPO_ROOT / "tests/e2e/configs/unified_cli/fixtures/traces/weka-g2-restore.jsonl")], + }, + "load": {"type": "trace_timestamps", "agentic_lanes": 1}, + } + for worker in config["engine"]["workers"].values(): + worker["kv_cache"]["block_size"] = 4 + report = _predict_conversations(tmp_path, config, "weka") + assert report["completed_requests"] == 4 + _assert_conversation_bindings(report, topology, mode, block_size=4) + + def source_request(suffix): + return next(record for record in report["per_request"] if record["request_id"].endswith(suffix)) + + parent = source_request(":outer:0")["agentic"] + child = source_request(":outer:1:inner:0")["agentic"] + resumed = source_request(":outer:2") + independent = source_request(":outer:3")["agentic"] + # The first and resumed inputs share hashes [1, 2, 3]. The explicit subagent + # has hashes [11, 12, 13]; the final, unrelated [21, 22, 23] input lowers to + # a separate worker stream, with its own conversation root. + assert resumed["agentic"]["conversation_id"] == parent["conversation_id"] + assert resumed["agentic"]["lineage"] == parent["lineage"] + assert resumed["reused_input_tokens"] == 8 + assert len({identity["conversation_id"] for identity in (parent, child, independent)}) == 3 + assert len({identity["play_id"] for identity in (parent, child, independent)}) == 1 + assert child["lineage"]["parent_conversation_id"] == parent["conversation_id"] + assert child["lineage"]["root_conversation_id"] == parent["conversation_id"] + for root in (parent, independent): + assert root["lineage"]["root_conversation_id"] == root["conversation_id"] + assert "parent_conversation_id" not in root["lineage"] + + +@pytest.mark.parametrize("dispatch", ["dispatch_committed", "dispatch_aborted"]) +def test_policy_releases_conversation_waiters_after_dispatch(dispatch, dynamo_policy): + from aisimulate.router import _DynamoPolicy + + policy = _DynamoPolicy( + json.dumps( + { + "schema_version": 1, + "role": "aggregated", + "dp_size": 1, + "topology": [{"worker_id": 0, "scheduler_ids": [0]}], + "engine": {"rank": {"block_size": 4, "num_gpu_blocks": 128, "max_num_batched_tokens": 32}}, + "placement": {"config": {"affinity": {"mode": "session", "ttl_seconds": 1}}}, + "selector_seed": 7, + "capture_decisions": True, + } + ) + ) + + def call(op, **fields): + return json.loads(policy.call_json(json.dumps({"op": op, **fields}))) + + def place(request_id): + return call( + "place", + now_ms=0, + input_length=4, + session_id="conversation", + request={"uuid": request_id, "tokens": [1, 2, 3, 4], "max_output_tokens": 2}, + ) + + # Authored tokens without precomputed hashes exercise native hashing too. + first = place("first")["result"]["placement"] + assert first is not None + assert place("second")["result"]["placement"] is None + call(dispatch, request_id="first", now_ms=0) + assert call("cancel_pending", request_id="unknown")["next_wakeup_ms"] == 0 + released = call("advance_clock", now_ms=0) + assert [row["request_id"] for row in released["result"]] == ["second"] + assert released["pending_count"] == 0 + evidence = call("snapshot")["evidence"] + assert evidence["decisions"][-1]["binding_reused"] == (dispatch == "dispatch_committed") + call("dispatch_committed", request_id="second", now_ms=0) + call("request_terminal", request_id="first", now_ms=0) + call("request_terminal", request_id="second", now_ms=0) + assert not policy.active and not policy.staged diff --git a/tests/test_host_resources.py b/tests/test_host_resources.py index 57312e9f5..b0aa6fed0 100644 --- a/tests/test_host_resources.py +++ b/tests/test_host_resources.py @@ -856,8 +856,12 @@ def test_snapshot_and_warmup_keep_previous_accounting(tmp_path, warmup, trace_fo stacks = ("dynamo",) if trace_format == "weka" else ("engine", "dynamo") for stack in stacks: estimate = estimate_workload(bounds, stack=stack) - assert estimate.allocation_model == "trace-json-metadata-v1" - assert estimate.estimated_peak_bytes is not None + if stack == "dynamo": + assert estimate.allocation_model == "dynamo-unqualified-v1" + assert estimate.estimated_peak_bytes is None + else: + assert estimate.allocation_model == "trace-json-metadata-v1" + assert estimate.estimated_peak_bytes is not None assert estimate_workload(concrete, stack=stack) == estimate ordinary = {key: value for key, value in bounds.items() if key not in {"agentic_snapshot", "agentic_warmup"}} if trace_format == "weka": @@ -1195,8 +1199,21 @@ def test_profile_resource_admission_requires_supervision_and_one_worker(tmp_path factory.admit_wave([spec, spec]) -@pytest.mark.parametrize("oversized", ["storage", "tokens"]) -def test_profile_keeps_initial_trace_materialization_guards(tmp_path, host, monkeypatch, oversized): +@pytest.mark.parametrize( + "stack,profile,oversized", + [ + ("engine", None, "storage"), + ("engine", None, "tokens"), + ("engine", {}, "storage"), + ("engine", {}, "tokens"), + ("dynamo", {}, "storage"), + ("dynamo", {}, "tokens"), + # Finite Dynamo traces stream metadata: array storage, not file size, + # supplies their lower bound. Oversized arrays must still be refused. + ("dynamo", None, "tokens"), + ], +) +def test_profile_keeps_initial_trace_materialization_guards(tmp_path, host, monkeypatch, oversized, profile, stack): trace = tmp_path / "oversized.jsonl" if oversized == "storage": with trace.open("wb") as stream: @@ -1215,8 +1232,145 @@ def test_profile_keeps_initial_trace_materialization_guards(tmp_path, host, monk ) + "\n" ) - plan = build_plan( - {"trace_path": str(trace), "trace_format": "weka", "agentic_profile": {}}, stack="engine", host=host + workload = {"trace_path": str(trace), "trace_format": "weka"} + if profile is not None: + workload["agentic_profile"] = profile + plan = build_plan(workload, stack=stack, host=host) + assert plan["status"] == "resource_limited" + if stack == "dynamo" and profile is None: + assert plan["estimate"]["estimated_peak_bytes"] is None + assert plan["estimate"]["lower_bound_bytes"] > plan["budget"]["memory_limit_bytes"] + else: + assert plan["estimate"]["estimated_peak_bytes"] > plan["budget"]["memory_limit_bytes"] + + +@pytest.mark.parametrize( + "profile,snapshot,trace_format,expected_model", + [ + (None, False, "weka", resources.OBSERVED_TRACE_MODEL), + (None, False, "mooncake", resources.OBSERVED_TRACE_MODEL), + (None, True, "weka", "dynamo-unqualified-v1"), + ({}, False, "weka", "agentic-profile-unqualified-v1"), + ], +) +@pytest.mark.parametrize("stack", ["dynamo", "engine"]) +def test_dynamo_policy_finite_and_profile_memory_remains_unqualified( + tmp_path, host, monkeypatch, profile, snapshot, trace_format, expected_model, stack +): + from aisimulate.config_adapter import PredictionAdapterContext + from aisimulate.router import EngineRouterConfigAdapter + + trace = tmp_path / "play.json" + trace.write_text( + json.dumps( + { + "id": "play", + "block_size": 64, + "hash_id_scope": "local", + "models": ["model"], + "requests": [{"t": 0, "type": "s", "model": "model", "in": 8, "out": 1}], + } + ) + + "\n" + ) + if trace_format == "mooncake": + trace.write_text(json.dumps({"timestamp": 0, "input_length": 8, "output_length": 1}) + "\n") + workload = {"trace_path": str(trace), "trace_format": trace_format} + if profile is not None: + workload["agentic_profile"] = profile + if snapshot: + workload["agentic_snapshot"] = {"seed": 42} + native_router = stack == "engine" + if native_router and profile is None: + expected_model = "dynamo-unqualified-v1" + monkeypatch.delenv("_AISIMULATE_SUPERVISED_BUDGET", raising=False) + assert build_plan(workload, stack=stack, host=host, native_router=native_router)["status"] == "resource_limited" + monkeypatch.setenv( + "_AISIMULATE_SUPERVISED_BUDGET", + json.dumps( + { + "supervisor_pid": os.getpid(), + "memory_limit_bytes": 8 * GB, + "cpu_limit": 4, + "reserved_host_memory_bytes": GB, + } + ), + ) + plan = build_plan(workload, stack=stack, host=host, requested_parallelism=4, native_router=native_router) + assert plan["status"] == "admitted" + assert plan["effective_parallelism"] == 1 + assert plan["estimate"]["estimated_peak_bytes"] is None + assert plan["estimate"]["allocation_model"] == expected_model + monkeypatch.setattr(resources, "discover_host", lambda: host) + adapter = EngineRouterConfigAdapter().compile_prediction( + {"policy": "kv_router"}, PredictionAdapterContext(engine={}, traffic={}, evaluation={}) + ) + spec = SimpleNamespace( + workload=workload, concurrency=None, runtime_hooks=adapter.runtime_hooks if native_router else () + ) + assert guard_replay(spec, stack=stack)["estimate"] == plan["estimate"] + factory = GuardedRunnerFactory(object(), stack, ResourceConfig()) + assert factory.admit_wave([spec])["status"] == "admitted" + if native_router: + # A completed run must not qualify a later native-router parallel wave. + factory.record_peak(spec, resources.WORKER_BASELINE_BYTES) + with pytest.raises(ResourceLimitError, match="wave"): + factory.admit_wave([spec, spec]) + + +@pytest.mark.parametrize("router", [None, {}, {"policy": "round_robin"}, {"policy": "kv_router"}]) +def test_compiled_router_policy_controls_synthetic_resource_admission(monkeypatch, host, router): + from aisimulate.config_adapter import PredictionAdapterContext + from aisimulate.router import EngineRouterConfigAdapter + + hooks = () + if router is not None: + hooks = ( + EngineRouterConfigAdapter() + .compile_prediction(router, PredictionAdapterContext(engine={}, traffic={}, evaluation={})) + .runtime_hooks + ) + spec = SimpleNamespace(workload={"request_count": 4, "isl": 8, "osl": 1}, concurrency=None, runtime_hooks=hooks) + monkeypatch.setattr(resources, "discover_host", lambda: host) + monkeypatch.delenv("_AISIMULATE_SUPERVISED_BUDGET", raising=False) + if hooks: + with pytest.raises(ResourceLimitError, match="Dynamo policy"): + guard_replay(spec, stack="engine") + else: + assert guard_replay(spec, stack="engine")["estimate"]["estimated_peak_bytes"] is not None + _supervised(monkeypatch) + factory = GuardedRunnerFactory(object(), "engine", ResourceConfig()) + if hooks: + with pytest.raises(ResourceLimitError, match="wave"): + factory.admit_wave([spec, spec]) + else: + assert factory.admit_wave([spec, spec])["status"] == "admitted" + + +def test_cli_native_router_override_requires_supervision_before_execution(tmp_path, monkeypatch, host): + raw = _config() + raw.pop("optimizer") + raw.pop("optimization") + raw["traffic"]["load"]["concurrency"] = 4 + raw["router"] = {"policy": "round_robin"} + config = tmp_path / "case.yaml" + config.write_text(yaml.safe_dump(raw)) + output = tmp_path / "output" + monkeypatch.setattr(cli, "resolve_runner_factory", lambda _: object()) + monkeypatch.setattr(resources, "discover_host", lambda: host) + monkeypatch.delenv("_AISIMULATE_SUPERVISED_BUDGET", raising=False) + assert ( + cli.main(["predict", "--config", str(config), "--output-dir", str(output), "--set", "router.policy=kv_router"]) + == 3 ) + plan = json.loads((output / "resource-plan.json").read_text()) + assert plan["stack"] == "engine" + assert plan["estimate"]["allocation_model"] == "dynamo-unqualified-v1" + assert plan["estimate"]["estimated_peak_bytes"] is None + + +def test_native_router_keeps_initial_materialization_refusal(host, monkeypatch): + _supervised(monkeypatch) + plan = build_plan({"request_count": 1, "isl": 10**12}, stack="engine", host=host, native_router=True) assert plan["status"] == "resource_limited" assert plan["estimate"]["estimated_peak_bytes"] > plan["budget"]["memory_limit_bytes"] diff --git a/tests/test_runner.py b/tests/test_runner.py index 8437faf4b..79d6b7564 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -136,7 +136,7 @@ def test_legacy_replay_config_preserves_execution_mode() -> None: assert config.to_replay_spec().execution_mode == "online" -def test_factory_is_pickleable_and_advertises_engine_only_capabilities(): +def test_factory_is_pickleable_and_advertises_native_router_capability(): factory = pickle.loads(pickle.dumps(EngineReplayRunnerFactory())) capabilities = factory.capabilities() @@ -145,7 +145,8 @@ def test_factory_is_pickleable_and_advertises_engine_only_capabilities(): assert capabilities.supports_backend_topology("trtllm", "disagg") assert capabilities.supports_disaggregated_attention_dp assert capabilities.supported_execution_modes == ("offline",) - assert capabilities.supported_hooks == () + assert capabilities.supports_hook(RuntimeHookSpec("engine.router", "placement_policy", 1)) + assert not capabilities.supports_hook(RuntimeHookSpec("dynamo.router", "placement_policy", 1)) assert capabilities.supports_trace_format("weka") assert capabilities.supports_trace_format("agentic_mooncake") assert capabilities.supports_agentic_lanes @@ -1806,6 +1807,20 @@ def test_engine_runner_does_not_silently_parse_a_dynamo_trace_as_mooncake(): ) +@pytest.mark.parametrize("provider", ["engine.router", "dynamo.router"]) +def test_engine_runner_rejects_unhandled_router_config_without_execution(provider): + runtime = RecordingRuntime() + spec = _spec(adapters={provider: AdapterReplaySpec(config={"policy": "kv_router"})}) + error, message = ( + (ValueError, "runtime hooks must match") + if provider == "engine.router" + else (InvalidRunnerError, "does not support component adapters") + ) + with pytest.raises(error, match=message): + EngineReplayRunnerFactory(runtime=runtime).create(0).run(spec) + assert runtime.execution_spec is None + + def test_engine_runner_rejects_dynamo_runtime_hooks(): hook = RuntimeHookSpec( provider="dynamo.router",