From 12ea1e7febd868d931cc3c38c618704f4337e755 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 10:27:54 -0700 Subject: [PATCH 01/13] feat: import reference AgentX snapshot selections Signed-off-by: hongkuanz --- THIRD_PARTY_NOTICES.md | 19 ++ crates/core/src/python.rs | 2 +- crates/core/src/replay/agg.rs | 4 +- crates/core/src/replay/disagg_tests.rs | 2 +- crates/core/src/replay/loadgen/driver.rs | 4 +- .../replay/loadgen/dynamo_snapshot_tests.rs | 4 +- crates/core/src/replay/loadgen/mod.rs | 2 +- crates/core/src/replay/loadgen/phase.rs | 4 +- crates/core/src/replay/loadgen/snapshot.rs | 75 +++++++- .../core/src/replay/loadgen/snapshot_tests.rs | 167 ++++++++++++++++-- crates/core/src/replay/loadgen/tests.rs | 4 +- crates/core/src/replay/snapshot_tests.rs | 16 +- crates/core/tests/agentx_qualification.rs | 4 +- crates/tests/public-api/src/lib.rs | 2 +- docs/replay/workloads.md | 61 ++++++- python/aisimulate/THIRD_PARTY_NOTICES.md | 19 ++ .../src/aisimulate/config/traffic.py | 31 +++- .../src/aisimulate/sweeper/replay.py | 12 +- .../export_agentx_snapshots.py | 131 ++++++++++++++ tests/test_agentx_snapshot_export.py | 73 ++++++++ tests/test_cli_config.py | 42 +++++ 21 files changed, 626 insertions(+), 52 deletions(-) create mode 100644 scripts/prediction_regression/export_agentx_snapshots.py create mode 100644 tests/test_agentx_snapshot_export.py diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 9553ef2ed..11ebdfd95 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -13,6 +13,25 @@ distributed with those packages. Unless otherwise stated, AISimulate file paths in this document are relative to `python/aisimulate/` in the repository source tree. +## AIPerf AgentX snapshot oracle + +The reference sampling vectors and recycling contract in repository-root +`crates/core/src/replay/loadgen/snapshot_tests.rs` and +`docs/replay/workloads.md`, plus the invocation contract in +`scripts/prediction_regression/export_agentx_snapshots.py`, are derived from +`src/aiperf/timing/trajectory_source.py` and +`src/aiperf/timing/strategies/agentic_replay.py` in SemiAnalysisAI/aiperf at +commit `754356e9a39acc6cc6afb242d123bb57c3fb6f75`: +https://github.com/SemiAnalysisAI/aiperf/tree/754356e9a39acc6cc6afb242d123bb57c3fb6f75 + +Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +Licensed under Apache-2.0. The upstream license is at +https://github.com/SemiAnalysisAI/aiperf/blob/754356e9a39acc6cc6afb242d123bb57c3fb6f75/LICENSE +and the full license is included in this repository as `LICENSE`. +Modified by NVIDIA: executed the pinned sampler on a small synthetic corpus +to freeze expected cuts, and documented its sequential recycling behavior. +No upstream implementation is vendored; AISimulate consumes explicit cuts. + ## AIConfigurator The repository-root `.coderabbit.yaml` is adapted and modified from diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index f5dd07f79..1c52cbf85 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -1501,7 +1501,7 @@ fn build_agentic_driver( .agentic_lanes .context("agentic_snapshot requires positive agentic_lanes")?; // Sample recorded time before applying speedup to remaining timers. - let prepared = graph.into_prepared_snapshots(lanes, *options)?; + let prepared = graph.into_prepared_snapshots(lanes, options.clone())?; let mut driver = if traffic.agentic_warmup { WorkloadDriver::new_agentic_warmup(prepared, engine_block_size, true, speedup) } else { diff --git a/crates/core/src/replay/agg.rs b/crates/core/src/replay/agg.rs index ade698a30..74338c268 100644 --- a/crates/core/src/replay/agg.rs +++ b/crates/core/src/replay/agg.rs @@ -1691,7 +1691,7 @@ mod agentic_warmup_tests { ) .unwrap(); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let play = prepared.context().prepare_play(0, 0, Some(50.0)).unwrap(); let driver = WorkloadDriver::new_agentic_warmup( @@ -1851,7 +1851,7 @@ mod agentic_warmup_tests { ) .unwrap(); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let play = prepared.context().prepare_play_from_start(0, 0).unwrap(); let driver = WorkloadDriver::new_agentic_snapshots( diff --git a/crates/core/src/replay/disagg_tests.rs b/crates/core/src/replay/disagg_tests.rs index fe4e45a27..8652d07fc 100644 --- a/crates/core/src/replay/disagg_tests.rs +++ b/crates/core/src/replay/disagg_tests.rs @@ -3116,7 +3116,7 @@ mod agentic_pd_qualification { fn prepared(graph: AgenticTrace, lanes: usize) -> PreparedAgenticSnapshots { let prepared = graph - .prepare_snapshots(lanes, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(lanes, AgenticSnapshotOptions::seeded(42)) .unwrap(); PreparedAgenticSnapshots::from_plays( (0..lanes) diff --git a/crates/core/src/replay/loadgen/driver.rs b/crates/core/src/replay/loadgen/driver.rs index 46a963b34..c05de5cab 100644 --- a/crates/core/src/replay/loadgen/driver.rs +++ b/crates/core/src/replay/loadgen/driver.rs @@ -905,7 +905,7 @@ impl WorkloadDriver { start_ms: None, last_advanced_ms: 0.0, cutoff: false, - cursor: u64::try_from(snapshots.len())?, + cursor: context.initial_corpus_cursor(), active_plays, next_ordinals, replacements_at_ms: (None, 0), @@ -4955,7 +4955,7 @@ mod tests { ) .unwrap(); let prepared = graph - .prepare_snapshots(lanes, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(lanes, AgenticSnapshotOptions::seeded(42)) .unwrap(); // Tests start at the source origin, independently of initial cut sampling. let plays = (0..lanes) diff --git a/crates/core/src/replay/loadgen/dynamo_snapshot_tests.rs b/crates/core/src/replay/loadgen/dynamo_snapshot_tests.rs index b76240b62..0ef8c5685 100644 --- a/crates/core/src/replay/loadgen/dynamo_snapshot_tests.rs +++ b/crates/core/src/replay/loadgen/dynamo_snapshot_tests.rs @@ -55,7 +55,7 @@ fn graph() -> ValidatedAgenticGraph { fn dynamo_snapshot_uses_original_intervals_for_history_and_frontier() { let graph = graph(); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let sampled = &prepared.snapshots()[0]; assert!((500.0..1500.0).contains(&sampled.t_star_ms)); @@ -125,7 +125,7 @@ fn dynamo_snapshot_uses_original_intervals_for_history_and_frontier() { #[test] fn dynamo_snapshot_live_dependencies_use_actual_completion_after_the_cut() { let prepared = graph() - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 0 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(0)) .unwrap(); for cut_ms in [600.0, 1000.0] { for speedup in [1.0, 2.0] { diff --git a/crates/core/src/replay/loadgen/mod.rs b/crates/core/src/replay/loadgen/mod.rs index b07986d30..8a9bc1c6b 100644 --- a/crates/core/src/replay/loadgen/mod.rs +++ b/crates/core/src/replay/loadgen/mod.rs @@ -31,7 +31,7 @@ pub use profile::{AgenticIdleShift, AgenticProfileOptions, AgenticProfileReport} pub use snapshot::{ AGENTIC_SNAPSHOT_SCHEMA_V1, AgenticPlaySnapshot, AgenticPrimer, AgenticReplayContext, AgenticSnapshotEvidence, AgenticSnapshotOptions, AgenticSnapshotRequest, - PreparedAgenticSnapshots, + AgenticSnapshotSelection, PreparedAgenticSnapshots, }; pub use steppable::{EngineEvent, StepOutcome, SteppableAgg, SteppableEngine, SteppableReplay}; pub use trace::{AgenticGraphBuilder, load_agentic_mooncake, validate_trace_files}; diff --git a/crates/core/src/replay/loadgen/phase.rs b/crates/core/src/replay/loadgen/phase.rs index 33e6db05f..23fd95729 100644 --- a/crates/core/src/replay/loadgen/phase.rs +++ b/crates/core/src/replay/loadgen/phase.rs @@ -680,7 +680,7 @@ mod tests { ) .unwrap(); let sampled = graph - .prepare_snapshots(lanes, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(lanes, AgenticSnapshotOptions::seeded(42)) .unwrap(); PreparedAgenticSnapshots::from_plays( (0..lanes) @@ -981,7 +981,7 @@ mod tests { ) .unwrap(); let snapshots = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 17 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(17)) .unwrap(); let original = snapshots.plays[0] .materialize_prefix("prefill-only", 126) diff --git a/crates/core/src/replay/loadgen/snapshot.rs b/crates/core/src/replay/loadgen/snapshot.rs index 98f3631eb..205c3372f 100644 --- a/crates/core/src/replay/loadgen/snapshot.rs +++ b/crates/core/src/replay/loadgen/snapshot.rs @@ -24,10 +24,32 @@ use crate::replay::AgenticRuntimeIdentity; pub const AGENTIC_SNAPSHOT_SCHEMA_V1: &str = "aisimulate.agentic.snapshot.v1"; -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AgenticSnapshotSelection { + pub source_play_id: String, + pub t_star_ms: f64, +} + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct AgenticSnapshotOptions { pub seed: u64, + /// Exact, externally sampled cuts in lane order. Unset preserves v1 draws. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub initial_selections: Option>, + /// Next draw from the shared sequential corpus sampler, including skipped draws. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub corpus_cursor: Option, +} + +impl AgenticSnapshotOptions { + pub fn seeded(seed: u64) -> Self { + Self { + seed, + ..Self::default() + } + } } /// Recorded provenance and the independent benchmark request boundary. @@ -80,6 +102,7 @@ pub struct AgenticReplayContext { graph_identity: AgenticGraphIdentity, pub(super) lanes: usize, seed: u64, + initial_corpus_cursor: u64, index_by_id: FxHashMap, node_to_play: Vec, topological_rank: Vec, @@ -129,6 +152,18 @@ impl ValidatedAgenticGraph { if lanes == 0 || self.plays.is_empty() { bail!("agentic snapshots require positive lanes and a nonempty corpus"); } + match (&options.initial_selections, options.corpus_cursor) { + (Some(selections), Some(cursor)) => { + if selections.len() != lanes { + bail!("initial_selections must contain exactly one selection per lane"); + } + if cursor < lanes as u64 { + bail!("corpus_cursor cannot precede the initial lane draws"); + } + } + (None, None) => {} + _ => bail!("initial_selections and corpus_cursor must be specified together"), + } let index_by_id = self .nodes .iter() @@ -222,6 +257,7 @@ impl ValidatedAgenticGraph { graph: self, lanes, seed: options.seed, + initial_corpus_cursor: options.corpus_cursor.unwrap_or(lanes as u64), outputs, index_by_id, node_to_play, @@ -229,9 +265,36 @@ impl ValidatedAgenticGraph { recorded_intervals, token_stride, }); - let plays = (0..lanes) - .map(|lane| context.prepare_play(lane, 0, None)) - .collect::>>()?; + let plays = if let Some(selections) = options.initial_selections { + let play_by_id = context + .graph + .plays + .iter() + .enumerate() + .map(|(index, play)| (play.play_id.as_str(), index)) + .collect::>(); + selections + .into_iter() + .enumerate() + .map(|(lane, selection)| { + let source = *play_by_id + .get(selection.source_play_id.as_str()) + .with_context(|| { + format!("unknown snapshot source play {}", selection.source_play_id) + })?; + context.prepare_play_with_cut( + lane, + 0, + PlayCut::Absolute(selection.t_star_ms), + Some(source), + ) + }) + .collect::>>()? + } else { + (0..lanes) + .map(|lane| context.prepare_play(lane, 0, None)) + .collect::>>()? + }; PreparedAgenticSnapshots::from_plays(plays) } } @@ -261,6 +324,10 @@ impl AgenticReplayContext { self.lanes } + pub fn initial_corpus_cursor(&self) -> u64 { + self.initial_corpus_cursor + } + /// Prepare a deterministic corpus-cycle instance. `None` samples an initial /// cut; `Some(cut_ms)` uses an absolute recorded timestamp, so `Some(0.0)` /// starts a play only when its recorded clock begins at zero. Use diff --git a/crates/core/src/replay/loadgen/snapshot_tests.rs b/crates/core/src/replay/loadgen/snapshot_tests.rs index 64cbb31e0..3808209b5 100644 --- a/crates/core/src/replay/loadgen/snapshot_tests.rs +++ b/crates/core/src/replay/loadgen/snapshot_tests.rs @@ -82,7 +82,7 @@ fn request<'a>(play: &'a AgenticPlaySnapshot, source: &str) -> &'a AgenticSnapsh fn prepare(graph: &ValidatedAgenticGraph, cut: f64) -> AgenticPlaySnapshot { graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap() .context() .prepare_play(0, 0, Some(cut)) @@ -141,7 +141,7 @@ fn full_play_start_handles_nonzero_source_clocks_without_changing_absolute_cuts( // Weka normalizes its nonzero source clock; Mooncake retains absolute starts. for (graph, first_ms) in [(mooncake, 10_000.0), (weka, 0.0)] { let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let context = prepared.context(); let play = context.prepare_play_from_start(0, 1).unwrap(); @@ -402,10 +402,10 @@ fn cycle_graph() -> ValidatedAgenticGraph { fn seeded_snapshots_are_repeatable_and_cycle_lanes_without_completion_order() { let graph = cycle_graph(); let first = graph - .prepare_snapshots(3, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(3, AgenticSnapshotOptions::seeded(42)) .unwrap(); let second = graph - .prepare_snapshots(3, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(3, AgenticSnapshotOptions::seeded(42)) .unwrap(); assert_eq!( serde_json::to_vec(first.snapshots()).unwrap(), @@ -426,7 +426,7 @@ fn seeded_snapshots_are_repeatable_and_cycle_lanes_without_completion_order() { .all(|snapshot| (25.0..75.0).contains(&snapshot.t_star_ms)) ); let changed = graph - .prepare_snapshots(3, AgenticSnapshotOptions { seed: 43 }) + .prepare_snapshots(3, AgenticSnapshotOptions::seeded(43)) .unwrap(); assert_ne!( first.snapshots()[0].t_star_ms, @@ -466,7 +466,7 @@ fn seeded_snapshots_are_repeatable_and_cycle_lanes_without_completion_order() { fn single_start_and_missing_or_zero_api_duration_preserve_provenance() { let single = graph(vec![row("only", "main", 0.0, None)]); let prepared = single - .prepare_snapshots(2, AgenticSnapshotOptions { seed: 3 }) + .prepare_snapshots(2, AgenticSnapshotOptions::seeded(3)) .unwrap(); for snapshot in prepared.snapshots() { assert_eq!(snapshot.t_star_ms, 0.0); @@ -490,11 +490,150 @@ fn single_start_and_missing_or_zero_api_duration_preserve_provenance() { } } +// Sampling oracle: the unmodified _seed_for_trace_lane and +// TrajectorySource._build_timestamped_trajectory methods in +// SemiAnalysisAI/aiperf@754356e9a39acc6cc6afb242d123bb57c3fb6f75, +// src/aiperf/timing/trajectory_source.py (Apache-2.0; NVIDIA copyright). +// These expected results were generated by executing those methods with +// NumPy, seed 42, request-start bounds [0, 100], and range [0.25, 0.75]. +#[test] +fn imported_reference_cuts_preserve_lane_frontiers_and_corpus_cursor() { + let graph = cycle_graph(); + let options: AgenticSnapshotOptions = serde_json::from_value(serde_json::json!({ + "seed": 42, + "initial_selections": [ + {"source_play_id": "a", "t_star_ms": 40.56078627883167}, + {"source_play_id": "b", "t_star_ms": 61.46074327942309}, + {"source_play_id": "a", "t_star_ms": 32.230547630728225} + ], + "corpus_cursor": 3 + })) + .unwrap(); + let prepared = graph.prepare_snapshots(3, options.clone()).unwrap(); + for (actual, selected) in prepared + .snapshots() + .iter() + .zip(options.initial_selections.unwrap()) + { + assert_eq!(actual.source_play_id, selected.source_play_id); + assert_eq!(actual.t_star_ms.to_bits(), selected.t_star_ms.to_bits()); + assert!( + actual + .requests + .iter() + .all(|request| request.historical + == (request.recorded_start_ms < selected.t_star_ms)) + ); + } + let context = prepared.context(); + assert_eq!(context.initial_corpus_cursor(), 3); + // Reference recycling uses the SAME sequential sampler and starts at turn + // zero: it does not draw a new randomized cut when a lane is reused. + let recycled = context + .prepare_source_play_from_start( + (context.initial_corpus_cursor() % context.source_play_count() as u64) as usize, + 0, + 1, + ) + .unwrap(); + assert_eq!(recycled.evidence().source_play_id, "b"); + assert_eq!(recycled.evidence().t_star_ms, 0.0); + assert!( + recycled + .evidence() + .requests + .iter() + .all(|request| !request.historical) + ); + assert_ne!( + recycled.evidence().cache_id, + prepared.snapshots()[0].cache_id + ); + // A recorded skipped draw must survive preparation rather than resetting + // the shared cursor to lane count. + let options = serde_json::from_value(serde_json::json!({ + "seed": 42, + "initial_selections": [{"source_play_id": "b", "t_star_ms": 100.0}], + "corpus_cursor": 4 + })) + .unwrap(); + let prepared = graph.prepare_snapshots(1, options).unwrap(); + assert_eq!(prepared.context().initial_corpus_cursor(), 4); + let mut driver = WorkloadDriver::new_agentic_snapshots(prepared, 64, true, 1.0).unwrap(); + driver + .enable_agentic_profile(AgenticProfileOptions::default()) + .unwrap(); + let first = driver.pop_ready(0.0, 1).pop().unwrap(); + assert_eq!(first.authored_request_id.as_deref(), Some("b1")); + driver + .on_terminal( + first.request_uuid, + 1.0, + crate::replay::ReplayTerminalStatus::Completed, + ) + .unwrap(); + assert!(driver.advance_agentic_profile(1.0).unwrap()); + let next = driver.pop_ready(1.0, 1).pop().unwrap(); + assert_eq!(next.authored_request_id.as_deref(), Some("a0")); + assert_eq!(driver.agentic_profile_report().unwrap().corpus_cursor, 5); + assert_ne!(first.request_uuid, next.request_uuid); +} + +#[test] +fn imported_snapshot_selection_is_strict_and_never_falls_back_to_sampling() { + let graph = cycle_graph(); + let valid = serde_json::json!({ + "seed": 42, + "initial_selections": [{"source_play_id": "a", "t_star_ms": 50.0}], + "corpus_cursor": 1 + }); + for field in ["initial_selections", "corpus_cursor"] { + let mut value = valid.clone(); + value.as_object_mut().unwrap().remove(field); + assert!( + graph + .prepare_snapshots(1, serde_json::from_value(value).unwrap()) + .is_err() + ); + } + for (path, value) in [ + ("source_play_id", serde_json::json!("missing")), + ("t_star_ms", serde_json::json!(-1.0)), + ("t_star_ms", serde_json::json!(101.0)), + ] { + let mut options = valid.clone(); + options["initial_selections"][0][path] = value; + assert!( + graph + .prepare_snapshots(1, serde_json::from_value(options).unwrap()) + .is_err() + ); + } + assert!( + graph + .prepare_snapshots(2, serde_json::from_value(valid.clone()).unwrap()) + .is_err() + ); + let mut options = valid; + options["corpus_cursor"] = serde_json::json!(0); + assert!( + graph + .prepare_snapshots(1, serde_json::from_value(options).unwrap()) + .is_err() + ); + let legacy: AgenticSnapshotOptions = + serde_json::from_value(serde_json::json!({"seed": 42})).unwrap(); + assert_eq!( + serde_json::to_value(legacy).unwrap(), + serde_json::json!({"seed": 42}) + ); +} + #[test] fn snapshots_retain_full_graph_outputs_and_logical_prompt_views() { let graph = boundary_graph(); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 8 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(8)) .unwrap(); let zero = prepared.context().prepare_play(0, 0, Some(0.0)).unwrap(); let cut = prepared.context().prepare_play(0, 0, Some(100.0)).unwrap(); @@ -583,7 +722,7 @@ fn encounter_allocated_prompt_preserves_full_width_source_ids_and_lane_offsets() authored.hash_ids = Some(vec![u64::MAX - 1, 0, 1_u64 << 32, u64::MAX, 7]); let graph = graph(vec![authored]); let prepared = graph - .prepare_snapshots(2, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(2, AgenticSnapshotOptions::seeded(42)) .unwrap(); let play = prepared.context().prepare_play(0, 0, Some(0.0)).unwrap(); let mut expected = [vec![0; 64], vec![1; 64], vec![2; 64], vec![3; 64]].concat(); @@ -601,11 +740,11 @@ fn invalid_snapshot_inputs_and_exhausted_identity_ranges_fail_before_execution() let graph = cycle_graph(); assert!( graph - .prepare_snapshots(0, AgenticSnapshotOptions { seed: 1 }) + .prepare_snapshots(0, AgenticSnapshotOptions::seeded(1)) .is_err() ); let prepared = graph - .prepare_snapshots(2, AgenticSnapshotOptions { seed: 1 }) + .prepare_snapshots(2, AgenticSnapshotOptions::seeded(1)) .unwrap(); let context = prepared.context(); for cut in [-1.0, 101.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { @@ -625,7 +764,7 @@ fn invalid_snapshot_inputs_and_exhausted_identity_ranges_fail_before_execution() assert!(PreparedAgenticSnapshots::from_plays(Vec::new()).is_err()); assert!(PreparedAgenticSnapshots::from_plays(vec![play.clone(), play.clone()]).is_err()); let independent = graph - .prepare_snapshots(2, AgenticSnapshotOptions { seed: 1 }) + .prepare_snapshots(2, AgenticSnapshotOptions::seeded(1)) .unwrap(); let other_context = independent.context().prepare_play(1, 0, Some(0.0)).unwrap(); assert!(PreparedAgenticSnapshots::from_plays(vec![play.clone(), other_context]).is_err()); @@ -648,7 +787,7 @@ fn overflowing_recorded_intervals_and_dependency_deadlines_are_rejected() { ]); assert!( interval - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 1 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(1)) .is_err() ); let root = row("first", "main", 0.0, Some(1.0e308)); @@ -662,7 +801,7 @@ fn overflowing_recorded_intervals_and_dependency_deadlines_are_rejected() { let deadline = graph(vec![root, last]); assert!( deadline - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 1 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(1)) .is_err() ); } @@ -679,7 +818,7 @@ fn historical_request_cannot_depend_on_a_request_not_started_at_the_cut() { )); let graph = graph(vec![source, historical]); let error = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 1 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(1)) .unwrap_err(); assert!( error diff --git a/crates/core/src/replay/loadgen/tests.rs b/crates/core/src/replay/loadgen/tests.rs index 77b4486f9..b7c715a86 100644 --- a/crates/core/src/replay/loadgen/tests.rs +++ b/crates/core/src/replay/loadgen/tests.rs @@ -356,7 +356,7 @@ fn legacy_snapshot_preserves_recorded_starts_without_adding_execution_gates() { assert!(cold.is_drained()); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 0 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(0)) .unwrap(); let snapshot = &prepared.snapshots()[0]; assert!((1250.0..1750.0).contains(&snapshot.t_star_ms)); @@ -419,7 +419,7 @@ fn legacy_snapshot_retains_missing_times_and_rejects_invalid_recorded_starts() { cold.on_complete(ready.request_uuid, now_ms + 10.0).unwrap(); } assert!(cold.is_drained()); - let prepared = graph.prepare_snapshots(1, AgenticSnapshotOptions { seed: 0 }); + let prepared = graph.prepare_snapshots(1, AgenticSnapshotOptions::seeded(0)); if timestamp.is_some() { assert!( prepared diff --git a/crates/core/src/replay/snapshot_tests.rs b/crates/core/src/replay/snapshot_tests.rs index 1fa49bb1d..a03d2a5be 100644 --- a/crates/core/src/replay/snapshot_tests.rs +++ b/crates/core/src/replay/snapshot_tests.rs @@ -189,7 +189,7 @@ fn qualify_prefix(case: PrefixCase) { .request_id() .to_string(); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let play = prepared.context().prepare_play(0, 0, Some(500.0)).unwrap(); let next_play = prepared.context().prepare_play(0, 1, Some(500.0)).unwrap(); @@ -539,7 +539,7 @@ fn frontier_graph() -> ValidatedAgenticGraph { fn snapshot_consumes_recorded_active_history_and_drains_live_child_join_without_primers() { for backend in [Backend::Vllm, Backend::Sglang] { let prepared = frontier_graph() - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let play = prepared.context().prepare_play(0, 0, Some(500.0)).unwrap(); let evidence = play.evidence(); @@ -641,7 +641,7 @@ fn snapshot_consumes_recorded_active_history_and_drains_live_child_join_without_ fn snapshot_seeded_initial_lanes_use_separate_cache_identities_in_the_same_engine() { for backend in [Backend::Vllm, Backend::Sglang] { let prepared = frontier_graph() - .prepare_snapshots(2, AgenticSnapshotOptions { seed: 2026 }) + .prepare_snapshots(2, AgenticSnapshotOptions::seeded(2026)) .unwrap(); let evidence = prepared.snapshots().to_vec(); assert_eq!(evidence.len(), 2); @@ -696,7 +696,7 @@ fn snapshot_seeded_initial_lanes_use_separate_cache_identities_in_the_same_engin fn snapshot_request_start_at_the_cut_stays_live_and_matches_turn_zero_tokens() { for backend in [Backend::Vllm, Backend::Sglang] { let prepared = frontier_graph() - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let zero = prepared.context().prepare_play(0, 0, Some(0.0)).unwrap(); let cut = prepared.context().prepare_play(0, 0, Some(800.0)).unwrap(); @@ -767,7 +767,7 @@ fn snapshot_original_prompt_has_fixed_token_and_native_hash_vectors() { .request_id() .to_string(); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let play = prepared.context().prepare_play(0, 0, Some(500.0)).unwrap(); // This one original prompt introduces four complete source units and one @@ -921,7 +921,7 @@ fn snapshot_new_incarnation_can_finish_before_old_incarnation_without_identity_o .request_id() .to_string(); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let old = prepared.context().prepare_play(0, 0, Some(0.0)).unwrap(); let new = prepared.context().prepare_play(0, 1, Some(0.0)).unwrap(); @@ -993,7 +993,7 @@ fn snapshot_new_incarnation_can_finish_before_old_incarnation_without_identity_o fn warmup_preserves_snapshot_spawn_join_timers_at_nonunit_speedup() { for backend in [Backend::Vllm, Backend::Sglang] { let prepared = frontier_graph() - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let play = prepared.context().prepare_play(0, 0, Some(500.0)).unwrap(); let make = || PreparedAgenticSnapshots::from_plays(vec![play.clone()]).unwrap(); @@ -1046,7 +1046,7 @@ fn warmup_preserves_snapshot_spawn_join_timers_at_nonunit_speedup() { fn warmup_initial_lanes_keep_native_cache_identities_separate() { for backend in [Backend::Vllm, Backend::Sglang] { let prepared = frontier_graph() - .prepare_snapshots(2, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(2, AgenticSnapshotOptions::seeded(42)) .unwrap(); let context = prepared.context(); let plays = (0..2) diff --git a/crates/core/tests/agentx_qualification.rs b/crates/core/tests/agentx_qualification.rs index bab8021be..9bf29e7b6 100644 --- a/crates/core/tests/agentx_qualification.rs +++ b/crates/core/tests/agentx_qualification.rs @@ -408,7 +408,7 @@ fn agentx_g2_restores_evicted_parent_after_tool_gap_and_join( fn g2_snapshot_driver(warmup: bool, from_start: bool) -> WorkloadDriver { let prepared = g2_graph() - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let play = if from_start { prepared.context().prepare_play_from_start(0, 0).unwrap() @@ -555,7 +555,7 @@ fn snapshot_pin_report( disaggregated: bool, ) -> ReplayReport { let prepared = graph - .prepare_snapshots(3, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(3, AgenticSnapshotOptions::seeded(42)) .unwrap(); // "profile" profiles after warmup; "profile-cold" profiles straight from // the snapshot without warmup's timer rebasing. diff --git a/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index 06540e07a..8fd2f23b4 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -536,7 +536,7 @@ mod tests { } let graph = builder.finish().unwrap(); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); assert_eq!(prepared.snapshots().len(), 1); let evidence = &prepared.snapshots()[0]; diff --git a/docs/replay/workloads.md b/docs/replay/workloads.md index 57825ef86..e09b074eb 100644 --- a/docs/replay/workloads.md +++ b/docs/replay/workloads.md @@ -71,7 +71,9 @@ the current SA convention. | `traffic.load.fraction` | `null` | `-` | `-` | Positive finite number; `kv_capacity_fraction` only and may exceed `1`. | | `traffic.load.speedup` | `1` | `-` | `-` | Positive; trace timestamp load only. | | `traffic.load.agentic_lanes` | `null` | `x` | `-` | Positive integer; `weka`, `agentic_mooncake`, or agentic `dynamo` timestamp replay only. | -| `traffic.load.agentic_snapshot` | `null` (unset) | `x` | `-` | Optional object `{seed: u64}`; required `seed` is an unsigned 64-bit integer (`0` through `2^64 - 1`). Requires `traffic.load.type: trace_timestamps` and positive `agentic_lanes`; supported formats are `weka`, `agentic_mooncake`, and agentic `dynamo`. Unset preserves turn-zero execution. See [snapshot](#agentic-snapshot). | +| `traffic.load.agentic_snapshot` | `null` (unset) | `x` | `-` | Optional object with required unsigned 64-bit `seed` (`0` through `2^64 - 1`). Requires `traffic.load.type: trace_timestamps` and positive `agentic_lanes`; supported formats are `weka`, `agentic_mooncake`, and agentic `dynamo`. Unset preserves turn-zero execution. See [snapshot](#agentic-snapshot). | +| `traffic.load.agentic_snapshot.initial_selections` | unset | `x` | `-` | Exact initial `{source_play_id, t_star_ms}` selections in lane order, one per lane. Requires `corpus_cursor`. Unset retains the original seed-based sampler. | +| `traffic.load.agentic_snapshot.corpus_cursor` | unset | `x` | `-` | Unsigned 64-bit next sequential corpus draw, including any skipped initial draws. Requires `initial_selections` and must be at least the lane count. Recycled plays start at turn zero. | | `traffic.load.agentic_warmup` | `false` | `x` | `-` | Optional boolean; `true` requires `agentic_snapshot` and positive `agentic_lanes`. Primes caches before measurement; see [warmup](#agentic-warmup). | | `traffic.load.agentic_profile` | `null` (unset) | `x` | `-` | Optional object; `{}` enables continuous lane replenishment with the defaults below. Requires `trace_timestamps`, positive `agentic_lanes`, and `agentic_snapshot`; cannot be combined with `traffic.stop.max_virtual_time_seconds`. Unset preserves finite replay. See [continuous profile](#agentic-profile). | | `traffic.load.agentic_profile.duration_seconds` | `3600` when enabled | `x` | `-` | Positive finite admission duration, starting at the preparation barrier or simulation start without warmup. No new workload requests or replacement plays are issued after the deadline. | @@ -319,6 +321,63 @@ A snapshot is cut at request boundaries. It does not model partial decode progress or restore an engine checkpoint: without warmup, the caches start empty. +For a reference run, matching the seed alone is insufficient: AISimulate's +versioned BLAKE3 sampler is not AIPerf's NumPy sampler. Import the actual cuts +and next sequential corpus position instead: + +```yaml +agentic_snapshot: + seed: 42 + initial_selections: + - source_play_id: "" + t_star_ms: 40560.78627883167 + corpus_cursor: 1 +``` + +Selections are in lane order. Each source ID must exist in the exact input +corpus and each cut must lie within that play's request-start bounds. Cuts +are in source milliseconds, before `speedup`; they are never re-sampled or +silently replaced. Repeated source plays still get distinct lane/cache +identities. Preserve the source corpus bytes/order, reference revision and +selection artifact alongside the replay configuration. + +The sequential cursor counts source draws, not completed requests or unique +plays; it can exceed the initial lane count when the reference skipped a +source. After a tree drains, the shared cursor selects the next source play +and starts it at turn zero, with a fresh cache identity. This matches the +sequential-sampler recycling contract of +[AIPerf `754356e9`](https://github.com/SemiAnalysisAI/aiperf/blob/754356e9a39acc6cc6afb242d123bb57c3fb6f75/src/aiperf/timing/strategies/agentic_replay.py). +Random/shuffle recycling is not implemented by this import. Initial cuts do +not import an engine checkpoint, tokenizer/chat-template bytes, or a +post-warmup frontier; the warmup policy below remains unchanged. + +To produce selections without manually computing cuts, run the optional exporter +in the **pinned AIPerf environment** (its existing dependencies are required): + +```bash +python scripts/prediction_regression/export_agentx_snapshots.py \ + --aiperf-source /path/to/aiperf-at-754356e9 \ + --metadata reference-dataset-metadata.json \ + --source-play-map source-play-map.json --lanes 15 --seed 42 \ + > snapshot-options.json 2> snapshot-provenance.jsonl +``` + +`reference-dataset-metadata.json` is the reference loader's +`DatasetMetadata.model_dump_json()` output, with every conversation, branch, +and timestamp retained. `source-play-map.json` maps each root +`conversation_id` from that metadata to its normalized AISimulate +`source_play_id` (as recorded in `agentic_snapshots`). The root order in the +metadata must match the replay corpus order and each root mapping must identify +the corresponding source play; do not sort either input. The importer validates +IDs and cut bounds, but cannot infer or verify reference corpus order from a +cursor alone. A wrong mapping/order is not a valid parity comparison. +Paste the emitted JSON object under `traffic.load.agentic_snapshot`. +The exporter verifies the source revision and clean source tree, invokes the +reference `TrajectorySource` directly, and records input hashes and corpus +order on stderr. It fails on unsupported samplers, timestamp-less trajectories, +or incomplete lane selection. Neither AIPerf nor NumPy becomes an AISimulate +runtime dependency. + ### Warmup diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index 9553ef2ed..11ebdfd95 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -13,6 +13,25 @@ distributed with those packages. Unless otherwise stated, AISimulate file paths in this document are relative to `python/aisimulate/` in the repository source tree. +## AIPerf AgentX snapshot oracle + +The reference sampling vectors and recycling contract in repository-root +`crates/core/src/replay/loadgen/snapshot_tests.rs` and +`docs/replay/workloads.md`, plus the invocation contract in +`scripts/prediction_regression/export_agentx_snapshots.py`, are derived from +`src/aiperf/timing/trajectory_source.py` and +`src/aiperf/timing/strategies/agentic_replay.py` in SemiAnalysisAI/aiperf at +commit `754356e9a39acc6cc6afb242d123bb57c3fb6f75`: +https://github.com/SemiAnalysisAI/aiperf/tree/754356e9a39acc6cc6afb242d123bb57c3fb6f75 + +Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +Licensed under Apache-2.0. The upstream license is at +https://github.com/SemiAnalysisAI/aiperf/blob/754356e9a39acc6cc6afb242d123bb57c3fb6f75/LICENSE +and the full license is included in this repository as `LICENSE`. +Modified by NVIDIA: executed the pinned sampler on a small synthetic corpus +to freeze expected cuts, and documented its sequential recycling behavior. +No upstream implementation is vendored; AISimulate consumes explicit cuts. + ## AIConfigurator The repository-root `.coderabbit.yaml` is adapted and modified from diff --git a/python/aisimulate/src/aisimulate/config/traffic.py b/python/aisimulate/src/aisimulate/config/traffic.py index 47b2d3b0a..8fa89333d 100644 --- a/python/aisimulate/src/aisimulate/config/traffic.py +++ b/python/aisimulate/src/aisimulate/config/traffic.py @@ -8,7 +8,7 @@ import math from typing import Annotated, Literal -from pydantic import Field, field_validator, model_validator +from pydantic import Field, field_validator, model_serializer, model_validator from .common import Choices, IntegerRange, NumericRange, StrictModel @@ -111,10 +111,33 @@ class TrafficStop(StrictModel): max_virtual_time_seconds: PositiveFloat | None = None +class AgenticSnapshotSelection(StrictModel): + """An exact reference-selected initial cut, in lane order.""" + + source_play_id: str = Field(strict=True, min_length=1) + t_star_ms: NonNegativeFloat + + class AgenticSnapshotOptions(StrictModel): """Seed for deterministic initial AgentX request-boundary snapshots.""" seed: int = Field(strict=True, ge=0, le=0xFFFF_FFFF_FFFF_FFFF) + initial_selections: list[AgenticSnapshotSelection] | None = Field(default=None, min_length=1) + corpus_cursor: int | None = Field(default=None, strict=True, ge=0, le=0xFFFF_FFFF_FFFF_FFFF) + + @model_validator(mode="after") + def _validate_reference_selection(self): + if (self.initial_selections is None) != (self.corpus_cursor is None): + raise ValueError("initial_selections and corpus_cursor must be specified together") + if self.initial_selections is not None and self.corpus_cursor < len(self.initial_selections): + raise ValueError("corpus_cursor cannot precede the initial lane draws") + return self + + @model_serializer(mode="wrap") + def _serialize(self, handler): + result = handler(self) + # Keep legacy seed-only saved configurations byte-for-byte compatible. + return {key: value for key, value in result.items() if value is not None} class AgenticProfileOptions(StrictModel): @@ -267,6 +290,12 @@ def _validate_source_load_stop(self, load) -> None: raise ValueError("agentic_warmup requires agentic_snapshot") if load.agentic_snapshot is not None and load.agentic_lanes is None: raise ValueError("agentic_snapshot requires positive agentic_lanes") + if ( + load.agentic_snapshot is not None + and load.agentic_snapshot.initial_selections is not None + and len(load.agentic_snapshot.initial_selections) != load.agentic_lanes + ): + raise ValueError("initial_selections must contain exactly one selection per lane") if isinstance(source, TraceSource): if load.type not in {"trace_timestamps", "concurrency"}: raise ValueError("trace traffic requires trace_timestamps or concurrency load") diff --git a/python/aisimulate/src/aisimulate/sweeper/replay.py b/python/aisimulate/src/aisimulate/sweeper/replay.py index 2f1c3b70b..bb8c28ab6 100644 --- a/python/aisimulate/src/aisimulate/sweeper/replay.py +++ b/python/aisimulate/src/aisimulate/sweeper/replay.py @@ -16,7 +16,7 @@ from pydantic import BaseModel from ..config.common import ENGINE_MODEL_CONTROL_FIELDS, is_active_engine_model_control -from ..config.traffic import AgenticProfileOptions +from ..config.traffic import AgenticProfileOptions, AgenticSnapshotOptions from ..fpm_profile import load_fpm_profile from ..power import POWER_FIELDS, normalize_power_summary from .provider import AdapterReplaySpec, JSONValue, RuntimeHookSpec @@ -512,15 +512,11 @@ def require_compatible(self, spec: ReplaySpec) -> None: if not self.supports_agentic_warmup: raise ValueError("runner does not support agentic warmup") if agentic_snapshot is not None: - if ( - not isinstance(agentic_snapshot, Mapping) - or set(agentic_snapshot) != {"seed"} - or type(agentic_snapshot["seed"]) is not int - or not 0 <= agentic_snapshot["seed"] <= 0xFFFF_FFFF_FFFF_FFFF - ): - raise ValueError("agentic_snapshot requires exactly one unsigned 64-bit integer seed") + snapshot = AgenticSnapshotOptions.model_validate(agentic_snapshot) if agentic_lanes is None or spec.workload.get("load_type") != "trace_timestamps": raise ValueError("agentic_snapshot requires trace_timestamps load with positive agentic_lanes") + if snapshot.initial_selections is not None and len(snapshot.initial_selections) != agentic_lanes: + raise ValueError("initial_selections must contain exactly one selection per lane") if spec.workload.get("source_type") != "trace" or spec.workload.get("replay_concurrency") is not None: raise ValueError("agentic_snapshot requires agentic trace input without replay_concurrency") if not self.supports_agentic_snapshots: diff --git a/scripts/prediction_regression/export_agentx_snapshots.py b/scripts/prediction_regression/export_agentx_snapshots.py new file mode 100644 index 000000000..2992bd96a --- /dev/null +++ b/scripts/prediction_regression/export_agentx_snapshots.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Export exact initial selections from an installed, pinned AIPerf reference. + +This tool invokes upstream code; it does not approximate NumPy's RNG in Rust. +Run in the reference AIPerf environment, not the AISimulate runtime environment. +Source: SemiAnalysisAI/aiperf@754356e9a39acc6cc6afb242d123bb57c3fb6f75, +src/aiperf/timing/trajectory_source.py (Apache-2.0, NVIDIA copyright). +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib +import json +import subprocess +import sys +from pathlib import Path + +REFERENCE_REVISION = "754356e9a39acc6cc6afb242d123bb57c3fb6f75" + + +class SequentialCursor: + """Record every sequential sampler draw, including unspawnable samples.""" + + def __init__(self, ids): + self.ids = tuple(ids) + if not self.ids or len(set(self.ids)) != len(self.ids): + raise ValueError("reference roots must be nonempty and uniquely identified") + self.cursor = 0 + + def next_conversation_id(self): + selected = self.ids[self.cursor % len(self.ids)] + self.cursor += 1 + return selected + + +def export_snapshot(metadata, source_play_map, *, lanes, seed, trajectory_type): + """Invoke the reference selector and retain its actual shared cursor.""" + if type(lanes) is not int or lanes <= 0 or type(seed) is not int or not 0 <= seed < 2**64: + raise ValueError("lanes must be positive and seed must be an unsigned 64-bit integer") + strategy = getattr(metadata.sampling_strategy, "value", metadata.sampling_strategy) + if str(strategy).lower() != "sequential": + raise ValueError("only the pinned reference's sequential sampler is supported") + root_ids = [item.conversation_id for item in metadata.conversations if getattr(item, "is_root", True) is not False] + if set(source_play_map) != set(root_ids) or any(not isinstance(v, str) or not v for v in source_play_map.values()): + raise ValueError("source-play-map must map every reference root to its normalized AISimulate source_play_id") + if len(set(source_play_map.values())) != len(root_ids): + raise ValueError("distinct reference roots must map to distinct source plays") + sampler = SequentialCursor(root_ids) + source = trajectory_type( + dataset_metadata=metadata, + dataset_sampler=sampler, + concurrency=lanes, + random_seed=seed, + start_min_ratio=0.25, + start_max_ratio=0.75, + allow_dataset_wrap=True, + ) + if len(source.trajectories) != lanes: + raise ValueError("reference could not construct every requested lane") + if any(item.snapshot is None for item in source.trajectories): + raise ValueError("timestamp-less reference snapshots cannot be imported as time cuts") + return { + "seed": seed, + "initial_selections": [ + {"source_play_id": source_play_map[item.conversation_id], "t_star_ms": item.snapshot.t_star_ms} + for item in source.trajectories + ], + "corpus_cursor": sampler.cursor, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--aiperf-source", type=Path, required=True) + parser.add_argument( + "--metadata", type=Path, required=True, help="DatasetMetadata.model_dump_json() from the reference loader" + ) + parser.add_argument( + "--source-play-map", + type=Path, + required=True, + help="JSON object: reference root ID -> normalized AIS source_play_id", + ) + parser.add_argument("--lanes", type=int, required=True) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + checkout = args.aiperf_source.resolve() + head = subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip() + if head != REFERENCE_REVISION: + parser.error(f"reference checkout must be exactly {REFERENCE_REVISION}, got {head}") + if subprocess.check_output( + ["git", "-C", str(checkout), "status", "--porcelain", "--untracked-files=all", "--", "src/aiperf"], text=True + ).strip(): + parser.error("reference AIPerf source has local modifications") + sys.path.insert(0, str(checkout / "src")) + module = importlib.import_module("aiperf.timing.trajectory_source") + if Path(module.__file__).resolve() != checkout / "src/aiperf/timing/trajectory_source.py": + parser.error("loaded trajectory source does not belong to the verified checkout") + metadata_type = importlib.import_module("aiperf.common.models").DatasetMetadata + metadata_bytes = args.metadata.read_bytes() + mapping_bytes = args.source_play_map.read_bytes() + metadata = metadata_type.model_validate_json(metadata_bytes) + result = export_snapshot( + metadata, json.loads(mapping_bytes), lanes=args.lanes, seed=args.seed, trajectory_type=module.TrajectorySource + ) + print(json.dumps(result, indent=2, allow_nan=False)) + print( + json.dumps( + { + "reference_revision": head, + "metadata_sha256": hashlib.sha256(metadata_bytes).hexdigest(), + "source_play_map_sha256": hashlib.sha256(mapping_bytes).hexdigest(), + "source_order": [ + item.conversation_id + for item in metadata.conversations + if getattr(item, "is_root", True) is not False + ], + "sampling_strategy": "sequential", + "start_ratio": [0.25, 0.75], + } + ), + file=sys.stderr, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_agentx_snapshot_export.py b/tests/test_agentx_snapshot_export.py new file mode 100644 index 000000000..3537cb632 --- /dev/null +++ b/tests/test_agentx_snapshot_export.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace as Ns + +import pytest + +from scripts.prediction_regression.export_agentx_snapshots import SequentialCursor, export_snapshot + + +def metadata(strategy="sequential"): + return Ns( + sampling_strategy=strategy, + conversations=[ + Ns(conversation_id="skip", is_root=True), + Ns(conversation_id="a", is_root=True), + Ns(conversation_id="child", is_root=False), + Ns(conversation_id="b", is_root=True), + ], + ) + + +def test_export_keeps_actual_draw_count_and_lane_order(): + def reference(**kwargs): + assert kwargs["random_seed"] == 42 + assert (kwargs["start_min_ratio"], kwargs["start_max_ratio"]) == (0.25, 0.75) + sampler = kwargs["dataset_sampler"] + assert sampler.next_conversation_id() == "skip" + assert sampler.next_conversation_id() == "a" + assert sampler.next_conversation_id() == "b" + return Ns( + trajectories=[ + Ns(conversation_id="a", snapshot=Ns(t_star_ms=12.5)), + Ns(conversation_id="b", snapshot=Ns(t_star_ms=40.5)), + ] + ) + + result = export_snapshot(metadata(), {"skip": "s", "a": "A", "b": "B"}, lanes=2, seed=42, trajectory_type=reference) + assert result == { + "seed": 42, + "initial_selections": [{"source_play_id": "A", "t_star_ms": 12.5}, {"source_play_id": "B", "t_star_ms": 40.5}], + "corpus_cursor": 3, + } + + +def test_sequential_cursor_wraps_without_resetting_consumed_draws(): + cursor = SequentialCursor(["a", "b"]) + assert [cursor.next_conversation_id() for _ in range(5)] == ["a", "b", "a", "b", "a"] + assert cursor.cursor == 5 + + +@pytest.mark.parametrize("strategy", ["random", "shuffle"]) +def test_export_rejects_unimplemented_recycle_samplers(strategy): + with pytest.raises(ValueError, match="sequential"): + export_snapshot(metadata(strategy), {}, lanes=2, seed=42, trajectory_type=None) + + +@pytest.mark.parametrize("mapping", [{}, {"skip": "S", "a": "A", "b": "A"}, {"skip": "S", "a": "A", "b": ""}]) +def test_export_requires_complete_unambiguous_normalized_play_mapping(mapping): + with pytest.raises(ValueError, match="source.play|distinct"): + export_snapshot(metadata(), mapping, lanes=2, seed=42, trajectory_type=None) + + +@pytest.mark.parametrize("trajectories", [[], [Ns(conversation_id="a", snapshot=None)]]) +def test_export_fails_closed_for_missing_lanes_or_timestamp_less_source(trajectories): + with pytest.raises(ValueError, match="lane|timestamp"): + export_snapshot( + metadata(), + {"skip": "S", "a": "A", "b": "B"}, + lanes=1, + seed=42, + trajectory_type=lambda **_: Ns(trajectories=trajectories), + ) diff --git a/tests/test_cli_config.py b/tests/test_cli_config.py index fbe0e2d99..df3f8bd55 100644 --- a/tests/test_cli_config.py +++ b/tests/test_cli_config.py @@ -1658,6 +1658,48 @@ def test_agentic_snapshot_is_opt_in() -> None: assert "agentic_profile" not in prediction_to_replay_spec(config).workload +def test_imported_agentic_snapshot_round_trips_without_changing_legacy_serialization(): + from aisimulate.config.traffic import AgenticSnapshotOptions, TrafficPredictionConfig + + snapshot = { + "seed": 42, + "initial_selections": [{"source_play_id": "reference-play", "t_star_ms": 40.56078627883167}], + "corpus_cursor": 2, + } + config = TrafficPredictionConfig.model_validate( + { + "source": {"type": "trace", "paths": ["corpus.jsonl"], "format": "weka"}, + "load": {"type": "trace_timestamps", "agentic_lanes": 1, "agentic_snapshot": snapshot}, + } + ) + assert config.load.agentic_snapshot.model_dump() == snapshot + assert ( + AgenticSnapshotOptions.model_validate(snapshot).model_dump_json() + == config.load.agentic_snapshot.model_dump_json() + ) + assert AgenticSnapshotOptions(seed=42).model_dump() == {"seed": 42} + + +@pytest.mark.parametrize( + "snapshot", + [ + {"seed": 42, "initial_selections": []}, + {"seed": 42, "corpus_cursor": 1}, + {"seed": 42, "initial_selections": [{"source_play_id": "a", "t_star_ms": 5.0}]}, + {"seed": 42, "initial_selections": [{"source_play_id": "", "t_star_ms": 5.0}], "corpus_cursor": 1}, + {"seed": 42, "initial_selections": [{"source_play_id": "a", "t_star_ms": float("nan")}], "corpus_cursor": 1}, + {"seed": 42, "initial_selections": [{"source_play_id": "a", "t_star_ms": -1.0}], "corpus_cursor": 1}, + {"seed": 42, "initial_selections": [{"source_play_id": "a", "t_star_ms": 5.0}], "corpus_cursor": 0}, + {"seed": 42, "initial_selections": [{"source_play_id": "a", "t_star_ms": 5.0}], "corpus_cursor": True}, + ], +) +def test_imported_agentic_snapshot_rejects_invalid_metadata(snapshot): + from aisimulate.config.traffic import AgenticSnapshotOptions + + with pytest.raises(ValidationError): + AgenticSnapshotOptions.model_validate(snapshot) + + @pytest.mark.parametrize("warmup", [False, True]) @pytest.mark.parametrize("profile", [{}, {"duration_seconds": 2.0, "response_grace_seconds": 0.0}]) def test_agentic_profile_roundtrip_and_compiler_preserve_controls(warmup: bool, profile: dict) -> None: From 929f51383ee619b4f49c9e8723082efe3d6b38e5 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 10:39:47 -0700 Subject: [PATCH 02/13] fix: retain snapshot context in runner validation errors Signed-off-by: hongkuanz --- python/aisimulate/src/aisimulate/sweeper/replay.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/aisimulate/src/aisimulate/sweeper/replay.py b/python/aisimulate/src/aisimulate/sweeper/replay.py index bb8c28ab6..4039c1366 100644 --- a/python/aisimulate/src/aisimulate/sweeper/replay.py +++ b/python/aisimulate/src/aisimulate/sweeper/replay.py @@ -512,7 +512,10 @@ def require_compatible(self, spec: ReplaySpec) -> None: if not self.supports_agentic_warmup: raise ValueError("runner does not support agentic warmup") if agentic_snapshot is not None: - snapshot = AgenticSnapshotOptions.model_validate(agentic_snapshot) + try: + snapshot = AgenticSnapshotOptions.model_validate(agentic_snapshot) + except ValueError as error: + raise ValueError(f"agentic_snapshot: {error}") from error if agentic_lanes is None or spec.workload.get("load_type") != "trace_timestamps": raise ValueError("agentic_snapshot requires trace_timestamps load with positive agentic_lanes") if snapshot.initial_selections is not None and len(snapshot.initial_selections) != agentic_lanes: From 9f6eafe6295caa2b6a765199c5c5330b229e3cd6 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 10:55:55 -0700 Subject: [PATCH 03/13] feat: advance AgentX trajectories during opt-in live warmup Signed-off-by: hongkuanz --- THIRD_PARTY_NOTICES.md | 12 + crates/core/src/python.rs | 107 ++- crates/core/src/replay/loadgen/driver.rs | 544 +++++++++++-- crates/core/src/replay/loadgen/mod.rs | 1 + crates/core/src/replay/loadgen/phase.rs | 734 ++++++++++++++++-- crates/core/src/replay/loadgen/snapshot.rs | 4 + crates/core/src/replay/report.rs | 1 + docs/replay/agentic/continuous-profiles.md | 9 +- docs/replay/agentic/warmup.md | 61 +- docs/replay/workloads.md | 34 +- python/aisimulate/THIRD_PARTY_NOTICES.md | 12 + python/aisimulate/src/aisimulate/compiler.py | 2 + .../src/aisimulate/config/traffic.py | 33 +- python/aisimulate/src/aisimulate/main.py | 8 + python/aisimulate/src/aisimulate/recommend.py | 2 + python/aisimulate/src/aisimulate/runner.py | 2 + .../aisimulate/src/aisimulate/supervision.py | 6 + .../src/aisimulate/sweeper/config.py | 12 + .../src/aisimulate/sweeper/replay.py | 12 + tests/test_agentx_live_warmup.py | 119 +++ tests/test_runner.py | 21 + 21 files changed, 1603 insertions(+), 133 deletions(-) create mode 100644 tests/test_agentx_live_warmup.py diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 11ebdfd95..59de8b86e 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -32,6 +32,18 @@ Modified by NVIDIA: executed the pinned sampler on a small synthetic corpus to freeze expected cuts, and documented its sequential recycling behavior. No upstream implementation is vendored; AISimulate consumes explicit cuts. +The live warmup policy in repository-root +`crates/core/src/replay/loadgen/phase.rs` and +`crates/core/src/replay/loadgen/driver.rs`, its regression tests in those files, +and `docs/replay/agentic/warmup.md`, `docs/replay/agentic/continuous-profiles.md` +and `docs/replay/workloads.md` are modified re-implementations of the same +pinned `agentic_replay.py` behavior: globally aligned primers, per-lane +admission quotas, compressed dependency waits, shared-cursor recycling, and +source-clock frontier handoff. NVIDIA modified these algorithms to use the +native deterministic dependency executor and its transactional server-settlement +barrier. The original NVIDIA copyright and Apache-2.0 license apply; no +upstream serving implementation is vendored. + ## AIConfigurator The repository-root `.coderabbit.yaml` is adapted and modified from diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 1c52cbf85..6d75ff91c 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -18,9 +18,9 @@ use crate::replay::{ ReplayOperationPowerDiagnostics, ReplayPhasePowerDiagnostics, ReplayPowerDiagnostics, ReplayRoleConfig, ReplayRuntimeInput, ReplaySpec, ReplayTopology, Replayer, TracePowerStats, loadgen::{ - AgenticProfileOptions, AgenticSnapshotOptions, ArrivalSpec, DelaySpec, DynamoRequestTrace, - LengthSpec, SyntheticTraceSpec, Trace, ValidatedAgenticGraph, WekaImportOptions, - WekaNestedTimestampBasis, WekaResolvedTimestampBasis, WorkloadDriver, + AgenticProfileOptions, AgenticSnapshotOptions, AgenticWarmupPolicy, ArrivalSpec, DelaySpec, + DynamoRequestTrace, LengthSpec, SyntheticTraceSpec, Trace, ValidatedAgenticGraph, + WekaImportOptions, WekaNestedTimestampBasis, WekaResolvedTimestampBasis, WorkloadDriver, load_agentic_mooncake, load_weka_agentic_graph_cached, }, }; @@ -83,6 +83,8 @@ struct RuntimeTraffic { #[serde(default)] agentic_warmup: bool, #[serde(default)] + agentic_warmup_policy: AgenticWarmupPolicy, + #[serde(default)] agentic_profile: Option, #[serde(default)] isl: Option, @@ -1503,7 +1505,17 @@ fn build_agentic_driver( // Sample recorded time before applying speedup to remaining timers. let prepared = graph.into_prepared_snapshots(lanes, options.clone())?; let mut driver = if traffic.agentic_warmup { - WorkloadDriver::new_agentic_warmup(prepared, engine_block_size, true, speedup) + WorkloadDriver::new_agentic_warmup_with_policy( + prepared, + engine_block_size, + true, + speedup, + traffic.agentic_warmup_policy, + traffic + .agentic_profile + .as_ref() + .map_or(10_000.0, |profile| profile.global_idle_cap_seconds * 1000.0), + ) } else { WorkloadDriver::new_agentic_snapshots(prepared, engine_block_size, true, speedup) }?; @@ -1526,6 +1538,11 @@ fn build_runtime_input( engine_block_size: usize, ) -> Result { ensure!(engine_block_size > 0, "engine block size must be positive"); + ensure!( + traffic.agentic_warmup_policy != AgenticWarmupPolicy::LiveTrajectoryV1 + || (traffic.agentic_warmup && traffic.agentic_profile.is_some()), + "live_trajectory_v1 requires agentic_warmup and agentic_profile" + ); if let Some(profile) = &traffic.agentic_profile { profile.validate()?; ensure!( @@ -3099,6 +3116,88 @@ mod tests { } } + #[test] + fn public_live_agentic_warmup_keeps_only_formal_requests_in_the_report() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("live-weka.jsonl"); + let rows = (0..3) + .map(|index| { + serde_json::json!({ + "t": index as f64 * 0.1, "type": "s", "model": "model", + "in": 4, "out": 3, "hash_ids": [index + 1], "api_time": 0.001 + }) + }) + .collect::>(); + std::fs::write( + &path, + serde_json::json!({ + "id": "play", "models": ["model"], "block_size": 4, + "hash_id_scope": "local", "requests": rows + }) + .to_string(), + ) + .unwrap(); + for backend in ["vllm", "sglang"] { + let payload = serde_json::json!({ + "spec": { + "version": 1, "topology": {"kind": "aggregated", "workers": {"initial_workers": 1}}, + "engine": {"rank": {"backend": backend, "block_size": 4, "num_gpu_blocks": 64, + "timing_model": {"type": "fixed", "prefill_ms": 1.0, "decode_ms": 1.0}}}, + "requests": [] + }, + "traffic": {"source_type": "trace", "load_type": "trace_timestamps", "trace_format": "weka", + "trace_path": path, "trace_block_size": 4, "execution_model": "model", + "agentic_lanes": 1, "agentic_snapshot": {"seed": 42}, "agentic_warmup": true, + "agentic_warmup_policy": "live_trajectory_v1", "agentic_profile": {"duration_seconds": 0.01}} + }); + let report: serde_json::Value = + serde_json::from_str(&execute_json(&payload.to_string(), false).unwrap()).unwrap(); + let phases = &report["agentic_phases"]; + assert_eq!(phases["phase"], "profile", "{backend}: {report}"); + assert_eq!(phases["warmup_policy"], "live_trajectory_v1"); + assert_eq!(phases["lanes"][0]["warmup_completed"], 10); + assert_eq!(phases["lanes"][0]["primers_completed"], 1); + let formal = report["per_request"].as_array().unwrap(); + assert!(!formal.is_empty()); + assert_eq!( + formal.len() as u64, + report["agentic_profile"]["issued_requests"] + .as_u64() + .unwrap() + ); + assert!(report["completed_requests"].as_u64().unwrap() > 0); + let outcomes = report["agentic_play_outcomes"].as_array().unwrap(); + assert_eq!( + outcomes.len() as u64, + report["total_trajectories"].as_u64().unwrap() + ); + assert_eq!( + outcomes.len() as u64, + report["agentic_profile"]["plays_started"].as_u64().unwrap() + ); + for outcome in outcomes { + for field in ["causal_terminal_ms", "settled_at_ms"] { + assert!( + outcome[field].as_f64().is_none_or(|value| value >= 0.0), + "{backend}: {outcome}" + ); + } + } + let mut invalid = payload.clone(); + invalid["traffic"] + .as_object_mut() + .unwrap() + .remove("agentic_profile"); + assert!( + format!( + "{:#}", + execute_json(&invalid.to_string(), false).unwrap_err() + ) + .contains("live_trajectory_v1 requires agentic_warmup and agentic_profile") + ); + } + } + #[test] fn public_agentic_json_rejects_unqualified_modes_without_restricting_standard_dynamo() { let directory = tempfile::tempdir().unwrap(); diff --git a/crates/core/src/replay/loadgen/driver.rs b/crates/core/src/replay/loadgen/driver.rs index c05de5cab..9ea51f7ee 100644 --- a/crates/core/src/replay/loadgen/driver.rs +++ b/crates/core/src/replay/loadgen/driver.rs @@ -1,6 +1,10 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +// Live warmup is a modified native re-implementation of NVIDIA AIPerf's +// Apache-2.0 agentic_replay.py at SemiAnalysisAI/aiperf commit +// 754356e9a39acc6cc6afb242d123bb57c3fb6f75; see THIRD_PARTY_NOTICES.md. + use std::cmp::Ordering; use std::collections::BinaryHeap; use std::sync::Arc; @@ -12,7 +16,9 @@ use rustc_hash::{FxHashMap, FxHashSet}; use serde::Serialize; use uuid::Uuid; -use super::phase::{AgenticPhaseEvidence, AgenticPreparation, AgenticPreparationTransition}; +use super::phase::{ + AgenticPhaseEvidence, AgenticPreparation, AgenticPreparationTransition, AgenticWarmupPolicy, +}; use super::profile::{ AgenticIdleShift, AgenticProfileOptions, AgenticProfileReport, AgenticProfileState, }; @@ -140,6 +146,7 @@ struct ConcurrencyState { #[derive(Debug)] struct AgenticState { + compressed_warmup_timing: bool, identities: Vec, node_states: Vec, remaining_dependencies: Vec, @@ -177,6 +184,7 @@ struct AgenticDependentEdge { #[derive(Debug)] struct AgenticPlayState { + preparation_only: bool, play_id: String, nodes: Vec, root_nodes: Vec, @@ -197,6 +205,7 @@ struct AgenticPlayState { impl AgenticPlayState { fn new(play_id: String, nodes: Vec, root_nodes: Vec) -> Self { Self { + preparation_only: false, play_id, nodes, root_nodes, @@ -249,6 +258,15 @@ struct TurnResolution { session_ended: bool, } +struct PreparedAgenticProfilePlay { + appended: WorkloadDriver, + lane: usize, + retired_play: usize, + next_cursor: u64, + next_ordinal: u64, + replacements: usize, +} + #[derive(Debug)] struct SessionRuntime { session_id: String, @@ -532,6 +550,7 @@ impl AgenticState { sessions, ready_sessions, now_ms, + self.compressed_warmup_timing, ); } @@ -550,6 +569,7 @@ impl AgenticState { sessions, ready_sessions, now_ms, + self.compressed_warmup_timing, ); } @@ -562,6 +582,7 @@ impl AgenticState { sessions: &mut [SessionRuntime], ready_sessions: &mut BinaryHeap, now_ms: f64, + compressed_warmup_timing: bool, ) { for edge in edges { let target = edge.target_node; @@ -571,7 +592,11 @@ impl AgenticState { remaining_dependencies[target] = remaining_dependencies[target] .checked_sub(1) .expect("compiled agentic dependency must fire exactly once"); - ready_after_ms[target] = ready_after_ms[target].max(now_ms + edge.delay_ms); + ready_after_ms[target] = if compressed_warmup_timing { + now_ms + } else { + ready_after_ms[target].max(now_ms + edge.delay_ms) + }; if remaining_dependencies[target] != 0 { continue; } @@ -796,7 +821,7 @@ impl AgenticState { fn trajectory_snapshot(&self) -> AgenticTrajectorySnapshot { let mut e2e_latencies_ms = Vec::new(); for play in &self.plays { - if play.failed || play.completed_nodes != play.nodes.len() { + if play.preparation_only || play.failed || play.completed_nodes != play.nodes.len() { continue; } let (Some(root_dispatch_ms), Some(max_terminal_ms)) = @@ -807,7 +832,11 @@ impl AgenticState { e2e_latencies_ms.push(max_terminal_ms - root_dispatch_ms); } AgenticTrajectorySnapshot { - total_trajectories: self.plays.len(), + total_trajectories: self + .plays + .iter() + .filter(|play| !play.preparation_only) + .count(), completed_trajectories: e2e_latencies_ms.len(), e2e_latencies_ms, } @@ -816,6 +845,7 @@ impl AgenticState { fn play_outcomes(&self, sessions: &[SessionRuntime]) -> Vec { self.plays .iter() + .filter(|play| !play.preparation_only) .map(|play| { let status = if !play.quiescent || play.censored { AgenticPlayStatus::Incomplete @@ -951,7 +981,7 @@ impl WorkloadDriver { /// queues and P/D handoffs. Preparation traffic and server-only cleanup are /// deliberately excluded. pub fn agentic_profile_pending_request_ids(&self) -> Vec { - if self.agentic_profile.is_none() { + if self.agentic_profile.is_none() || self.is_agentic_preparing() { return Vec::new(); } let mut entries = self.in_flight.iter().collect::>(); @@ -968,6 +998,9 @@ impl WorkloadDriver { /// Apply client control events at a runtime-owned timestamp. The runtime /// must deliver existing responses first, then call this before admission. pub fn advance_agentic_profile(&mut self, now_ms: f64) -> Result { + if self.live_warmup_active() { + return self.advance_live_warmup(now_ms); + } let Some(profile) = &self.agentic_profile else { return Ok(false); }; @@ -1048,26 +1081,52 @@ impl WorkloadDriver { } fn append_agentic_profile_play(&mut self, lane: usize, now_ms: f64) -> Result<()> { + let prepared = self.prepare_agentic_profile_play(lane, now_ms, 0)?; + self.commit_agentic_profile_play(prepared, now_ms); + Ok(()) + } + + fn prepare_agentic_profile_play( + &self, + lane: usize, + now_ms: f64, + pending: usize, + ) -> Result { let context = self.agentic_replay_context.as_ref().unwrap(); let profile = self.agentic_profile.as_ref().unwrap(); - let cursor = profile.cursor; + let cursor = profile + .cursor + .checked_add(u64::try_from(pending)?) + .context("agentic corpus cursor overflow")?; let ordinal = profile.next_ordinals[lane]; let (last_replacement_ms, replacements) = profile.replacements_at_ms; let replacements = if last_replacement_ms == Some(now_ms) { replacements - .checked_add(1) + .checked_add( + pending + .checked_add(1) + .context("same-time replacement count overflow")?, + ) .context("same-time replacement count overflow")? } else { - 1 + pending + .checked_add(1) + .context("same-time replacement count overflow")? }; // Source selection is shared across lanes: another lane may consume // the positive-duration source while this lane draws a zero-duration // source repeatedly. Each full corpus cycle can occupy one more lane // with future work, so allow a full cycle for every configured lane. - let same_time_budget = context - .source_play_count() - .checked_mul(context.lane_count()) - .context("same-time replacement budget overflow")?; + let same_time_budget = if self.live_warmup_active() { + // Unlike an unbounded profile, warmup is bounded by ten admissions + // per lane plus one fresh handoff frontier. Zero-time one-token + // passes are therefore finite even when full requests take time. + super::AGENTIC_WARMUP_REQUESTS_PER_LANE + 1 + } else { + context.source_play_count() + } + .checked_mul(context.lane_count()) + .context("same-time replacement budget overflow")?; if replacements > same_time_budget { bail!( "agentic profile repeated a corpus cycle for every lane without virtual-time progress at {now_ms} ms" @@ -1087,6 +1146,28 @@ impl WorkloadDriver { self.include_replay_hashes, self.agentic_speedup, )?; + if self.live_warmup_active() { + appended.compress_live_warmup_timers(0.0); + } + Ok(PreparedAgenticProfilePlay { + appended, + lane, + retired_play, + next_cursor, + next_ordinal, + replacements, + }) + } + + fn commit_agentic_profile_play(&mut self, prepared: PreparedAgenticProfilePlay, now_ms: f64) { + let PreparedAgenticProfilePlay { + mut appended, + lane, + retired_play, + next_cursor, + next_ordinal, + replacements, + } = prepared; let SchedulingPolicy::Agentic(mut added) = appended.policy else { unreachable!() }; @@ -1165,7 +1246,6 @@ impl WorkloadDriver { profile.replacements_at_ms = (Some(now_ms), replacements); profile.active_plays[lane] = play_index; profile.tree_idle_since_ms.push(Some(now_ms)); - Ok(()) } fn shift_agentic_timers(&mut self, only_play: Option, now_ms: f64) -> bool { @@ -1232,14 +1312,22 @@ impl WorkloadDriver { finished_at_ms: None, cancel_drain_timed_out: false, unsettled_server_requests: self.in_flight.len() + self.agentic_settling.len(), - plays_started: state.plays.len(), + plays_started: state.plays.iter().filter(|p| !p.preparation_only).count(), client_completed_plays: state .plays .iter() - .filter(|p| p.client_finished && !p.censored) + .filter(|p| !p.preparation_only && p.client_finished && !p.censored) + .count(), + retired_plays: state + .plays + .iter() + .filter(|p| !p.preparation_only && p.client_finished) + .count(), + server_quiescent_plays: state + .plays + .iter() + .filter(|p| !p.preparation_only && p.quiescent) .count(), - retired_plays: state.plays.iter().filter(|p| p.client_finished).count(), - server_quiescent_plays: state.plays.iter().filter(|p| p.quiescent).count(), issued_requests: profile.issued_requests, successful_responses: profile.successful_responses, canceled_requests: profile.canceled_requests, @@ -1259,6 +1347,21 @@ impl WorkloadDriver { fn agentic_profile_next_control_ms(&self) -> Option { let profile = self.agentic_profile.as_ref()?; + if self.live_warmup_active() { + let SchedulingPolicy::Agentic(state) = &self.policy else { + return None; + }; + let preparation = self.agentic_preparation.as_ref()?; + return profile + .active_plays + .iter() + .enumerate() + .filter(|(lane, play)| { + preparation.lane_has_budget(*lane) && state.plays[**play].client_finished + }) + .filter_map(|(_, play)| state.plays[*play].max_terminal_ms) + .min_by(f64::total_cmp); + } let (cutoff, _, _) = profile.deadlines()?; if profile.cutoff { return None; @@ -1321,11 +1424,42 @@ impl WorkloadDriver { include_replay_hashes: bool, speedup: f64, ) -> Result { - let preparation = AgenticPreparation::new( - &prepared, - u32::try_from(engine_block_size).context("engine block size exceeds u32")?, + Self::new_agentic_warmup_with_policy( + prepared, + engine_block_size, include_replay_hashes, - )?; + speedup, + AgenticWarmupPolicy::RepeatedPrefixes, + 10_000.0, + ) + } + + pub fn new_agentic_warmup_with_policy( + prepared: PreparedAgenticSnapshots, + engine_block_size: usize, + include_replay_hashes: bool, + speedup: f64, + policy: AgenticWarmupPolicy, + primer_idle_cap_ms: f64, + ) -> Result { + if !primer_idle_cap_ms.is_finite() || primer_idle_cap_ms < 0.0 { + bail!("warmup primer idle cap must be finite and nonnegative"); + } + let preparation = if policy == AgenticWarmupPolicy::RepeatedPrefixes { + AgenticPreparation::new( + &prepared, + u32::try_from(engine_block_size).context("engine block size exceeds u32")?, + include_replay_hashes, + )? + } else { + AgenticPreparation::with_policy( + &prepared, + u32::try_from(engine_block_size).context("engine block size exceeds u32")?, + include_replay_hashes, + policy, + primer_idle_cap_ms, + )? + }; let mut driver = Self::new_agentic_snapshots( prepared, engine_block_size, @@ -1336,6 +1470,112 @@ impl WorkloadDriver { Ok(driver) } + fn live_warmup_active(&self) -> bool { + self.agentic_preparation + .as_ref() + .is_some_and(AgenticPreparation::live_active) + } + + fn is_live_preparation_request(&self, uuid: Uuid) -> bool { + self.agentic_preparation + .as_ref() + .is_some_and(|phase| phase.is_live_request(uuid)) + } + + fn compress_live_warmup_timers(&mut self, now_ms: f64) { + let SchedulingPolicy::Agentic(state) = &mut self.policy else { + return; + }; + state.compressed_warmup_timing = true; + state.ready_after_ms.fill(now_ms); + for session in &mut self.sessions { + if let Some(at) = &mut session.next_ready_at_ms { + *at = now_ms; + } + } + self.ready_sessions = std::mem::take(&mut self.ready_sessions) + .into_iter() + .map(|mut ready| { + ready.ready_at_ms = now_ms; + ready + }) + .collect(); + } + + fn advance_live_warmup(&mut self, now_ms: f64) -> Result { + let profile = self + .agentic_profile + .as_ref() + .context("live warmup requires agentic_profile")?; + if !now_ms.is_finite() || now_ms < profile.last_advanced_ms { + bail!("live warmup clock must be finite and nondecreasing"); + } + let active = profile.active_plays.clone(); + let mut changed = false; + for (lane, play) in active.into_iter().enumerate() { + let SchedulingPolicy::Agentic(state) = &self.policy else { + unreachable!() + }; + if self + .agentic_preparation + .as_ref() + .is_some_and(|phase| phase.lane_has_budget(lane)) + && state.plays[play].client_finished + { + self.append_agentic_profile_play(lane, now_ms)?; + changed = true; + } + } + self.agentic_profile.as_mut().unwrap().last_advanced_ms = now_ms; + Ok(changed) + } + + // Reconstruct source-clock offsets, not warmup wall time. This is computed + // before reset_measurements so any error leaves the phase retryable. + fn live_handoff_times(&self, now_ms: f64, fresh_origin: bool) -> Result> { + let SchedulingPolicy::Agentic(state) = &self.policy else { + unreachable!() + }; + let context = self + .agentic_replay_context + .as_ref() + .context("missing live replay context")?; + let snapshots = self + .agentic_snapshots + .as_ref() + .context("missing live snapshots")?; + let mut offsets = Vec::new(); + let mut first = if fresh_origin { 0.0 } else { f64::INFINITY }; + for (index, node) in state.node_states.iter().enumerate() { + if !matches!(node, AgenticNodeState::Ready | AgenticNodeState::Blocked) { + continue; + } + let PromptTokens::Snapshot { source_index, .. } = + self.sessions[index].turns[0].prompt_tokens + else { + bail!("live warmup requires immutable snapshot prompts"); + }; + let offset = (context.recorded_start_ms(source_index) + - snapshots[state.node_to_play[index]].t_star_ms) + .max(0.0) + / self.agentic_speedup; + if *node == AgenticNodeState::Ready { + first = first.min(offset); + } + offsets.push((index, offset)); + } + if !first.is_finite() { + first = 0.0; + } + for (_, at) in &mut offsets { + *at = now_ms + (*at - first).max(0.0); + if !at.is_finite() { + bail!("live warmup handoff timer overflow"); + } + } + Ok(offsets) + } + pub fn is_agentic_preparing(&self) -> bool { self.agentic_preparation .as_ref() @@ -1375,7 +1615,8 @@ impl WorkloadDriver { /// Open profiling only after all preparation requests succeed and settle, /// or finish an aborted preparation after its submitted server work drains. - /// This does not change the saved graph's frontier or dependency delays. + /// The default keeps the saved frontier; live warmup retains its advanced + /// dependency state and restores source-clock delays at the barrier. pub fn finish_agentic_preparation( &mut self, now_ms: f64, @@ -1390,12 +1631,74 @@ impl WorkloadDriver { now_ms: f64, reset_measurements: impl FnOnce() -> Result<()>, ) -> Result> { + let pending_transition = self + .agentic_preparation + .as_ref() + .map(|phase| phase.ready_transition(now_ms)) + .transpose()? + .flatten(); + let mut handoff_plays = Vec::new(); + if self.live_warmup_active() + && pending_transition == Some(AgenticPreparationTransition::OpenProfile) + { + let SchedulingPolicy::Agentic(state) = &self.policy else { + unreachable!() + }; + for (lane, &play) in self + .agentic_profile + .as_ref() + .unwrap() + .active_plays + .iter() + .enumerate() + { + if state.plays[play].client_finished { + handoff_plays.push(self.prepare_agentic_profile_play( + lane, + now_ms, + handoff_plays.len(), + )?); + } + } + } let Some(preparation) = &self.agentic_preparation else { return Ok(None); }; + if preparation.ready_to_start_live() { + if self.agentic_profile.is_none() { + bail!("live warmup requires agentic_profile"); + } + preparation.validate_boundary_time(now_ms)?; + self.compress_live_warmup_timers(now_ms); + self.agentic_preparation + .as_mut() + .unwrap() + .start_live(now_ms); + return Ok(None); + } let Some(transition) = preparation.ready_transition(now_ms)? else { return Ok(None); }; + let mut handoff_times = (transition == AgenticPreparationTransition::OpenProfile + && preparation.live_active()) + .then(|| self.live_handoff_times(now_ms, !handoff_plays.is_empty())) + .transpose()?; + if let Some(times) = &mut handoff_times { + let mut node_offset = self.sessions.len(); + for plan in &handoff_plays { + for (index, at) in plan.appended.live_handoff_times(now_ms, true)? { + times.push(( + node_offset + .checked_add(index) + .context("live handoff node index overflow")?, + at, + )); + } + node_offset = node_offset + .checked_add(plan.appended.sessions.len()) + .context("live handoff node index overflow")?; + } + } if let SchedulingPolicy::Agentic(state) = &self.policy && state .last_runtime_feedback_at_ms @@ -1403,7 +1706,7 @@ impl WorkloadDriver { { bail!("agentic preparation boundary precedes runtime feedback"); } - if transition == AgenticPreparationTransition::OpenProfile { + if transition == AgenticPreparationTransition::OpenProfile && handoff_times.is_none() { let SchedulingPolicy::Agentic(state) = &self.policy else { unreachable!() }; @@ -1437,7 +1740,55 @@ impl WorkloadDriver { bail!("agentic profile deadline overflow at activation"); } reset_measurements()?; - if transition == AgenticPreparationTransition::OpenProfile { + for plan in handoff_plays { + self.commit_agentic_profile_play(plan, now_ms); + } + if transition == AgenticPreparationTransition::Aborted + && self + .agentic_preparation + .as_ref() + .is_some_and(AgenticPreparation::uses_live_policy) + { + let SchedulingPolicy::Agentic(state) = &mut self.policy else { + unreachable!() + }; + for play in &mut state.plays { + play.preparation_only = true; + } + state.lifecycle.clear(); + state.next_lifecycle_ordinal = 0; + } + if let Some(times) = handoff_times { + let SchedulingPolicy::Agentic(state) = &mut self.policy else { + unreachable!() + }; + state.compressed_warmup_timing = false; + // Formal metrics describe only the post-barrier cohort. Keep the + // compact old plays for dependency/identity audit, but not outcome + // latency samples. Preparation requests remain in the phase ledger. + for play in &mut state.plays { + if play.client_finished { + play.preparation_only = true; + } else { + play.root_dispatch_ms = Some(now_ms); + play.max_terminal_ms = None; + } + } + state.lifecycle.clear(); + state.next_lifecycle_ordinal = 0; + self.ready_sessions.clear(); + for (index, at) in times { + state.ready_after_ms[index] = at; + if state.node_states[index] == AgenticNodeState::Ready { + self.sessions[index].next_ready_at_ms = Some(at); + self.ready_sessions.push(ReadySession { + ready_at_ms: at, + session_index: index, + turn_index: 0, + }); + } + } + } else if transition == AgenticPreparationTransition::OpenProfile { let SchedulingPolicy::Agentic(state) = &mut self.policy else { unreachable!() }; @@ -1626,6 +1977,7 @@ impl WorkloadDriver { agentic_plays.push(play); } let mut state = AgenticState { + compressed_warmup_timing: false, identities, node_states: vec![AgenticNodeState::Blocked; sessions.len()], remaining_dependencies, @@ -1964,6 +2316,7 @@ impl WorkloadDriver { } let mut state = AgenticState { + compressed_warmup_timing: false, identities, node_states: vec![AgenticNodeState::Blocked; sessions.len()], remaining_dependencies, @@ -2282,9 +2635,13 @@ impl WorkloadDriver { #[doc(hidden)] pub fn pop_ready_compact(&mut self, now_ms: f64, limit: usize) -> Vec { + let live_warmup = self.live_warmup_active(); if let Some(preparation) = &mut self.agentic_preparation { match preparation.transition() { - None => return preparation.pop_ready(now_ms, limit, self.emit_session_metadata), + None if !live_warmup => { + return preparation.pop_ready(now_ms, limit, self.emit_session_metadata); + } + None => {} Some(AgenticPreparationTransition::Aborted) => return Vec::new(), Some(AgenticPreparationTransition::OpenProfile) => {} } @@ -2302,6 +2659,7 @@ impl WorkloadDriver { } let mut emitted = Vec::new(); + let mut quota_deferred = Vec::new(); while emitted.len() < effective_limit { let Some(ready_session) = self.ready_sessions.pop() else { break; @@ -2312,6 +2670,26 @@ impl WorkloadDriver { } let session_index = ready_session.session_index; + let lane = if live_warmup { + let SchedulingPolicy::Agentic(state) = &self.policy else { + unreachable!() + }; + let lane = self.agentic_snapshots.as_ref().unwrap() + [state.node_to_play[session_index]] + .lane_id; + if !self + .agentic_preparation + .as_ref() + .unwrap() + .lane_has_budget(lane) + { + quota_deferred.push(ready_session); + continue; + } + Some(lane) + } else { + None + }; let Some((turn_index, scheduled_ready_at_ms)) = self .sessions .get(session_index) @@ -2343,7 +2721,11 @@ impl WorkloadDriver { authored_id: request_id.clone(), session_id: Some(session.session_id.clone()), turn_index: None, - metadata: serde_json::Value::Null, + metadata: if live_warmup { + serde_json::json!({"agentic_phase": "warmup"}) + } else { + serde_json::Value::Null + }, prompt_token_source: ReplayPromptTokenSource::Materialized, agentic: agentic_identity, }, @@ -2365,7 +2747,13 @@ impl WorkloadDriver { ( input_length, context.token_ids(slot, source_index), - Some(context.outputs[source_index].clone()), + Some(if live_warmup { + vec![ + context.outputs[source_index].first().copied().unwrap_or(0), + ] + } else { + context.outputs[source_index].clone() + }), ) } else { let (input_length, hash_ids) = turn.prompt_tokens.take_deferred(); @@ -2373,7 +2761,11 @@ impl WorkloadDriver { }; let request_metadata = DirectRequest { tokens: Vec::new(), - max_output_tokens: turn.max_output_tokens, + max_output_tokens: if live_warmup { + 1 + } else { + turn.max_output_tokens + }, output_token_ids, uuid: Some(request_uuid), dp_rank: 0, @@ -2428,7 +2820,7 @@ impl WorkloadDriver { emitted_output_tokens: 0, }, ); - emitted.push(CompactReadyTurn { + let ready = CompactReadyTurn { request_uuid, authored_request_id: turn.request_id.clone(), play_id: turn.play_id.clone(), @@ -2441,10 +2833,29 @@ impl WorkloadDriver { emit_session_metadata: self.emit_session_metadata, synthetic_session_id: turn.synthetic_session_id, request, - }); + }; + if let Some(lane) = lane { + let input_length = ready.request.input_length(); + let identity = ready + .request + .metadata() + .replay_context + .as_ref() + .unwrap() + .agentic + .clone() + .unwrap(); + self.agentic_preparation + .as_mut() + .unwrap() + .record_live_dispatch(&ready, lane, input_length, identity); + } + emitted.push(ready); if let SchedulingPolicy::Agentic(state) = &mut self.policy { state.on_node_emitted(session_index, now_ms); - if let Some(profile) = &mut self.agentic_profile { + if let Some(profile) = &mut self.agentic_profile + && profile.start_ms.is_some() + { profile.issued_requests += 1; profile.first_request_ms.get_or_insert(now_ms); profile.global_idle_since_ms = None; @@ -2458,6 +2869,7 @@ impl WorkloadDriver { ); } } + self.ready_sessions.extend(quota_deferred); emitted } @@ -2511,11 +2923,13 @@ impl WorkloadDriver { } fn agentic_node_ordinal(&self, request_uuid: Uuid) -> Result { - if let Some(ordinal) = self - .agentic_preparation - .as_ref() - .and_then(|phase| phase.ordinal(request_uuid)) - { + if let Some(ordinal) = self.agentic_preparation.as_ref().and_then(|phase| { + if phase.is_live_request(request_uuid) { + None + } else { + phase.ordinal(request_uuid) + } + }) { return self .sessions .len() @@ -2620,7 +3034,9 @@ impl WorkloadDriver { debug_assert_eq!(self.prompt_mode, PromptMode::Full); let mut becoming_terminal = FxHashSet::default(); for terminal in &feedback.causal_terminals { - if self.knows_preparation_request(terminal.request_uuid) { + if self.knows_preparation_request(terminal.request_uuid) + && !self.is_live_preparation_request(terminal.request_uuid) + { continue; } let outcome = match terminal.status { @@ -2642,7 +3058,9 @@ impl WorkloadDriver { becoming_terminal.insert(terminal.request_uuid); } for request_uuid in &feedback.quiescent_requests { - if self.knows_preparation_request(*request_uuid) { + if self.knows_preparation_request(*request_uuid) + && !self.is_live_preparation_request(*request_uuid) + { continue; } if !self.agentic_settling.contains_key(request_uuid) @@ -2699,11 +3117,13 @@ impl WorkloadDriver { status: ReplayTerminalStatus, ) -> Result<()> { if self.knows_preparation_request(request_uuid) { - return self.agentic_preparation.as_mut().unwrap().on_terminal( - request_uuid, - now_ms, - status, - ); + self.agentic_preparation + .as_mut() + .unwrap() + .on_terminal(request_uuid, now_ms, status)?; + if !self.is_live_preparation_request(request_uuid) { + return Ok(()); + } } let outcome = match status { ReplayTerminalStatus::Completed => TurnOutcome::Completed, @@ -2744,10 +3164,10 @@ impl WorkloadDriver { if self.in_flight.is_empty() { profile.global_idle_since_ms = Some(now_ms); } - if status == ReplayTerminalStatus::Completed { + if profile.start_ms.is_some() && status == ReplayTerminalStatus::Completed { profile.successful_responses += 1; profile.last_successful_response_ms = Some(now_ms); - } else if status == ReplayTerminalStatus::Canceled { + } else if profile.start_ms.is_some() && status == ReplayTerminalStatus::Canceled { profile.canceled_requests += 1; } } @@ -2757,11 +3177,13 @@ impl WorkloadDriver { /// Mark an already-terminal agentic request free of runtime-owned state. pub fn on_quiescent(&mut self, request_uuid: Uuid, now_ms: f64) -> Result<()> { if self.knows_preparation_request(request_uuid) { - return self - .agentic_preparation + self.agentic_preparation .as_mut() .unwrap() - .on_quiescent(request_uuid, now_ms); + .on_quiescent(request_uuid, now_ms)?; + if !self.is_live_preparation_request(request_uuid) { + return Ok(()); + } } let SchedulingPolicy::Agentic(state) = &mut self.policy else { return Ok(()); @@ -2914,9 +3336,11 @@ impl WorkloadDriver { } pub fn next_ready_time_ms(&mut self) -> Option { + let live_warmup = self.live_warmup_active(); if let Some(preparation) = &self.agentic_preparation { match preparation.transition() { - None => return preparation.next_ready_time_ms(), + None if !live_warmup => return preparation.next_ready_time_ms(), + None => {} Some(AgenticPreparationTransition::Aborted) => return None, Some(AgenticPreparationTransition::OpenProfile) => {} } @@ -2925,6 +3349,28 @@ impl WorkloadDriver { return None; } let control_ms = self.agentic_profile_next_control_ms(); + if live_warmup { + let SchedulingPolicy::Agentic(state) = &self.policy else { + unreachable!() + }; + let preparation = self.agentic_preparation.as_ref().unwrap(); + return self + .ready_sessions + .iter() + .filter(|ready| { + let session = &self.sessions[ready.session_index]; + let lane = self.agentic_snapshots.as_ref().unwrap() + [state.node_to_play[ready.session_index]] + .lane_id; + preparation.lane_has_budget(lane) + && session.in_flight.is_none() + && session.next_ready_at_ms == Some(ready.ready_at_ms) + && session.next_turn_index == ready.turn_index + }) + .map(|ready| ready.ready_at_ms) + .chain(control_ms) + .min_by(f64::total_cmp); + } loop { let Some(ready_session) = self.ready_sessions.peek().copied() else { return control_ms; diff --git a/crates/core/src/replay/loadgen/mod.rs b/crates/core/src/replay/loadgen/mod.rs index 8a9bc1c6b..43551da53 100644 --- a/crates/core/src/replay/loadgen/mod.rs +++ b/crates/core/src/replay/loadgen/mod.rs @@ -26,6 +26,7 @@ pub use generated::GeneratedRequests; pub use phase::{ AGENTIC_PHASE_SCHEMA_V1, AGENTIC_WARMUP_REQUESTS_PER_LANE, AgenticPhaseEvidence, AgenticPhaseLane, AgenticPhaseRequest, AgenticPreparationTransition, AgenticReplayPhase, + AgenticWarmupPolicy, }; pub use profile::{AgenticIdleShift, AgenticProfileOptions, AgenticProfileReport}; pub use snapshot::{ diff --git a/crates/core/src/replay/loadgen/phase.rs b/crates/core/src/replay/loadgen/phase.rs index 23fd95729..10cb0f305 100644 --- a/crates/core/src/replay/loadgen/phase.rs +++ b/crates/core/src/replay/loadgen/phase.rs @@ -1,8 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Preparation traffic is separate from the saved graph frontier. Warmup repeats -//! full snapshot prefixes; it never advances a source request or its output RNG. +//! Preparation traffic is excluded from the formal measurement cohort. The +//! default repeats saved prefixes; the explicit live policy advances the DAG. +//! +//! Live policy adapted from the Apache-2.0 NVIDIA AIPerf implementation: +//! https://github.com/SemiAnalysisAI/aiperf/blob/754356e9a39acc6cc6afb242d123bb57c3fb6f75/src/aiperf/timing/strategies/agentic_replay.py +//! Modified for native deterministic scheduling and server-settlement barriers. //! //! Methodology reference for one-token primers and ten warmups per lane: //! https://github.com/SemiAnalysisAI/InferenceX-app/blob/9bb7b13eb4985217a6282f340459fd5948613276/packages/app/src/components/datasets/agentx-methodology-article.tsx @@ -13,7 +17,7 @@ use anyhow::{Context, Result, ensure}; use rustc_hash::{FxHashMap, FxHashSet}; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use uuid::Uuid; use super::driver::{AgenticRuntimeFeedback, deferred_request_with_hashes}; @@ -25,6 +29,20 @@ pub const AGENTIC_PHASE_SCHEMA_V1: &str = "aisimulate.agentic.phases.v1"; pub const AGENTIC_WARMUP_REQUESTS_PER_LANE: usize = 10; const PREPARATION_UUID_BASE: u128 = 1_u128 << 127; +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AgenticWarmupPolicy { + #[default] + RepeatedPrefixes, + LiveTrajectoryV1, +} + +impl AgenticWarmupPolicy { + fn is_repeated(&self) -> bool { + *self == Self::RepeatedPrefixes + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] #[serde(rename_all = "snake_case")] pub enum AgenticReplayPhase { @@ -86,6 +104,8 @@ pub struct AgenticPhaseEvidence { pub failure_reason: Option, pub lanes: Vec, pub requests: Vec, + #[serde(skip_serializing_if = "AgenticWarmupPolicy::is_repeated")] + pub warmup_policy: AgenticWarmupPolicy, } #[derive(Debug)] @@ -96,6 +116,7 @@ struct PreparationLane { next: usize, in_flight: Option, ready_at_ms: f64, + primer_times_ms: Vec, } #[derive(Debug)] @@ -110,6 +131,10 @@ pub(super) struct AgenticPreparation { last_event_at_ms: f64, engine_block_size: u32, include_replay_hashes: bool, + policy: AgenticWarmupPolicy, + live_started: bool, + live_counts: Vec, + live_requests: FxHashSet, } impl AgenticPreparation { @@ -117,6 +142,22 @@ impl AgenticPreparation { prepared: &PreparedAgenticSnapshots, engine_block_size: u32, include_replay_hashes: bool, + ) -> Result { + Self::with_policy( + prepared, + engine_block_size, + include_replay_hashes, + AgenticWarmupPolicy::RepeatedPrefixes, + 10_000.0, + ) + } + + pub(super) fn with_policy( + prepared: &PreparedAgenticSnapshots, + engine_block_size: u32, + include_replay_hashes: bool, + policy: AgenticWarmupPolicy, + primer_idle_cap_ms: f64, ) -> Result { ensure!( engine_block_size > 0, @@ -133,6 +174,10 @@ impl AgenticPreparation { last_event_at_ms: 0.0, engine_block_size, include_replay_hashes, + policy, + live_started: false, + live_counts: vec![0; prepared.context().lane_count()], + live_requests: FxHashSet::default(), }; for snapshot in &prepared.plays { // Snapshot UUIDs use checked arithmetic. Check the separate domain @@ -165,10 +210,12 @@ impl AgenticPreparation { primers.clone() }; let mut source_indices = primers.clone(); - source_indices.extend( - (0..AGENTIC_WARMUP_REQUESTS_PER_LANE) - .map(|ordinal| warmup_sources[ordinal % warmup_sources.len()]), - ); + if policy == AgenticWarmupPolicy::RepeatedPrefixes { + source_indices.extend( + (0..AGENTIC_WARMUP_REQUESTS_PER_LANE) + .map(|ordinal| warmup_sources[ordinal % warmup_sources.len()]), + ); + } let mut request_indices = Vec::with_capacity(source_indices.len()); for (ordinal, &source_index) in source_indices.iter().enumerate() { let index = this.requests.len(); @@ -224,11 +271,97 @@ impl AgenticPreparation { next: 0, in_flight: None, ready_at_ms: 0.0, + primer_times_ms: Vec::new(), }); } + if policy == AgenticWarmupPolicy::LiveTrajectoryV1 { + let mut max_lead_ms = 0.0_f64; + for lane in &mut this.lanes { + lane.primer_times_ms = lane + .source_indices + .iter() + .map(|&index| { + let start = lane.snapshot.context.recorded_start_ms(index); + let lead = (lane.snapshot.evidence.t_star_ms - start) + .max(0.0) + .min(primer_idle_cap_ms); + max_lead_ms = max_lead_ms.max(lead); + lead + }) + .collect(); + } + for lane in &mut this.lanes { + for lead in &mut lane.primer_times_ms { + *lead = max_lead_ms - *lead; + } + } + } Ok(this) } + pub(super) fn live_active(&self) -> bool { + self.live_started && self.transition.is_none() && self.failure_request_id.is_none() + } + + pub(super) fn uses_live_policy(&self) -> bool { + self.policy == AgenticWarmupPolicy::LiveTrajectoryV1 + } + + pub(super) fn is_live_request(&self, uuid: Uuid) -> bool { + self.live_requests.contains(&uuid) + } + + pub(super) fn lane_has_budget(&self, lane: usize) -> bool { + self.live_active() && self.live_counts[lane] < AGENTIC_WARMUP_REQUESTS_PER_LANE + } + + pub(super) fn ready_to_start_live(&self) -> bool { + self.policy == AgenticWarmupPolicy::LiveTrajectoryV1 + && !self.live_started + && self.failure_request_id.is_none() + && self + .requests + .iter() + .all(|request| request.quiescent_at_ms.is_some()) + } + + pub(super) fn start_live(&mut self, now_ms: f64) { + self.live_started = true; + self.last_event_at_ms = now_ms; + } + + pub(super) fn record_live_dispatch( + &mut self, + turn: &CompactReadyTurn, + lane: usize, + input_length: usize, + identity: AgenticRuntimeIdentity, + ) { + debug_assert!(self.lane_has_budget(lane)); + let index = self.requests.len(); + self.index_by_uuid.insert(turn.request_uuid, index); + self.live_requests.insert(turn.request_uuid); + self.live_counts[lane] += 1; + self.requests.push(AgenticPhaseRequest { + uuid: turn.request_uuid, + phase: AgenticReplayPhase::Warmup, + lane_id: lane, + source_request_id: turn.authored_request_id.clone().unwrap_or_default(), + identity, + input_length, + max_output_tokens: 1, + expected_full_block_tokens: input_length / self.engine_block_size as usize + * self.engine_block_size as usize, + dispatched_at_ms: Some(turn.dispatched_at_ms), + terminal_status: None, + causal_terminal_at_ms: None, + quiescent_at_ms: None, + first_admit_ms: None, + first_admission_reused_input_tokens: None, + observed_output_tokens: 0, + }); + } + pub(super) fn contains(&self, uuid: Uuid) -> bool { self.index_by_uuid.contains_key(&uuid) } @@ -248,7 +381,17 @@ impl AgenticPreparation { self.lanes .iter() .filter(|lane| lane.in_flight.is_none() && lane.next < lane.request_indices.len()) - .map(|lane| lane.ready_at_ms) + .map(|lane| { + if self.policy == AgenticWarmupPolicy::LiveTrajectoryV1 { + lane.primer_times_ms[lane.next..] + .iter() + .copied() + .min_by(f64::total_cmp) + .unwrap() + } else { + lane.ready_at_ms + } + }) .min_by(f64::total_cmp) } @@ -262,6 +405,38 @@ impl AgenticPreparation { return Vec::new(); } let mut ready = Vec::new(); + if self.policy == AgenticWarmupPolicy::LiveTrajectoryV1 && !self.live_started { + // Independent predecessor streams on one lane may prime concurrently. + // Select each due primer once; a completion is not a lane-serial gate. + for lane in &mut self.lanes { + while lane.next < lane.request_indices.len() && ready.len() < limit { + // Per-lane source order is not time order. Find the earliest + // remaining primer while keeping deterministic ties. + let best = (lane.next..lane.request_indices.len()) + .min_by(|&a, &b| { + lane.primer_times_ms[a].total_cmp(&lane.primer_times_ms[b]) + }) + .unwrap(); + if lane.primer_times_ms[best] > now_ms { + break; + } + lane.source_indices.swap(lane.next, best); + lane.request_indices.swap(lane.next, best); + lane.primer_times_ms.swap(lane.next, best); + lane.ready_at_ms = lane.primer_times_ms[lane.next]; + ready.push(Self::dispatch_preparation( + lane, + &mut self.requests, + now_ms, + emit_session_metadata, + self.include_replay_hashes, + self.engine_block_size, + )); + lane.in_flight = None; + } + } + return ready; + } for lane in &mut self.lanes { if ready.len() == limit { break; @@ -272,65 +447,83 @@ impl AgenticPreparation { { continue; } - let ordinal = lane.next; - let request_index = lane.request_indices[ordinal]; - let source_index = lane.source_indices[ordinal]; - let evidence = &mut self.requests[request_index]; - let source = &lane.snapshot.context.graph.nodes[source_index]; - let metadata = DirectRequest { - uuid: Some(evidence.uuid), - max_output_tokens: 1, - // A source prefill-only request has no output plan. Its warmup - // emits a fixed execution-only token without consuming the - // corpus RNG or appending output to another request's input. - output_token_ids: Some(vec![ - lane.snapshot.context.outputs[source_index] - .first() - .copied() - .unwrap_or(0), - ]), - arrival_timestamp_ms: Some(lane.ready_at_ms), - priority: source.priority, - strict_priority: source.strict_priority, - policy_class: source.policy_class.clone(), - replay_context: Some(ReplayRequestContext { - authored_id: evidence.source_request_id.clone(), - session_id: Some(evidence.identity.conversation_id.clone()), - turn_index: Some(ordinal), - metadata: serde_json::json!({"agentic_phase": evidence.phase}), - prompt_token_source: ReplayPromptTokenSource::Materialized, - agentic: Some(evidence.identity.clone()), - }), - ..Default::default() - }; - let (request, replay_hashes) = deferred_request_with_hashes( - metadata, - source.input_length, - lane.snapshot.token_ids(source_index), - lane.snapshot.context.graph.block_size, - self.include_replay_hashes.then_some(self.engine_block_size), - ); - evidence.dispatched_at_ms = Some(now_ms); - lane.in_flight = Some(evidence.uuid); - lane.next += 1; - ready.push(CompactReadyTurn { - request_uuid: evidence.uuid, - authored_request_id: Some(evidence.source_request_id.clone()), - play_id: Some(evidence.identity.play_id.clone()), - dispatched_at_ms: now_ms, - session_id: evidence.identity.conversation_id.clone(), - turn_index: ordinal, - replay_key: None, - scheduled_ready_at_ms: lane.ready_at_ms, - replay_hashes, + ready.push(Self::dispatch_preparation( + lane, + &mut self.requests, + now_ms, emit_session_metadata, - synthetic_session_id: false, - request, - }); + self.include_replay_hashes, + self.engine_block_size, + )); } ready } + fn dispatch_preparation( + lane: &mut PreparationLane, + requests: &mut [AgenticPhaseRequest], + now_ms: f64, + emit_session_metadata: bool, + include_replay_hashes: bool, + engine_block_size: u32, + ) -> CompactReadyTurn { + let ordinal = lane.next; + let request_index = lane.request_indices[ordinal]; + let source_index = lane.source_indices[ordinal]; + let evidence = &mut requests[request_index]; + let source = &lane.snapshot.context.graph.nodes[source_index]; + let metadata = DirectRequest { + uuid: Some(evidence.uuid), + max_output_tokens: 1, + // A source prefill-only request has no output plan. Its warmup + // emits a fixed execution-only token without consuming the + // corpus RNG or appending output to another request's input. + output_token_ids: Some(vec![ + lane.snapshot.context.outputs[source_index] + .first() + .copied() + .unwrap_or(0), + ]), + arrival_timestamp_ms: Some(lane.ready_at_ms), + priority: source.priority, + strict_priority: source.strict_priority, + policy_class: source.policy_class.clone(), + replay_context: Some(ReplayRequestContext { + authored_id: evidence.source_request_id.clone(), + session_id: Some(evidence.identity.conversation_id.clone()), + turn_index: Some(ordinal), + metadata: serde_json::json!({"agentic_phase": evidence.phase}), + prompt_token_source: ReplayPromptTokenSource::Materialized, + agentic: Some(evidence.identity.clone()), + }), + ..Default::default() + }; + let (request, replay_hashes) = deferred_request_with_hashes( + metadata, + source.input_length, + lane.snapshot.token_ids(source_index), + lane.snapshot.context.graph.block_size, + include_replay_hashes.then_some(engine_block_size), + ); + evidence.dispatched_at_ms = Some(now_ms); + lane.in_flight = Some(evidence.uuid); + lane.next += 1; + CompactReadyTurn { + request_uuid: evidence.uuid, + authored_request_id: Some(evidence.source_request_id.clone()), + play_id: Some(evidence.identity.play_id.clone()), + dispatched_at_ms: now_ms, + session_id: evidence.identity.conversation_id.clone(), + turn_index: ordinal, + replay_key: None, + scheduled_ready_at_ms: lane.ready_at_ms, + replay_hashes, + emit_session_metadata, + synthetic_session_id: false, + request, + } + } + fn request(&self, uuid: Uuid) -> Result<&AgenticPhaseRequest> { let request = &self.requests[*self .index_by_uuid @@ -351,6 +544,10 @@ impl AgenticPreparation { Ok(()) } + pub(super) fn validate_boundary_time(&self, now_ms: f64) -> Result<()> { + self.validate_time(now_ms) + } + pub(super) fn validate_feedback(&self, feedback: &AgenticRuntimeFeedback) -> Result<()> { self.validate_time(feedback.at_ms)?; let mut terminal = FxHashSet::default(); @@ -449,6 +646,9 @@ impl AgenticPreparation { Some("preparation request completed without exactly one output token".into()); } } + if self.policy == AgenticWarmupPolicy::LiveTrajectoryV1 { + return Ok(()); + } let lane = self .lanes .iter_mut() @@ -493,10 +693,16 @@ impl AgenticPreparation { if self.failure_request_id.is_some() { return Ok(Some(AgenticPreparationTransition::Aborted)); } - if self - .lanes - .iter() - .all(|lane| lane.next == lane.request_indices.len()) + if (self.policy == AgenticWarmupPolicy::RepeatedPrefixes + && self + .lanes + .iter() + .all(|lane| lane.next == lane.request_indices.len())) + || (self.live_started + && self + .live_counts + .iter() + .all(|&count| count == AGENTIC_WARMUP_REQUESTS_PER_LANE)) { return Ok(Some(AgenticPreparationTransition::OpenProfile)); } @@ -549,6 +755,7 @@ impl AgenticPreparation { None => AgenticReplayPhase::Warmup, }; AgenticPhaseEvidence { + warmup_policy: self.policy, schema: AGENTIC_PHASE_SCHEMA_V1, phase, barrier_condition: "all_preparation_succeeded_and_server_quiescent", @@ -562,10 +769,10 @@ impl AgenticPreparation { .lanes .iter() .map(|lane| { - let requests = lane - .request_indices + let requests = self + .requests .iter() - .map(|&i| &self.requests[i]) + .filter(|request| request.lane_id == lane.snapshot.evidence.lane_id) .collect::>(); AgenticPhaseLane { lane_id: lane.snapshot.evidence.lane_id, @@ -690,6 +897,395 @@ mod tests { .unwrap() } + fn live_driver(lanes: usize, cut: f64) -> WorkloadDriver { + let mut driver = WorkloadDriver::new_agentic_warmup_with_policy( + prepared(lanes, cut), + 64, + true, + 1.0, + AgenticWarmupPolicy::LiveTrajectoryV1, + 10_000.0, + ) + .unwrap(); + driver + .enable_agentic_profile(super::super::AgenticProfileOptions::default()) + .unwrap(); + driver + } + + #[test] + fn live_warmup_advances_dag_recycles_and_preserves_the_formal_frontier() { + let mut driver = live_driver(1, 50.0); + let mut now = 0.0; + let mut source_ids = Vec::new(); + let mut late_quiescence = Vec::new(); + for _ in 0..30 { + driver.finish_agentic_preparation(now).unwrap(); + driver.advance_agentic_profile(now).unwrap(); + let next = driver + .next_ready_time_ms() + .expect("preparation owns a concrete next dispatch"); + now = now.max(next); + let ready = driver.pop_ready_compact(now, 100); + assert!(!ready.is_empty()); + for turn in &ready { + assert!(turn.request.materialized_tokens().is_none()); + assert_eq!(turn.request.metadata().max_output_tokens, 1); + assert_eq!( + turn.request + .metadata() + .output_token_ids + .as_ref() + .unwrap() + .len(), + 1 + ); + assert!(driver.knows_preparation_request(turn.request_uuid)); + source_ids.push(turn.authored_request_id.clone().unwrap()); + driver.on_output_token(turn.request_uuid, 0).unwrap(); + } + now += 1.0; + let counts = driver.agentic_phase_evidence().unwrap(); + let last = counts + .requests + .iter() + .filter(|r| r.phase == AgenticReplayPhase::Warmup) + .count() + == 10; + for turn in ready { + driver + .on_causal_terminal(turn.request_uuid, now, ReplayTerminalStatus::Completed) + .unwrap(); + if last { + late_quiescence.push(turn.request_uuid); + } else { + driver.on_quiescent(turn.request_uuid, now).unwrap(); + } + } + if last { + break; + } + } + assert_eq!( + source_ids, + [ + "a-history", + "b-history", // globally spread primers + "a-future", + "b-future", + "join", // live sampled suffix + "a-history", + "b-history", + "a-future", + "b-future", + "join", // first recycle + "a-history", + "b-history", // second recycle reaches quota + ] + ); + assert!(!late_quiescence.is_empty()); + assert_eq!( + driver.next_ready_time_ms(), + None, + "quota-paused requests must not spin" + ); + assert_eq!(driver.finish_agentic_preparation(now).unwrap(), None); + assert!(driver.pop_ready_compact(now, 100).is_empty()); + now += 100.0; + for uuid in late_quiescence { + driver.on_quiescent(uuid, now).unwrap(); + } + assert_eq!( + driver.finish_agentic_preparation(now).unwrap(), + Some(AgenticPreparationTransition::OpenProfile) + ); + let evidence = driver.agentic_phase_evidence().unwrap(); + assert_eq!( + evidence.warmup_policy, + AgenticWarmupPolicy::LiveTrajectoryV1 + ); + assert_eq!(evidence.lanes[0].primers_completed, 2); + assert_eq!(evidence.lanes[0].warmup_completed, 10); + assert_eq!(evidence.lanes[0].requests_quiescent, 12); + let profile = driver.agentic_profile_report().unwrap(); + assert_eq!(profile.corpus_cursor, 3); + assert_eq!(profile.issued_requests, 0); + assert_eq!(profile.plays_started, 1); + let first = driver.pop_ready(now, 100); + assert_eq!(first.len(), 1); + assert_eq!(first[0].authored_request_id.as_deref(), Some("a-future")); + assert_eq!(first[0].request.max_output_tokens, 3); + assert!(!driver.knows_preparation_request(first[0].request_uuid)); + assert_eq!( + driver.next_ready_time_ms(), + Some(now + 40.0), + "barrier waiting must not consume source-clock spacing" + ); + let identity = first[0] + .request + .replay_context + .as_ref() + .unwrap() + .agentic + .as_ref() + .unwrap(); + assert_eq!( + identity.cache_id, + evidence.requests.last().unwrap().identity.cache_id + ); + } + + #[test] + fn live_primers_are_concurrent_and_globally_aligned() { + let mut driver = live_driver(2, 50.0); + let first = driver.pop_ready(0.0, 100); + assert_eq!(first.len(), 2); + assert!( + first + .iter() + .all(|turn| turn.authored_request_id.as_deref() == Some("a-history")) + ); + assert_eq!(driver.next_ready_time_ms(), Some(20.0)); + let second = driver.pop_ready(20.0, 100); + assert_eq!( + second.len(), + 2, + "child primers need not wait for their lane's root primer response" + ); + assert!( + second + .iter() + .all(|turn| turn.authored_request_id.as_deref() == Some("b-history")) + ); + for turn in first.into_iter().chain(second) { + driver.on_output_token(turn.request_uuid, 0).unwrap(); + driver + .on_terminal(turn.request_uuid, 30.0, ReplayTerminalStatus::Completed) + .unwrap(); + } + assert_eq!(driver.finish_agentic_preparation(30.0).unwrap(), None); + let live = driver.pop_ready(30.0, 100); + assert_eq!(live.len(), 4); + assert!(live.iter().all(|turn| matches!( + turn.authored_request_id.as_deref(), + Some("a-future" | "b-future") + ))); + } + + #[test] + fn failed_live_warmup_drains_and_aborts_before_profile() { + let mut driver = live_driver(1, 0.0); + driver.finish_agentic_preparation(0.0).unwrap(); // no primers + let ready = driver.pop_ready(0.0, 100); + assert_eq!(ready.len(), 2); + driver + .on_causal_terminal(ready[0].request_uuid, 1.0, ReplayTerminalStatus::Failed) + .unwrap(); + driver.on_output_token(ready[1].request_uuid, 0).unwrap(); + driver + .on_causal_terminal(ready[1].request_uuid, 1.0, ReplayTerminalStatus::Completed) + .unwrap(); + assert_eq!(driver.finish_agentic_preparation(1.0).unwrap(), None); + assert!(driver.pop_ready(1.0, 100).is_empty()); + for turn in ready { + driver.on_quiescent(turn.request_uuid, 5.0).unwrap(); + } + assert_eq!( + driver.finish_agentic_preparation(5.0).unwrap(), + Some(AgenticPreparationTransition::Aborted) + ); + assert_eq!( + driver.agentic_profile_report().unwrap().profile_start_ms, + None + ); + } + + fn linear_live_driver(nodes: usize, rootless: bool) -> WorkloadDriver { + let rows = (0..nodes) + .map(|index| AgenticMooncakeRow { + request_id: format!("n{index}"), + play_id: "linear".into(), + session_id: if rootless && index > 0 { + "child" + } else { + "main" + } + .into(), + model: "model".into(), + input_length: Some(64), + output_length: Some(3), + hash_ids: Some(vec![index as u64 + 1]), + not_before_ms: index as f64, + recorded_api_time_ms: Some(0.1), + dependencies: if index == 0 { + vec![] + } else { + vec![AgenticDependency { + request_id: format!("n{}", index - 1), + delay_ms: 100.0, + relation: if rootless && index == 1 { + AgenticDependencyRelation::Spawn + } else { + AgenticDependencyRelation::Sequence + }, + trigger: if rootless && index == 1 { + AgenticDependencyTrigger::Dispatch + } else { + AgenticDependencyTrigger::Completion + }, + }] + }, + ..Default::default() + }) + .collect(); + let graph = ValidatedAgenticGraph::from_agentic_mooncake_rows( + AgenticMooncakeHeader { + schema: AGENTIC_MOONCAKE_SCHEMA.into(), + version: AGENTIC_MOONCAKE_VERSION, + block_size: 64, + hash_id_scope: AgenticHashIdScope::Local, + source: AgenticSourceProvenance { + format: "self-authored".into(), + digest: "live-linear".into(), + }, + }, + rows, + ) + .unwrap(); + let prepared = graph + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) + .unwrap(); + let play = prepared + .context() + .prepare_play(0, 0, Some(if rootless { 0.5 } else { 0.0 })) + .unwrap(); + let mut driver = WorkloadDriver::new_agentic_warmup_with_policy( + PreparedAgenticSnapshots::from_plays(vec![play]).unwrap(), + 64, + true, + 1.0, + AgenticWarmupPolicy::LiveTrajectoryV1, + 10_000.0, + ) + .unwrap(); + driver + .enable_agentic_profile(super::super::AgenticProfileOptions::default()) + .unwrap(); + driver + } + + #[test] + fn zero_time_live_warmup_is_bounded_and_failed_reset_is_transactional() { + let mut driver = linear_live_driver(1, false); + driver.finish_agentic_preparation(0.0).unwrap(); + for _ in 0..10 { + driver.advance_agentic_profile(0.0).unwrap(); + let ready = driver.pop_ready(0.0, 10); + assert_eq!(ready.len(), 1); + driver.on_output_token(ready[0].request_uuid, 0).unwrap(); + driver + .on_terminal(ready[0].request_uuid, 0.0, ReplayTerminalStatus::Completed) + .unwrap(); + } + let snapshots = serde_json::to_value(driver.agentic_snapshot_evidence()).unwrap(); + let profile = serde_json::to_value(driver.agentic_profile_report()).unwrap(); + let lifecycle = driver.agentic_lifecycle_transcript(); + let phase = driver.agentic_phase_evidence().unwrap(); + assert!( + driver + .finish_agentic_preparation_with(0.0, || anyhow::bail!("reset failed")) + .is_err() + ); + assert_eq!( + serde_json::to_value(driver.agentic_snapshot_evidence()).unwrap(), + snapshots + ); + assert_eq!( + serde_json::to_value(driver.agentic_profile_report()).unwrap(), + profile + ); + assert_eq!(driver.agentic_lifecycle_transcript(), lifecycle); + assert_eq!(driver.agentic_phase_evidence().unwrap(), phase); + assert_eq!( + driver.finish_agentic_preparation(0.0).unwrap(), + Some(AgenticPreparationTransition::OpenProfile) + ); + assert_eq!(driver.agentic_profile_report().unwrap().corpus_cursor, 11); + let formal = driver.pop_ready(0.0, 10); + assert_eq!(formal.len(), 1); + assert_eq!(formal[0].request.max_output_tokens, 3); + assert!(!driver.knows_preparation_request(formal[0].request_uuid)); + } + + #[test] + fn rootless_live_frontier_advances_and_recycles_without_replaying_history() { + let mut driver = linear_live_driver(3, true); + driver.finish_agentic_preparation(0.0).unwrap(); + let mut ids = Vec::new(); + for iteration in 0..10 { + let now = iteration as f64; + driver.advance_agentic_profile(now).unwrap(); + let ready = driver.pop_ready(now, 1).pop().unwrap(); + ids.push(ready.authored_request_id.clone().unwrap()); + driver.on_output_token(ready.request_uuid, 0).unwrap(); + driver + .on_terminal( + ready.request_uuid, + now + 1.0, + ReplayTerminalStatus::Completed, + ) + .unwrap(); + } + assert_eq!( + ids, + ["n1", "n2", "n0", "n1", "n2", "n0", "n1", "n2", "n0", "n1"] + ); + assert_eq!( + driver.finish_agentic_preparation(10.0).unwrap(), + Some(AgenticPreparationTransition::OpenProfile) + ); + let ready = driver.pop_ready(10.0, 1).pop().unwrap(); + assert_eq!(ready.authored_request_id.as_deref(), Some("n2")); + assert_eq!(ready.request.max_output_tokens, 3); + } + + #[test] + fn aborted_live_warmup_has_no_formal_trajectory_cohort() { + let mut driver = linear_live_driver(1, false); + driver.finish_agentic_preparation(0.0).unwrap(); + for now in [0.0, 1.0] { + driver.advance_agentic_profile(now).unwrap(); + let turn = driver.pop_ready(now, 1).pop().unwrap(); + driver.on_output_token(turn.request_uuid, 0).unwrap(); + driver + .on_terminal( + turn.request_uuid, + now + 1.0, + ReplayTerminalStatus::Completed, + ) + .unwrap(); + } + driver.advance_agentic_profile(2.0).unwrap(); + let turn = driver.pop_ready(2.0, 1).pop().unwrap(); + driver + .on_terminal(turn.request_uuid, 3.0, ReplayTerminalStatus::Failed) + .unwrap(); + assert_eq!( + driver.finish_agentic_preparation(3.0).unwrap(), + Some(AgenticPreparationTransition::Aborted) + ); + assert_eq!( + driver + .agentic_trajectory_snapshot() + .unwrap() + .total_trajectories, + 0 + ); + assert!(driver.agentic_play_outcomes().unwrap().is_empty()); + assert_eq!(driver.agentic_profile_report().unwrap().plays_started, 0); + assert_eq!(driver.agentic_phase_evidence().unwrap().requests.len(), 3); + } + #[test] fn full_snapshot_prefixes_repeat_ten_times_without_advancing_the_saved_frontier() { let prepared = prepared(2, 50.0); diff --git a/crates/core/src/replay/loadgen/snapshot.rs b/crates/core/src/replay/loadgen/snapshot.rs index 205c3372f..c8a6a1fe9 100644 --- a/crates/core/src/replay/loadgen/snapshot.rs +++ b/crates/core/src/replay/loadgen/snapshot.rs @@ -328,6 +328,10 @@ impl AgenticReplayContext { self.initial_corpus_cursor } + pub(super) fn recorded_start_ms(&self, index: usize) -> f64 { + self.recorded_intervals[index].0 + } + /// Prepare a deterministic corpus-cycle instance. `None` samples an initial /// cut; `Some(cut_ms)` uses an absolute recorded timestamp, so `Some(0.0)` /// starts a play only when its recorded clock begins at zero. Use diff --git a/crates/core/src/replay/report.rs b/crates/core/src/replay/report.rs index a92fff1db..0a387d8b9 100644 --- a/crates/core/src/replay/report.rs +++ b/crates/core/src/replay/report.rs @@ -2312,6 +2312,7 @@ mod tests { 1 ); collector.set_agentic_phases(AgenticPhaseEvidence { + warmup_policy: super::super::loadgen::AgenticWarmupPolicy::RepeatedPrefixes, schema: super::super::loadgen::AGENTIC_PHASE_SCHEMA_V1, phase: AgenticReplayPhase::Profile, barrier_condition: "all_preparation_quiescent", diff --git a/docs/replay/agentic/continuous-profiles.md b/docs/replay/agentic/continuous-profiles.md index 35ed85409..f05bbb1b9 100644 --- a/docs/replay/agentic/continuous-profiles.md +++ b/docs/replay/agentic/continuous-profiles.md @@ -87,10 +87,11 @@ is zero. The report keeps the admission cutoff separately. A configured one-hour run can have a request observation interval shorter or longer than one hour. -The warmup policy is to repeat saved primer inputs and keep -the original snapshot frontier. The referenced AgentX executor advances its live -trajectory during warmup. Initial snapshot sampling also retains AISimulate's -existing deterministic algorithm. These differences remain explicit; this is +The default warmup policy repeats saved primer inputs and keeps the original +snapshot frontier. The opt-in [live trajectory policy](warmup.md#live-trajectory-policy) +advances the dependency graph during warmup and preserves its surviving frontier. +Initial sampling retains AISimulate's existing deterministic algorithm unless +reference cuts and the shared sequential cursor are explicitly imported. This is functional qualification of profile control, not complete AgentX parity or hardware performance accuracy. Long profiles retain lifecycle evidence and can use more host memory than a finite corpus replay. Resource admission still checks diff --git a/docs/replay/agentic/warmup.md b/docs/replay/agentic/warmup.md index 188b65747..f4e9a9ec4 100644 --- a/docs/replay/agentic/warmup.md +++ b/docs/replay/agentic/warmup.md @@ -8,8 +8,11 @@ SPDX-License-Identifier: Apache-2.0 The warmup run fills the same native engine cache that will serve the measured suffix. Its input is a prepared snapshot, its original prompt identity context, and the selected -engine configuration. Its output is retained KV state, an unchanged logical -snapshot frontier, and separate preparation evidence. +engine configuration. Its output is retained KV state and separate preparation +evidence. The default `repeated_prefixes` policy preserves the snapshot frontier; +the opt-in `live_trajectory_v1` policy advances it. The next sections describe +the backward-compatible default; see [live trajectory policy](#live-trajectory-policy) +for the alternative. ## Request inputs and outputs @@ -167,6 +170,54 @@ request-instance identities; primer and profile prefix views share their play's mapping. Identity capacity exhaustion fails instead of reusing an old range. These APIs reuse the existing dependency executor and runtime feedback contract. +## Live trajectory policy + +Set `traffic.load.agentic_warmup_policy: live_trajectory_v1` together with +`agentic_warmup: true` and `agentic_profile`. Finite warmup without a continuous +profile is not supported by this policy, because shared-cursor lane recycling +belongs to profile orchestration. The CLI currently requires offline +`--stack engine`; it rejects unsupported execution before output-directory +cleanup. SDK integrations need the explicit `supports_agentic_live_warmup` +capability and the new native JSON field; passing only the warmup boolean is +not equivalent. Existing configuration serialization and seed sampling remain +unchanged when the policy is omitted. + +The policy follows the accelerated request-count warmup contract in +[AIPerf `754356e9`](https://github.com/SemiAnalysisAI/aiperf/blob/754356e9a39acc6cc6afb242d123bb57c3fb6f75/src/aiperf/timing/strategies/agentic_replay.py): + +1. Prime each continuing conversation's last historical full input. Primer + leads are globally aligned and capped by the configured global idle cap. + Independent streams on the same lane may prime concurrently. No live + request starts until every primer succeeds and settles. +2. Advance the existing dependency graph, including dispatch-triggered spawns, + completion-triggered joins, rootless children and lane recycling. Compress + authored idle delays to zero and request one output token. Admit exactly ten + additional requests per lane, excluding primers; preserve deferred requests + and cache identities when a lane reaches its quota. +3. Wait for all submitted server work to settle. Prepare the surviving frontier + using each tree's source clock; the earliest eligible offset becomes formal + time zero. An empty lane draws a fresh turn-zero tree from the shared corpus + cursor. Restore original output lengths and dependency delays. A failed + measurement reset commits neither this new frontier nor the cursor update. + +The formal request and trajectory cohorts exclude all preparation-only plays. +Partially warmed active trees retain their cache/request identity but start +their formal trajectory observation at the barrier. Formal lifecycle ordinals +restart there; complete warmup request provenance stays in `agentic_phases`. +Failure drains issued work, keeps the phase ledger, and produces no formal +cohort. The finite warmup quota permits zero-time one-token iterations; the +unbounded continuous-profile livelock guard remains enabled afterward. + +This is native functional qualification, not a claim of byte-exact HTTP prompt +or hardware-performance parity. The reference's sampler must be aligned with +[imported initial selections](../workloads.md#agentic-snapshot); identical +seeds alone are insufficient. AISimulate still uses its normalized dependency +graph and logical prefix identities, not the reference tokenizer/chat-template +payloads. Request scheduling ties use stable native graph order, not wall-clock +callback timing or random UUID ordering. CPU fixtures cover branching, +recycling, barriers, failure and formal-cohort exclusion; hardware/HTTP accuracy +requires an explicitly matched experiment. + The native report's `agentic_snapshots` collection records source/graph identity, seed, lane/play/cache identity, sampled cut, recorded request intervals, retained frontier and remaining dependency timers, and primer descriptions. Default Python @@ -191,9 +242,9 @@ The [retained source audit](../../../benchmarks/evidence/accuracy/replay-evidenc records the original comparison and review boundary. AISimulate uses its native dependency graph, a versioned BLAKE3-based snapshot -sample, and disjoint per-play token identities. Warmup repeats saved primer -inputs and preserves the frontier; the reference executor's warmup advances -its live trajectory. Matching preparation request counts does not establish +sample (or imported cuts), and disjoint per-play token identities. The default +warmup repeats saved primer inputs; `live_trajectory_v1` advances the graph as +described above. Matching preparation request counts does not establish full AgentX parity. Qualification must cover exact inputs, one output token, ten requests per lane, barrier ordering, failure handling, timer preservation, cache reuse with caching enabled/disabled, and deterministic cold/warm runs. diff --git a/docs/replay/workloads.md b/docs/replay/workloads.md index e09b074eb..132bbbc57 100644 --- a/docs/replay/workloads.md +++ b/docs/replay/workloads.md @@ -75,6 +75,7 @@ the current SA convention. | `traffic.load.agentic_snapshot.initial_selections` | unset | `x` | `-` | Exact initial `{source_play_id, t_star_ms}` selections in lane order, one per lane. Requires `corpus_cursor`. Unset retains the original seed-based sampler. | | `traffic.load.agentic_snapshot.corpus_cursor` | unset | `x` | `-` | Unsigned 64-bit next sequential corpus draw, including any skipped initial draws. Requires `initial_selections` and must be at least the lane count. Recycled plays start at turn zero. | | `traffic.load.agentic_warmup` | `false` | `x` | `-` | Optional boolean; `true` requires `agentic_snapshot` and positive `agentic_lanes`. Primes caches before measurement; see [warmup](#agentic-warmup). | +| `traffic.load.agentic_warmup_policy` | `repeated_prefixes` | `x` | `-` | `repeated_prefixes` preserves the original frontier; `live_trajectory_v1` advances the live dependency graph during ten additional one-token requests per lane. The live policy requires `agentic_warmup: true` and `agentic_profile`; initially offline `--stack engine` only. | | `traffic.load.agentic_profile` | `null` (unset) | `x` | `-` | Optional object; `{}` enables continuous lane replenishment with the defaults below. Requires `trace_timestamps`, positive `agentic_lanes`, and `agentic_snapshot`; cannot be combined with `traffic.stop.max_virtual_time_seconds`. Unset preserves finite replay. See [continuous profile](#agentic-profile). | | `traffic.load.agentic_profile.duration_seconds` | `3600` when enabled | `x` | `-` | Positive finite admission duration, starting at the preparation barrier or simulation start without warmup. No new workload requests or replacement plays are issued after the deadline. | | `traffic.load.agentic_profile.response_grace_seconds` | `30` when enabled | `x` | `-` | Nonnegative finite time for already submitted requests to respond after the admission deadline; remaining client requests are then canceled. | @@ -382,7 +383,8 @@ runtime dependency. ### Warmup -`agentic_warmup: true` fills the caches before measurement, in three stages: +`agentic_warmup: true` fills the caches before measurement. The default +`agentic_warmup_policy: repeated_prefixes` has three stages: | Stage | Requests | Output | | --- | --- | --- | @@ -402,6 +404,36 @@ prefixes can still be evicted or placed on another worker, so warmup does not guarantee cache hits. See [agentic warmup](agentic/warmup.md) for the full preparation, barrier and reporting contract. +For AgentX-style live advancement, select the versioned policy explicitly: + +```yaml +agentic_lanes: 15 +agentic_snapshot: {seed: 42} # import reference selections for matching cuts +agentic_warmup: true +agentic_warmup_policy: live_trajectory_v1 +agentic_profile: {duration_seconds: 3600} +``` + +This policy primes predecessor streams with globally aligned lead times, capped +by `agentic_profile.global_idle_cap_seconds`. After all primers settle, it runs +the actual dependency graph with zero authored idle delay and one output token +per request. Primers do not consume the additional quota of ten admitted +requests per lane. Spawns, joins, rootless children and sequential corpus +recycling continue during warmup; an exhausted lane pauses without discarding +its next request. At the all-settled barrier, the surviving frontier resumes +with original output lengths and source-clock offsets. One global minimum is +subtracted from eligible frontier offsets, preserving cross-lane spacing; +waiting at the barrier does not consume those delays. + +Successful profiles exclude preparation-only requests and retired plays from +counts and latency distributions. Active partially warmed plays retain their +identities and get a formal measurement origin at the barrier. An aborted live +warmup has an empty formal cohort; its request ledger remains in +`agentic_phases`. See [live trajectory warmup](agentic/warmup.md#live-trajectory-policy) +for qualification boundaries. Unsupported CLI stacks are rejected before +`--overwrite` prepares the output directory; external SDK adapters must +explicitly advertise `supports_agentic_live_warmup` and forward the policy. + ### Continuous profile diff --git a/python/aisimulate/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index 11ebdfd95..59de8b86e 100644 --- a/python/aisimulate/THIRD_PARTY_NOTICES.md +++ b/python/aisimulate/THIRD_PARTY_NOTICES.md @@ -32,6 +32,18 @@ Modified by NVIDIA: executed the pinned sampler on a small synthetic corpus to freeze expected cuts, and documented its sequential recycling behavior. No upstream implementation is vendored; AISimulate consumes explicit cuts. +The live warmup policy in repository-root +`crates/core/src/replay/loadgen/phase.rs` and +`crates/core/src/replay/loadgen/driver.rs`, its regression tests in those files, +and `docs/replay/agentic/warmup.md`, `docs/replay/agentic/continuous-profiles.md` +and `docs/replay/workloads.md` are modified re-implementations of the same +pinned `agentic_replay.py` behavior: globally aligned primers, per-lane +admission quotas, compressed dependency waits, shared-cursor recycling, and +source-clock frontier handoff. NVIDIA modified these algorithms to use the +native deterministic dependency executor and its transactional server-settlement +barrier. The original NVIDIA copyright and Apache-2.0 license apply; no +upstream serving implementation is vendored. + ## AIConfigurator The repository-root `.coderabbit.yaml` is adapted and modified from diff --git a/python/aisimulate/src/aisimulate/compiler.py b/python/aisimulate/src/aisimulate/compiler.py index da6445bcd..d23268159 100644 --- a/python/aisimulate/src/aisimulate/compiler.py +++ b/python/aisimulate/src/aisimulate/compiler.py @@ -837,6 +837,8 @@ def _traffic( workload["agentic_snapshot"] = load.agentic_snapshot.model_dump(mode="json") if load.agentic_warmup: workload["agentic_warmup"] = True + if load.agentic_warmup_policy != "repeated_prefixes": + workload["agentic_warmup_policy"] = load.agentic_warmup_policy if load.agentic_profile is not None: workload["agentic_profile"] = load.agentic_profile.model_dump(mode="json") if stop is not None and stop.max_virtual_time_seconds is not None: diff --git a/python/aisimulate/src/aisimulate/config/traffic.py b/python/aisimulate/src/aisimulate/config/traffic.py index 8fa89333d..da7d3c913 100644 --- a/python/aisimulate/src/aisimulate/config/traffic.py +++ b/python/aisimulate/src/aisimulate/config/traffic.py @@ -160,6 +160,15 @@ def _validate_millisecond_deadlines(self) -> AgenticProfileOptions: class TrafficPredictionLoad(StrictModel): + agentic_warmup_policy: Literal["repeated_prefixes", "live_trajectory_v1"] = "repeated_prefixes" + + @model_serializer(mode="wrap") + def _omit_default_warmup_policy(self, handler): + result = handler(self) + if self.agentic_warmup_policy == "repeated_prefixes": + result.pop("agentic_warmup_policy", None) + return result + type: Literal[ "concurrency", "poisson", @@ -191,6 +200,15 @@ def _validate_fields_for_type(self) -> TrafficPredictionLoad: class TrafficRecommendationLoad(StrictModel): + agentic_warmup_policy: Literal["repeated_prefixes", "live_trajectory_v1"] = "repeated_prefixes" + + @model_serializer(mode="wrap") + def _omit_default_warmup_policy(self, handler): + result = handler(self) + if self.agentic_warmup_policy == "repeated_prefixes": + result.pop("agentic_warmup_policy", None) + return result + type: Literal[ "concurrency", "poisson", @@ -252,12 +270,21 @@ def _validate_load_fields(load) -> None: } if load.agentic_warmup: used.add("agentic_warmup") + if load.agentic_warmup_policy != "repeated_prefixes": + used.add("agentic_warmup_policy") allowed = { "concurrency": {"concurrency"}, "poisson": {"requests_per_second", "sessions_per_second", "seed"}, "constant_rate": {"requests_per_second", "sessions_per_second"}, "kv_capacity_fraction": {"fraction"}, - "trace_timestamps": {"speedup", "agentic_lanes", "agentic_snapshot", "agentic_warmup", "agentic_profile"}, + "trace_timestamps": { + "speedup", + "agentic_lanes", + "agentic_snapshot", + "agentic_warmup", + "agentic_warmup_policy", + "agentic_profile", + }, }[load.type] unexpected = used - allowed if unexpected: @@ -288,6 +315,10 @@ def _validate_source_load_stop(self, load) -> None: raise ValueError("agentic_profile cannot be combined with max_virtual_time_seconds") if load.agentic_warmup and load.agentic_snapshot is None: raise ValueError("agentic_warmup requires agentic_snapshot") + if load.agentic_warmup_policy == "live_trajectory_v1" and ( + not load.agentic_warmup or load.agentic_profile is None + ): + raise ValueError("live_trajectory_v1 requires agentic_warmup and agentic_profile") if load.agentic_snapshot is not None and load.agentic_lanes is None: raise ValueError("agentic_snapshot requires positive agentic_lanes") if ( diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py index a517f1455..989d7b39f 100644 --- a/python/aisimulate/src/aisimulate/main.py +++ b/python/aisimulate/src/aisimulate/main.py @@ -157,6 +157,8 @@ def _compile_prediction_adapters( def _predict(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: core_raw, adapter_raw = split_config_sections(raw, command="predict") config = CorePredictionConfig.model_validate(core_raw) + if config.traffic.load.agentic_warmup_policy == "live_trajectory_v1" and (args.stack != "engine" or args.online): + raise ValueError("live_trajectory_v1 CLI requires offline --stack engine") if config.engine.speculation is not None and (args.stack != "engine" or args.online or adapter_raw): raise ValueError("ngram speculation requires offline --stack engine without adapters") plan = _resource_plan(args, config, factory) @@ -280,6 +282,12 @@ def _recommend(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: output_adapters = resolve_output_adapters(output_configs) core_raw, adapter_raw = split_config_sections(raw, command="recommend") config = CoreRecommendationConfig.model_validate(core_raw) + if ( + config.traffic is not None + and config.traffic.load.agentic_warmup_policy == "live_trajectory_v1" + and args.stack != "engine" + ): + raise ValueError("live_trajectory_v1 CLI requires offline --stack engine") adapters = _resolve_section_adapters(adapter_raw, args.stack) result = run_recommendation( config, diff --git a/python/aisimulate/src/aisimulate/recommend.py b/python/aisimulate/src/aisimulate/recommend.py index 2ba36b107..0985fad1b 100644 --- a/python/aisimulate/src/aisimulate/recommend.py +++ b/python/aisimulate/src/aisimulate/recommend.py @@ -732,6 +732,8 @@ def _recommendation_workload(raw: dict[str, Any] | None) -> dict[str, Any]: result["agentic_snapshot"] = deepcopy(load["agentic_snapshot"]) if load.get("agentic_warmup"): result["agentic_warmup"] = True + if load.get("agentic_warmup_policy", "repeated_prefixes") != "repeated_prefixes": + result["agentic_warmup_policy"] = load["agentic_warmup_policy"] if load.get("agentic_profile") is not None: result["agentic_profile"] = deepcopy(load["agentic_profile"]) if isinstance(stop, dict) and stop.get("max_virtual_time_seconds") is not None: diff --git a/python/aisimulate/src/aisimulate/runner.py b/python/aisimulate/src/aisimulate/runner.py index a408a28d4..78099eae5 100644 --- a/python/aisimulate/src/aisimulate/runner.py +++ b/python/aisimulate/src/aisimulate/runner.py @@ -83,6 +83,7 @@ "agentic_lanes", "agentic_snapshot", "agentic_warmup", + "agentic_warmup_policy", "agentic_profile", } ) @@ -398,6 +399,7 @@ def capabilities(self) -> RunnerCapabilities: supports_agentic_lanes=True, supports_agentic_snapshots=True, supports_agentic_warmup=True, + supports_agentic_live_warmup=True, supports_agentic_profile=True, supported_agentic_topologies=("agg", "disagg"), supported_agentic_backends=("vllm", "sglang"), diff --git a/python/aisimulate/src/aisimulate/supervision.py b/python/aisimulate/src/aisimulate/supervision.py index 5de1939a3..c8c0a3241 100644 --- a/python/aisimulate/src/aisimulate/supervision.py +++ b/python/aisimulate/src/aisimulate/supervision.py @@ -428,6 +428,12 @@ def main(argv: Sequence[str] | None = None) -> int: core_raw, adapter_raw = split_config_sections(validation_raw, command=args.command) config_type = CorePredictionConfig if args.command == "predict" else CoreRecommendationConfig config = config_type.model_validate(core_raw) + if ( + config.traffic is not None + and config.traffic.load.agentic_warmup_policy == "live_trajectory_v1" + and (args.stack != "engine" or (args.command == "predict" and args.online)) + ): + raise ValueError("live_trajectory_v1 CLI requires offline --stack engine") if config.engine.workers.encoder is not None: if args.command == "predict" and ( args.stack != "engine" or args.online or args.capture_per_request or adapter_raw diff --git a/python/aisimulate/src/aisimulate/sweeper/config.py b/python/aisimulate/src/aisimulate/sweeper/config.py index 85596e20c..bb9fa699c 100644 --- a/python/aisimulate/src/aisimulate/sweeper/config.py +++ b/python/aisimulate/src/aisimulate/sweeper/config.py @@ -281,6 +281,7 @@ class Workload(BaseModel): agentic_lanes: int | None = Field(default=None, strict=True, gt=0) agentic_snapshot: AgenticSnapshotOptions | None = None agentic_warmup: bool = Field(default=False, strict=True) + agentic_warmup_policy: Literal["repeated_prefixes", "live_trajectory_v1"] = "repeated_prefixes" agentic_profile: AgenticProfileOptions | None = None # Closed-loop replay over a *trace*: cap in-flight requests at this many (the # trace's timestamps are ignored; a new request starts as one finishes). For a @@ -339,6 +340,13 @@ def _validate_cached_prefix_tokens_type(cls, value: Any) -> Any: def is_trace_based(self) -> bool: return self.trace_path is not None + @model_serializer(mode="wrap") + def _omit_default_warmup_policy(self, handler): + result = handler(self) + if self.agentic_warmup_policy == "repeated_prefixes": + result.pop("agentic_warmup_policy", None) + return result + @property def is_synthetic(self) -> bool: return self.trace_path is None @@ -402,6 +410,10 @@ def _validate_workload(self) -> Workload: raise ValueError("weka_nested_timestamp_basis requires Weka trace input") if self.agentic_warmup and self.agentic_snapshot is None: raise ValueError("agentic_warmup requires agentic_snapshot") + if self.agentic_warmup_policy == "live_trajectory_v1" and ( + not self.agentic_warmup or self.agentic_profile is None + ): + raise ValueError("live_trajectory_v1 requires agentic_warmup and agentic_profile") if self.agentic_snapshot is not None and self.agentic_lanes is None: raise ValueError("agentic_snapshot requires positive agentic_lanes") if self.agentic_lanes is not None: diff --git a/python/aisimulate/src/aisimulate/sweeper/replay.py b/python/aisimulate/src/aisimulate/sweeper/replay.py index bb8c28ab6..81c11eba2 100644 --- a/python/aisimulate/src/aisimulate/sweeper/replay.py +++ b/python/aisimulate/src/aisimulate/sweeper/replay.py @@ -329,6 +329,7 @@ class RunnerCapabilities: supports_agentic_speculative_decoding: bool = True supports_agentic_snapshots: bool = False supports_agentic_warmup: bool = False + supports_agentic_live_warmup: bool = False supports_cached_prefix_tokens: bool = False supported_engine_model_controls: tuple[str, ...] = () supports_mtp_expected_acceptance: bool = False @@ -511,6 +512,17 @@ def require_compatible(self, spec: ReplaySpec) -> None: raise ValueError("agentic_warmup requires agentic_snapshot") if not self.supports_agentic_warmup: raise ValueError("runner does not support agentic warmup") + warmup_policy = spec.workload.get("agentic_warmup_policy", "repeated_prefixes") + if not isinstance(warmup_policy, str) or warmup_policy not in {"repeated_prefixes", "live_trajectory_v1"}: + raise ValueError("unknown agentic_warmup_policy") + if warmup_policy == "live_trajectory_v1": + if not agentic_warmup or agentic_profile is None: + raise ValueError("live_trajectory_v1 requires agentic_warmup and agentic_profile") + if not self.supports_agentic_live_warmup: + raise ValueError( + "runner does not support live agentic warmup; " + "update the adapter before selecting live_trajectory_v1" + ) if agentic_snapshot is not None: snapshot = AgenticSnapshotOptions.model_validate(agentic_snapshot) if agentic_lanes is None or spec.workload.get("load_type") != "trace_timestamps": diff --git a/tests/test_agentx_live_warmup.py b/tests/test_agentx_live_warmup.py new file mode 100644 index 000000000..ca56d9590 --- /dev/null +++ b/tests/test_agentx_live_warmup.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from pydantic import ValidationError + +from aisimulate.config.traffic import TrafficPredictionConfig, TrafficRecommendationConfig +from aisimulate.sweeper.config import Workload + + +def traffic(): + return { + "source": {"type": "trace", "paths": ["corpus"], "format": "weka"}, + "load": { + "type": "trace_timestamps", + "agentic_lanes": 2, + "agentic_snapshot": {"seed": 42}, + "agentic_warmup": True, + "agentic_warmup_policy": "live_trajectory_v1", + "agentic_profile": {}, + }, + } + + +@pytest.mark.parametrize("schema", [TrafficPredictionConfig, TrafficRecommendationConfig]) +def test_live_warmup_policy_roundtrip_and_legacy_default_omission(schema): + source = traffic() + parsed = schema.model_validate(source) + assert parsed.load.agentic_warmup_policy == "live_trajectory_v1" + assert parsed.model_dump()["load"]["agentic_warmup_policy"] == "live_trajectory_v1" + source["load"].pop("agentic_warmup_policy") + assert "agentic_warmup_policy" not in schema.model_validate(source).model_dump()["load"] + + +@pytest.mark.parametrize("missing", ["agentic_warmup", "agentic_profile"]) +def test_live_warmup_requires_preparation_and_continuous_profile(missing): + source = traffic() + source["load"].pop(missing) + with pytest.raises(ValidationError, match="live_trajectory_v1 requires"): + TrafficPredictionConfig.model_validate(source) + + +@pytest.mark.parametrize("policy", [True, "live", {}, None]) +def test_live_warmup_rejects_unknown_policy(policy): + source = traffic() + source["load"]["agentic_warmup_policy"] = policy + with pytest.raises(ValidationError): + TrafficPredictionConfig.model_validate(source) + + +def test_live_warmup_workload_serialization_keeps_the_policy_explicit(): + data = { + "source_type": "trace", + "load_type": "trace_timestamps", + "trace_path": "corpus", + "trace_format": "weka", + "agentic_lanes": 2, + "agentic_snapshot": {"seed": 42}, + "agentic_warmup": True, + "agentic_profile": {}, + } + assert "agentic_warmup_policy" not in Workload.model_validate(data).model_dump() + data["agentic_warmup_policy"] = "live_trajectory_v1" + assert Workload.model_validate(data).model_dump()["agentic_warmup_policy"] == "live_trajectory_v1" + + +@pytest.mark.parametrize("stack,online", [("legacy", False), ("engine", True)]) +def test_supervisor_rejects_unsupported_live_warmup_before_overwrite(tmp_path, monkeypatch, capsys, stack, online): + from aisimulate import cli_args, output, supervision + + prior = tmp_path / "results.json" + prior.write_bytes(b"previous valid results\n") + raw = { + "engine": {"model": "example/model", "hardware": "h200_sxm", "workers": {"aggregated": {}}}, + "traffic": traffic(), + } + monkeypatch.setattr(cli_args, "_load_mapping", lambda _: raw) + monkeypatch.setattr(supervision, "_installed_names", lambda _: ["legacy"]) + calls = [] + + def forbidden_prepare(*args, **kwargs): + calls.append((args, kwargs)) + pytest.fail("unsupported policy reached output cleanup") + + monkeypatch.setattr(output, "prepare_output_directory", forbidden_prepare) + arguments = ["predict", "--config", "unused.yaml", "--stack", stack, "--output-dir", str(tmp_path), "--overwrite"] + if online: + arguments.append("--online") + with pytest.raises(SystemExit) as exc: + supervision.main(arguments) + assert exc.value.code == 2 + assert "live_trajectory_v1 CLI requires offline --stack engine" in capsys.readouterr().err + assert not calls + assert prior.read_bytes() == b"previous valid results\n" + + +def test_supervisor_default_recommendation_without_traffic_keeps_existing_path(monkeypatch): + from aisimulate import cli_args, output, supervision + + raw = { + "engine": { + "model": "example/model", + "hardware": "h200_sxm", + "mode": "aggregated", + "workers": {"aggregated": {}}, + }, + "optimization": {"target": "throughput_per_gpu", "constraints": {"max_candidate_gpus": 1}}, + } + monkeypatch.setattr(cli_args, "_load_mapping", lambda _: raw) + + class ReachedPreparation(Exception): + pass + + def observe_prepare(*args, **kwargs): + raise ReachedPreparation + + monkeypatch.setattr(output, "prepare_output_directory", observe_prepare) + with pytest.raises(ReachedPreparation): + supervision.main(["recommend", "--config", "unused.yaml", "--stack", "engine", "--output-dir", "/unused"]) diff --git a/tests/test_runner.py b/tests/test_runner.py index cb67af018..9005fdf09 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -2191,6 +2191,27 @@ def test_agentic_profile_capability_is_explicit() -> None: ) +def test_live_agentic_warmup_requires_explicit_runner_capability_and_reaches_native(): + workload = { + "source_type": "trace", + "load_type": "trace_timestamps", + "trace_path": "corpus", + "trace_format": "weka", + "agentic_lanes": 1, + "agentic_snapshot": {"seed": 42}, + "agentic_warmup": True, + "agentic_profile": {}, + "agentic_warmup_policy": "live_trajectory_v1", + } + capabilities = EngineReplayRunnerFactory().capabilities() + assert capabilities.supports_agentic_live_warmup + with pytest.raises(ValueError, match="does not support live agentic warmup"): + replace(capabilities, supports_agentic_live_warmup=False).require_compatible(_spec(workload=workload)) + runtime = RecordingRuntime() + EngineReplayRunnerFactory(runtime=runtime).create(0).run(_spec(workload=workload)) + assert runtime.execution_spec["traffic"]["agentic_warmup_policy"] == "live_trajectory_v1" + + @pytest.mark.parametrize("sampler", ["numpy_random_state", "python_random"]) def test_runner_honors_length_sampler(sampler): # Fixed vectors from the benchmark's NumPy RandomState contract and the From 17356c1328ea5939ec60be83ff590974585060ce Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 11:05:45 -0700 Subject: [PATCH 04/13] chore: route snapshot exporter tests to replay owners Signed-off-by: hongkuanz --- .github/codeowners/areas.yaml | 3 +++ CODEOWNERS | 2 ++ 2 files changed, 5 insertions(+) diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 3fe35b246..503e8d8f4 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -72,6 +72,7 @@ areas: - tests/test_g3_offload_cli.py - tests/test_shared_g2_cli.py - tests/test_agentx_g2.py + - tests/test_agentx_snapshot_export.py - tests/test_runner.py - tests/test_traffic.py - tests/e2e/ @@ -271,6 +272,8 @@ shared: owners: [replay, maintainers] - glob: tests/test_agentx_g2.py owners: [replay, maintainers] + - glob: tests/test_agentx_snapshot_export.py + owners: [replay, maintainers] - glob: tests/test_runner.py owners: [replay, maintainers] - glob: tests/test_traffic.py diff --git a/CODEOWNERS b/CODEOWNERS index a09316d01..9e2732a36 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -78,6 +78,7 @@ /tests/test_host_offload_cli.py @ai-dynamo/aisimulate-replay-codeowners /tests/test_power_diagnostics.py @ai-dynamo/aisimulate-replay-codeowners /tests/test_state_cache_config.py @ai-dynamo/aisimulate-replay-codeowners +/tests/test_agentx_snapshot_export.py @ai-dynamo/aisimulate-replay-codeowners /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 @@ -215,6 +216,7 @@ /crates/core/parity_tests/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_ci_workflow_contracts.py @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_e2e_accuracy_overview.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain +/tests/test_agentx_snapshot_export.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_e2e_accuracy_workflow.mjs @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_fpm_accuracy_workflow.mjs @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_output_adapter_context.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain From 953711ce34b7b67131048fffab0726ee63f1dcc8 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 11:06:36 -0700 Subject: [PATCH 05/13] fix: isolate warmup progress counters and preserve linear reporting Signed-off-by: hongkuanz --- crates/core/src/python.rs | 29 ++++++++++++++++++++++++ crates/core/src/replay/loadgen/driver.rs | 3 +++ crates/core/src/replay/loadgen/phase.rs | 17 ++++++++++---- 3 files changed, 44 insertions(+), 5 deletions(-) diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index 6d75ff91c..83bc00183 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -3198,6 +3198,35 @@ mod tests { } } + #[test] + fn live_zero_time_warmup_does_not_poison_the_formal_progress_guard() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("mixed-duration.jsonl"); + let rows = [("zero", 1), ("positive", 3)].into_iter().map(|(id, output)| serde_json::json!({ + "id": id, "models": ["model"], "block_size": 4, "hash_id_scope": "local", + "requests": [{"t": 0.0, "type": "s", "model": "model", "in": 4, "out": output, "hash_ids": [1]}] + }).to_string()).collect::>().join("\n"); + std::fs::write(&path, rows).unwrap(); + let payload = serde_json::json!({ + "spec": { + "version": 1, "topology": {"kind": "aggregated", "workers": {"initial_workers": 1}}, + "engine": {"rank": {"backend": "vllm", "block_size": 4, "num_gpu_blocks": 64, + "timing_model": {"type": "fixed", "prefill_ms": 0.0, "decode_ms": 1.0}}}, + "requests": [] + }, + "traffic": {"source_type": "trace", "load_type": "trace_timestamps", "trace_format": "weka", + "trace_path": path, "trace_block_size": 4, "execution_model": "model", + "agentic_lanes": 1, "agentic_snapshot": {"seed": 42}, "agentic_warmup": true, + "agentic_warmup_policy": "live_trajectory_v1", "agentic_profile": {"duration_seconds": 0.01}} + }); + let report: serde_json::Value = + serde_json::from_str(&execute_json(&payload.to_string(), false).unwrap()).unwrap(); + assert_eq!(report["agentic_phases"]["profile_start_ms"], 0.0); + assert_eq!(report["agentic_phases"]["lanes"][0]["warmup_completed"], 10); + assert!(report["duration_ms"].as_f64().unwrap() >= 10.0); + assert!(report["completed_requests"].as_u64().unwrap() >= 2); + } + #[test] fn public_agentic_json_rejects_unqualified_modes_without_restricting_standard_dynamo() { let directory = tempfile::tempdir().unwrap(); diff --git a/crates/core/src/replay/loadgen/driver.rs b/crates/core/src/replay/loadgen/driver.rs index 9ea51f7ee..d4dae1df8 100644 --- a/crates/core/src/replay/loadgen/driver.rs +++ b/crates/core/src/replay/loadgen/driver.rs @@ -967,6 +967,9 @@ impl WorkloadDriver { bail!("agentic profile deadline overflow at activation"); } profile.start_ms = Some(now_ms); + // Warmup has a finite admission bound; the formal profile starts its + // own same-time progress window after measurement reset succeeds. + profile.replacements_at_ms = (None, 0); profile.last_advanced_ms = now_ms; profile.global_idle_since_ms = Some(now_ms); profile.tree_idle_since_ms.fill(Some(now_ms)); diff --git a/crates/core/src/replay/loadgen/phase.rs b/crates/core/src/replay/loadgen/phase.rs index 10cb0f305..8212e2171 100644 --- a/crates/core/src/replay/loadgen/phase.rs +++ b/crates/core/src/replay/loadgen/phase.rs @@ -275,6 +275,12 @@ impl AgenticPreparation { }); } if policy == AgenticWarmupPolicy::LiveTrajectoryV1 { + ensure!( + this.lanes.len() == prepared.context().lane_count(), + "live warmup requires a complete initial lane cohort" + ); + this.lanes + .sort_by_key(|lane| lane.snapshot.evidence.lane_id); let mut max_lead_ms = 0.0_f64; for lane in &mut this.lanes { lane.primer_times_ms = lane @@ -339,6 +345,7 @@ impl AgenticPreparation { ) { debug_assert!(self.lane_has_budget(lane)); let index = self.requests.len(); + self.lanes[lane].request_indices.push(index); self.index_by_uuid.insert(turn.request_uuid, index); self.live_requests.insert(turn.request_uuid); self.live_counts[lane] += 1; @@ -375,7 +382,7 @@ impl AgenticPreparation { } pub(super) fn next_ready_time_ms(&self) -> Option { - if self.failure_request_id.is_some() || self.transition.is_some() { + if self.live_started || self.failure_request_id.is_some() || self.transition.is_some() { return None; } self.lanes @@ -401,7 +408,7 @@ impl AgenticPreparation { limit: usize, emit_session_metadata: bool, ) -> Vec { - if self.failure_request_id.is_some() || self.transition.is_some() { + if self.live_started || self.failure_request_id.is_some() || self.transition.is_some() { return Vec::new(); } let mut ready = Vec::new(); @@ -769,10 +776,10 @@ impl AgenticPreparation { .lanes .iter() .map(|lane| { - let requests = self - .requests + let requests = lane + .request_indices .iter() - .filter(|request| request.lane_id == lane.snapshot.evidence.lane_id) + .map(|&index| &self.requests[index]) .collect::>(); AgenticPhaseLane { lane_id: lane.snapshot.evidence.lane_id, From 9d63c91089db4dde7b7cf6cd159b955056fb3954 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 11:06:39 -0700 Subject: [PATCH 06/13] chore: route AgentX warmup tests to replay owners Signed-off-by: hongkuanz --- .github/codeowners/areas.yaml | 6 ++++++ CODEOWNERS | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 3fe35b246..8fa69ae53 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -72,6 +72,8 @@ areas: - tests/test_g3_offload_cli.py - tests/test_shared_g2_cli.py - tests/test_agentx_g2.py + - tests/test_agentx_snapshot_export.py + - tests/test_agentx_live_warmup.py - tests/test_runner.py - tests/test_traffic.py - tests/e2e/ @@ -271,6 +273,10 @@ shared: owners: [replay, maintainers] - glob: tests/test_agentx_g2.py owners: [replay, maintainers] + - glob: tests/test_agentx_snapshot_export.py + owners: [replay, maintainers] + - glob: tests/test_agentx_live_warmup.py + owners: [replay, maintainers] - glob: tests/test_runner.py owners: [replay, maintainers] - glob: tests/test_traffic.py diff --git a/CODEOWNERS b/CODEOWNERS index a09316d01..5898da1e1 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -77,7 +77,9 @@ /tests/test_stack_discovery.py @ai-dynamo/aisimulate-replay-codeowners /tests/test_host_offload_cli.py @ai-dynamo/aisimulate-replay-codeowners /tests/test_power_diagnostics.py @ai-dynamo/aisimulate-replay-codeowners +/tests/test_agentx_live_warmup.py @ai-dynamo/aisimulate-replay-codeowners /tests/test_state_cache_config.py @ai-dynamo/aisimulate-replay-codeowners +/tests/test_agentx_snapshot_export.py @ai-dynamo/aisimulate-replay-codeowners /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 @@ -200,6 +202,7 @@ /docs/aic-backward-compatibility/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/.gitattributes @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/pyproject.toml @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-mocker-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain +/tests/test_agentx_live_warmup.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_fpm_query_evidence.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_resource_scheduler.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_state_cache_config.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-mocker-codeowners @ai-dynamo/access-aisimulate-maintain @@ -215,6 +218,7 @@ /crates/core/parity_tests/perfmodel/ @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_ci_workflow_contracts.py @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_e2e_accuracy_overview.py @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain +/tests/test_agentx_snapshot_export.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_e2e_accuracy_workflow.mjs @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_fpm_accuracy_workflow.mjs @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_output_adapter_context.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain From 93517a7ab99fa93f9b06c251c6fd96194bd6715c Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 13:26:35 -0700 Subject: [PATCH 07/13] fix: validate and document imported snapshots across replay entrypoints Signed-off-by: hongkuanz --- docs/replay/workloads.md | 6 ++- docs/sweeper/sdk.md | 8 ++- .../src/aisimulate/sweeper/config.py | 6 +++ tests/test_cli_config.py | 51 ++++++++++++++++-- tests/test_runner.py | 53 ++++++++++++------- 5 files changed, 96 insertions(+), 28 deletions(-) diff --git a/docs/replay/workloads.md b/docs/replay/workloads.md index e09b074eb..b04ec7194 100644 --- a/docs/replay/workloads.md +++ b/docs/replay/workloads.md @@ -263,7 +263,7 @@ traffic: | Control | Without it | With it | | --- | --- | --- | | `agentic_lanes` | Every play starts at its authored timestamp. | At most `N` plays run at once. | -| `agentic_snapshot` | Each play starts at turn zero. | Each lane starts from a sampled point inside its play. | +| `agentic_snapshot` | Each play starts at turn zero. | Each lane starts from a sampled or imported exact cut inside its play. | | `agentic_warmup` | Measurement starts with cold caches. | Caches are primed before measurement starts. | | `agentic_profile` | The run ends when the selected plays finish. | Lanes keep taking new plays until a deadline. | @@ -334,6 +334,8 @@ agentic_snapshot: corpus_cursor: 1 ``` +`initial_selections` and `corpus_cursor` must be supplied together, with exactly +one selection per lane and a cursor at least as large as the selection count. Selections are in lane order. Each source ID must exist in the exact input corpus and each cut must lie within that play's request-start bounds. Cuts are in source milliseconds, before `speedup`; they are never re-sampled or @@ -448,7 +450,7 @@ aisimulate predict --config examples/cli/agentic-profile.yaml \ | Section | Contents | | --- | --- | -| `agentic_snapshots` | Seed, sampled cut, history and remaining timers for each initial lane | +| `agentic_snapshots` | Seed, sampled or imported exact cut, history and remaining timers for each initial lane | | `agentic_phases` | Preparation requests, barrier state and time (`profile_start_ms`) | | `agentic_play_outcomes` | Each play's terminal status (`completed`, `failed` or `incomplete`) and times | | `agentic_profile` | Resolved profile options, deadline, recycled plays, canceled and never-issued requests, unsettled server work | diff --git a/docs/sweeper/sdk.md b/docs/sweeper/sdk.md index 3372aabd2..49995fac4 100644 --- a/docs/sweeper/sdk.md +++ b/docs/sweeper/sdk.md @@ -212,7 +212,9 @@ Every `Workload` field: | `trace_format` | `str` | `"mooncake"` | Replay-ready trace schema. A runner may validate supported formats. | | `arrival_speedup_ratio` | `float` | `1.0` | Scales the trace's inter-arrival times (open-loop trace only). `>1` speeds arrivals up. | | `agentic_lanes` | `int \| None` | `None` | Positive number of client play lanes for `weka`, `agentic_mooncake`, or agentic `dynamo` timestamp replay; does not cap concurrent child requests. Requires `source_type: trace`, `trace_path`, and `load_type: trace_timestamps`, with no `replay_concurrency`. | -| `agentic_snapshot` | `AgenticSnapshotOptions \| None` | `None` (unset) | Optional object `{seed: u64}`; required `seed` is an unsigned 64-bit integer (`0` through `2^64 - 1`). Requires positive `agentic_lanes`, `source_type: trace`, `trace_path`, and `load_type: trace_timestamps`, with no `replay_concurrency`; supported formats are `weka`, `agentic_mooncake`, and agentic `dynamo`. Unset preserves turn-zero execution. | +| `agentic_snapshot` | `AgenticSnapshotOptions \| None` | `None` (unset) | Optional object with required unsigned 64-bit `seed` (`0` through `2^64 - 1`) and the optional import fields below. Requires positive `agentic_lanes`, `source_type: trace`, `trace_path`, and `load_type: trace_timestamps`, with no `replay_concurrency`; supported formats are `weka`, `agentic_mooncake`, and agentic `dynamo`. Unset preserves turn-zero execution. | +| `agentic_snapshot.initial_selections` | `list[AgenticSnapshotSelection] \| None` | `None` (unset) | Exact imported `{source_play_id, t_star_ms}` cuts in lane order, with exactly one selection per `agentic_lanes`. Requires `corpus_cursor`; source IDs must exist in the input corpus and finite source-millisecond cuts must lie within the selected play's request-start bounds. Cuts are not re-sampled. | +| `agentic_snapshot.corpus_cursor` | `int \| None` | `None` (unset) | Unsigned 64-bit next sequential source draw, including skipped draws. Requires `initial_selections` and must be at least their count; it may exceed the lane count. | | `agentic_warmup` | `bool` | `False` | Requires `agentic_snapshot` and positive lanes when enabled. Runs primers and ten warmup requests per lane before measured replay, preserving native cache state; see [warmup behavior](../replay/workloads.md#agentic-warmup). Failed preparation is retained as candidate evidence and excluded from ranking. | | `agentic_profile` | `AgenticProfileOptions \| None` | `None` (unset) | Optional object; `{}` enables continuous lane replenishment with the defaults below. Requires `agentic_snapshot` and its agentic trace/lane controls; cannot be combined with `max_sim_time_ms`. Unset preserves finite replay. See [continuous profile](../replay/workloads.md#agentic-profile) for the complete configuration, results, and limits. | | `agentic_profile.duration_seconds` | `float` | `3600.0` when enabled | Positive finite admission duration, starting at the preparation barrier or simulation start without warmup. Stops new workload requests and replacement plays at the deadline. | @@ -222,6 +224,10 @@ Every `Workload` field: | `agentic_profile.global_idle_cap_seconds` | `float` | `10.0` when enabled | Positive finite idle cap for advancing pending workload timers when the entire client workload has no outstanding requests. Engine completions and server cleanup keep their actual timestamps. | | `replay_concurrency` | `int \| None` | `None` | Closed-loop in-flight cap **for a trace**; when set, trace timestamps are ignored. For synthetic closed-loop use `concurrency` instead. | +Seed-only snapshots retain the existing sampling behavior and serialize as +`{seed: S}`; unset import fields are omitted. See [snapshot import](../replay/workloads.md#snapshot) +for the reference exporter and exact-cut semantics. + The synthetic fields are `isl`, `osl`, `request_rate`, `concurrency`, `kv_load_ratio`, `num_request_ratio`, `random_range_ratio`, and `random_seed`; `shared_prefix_ratio`, `num_prefix_groups`, `turns_per_session`, `inter_turn_delay_ms` are diff --git a/python/aisimulate/src/aisimulate/sweeper/config.py b/python/aisimulate/src/aisimulate/sweeper/config.py index 85596e20c..b03601d32 100644 --- a/python/aisimulate/src/aisimulate/sweeper/config.py +++ b/python/aisimulate/src/aisimulate/sweeper/config.py @@ -404,6 +404,12 @@ def _validate_workload(self) -> Workload: raise ValueError("agentic_warmup requires agentic_snapshot") if self.agentic_snapshot is not None and self.agentic_lanes is None: raise ValueError("agentic_snapshot requires positive agentic_lanes") + if ( + self.agentic_snapshot is not None + and self.agentic_snapshot.initial_selections is not None + and len(self.agentic_snapshot.initial_selections) != self.agentic_lanes + ): + raise ValueError("initial_selections must contain exactly one selection per lane") if self.agentic_lanes is not None: if ( self.trace_path is None diff --git a/tests/test_cli_config.py b/tests/test_cli_config.py index df3f8bd55..4749c6558 100644 --- a/tests/test_cli_config.py +++ b/tests/test_cli_config.py @@ -1570,18 +1570,30 @@ def test_vllm_disaggregated_pcp_plus_dcp_prefill_pairs_with_a_dividing_decode_dc @pytest.mark.parametrize("seed", [0, 42, 2**64 - 1]) @pytest.mark.parametrize("warmup", [False, True]) -def test_agentic_snapshot_seed_compiles_for_prediction_and_recommendation(seed: int, warmup: bool) -> None: +@pytest.mark.parametrize("imported", [False, True]) +def test_agentic_snapshot_seed_compiles_for_prediction_and_recommendation( + seed: int, warmup: bool, imported: bool +) -> None: from aisimulate.compiler import prediction_to_replay_spec from aisimulate.config import TrafficPredictionConfig, TrafficRecommendationConfig from aisimulate.recommend import _recommendation_workload from aisimulate.sweeper.config import Workload + snapshot = {"seed": seed} + if imported: + snapshot.update( + initial_selections=[ + {"source_play_id": "reference-a", "t_star_ms": 40.56078627883167}, + {"source_play_id": "reference-b", "t_star_ms": 50.0}, + ], + corpus_cursor=3, + ) traffic = { "source": {"type": "trace", "format": "weka", "paths": ["corpus"]}, "load": { "type": "trace_timestamps", "agentic_lanes": 2, - "agentic_snapshot": {"seed": seed}, + "agentic_snapshot": snapshot, "agentic_warmup": warmup, }, } @@ -1591,12 +1603,12 @@ def test_agentic_snapshot_seed_compiles_for_prediction_and_recommendation(seed: {"engine": _engine() | {"context_length": 1024}, "traffic": traffic} ) workload = prediction_to_replay_spec(prediction).workload - assert workload["agentic_snapshot"] == {"seed": seed} + assert workload["agentic_snapshot"] == snapshot assert workload.get("agentic_warmup", False) is warmup recommended = _recommendation_workload(traffic) - assert recommended["agentic_snapshot"] == {"seed": seed} + assert recommended["agentic_snapshot"] == snapshot assert recommended.get("agentic_warmup", False) is warmup - assert Workload.model_validate(recommended).model_dump()["agentic_snapshot"] == {"seed": seed} + assert Workload.model_validate(recommended).model_dump()["agentic_snapshot"] == snapshot assert Workload.model_validate(recommended).agentic_warmup is warmup @@ -1700,6 +1712,35 @@ def test_imported_agentic_snapshot_rejects_invalid_metadata(snapshot): AgenticSnapshotOptions.model_validate(snapshot) +@pytest.mark.parametrize("selection_count", [1, 3]) +def test_imported_agentic_snapshot_requires_one_selection_per_lane(selection_count: int) -> None: + from aisimulate.config import TrafficPredictionConfig, TrafficRecommendationConfig + from aisimulate.sweeper.config import Workload + + snapshot = { + "seed": 42, + "initial_selections": [ + {"source_play_id": f"reference-{index}", "t_star_ms": 50.0} for index in range(selection_count) + ], + "corpus_cursor": selection_count, + } + load = {"type": "trace_timestamps", "agentic_lanes": 2, "agentic_snapshot": snapshot} + for schema in (TrafficPredictionConfig, TrafficRecommendationConfig): + with pytest.raises(ValidationError, match="exactly one selection per lane"): + schema.model_validate({"source": {"type": "trace", "format": "weka", "paths": ["corpus"]}, "load": load}) + with pytest.raises(ValidationError, match="exactly one selection per lane"): + Workload.model_validate( + { + "source_type": "trace", + "trace_path": "corpus", + "trace_format": "weka", + "load_type": "trace_timestamps", + "agentic_lanes": 2, + "agentic_snapshot": snapshot, + } + ) + + @pytest.mark.parametrize("warmup", [False, True]) @pytest.mark.parametrize("profile", [{}, {"duration_seconds": 2.0, "response_grace_seconds": 0.0}]) def test_agentic_profile_roundtrip_and_compiler_preserve_controls(warmup: bool, profile: dict) -> None: diff --git a/tests/test_runner.py b/tests/test_runner.py index cb67af018..34438b60e 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -2018,9 +2018,38 @@ def test_runner_rejects_overflowing_ordinary_metric(): @pytest.mark.parametrize("trace_format", ["weka", "agentic_mooncake", "dynamo"]) @pytest.mark.parametrize("warmup", [False, True]) -def test_agentic_snapshot_reaches_native_payload_and_default_python_evidence(trace_format: str, warmup: bool) -> None: +@pytest.mark.parametrize("imported", [False, True]) +def test_agentic_snapshot_reaches_native_payload_and_default_python_evidence( + trace_format: str, warmup: bool, imported: bool +) -> None: evidence = [{"seed": 2**64 - 1, "lane_id": "lane:0", "t_star_ms": 50.0}] phases = {"schema": "aisimulate.agentic.phases.v1", "phase": "profile"} + snapshot = {"seed": 2**64 - 1} + if imported: + snapshot.update( + initial_selections=[{"source_play_id": "reference-play", "t_star_ms": 40.56078627883167}], + corpus_cursor=3, + ) + config = CorePredictionConfig.model_validate( + { + "engine": { + "model": "example/model", + "hardware": "h200_sxm", + "context_length": 1024, + "workers": {"aggregated": {}}, + }, + "traffic": { + "source": {"type": "trace", "format": trace_format, "paths": ["corpus"]}, + "load": { + "type": "trace_timestamps", + "agentic_lanes": 1, + "agentic_snapshot": snapshot, + "agentic_warmup": warmup, + }, + }, + } + ) + compiled = prediction_to_replay_spec(config) class SnapshotRuntime(RecordingRuntime): def run_replay_json(self, execution_spec_json): @@ -2028,26 +2057,10 @@ def run_replay_json(self, execution_spec_json): return json.dumps(report | {"agentic_snapshots": evidence} | ({"agentic_phases": phases} if warmup else {})) runtime = SnapshotRuntime() - report = ( - EngineReplayRunnerFactory(runtime=runtime) - .create(0) - .run( - _spec( - workload={ - "source_type": "trace", - "load_type": "trace_timestamps", - "trace_path": "corpus", - "trace_format": trace_format, - "agentic_lanes": 1, - "agentic_snapshot": {"seed": 2**64 - 1}, - "agentic_warmup": warmup, - } - ) - ) - ) - assert runtime.execution_spec["traffic"]["agentic_snapshot"] == {"seed": 2**64 - 1} + report = EngineReplayRunnerFactory(runtime=runtime).create(0).run(_spec(workload=compiled.workload)) + assert runtime.execution_spec["traffic"]["agentic_snapshot"] == snapshot assert report.metadata["agentic_snapshots"] == evidence - assert runtime.execution_spec["traffic"]["agentic_warmup"] is warmup + assert runtime.execution_spec["traffic"].get("agentic_warmup", False) is warmup assert report.metadata.get("agentic_phases") == (phases if warmup else None) assert "native_report" not in report.metadata From e74427879b89e36e08db6ee19ff6e9aba8a8c3a2 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 13:27:15 -0700 Subject: [PATCH 08/13] fix: preserve capability constructor compatibility and isolate warmup tests Signed-off-by: hongkuanz --- .../src/aisimulate/sweeper/replay.py | 2 +- tests/test_agentx_live_warmup.py | 28 ++++++++++--------- tests/test_unified_cli.py | 2 +- 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/python/aisimulate/src/aisimulate/sweeper/replay.py b/python/aisimulate/src/aisimulate/sweeper/replay.py index b465a2902..eae610306 100644 --- a/python/aisimulate/src/aisimulate/sweeper/replay.py +++ b/python/aisimulate/src/aisimulate/sweeper/replay.py @@ -329,13 +329,13 @@ class RunnerCapabilities: supports_agentic_speculative_decoding: bool = True supports_agentic_snapshots: bool = False supports_agentic_warmup: bool = False - supports_agentic_live_warmup: bool = False supports_cached_prefix_tokens: bool = False supported_engine_model_controls: tuple[str, ...] = () supports_mtp_expected_acceptance: bool = False supports_state_cache: bool = False supports_agentic_profile: bool = False supports_grouped_kv_cache: bool = False + supports_agentic_live_warmup: bool = False def supports_backend_topology(self, backend: str, topology: str) -> bool: """Return whether a backend/topology pair is supported. diff --git a/tests/test_agentx_live_warmup.py b/tests/test_agentx_live_warmup.py index ca56d9590..c08e14aa5 100644 --- a/tests/test_agentx_live_warmup.py +++ b/tests/test_agentx_live_warmup.py @@ -66,9 +66,9 @@ def test_live_warmup_workload_serialization_keeps_the_policy_explicit(): @pytest.mark.parametrize("stack,online", [("legacy", False), ("engine", True)]) def test_supervisor_rejects_unsupported_live_warmup_before_overwrite(tmp_path, monkeypatch, capsys, stack, online): - from aisimulate import cli_args, output, supervision + from aisimulate import cli_args, supervision - prior = tmp_path / "results.json" + prior = tmp_path / "prediction.json" prior.write_bytes(b"previous valid results\n") raw = { "engine": {"model": "example/model", "hardware": "h200_sxm", "workers": {"aggregated": {}}}, @@ -78,11 +78,11 @@ def test_supervisor_rejects_unsupported_live_warmup_before_overwrite(tmp_path, m monkeypatch.setattr(supervision, "_installed_names", lambda _: ["legacy"]) calls = [] - def forbidden_prepare(*args, **kwargs): + def forbidden_run(*args, **kwargs): calls.append((args, kwargs)) - pytest.fail("unsupported policy reached output cleanup") + pytest.fail("unsupported policy reached execution") - monkeypatch.setattr(output, "prepare_output_directory", forbidden_prepare) + monkeypatch.setattr(supervision, "run_process", forbidden_run) arguments = ["predict", "--config", "unused.yaml", "--stack", stack, "--output-dir", str(tmp_path), "--overwrite"] if online: arguments.append("--online") @@ -94,8 +94,8 @@ def forbidden_prepare(*args, **kwargs): assert prior.read_bytes() == b"previous valid results\n" -def test_supervisor_default_recommendation_without_traffic_keeps_existing_path(monkeypatch): - from aisimulate import cli_args, output, supervision +def test_supervisor_default_recommendation_without_traffic_keeps_existing_path(tmp_path, monkeypatch): + from aisimulate import cli_args, supervision raw = { "engine": { @@ -108,12 +108,14 @@ def test_supervisor_default_recommendation_without_traffic_keeps_existing_path(m } monkeypatch.setattr(cli_args, "_load_mapping", lambda _: raw) - class ReachedPreparation(Exception): + class ReachedExecution(Exception): pass - def observe_prepare(*args, **kwargs): - raise ReachedPreparation + def observe_run(*args, **kwargs): + raise ReachedExecution - monkeypatch.setattr(output, "prepare_output_directory", observe_prepare) - with pytest.raises(ReachedPreparation): - supervision.main(["recommend", "--config", "unused.yaml", "--stack", "engine", "--output-dir", "/unused"]) + monkeypatch.setattr(supervision, "run_process", observe_run) + output = tmp_path / "output" + with pytest.raises(ReachedExecution): + supervision.main(["recommend", "--config", "unused.yaml", "--stack", "engine", "--output-dir", str(output)]) + assert output.is_dir() diff --git a/tests/test_unified_cli.py b/tests/test_unified_cli.py index 2cc7d1280..868ec8efc 100644 --- a/tests/test_unified_cli.py +++ b/tests/test_unified_cli.py @@ -1686,7 +1686,7 @@ def test_recommendation_writes_an_empty_result_before_returning_failure( monkeypatch.setattr( cli.CoreRecommendationConfig, "model_validate", - staticmethod(lambda raw: object()), + staticmethod(lambda raw: SimpleNamespace(traffic=None)), ) monkeypatch.setattr(cli, "_resolve_section_adapters", lambda sections, stack: {}) result = _RecommendationResult([], failed=1) From fdf0f4b8f8c370b63a9bb1f088d3c121428d5f2b Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Thu, 15 Oct 2026 17:58:48 -0700 Subject: [PATCH 09/13] fix: preserve causal delays after AgentX live warmup Signed-off-by: hongkuanz --- crates/core/src/replay/loadgen/driver.rs | 45 +++++ crates/core/src/replay/loadgen/phase.rs | 202 +++++++++++++++++++++-- 2 files changed, 237 insertions(+), 10 deletions(-) diff --git a/crates/core/src/replay/loadgen/driver.rs b/crates/core/src/replay/loadgen/driver.rs index d4dae1df8..67e315c6f 100644 --- a/crates/core/src/replay/loadgen/driver.rs +++ b/crates/core/src/replay/loadgen/driver.rs @@ -1547,6 +1547,23 @@ impl WorkloadDriver { .agentic_snapshots .as_ref() .context("missing live snapshots")?; + let states_by_request: FxHashMap<_, _> = self + .sessions + .iter() + .enumerate() + .map(|(index, session)| { + Ok(( + ( + state.node_to_play[index], + session.turns[0] + .request_id + .as_deref() + .context("live warmup requires source request IDs")?, + ), + state.node_states[index], + )) + }) + .collect::>()?; let mut offsets = Vec::new(); let mut first = if fresh_origin { 0.0 } else { f64::INFINITY }; for (index, node) in state.node_states.iter().enumerate() { @@ -1558,6 +1575,34 @@ impl WorkloadDriver { else { bail!("live warmup requires immutable snapshot prompts"); }; + // Only the immediate conversation frontier carries a source-clock + // deadline across handoff. A pending join/barrier retains its gate, + // but future continuations and unspawned children must create their + // timers from actual dependency events, including after idle caps. + if *node == AgenticNodeState::Blocked + && context.graph.nodes[source_index] + .dependencies + .iter() + .filter(|edge| { + matches!( + edge.relation, + AgenticDependencyRelation::Sequence | AgenticDependencyRelation::Spawn + ) + }) + .any(|edge| { + states_by_request + .get(&(state.node_to_play[index], edge.request_id.as_str())) + .is_some_and(|source| { + matches!( + source, + AgenticNodeState::Blocked | AgenticNodeState::Ready + ) || (edge.trigger == AgenticDependencyTrigger::Completion + && *source == AgenticNodeState::Emitted) + }) + }) + { + continue; + } let offset = (context.recorded_start_ms(source_index) - snapshots[state.node_to_play[index]].t_star_ms) .max(0.0) diff --git a/crates/core/src/replay/loadgen/phase.rs b/crates/core/src/replay/loadgen/phase.rs index 8212e2171..f55751622 100644 --- a/crates/core/src/replay/loadgen/phase.rs +++ b/crates/core/src/replay/loadgen/phase.rs @@ -1107,8 +1107,8 @@ mod tests { ); } - fn linear_live_driver(nodes: usize, rootless: bool) -> WorkloadDriver { - let rows = (0..nodes) + fn linear_live_rows(nodes: usize, rootless: bool) -> Vec { + (0..nodes) .map(|index| AgenticMooncakeRow { request_id: format!("n{index}"), play_id: "linear".into(), @@ -1144,7 +1144,15 @@ mod tests { }, ..Default::default() }) - .collect(); + .collect() + } + + fn rows_live_driver( + rows: Vec, + lanes: usize, + cut: f64, + speedup: f64, + ) -> WorkloadDriver { let graph = ValidatedAgenticGraph::from_agentic_mooncake_rows( AgenticMooncakeHeader { schema: AGENTIC_MOONCAKE_SCHEMA.into(), @@ -1160,17 +1168,16 @@ mod tests { ) .unwrap(); let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) - .unwrap(); - let play = prepared - .context() - .prepare_play(0, 0, Some(if rootless { 0.5 } else { 0.0 })) + .prepare_snapshots(lanes, AgenticSnapshotOptions::seeded(42)) .unwrap(); + let plays = (0..lanes) + .map(|lane| prepared.context().prepare_play(lane, 0, Some(cut)).unwrap()) + .collect(); let mut driver = WorkloadDriver::new_agentic_warmup_with_policy( - PreparedAgenticSnapshots::from_plays(vec![play]).unwrap(), + PreparedAgenticSnapshots::from_plays(plays).unwrap(), 64, true, - 1.0, + speedup, AgenticWarmupPolicy::LiveTrajectoryV1, 10_000.0, ) @@ -1181,6 +1188,181 @@ mod tests { driver } + fn linear_live_driver(nodes: usize, rootless: bool) -> WorkloadDriver { + rows_live_driver( + linear_live_rows(nodes, rootless), + 1, + if rootless { 0.5 } else { 0.0 }, + 1.0, + ) + } + + fn finish_ten_live_rounds(driver: &mut WorkloadDriver, lanes: usize) -> f64 { + driver.finish_agentic_preparation(0.0).unwrap(); + for round in 0..10 { + let now = round as f64; + driver.advance_agentic_profile(now).unwrap(); + let ready = driver.pop_ready(now, lanes); + assert_eq!(ready.len(), lanes); + for turn in ready { + driver.on_output_token(turn.request_uuid, 0).unwrap(); + driver + .on_terminal( + turn.request_uuid, + now + 1.0, + ReplayTerminalStatus::Completed, + ) + .unwrap(); + } + } + assert_eq!( + driver.finish_agentic_preparation(10.0).unwrap(), + Some(AgenticPreparationTransition::OpenProfile) + ); + 10.0 + } + + #[test] + fn live_handoff_idle_cap_does_not_restore_future_continuation_timestamps() { + let mut rows = linear_live_rows(13, false); + rows[11].not_before_ms = 601_000.0; + rows[11].dependencies[0].delay_ms = 600_000.0; + rows[12].not_before_ms = 605_001.0; + rows[12].dependencies[0].delay_ms = 4_000.0; + let mut driver = rows_live_driver(rows, 2, 0.0, 1.0); + let start = finish_ten_live_rounds(&mut driver, 2); + let frontier = driver.pop_ready(start, 2); + assert_eq!(frontier.len(), 2); + assert!( + frontier + .iter() + .all(|turn| turn.authored_request_id.as_deref() == Some("n10")) + ); + // Keep the second lane busy so only the 300-second per-tree guard fires. + driver + .on_terminal( + frontier[0].request_uuid, + start + 1.0, + ReplayTerminalStatus::Completed, + ) + .unwrap(); + let capped = start + 1.0 + 300_000.0; + assert_eq!(driver.next_ready_time_ms(), Some(capped)); + driver.advance_agentic_profile(capped).unwrap(); + let delayed = driver.pop_ready(capped, 1).pop().unwrap(); + assert_eq!(delayed.authored_request_id.as_deref(), Some("n11")); + driver + .on_terminal( + delayed.request_uuid, + capped + 1.0, + ReplayTerminalStatus::Completed, + ) + .unwrap(); + let continuation = capped + 1.0 + 4_000.0; + assert_eq!(driver.next_ready_time_ms(), Some(continuation)); + let ready = driver.pop_ready(continuation, 1).pop().unwrap(); + assert_eq!(ready.authored_request_id.as_deref(), Some("n12")); + assert_eq!( + driver.agentic_profile_report().unwrap().idle_shifts.len(), + 1 + ); + } + + #[rstest::rstest] + #[case::join(AgenticDependencyRelation::Join)] + #[case::replay_barrier(AgenticDependencyRelation::ReplayBarrier)] + fn live_handoff_preserves_blocked_frontier_but_not_later_turns( + #[case] relation: AgenticDependencyRelation, + ) { + let mut rows = linear_live_rows(12, false); + rows[10].not_before_ms = 5_000.0; + rows[10].recorded_api_time_ms = Some(10_000.0); + rows[10].dependencies.push(AgenticDependency { + request_id: "child".into(), + delay_ms: 0.0, + relation, + trigger: AgenticDependencyTrigger::Completion, + }); + rows[11].not_before_ms = 19_000.0; + rows[11].dependencies[0].delay_ms = 4_000.0; + let mut child = rows[0].clone(); + child.request_id = "child".into(); + child.session_id = "child".into(); + child.not_before_ms = 1_000.0; + child.dependencies = vec![AgenticDependency { + request_id: "n9".into(), + delay_ms: 100.0, + relation: AgenticDependencyRelation::Spawn, + trigger: AgenticDependencyTrigger::Dispatch, + }]; + rows.push(child.clone()); + child.request_id = "future-child".into(); + child.session_id = "future-child".into(); + child.not_before_ms = 21_000.0; + child.dependencies[0].request_id = "n11".into(); + child.dependencies[0].delay_ms = 2_000.0; + rows.push(child); + + let mut driver = rows_live_driver(rows, 1, 0.0, 1.0); + let start = finish_ten_live_rounds(&mut driver, 1); + let ready = driver.pop_ready(start, 1).pop().unwrap(); + assert_eq!(ready.authored_request_id.as_deref(), Some("child")); + driver + .on_terminal( + ready.request_uuid, + start + 1.0, + ReplayTerminalStatus::Completed, + ) + .unwrap(); + // The parent was already waiting at a join/barrier when warmup ended: + // keep its recorded deadline even if the remaining child finishes early. + let parent_at = start + 4_000.0; + assert_eq!(driver.next_ready_time_ms(), Some(parent_at)); + let parent = driver.pop_ready(parent_at, 1).pop().unwrap(); + assert_eq!(parent.authored_request_id.as_deref(), Some("n10")); + driver + .on_terminal( + parent.request_uuid, + parent_at + 1.0, + ReplayTerminalStatus::Completed, + ) + .unwrap(); + let next_at = parent_at + 1.0 + 4_000.0; + assert_eq!(driver.next_ready_time_ms(), Some(next_at)); + let next = driver.pop_ready(next_at, 1).pop().unwrap(); + assert_eq!(next.authored_request_id.as_deref(), Some("n11")); + // A child whose spawning turn was still blocked at handoff is not a + // pending handoff stream. Its timer starts only on actual dispatch. + assert_eq!(driver.next_ready_time_ms(), Some(next_at + 2_000.0)); + let spawned = driver.pop_ready(next_at + 2_000.0, 1).pop().unwrap(); + assert_eq!(spawned.authored_request_id.as_deref(), Some("future-child")); + } + + #[test] + fn live_handoff_fresh_root_preserves_other_lanes_scaled_frontier_offset() { + let mut rows = linear_live_rows(10, false); + let mut other = linear_live_rows(12, false); + for row in &mut other { + row.play_id = "other".into(); + row.request_id = format!("other-{}", row.request_id); + for edge in &mut row.dependencies { + edge.request_id = format!("other-{}", edge.request_id); + } + } + rows.extend(other); + let mut driver = rows_live_driver(rows, 2, 0.0, 2.0); + let start = finish_ten_live_rounds(&mut driver, 2); + // One lane finished exactly at its quota. Its new incarnation anchors + // the phase at zero, without also bursting the other lane's frontier. + let root = driver.pop_ready(start, 2); + assert_eq!(root.len(), 1); + assert_eq!(root[0].authored_request_id.as_deref(), Some("n0")); + assert_eq!(driver.next_ready_time_ms(), Some(start + 5.0)); + let other = driver.pop_ready(start + 5.0, 2); + assert_eq!(other.len(), 1); + assert_eq!(other[0].authored_request_id.as_deref(), Some("other-n10")); + } + #[test] fn zero_time_live_warmup_is_bounded_and_failed_reset_is_transactional() { let mut driver = linear_live_driver(1, false); From c218f118b180a872e3d4ce6c2f9bf4c83b5260b0 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Mon, 19 Oct 2026 17:02:24 -0700 Subject: [PATCH 10/13] fix: normalize exported AgentX snapshot clocks and retain provenance Signed-off-by: hongkuanz --- crates/core/src/replay/loadgen/snapshot.rs | 1 + crates/core/tests/replay_tokenization.rs | 2 +- docs/replay/workloads.md | 34 ++- .../export_agentx_snapshots.py | 132 ++++++++-- tests/test_agentx_snapshot_export.py | 225 +++++++++++++++++- 5 files changed, 355 insertions(+), 39 deletions(-) diff --git a/crates/core/src/replay/loadgen/snapshot.rs b/crates/core/src/replay/loadgen/snapshot.rs index 205c3372f..2f0040ebc 100644 --- a/crates/core/src/replay/loadgen/snapshot.rs +++ b/crates/core/src/replay/loadgen/snapshot.rs @@ -28,6 +28,7 @@ pub const AGENTIC_SNAPSHOT_SCHEMA_V1: &str = "aisimulate.agentic.snapshot.v1"; #[serde(deny_unknown_fields)] pub struct AgenticSnapshotSelection { pub source_play_id: String, + /// Loaded graph clock: per-play-relative for Weka, source clock for Mooncake. pub t_star_ms: f64, } diff --git a/crates/core/tests/replay_tokenization.rs b/crates/core/tests/replay_tokenization.rs index 3a8820b85..8790f769a 100644 --- a/crates/core/tests/replay_tokenization.rs +++ b/crates/core/tests/replay_tokenization.rs @@ -305,7 +305,7 @@ fn agentic_driver( return WorkloadDriver::new_agentic_trace_with_lanes(graph, 4, 1).unwrap(); }; let prepared = graph - .prepare_snapshots(1, AgenticSnapshotOptions { seed: 42 }) + .prepare_snapshots(1, AgenticSnapshotOptions::seeded(42)) .unwrap(); let play = prepared.context().prepare_play_from_start(0, 0).unwrap(); let snapshots = PreparedAgenticSnapshots::from_plays(vec![play]).unwrap(); diff --git a/docs/replay/workloads.md b/docs/replay/workloads.md index b2525e190..ed5c9efaa 100644 --- a/docs/replay/workloads.md +++ b/docs/replay/workloads.md @@ -72,7 +72,7 @@ the current SA convention. | `traffic.load.speedup` | `1` | `-` | `-` | Positive; trace timestamp load only. | | `traffic.load.agentic_lanes` | `null` | `x` | `-` | Positive integer; `weka`, `agentic_mooncake`, or agentic `dynamo` timestamp replay only. | | `traffic.load.agentic_snapshot` | `null` (unset) | `x` | `-` | Optional object with required unsigned 64-bit `seed` (`0` through `2^64 - 1`). Requires `traffic.load.type: trace_timestamps` and positive `agentic_lanes`; supported formats are `weka`, `agentic_mooncake`, and agentic `dynamo`. Unset preserves turn-zero execution. See [snapshot](#agentic-snapshot). | -| `traffic.load.agentic_snapshot.initial_selections` | unset | `x` | `-` | Exact initial `{source_play_id, t_star_ms}` selections in lane order, one per lane. Requires `corpus_cursor`. Unset retains the original seed-based sampler. | +| `traffic.load.agentic_snapshot.initial_selections` | unset | `x` | `-` | Exact initial `{source_play_id, t_star_ms}` selections in lane order, one per lane, in the loaded graph's clock (Weka: per-play relative milliseconds). Requires `corpus_cursor`. Unset retains the original seed-based sampler. | | `traffic.load.agentic_snapshot.corpus_cursor` | unset | `x` | `-` | Unsigned 64-bit next sequential corpus draw, including any skipped initial draws. Requires `initial_selections` and must be at least the lane count. Recycled plays start at turn zero. | | `traffic.load.agentic_warmup` | `false` | `x` | `-` | Optional boolean; `true` requires `agentic_snapshot` and positive `agentic_lanes`. Primes caches before measurement; see [warmup](#agentic-warmup). | | `traffic.load.agentic_profile` | `null` (unset) | `x` | `-` | Optional object; `{}` enables continuous lane replenishment with the defaults below. Requires `trace_timestamps`, positive `agentic_lanes`, and `agentic_snapshot`; cannot be combined with `traffic.stop.max_virtual_time_seconds`. Unset preserves finite replay. See [continuous profile](#agentic-profile). | @@ -352,7 +352,11 @@ agentic_snapshot: one selection per lane and a cursor at least as large as the selection count. Selections are in lane order. Each source ID must exist in the exact input corpus and each cut must lie within that play's request-start bounds. Cuts -are in source milliseconds, before `speedup`; they are never re-sampled or +are in the **loaded graph's clock**, before `speedup`: Weka subtracts the +earliest request timestamp in each complete play, including child streams, +whereas Agentic Mooncake retains its recorded source clock. Branch/spawn marker +times are not request starts and do not define the origin. The importer does +not guess a clock basis or convert untagged cuts; they are never re-sampled or silently replaced. Repeated source plays still get distinct lane/cache identities. Preserve the source corpus bytes/order, reference revision and selection artifact alongside the replay configuration. @@ -374,8 +378,8 @@ in the **pinned AIPerf environment** (its existing dependencies are required): python scripts/prediction_regression/export_agentx_snapshots.py \ --aiperf-source /path/to/aiperf-at-754356e9 \ --metadata reference-dataset-metadata.json \ - --source-play-map source-play-map.json --lanes 15 --seed 42 \ - > snapshot-options.json 2> snapshot-provenance.jsonl + --source-play-map source-play-map.json --trace-format weka \ + --lanes 15 --seed 42 --output snapshot-options.json ``` `reference-dataset-metadata.json` is the reference loader's @@ -387,10 +391,26 @@ metadata must match the replay corpus order and each root mapping must identify the corresponding source play; do not sort either input. The importer validates IDs and cut bounds, but cannot infer or verify reference corpus order from a cursor alone. A wrong mapping/order is not a valid parity comparison. -Paste the emitted JSON object under `traffic.load.agentic_snapshot`. +The required `--trace-format` selects the target loader clock explicitly: +`weka` subtracts the minimum request start across each selected reference tree; +`agentic_mooncake` leaves source-clock cuts unchanged. It uses the pinned +metadata's millisecond turn timestamps, not seconds or branch markers. The +reference and replay loaders must represent the same complete request tree +and nested timestamp interpretation. Missing children, ambiguous trees or +missing/nonfinite request timestamps in a selected tree fail explicitly; +do not invent an origin for an unmappable corpus. Existing exports made before +this conversion fix must be regenerated for nonzero-origin Weka plays, even +if their old cut happens to fall within the normalized bounds. + +Paste the emitted JSON object under `traffic.load.agentic_snapshot`. Keep the +automatically written `snapshot-options.json.provenance.json` beside it: it +records the pinned revision, metadata/map/output SHA-256 hashes, root order, +clock basis and each lane's raw cut, origin and converted cut. Both output +paths must be new. For importable JSON on stdout, omit `--output` and provide +`--provenance-out snapshot-provenance.json`; the sidecar is persisted before +anything is printed to stdout. The exporter verifies the source revision and clean source tree, invokes the -reference `TrajectorySource` directly, and records input hashes and corpus -order on stderr. It fails on unsupported samplers, timestamp-less trajectories, +reference `TrajectorySource` directly. It fails on unsupported samplers, timestamp-less trajectories, or incomplete lane selection. Neither AIPerf nor NumPy becomes an AISimulate runtime dependency. diff --git a/scripts/prediction_regression/export_agentx_snapshots.py b/scripts/prediction_regression/export_agentx_snapshots.py index 2992bd96a..74372a582 100644 --- a/scripts/prediction_regression/export_agentx_snapshots.py +++ b/scripts/prediction_regression/export_agentx_snapshots.py @@ -15,6 +15,7 @@ import hashlib import importlib import json +import math import subprocess import sys from pathlib import Path @@ -37,10 +38,42 @@ def next_conversation_id(self): return selected -def export_snapshot(metadata, source_play_map, *, lanes, seed, trajectory_type): +def _source_time_bounds(metadata, root_id): + """Require complete root/child request clocks; branch markers are not requests.""" + by_id = {item.conversation_id: item for item in metadata.conversations} + if len(by_id) != len(metadata.conversations): + raise ValueError("duplicate reference conversation IDs") + pending = [root_id] + seen = set() + timestamps = [] + while pending: + conversation_id = pending.pop() + if conversation_id in seen: + raise ValueError("reference child graph is cyclic or has shared children") + seen.add(conversation_id) + item = by_id.get(conversation_id) + if item is None: + raise ValueError(f"missing reference child metadata: {conversation_id}") + if conversation_id != root_id and getattr(item, "is_root", True): + raise ValueError("reference child also appears as a sampleable root") + for turn in item.turns: + timestamp = turn.timestamp_ms + if type(timestamp) not in (int, float) or not math.isfinite(timestamp) or timestamp < 0: + raise ValueError("every reference request must have a finite nonnegative timestamp") + timestamps.append(timestamp) + for branch in item.branches: + pending.extend(branch.child_conversation_ids) + if not timestamps: + raise ValueError("reference play has no request timestamps") + return min(timestamps), max(timestamps) + + +def export_snapshot(metadata, source_play_map, *, lanes, seed, trajectory_type, trace_format, clock_evidence=None): """Invoke the reference selector and retain its actual shared cursor.""" if type(lanes) is not int or lanes <= 0 or type(seed) is not int or not 0 <= seed < 2**64: raise ValueError("lanes must be positive and seed must be an unsigned 64-bit integer") + if trace_format not in ("weka", "agentic_mooncake"): + raise ValueError("trace_format must explicitly select weka or agentic_mooncake") strategy = getattr(metadata.sampling_strategy, "value", metadata.sampling_strategy) if str(strategy).lower() != "sequential": raise ValueError("only the pinned reference's sequential sampler is supported") @@ -63,16 +96,56 @@ def export_snapshot(metadata, source_play_map, *, lanes, seed, trajectory_type): raise ValueError("reference could not construct every requested lane") if any(item.snapshot is None for item in source.trajectories): raise ValueError("timestamp-less reference snapshots cannot be imported as time cuts") + selections = [] + for item in source.trajectories: + start, end = _source_time_bounds(metadata, item.conversation_id) + cut = item.snapshot.t_star_ms + if type(cut) not in (int, float) or not math.isfinite(cut) or not start <= cut <= end: + raise ValueError("reference cut is outside its source request-start bounds") + # Weka lower_trace subtracts the earliest request in the entire play, + # including child streams. Agentic Mooncake retains its source clock. + origin = start if trace_format == "weka" else 0 + selections.append({"source_play_id": source_play_map[item.conversation_id], "t_star_ms": cut - origin}) + if clock_evidence is not None: + clock_evidence.append( + { + "reference_root_id": item.conversation_id, + "source_play_id": source_play_map[item.conversation_id], + "source_origin_ms": start, + "subtracted_origin_ms": origin, + "reference_t_star_ms": cut, + "imported_t_star_ms": cut - origin, + } + ) return { "seed": seed, - "initial_selections": [ - {"source_play_id": source_play_map[item.conversation_id], "t_star_ms": item.snapshot.t_star_ms} - for item in source.trajectories - ], + "initial_selections": selections, "corpus_cursor": sampler.cursor, } +def write_export(result, provenance, *, output, provenance_out): + """Persist provenance before publishing usable config; never overwrite artifacts.""" + if output is not None: + if provenance_out is None: + provenance_out = output.with_name(output.name + ".provenance.json") + if output.resolve() == provenance_out.resolve(): + raise ValueError("output and provenance paths must differ") + if output.exists(): + raise FileExistsError(output) + elif provenance_out is None: + raise ValueError("stdout export requires --provenance-out") + result_json = json.dumps(result, indent=2, allow_nan=False) + "\n" + provenance = {**provenance, "snapshot_options_sha256": hashlib.sha256(result_json.encode()).hexdigest()} + with provenance_out.open("x") as stream: + stream.write(json.dumps(provenance, indent=2, allow_nan=False) + "\n") + if output is None: + print(result_json, end="") + else: + with output.open("x") as stream: + stream.write(result_json) + + def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--aiperf-source", type=Path, required=True) @@ -87,6 +160,11 @@ def main(): ) parser.add_argument("--lanes", type=int, required=True) parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--trace-format", choices=("weka", "agentic_mooncake"), required=True) + parser.add_argument( + "--output", type=Path, help="Write config and a sibling .provenance.json sidecar; never overwrite" + ) + parser.add_argument("--provenance-out", type=Path, help="Required for stdout; overrides the --output sidecar path") args = parser.parse_args() checkout = args.aiperf_source.resolve() head = subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip() @@ -104,26 +182,34 @@ def main(): metadata_bytes = args.metadata.read_bytes() mapping_bytes = args.source_play_map.read_bytes() metadata = metadata_type.model_validate_json(metadata_bytes) + clock_evidence = [] result = export_snapshot( - metadata, json.loads(mapping_bytes), lanes=args.lanes, seed=args.seed, trajectory_type=module.TrajectorySource + metadata, + json.loads(mapping_bytes), + lanes=args.lanes, + seed=args.seed, + trajectory_type=module.TrajectorySource, + trace_format=args.trace_format, + clock_evidence=clock_evidence, ) - print(json.dumps(result, indent=2, allow_nan=False)) - print( - json.dumps( - { - "reference_revision": head, - "metadata_sha256": hashlib.sha256(metadata_bytes).hexdigest(), - "source_play_map_sha256": hashlib.sha256(mapping_bytes).hexdigest(), - "source_order": [ - item.conversation_id - for item in metadata.conversations - if getattr(item, "is_root", True) is not False - ], - "sampling_strategy": "sequential", - "start_ratio": [0.25, 0.75], - } - ), - file=sys.stderr, + write_export( + result, + { + "reference_revision": head, + "metadata_sha256": hashlib.sha256(metadata_bytes).hexdigest(), + "source_play_map_sha256": hashlib.sha256(mapping_bytes).hexdigest(), + "source_order": [ + item.conversation_id for item in metadata.conversations if getattr(item, "is_root", True) is not False + ], + "sampling_strategy": "sequential", + "start_ratio": [0.25, 0.75], + "trace_format": args.trace_format, + "input_clock": "reference_source_ms", + "output_clock": "per_play_relative_ms" if args.trace_format == "weka" else "source_ms", + "clock_conversions": clock_evidence, + }, + output=args.output, + provenance_out=args.provenance_out, ) diff --git a/tests/test_agentx_snapshot_export.py b/tests/test_agentx_snapshot_export.py index 3537cb632..b13b5774d 100644 --- a/tests/test_agentx_snapshot_export.py +++ b/tests/test_agentx_snapshot_export.py @@ -1,21 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json from types import SimpleNamespace as Ns import pytest -from scripts.prediction_regression.export_agentx_snapshots import SequentialCursor, export_snapshot +from scripts.prediction_regression.export_agentx_snapshots import SequentialCursor, export_snapshot, write_export def metadata(strategy="sequential"): return Ns( sampling_strategy=strategy, conversations=[ - Ns(conversation_id="skip", is_root=True), - Ns(conversation_id="a", is_root=True), - Ns(conversation_id="child", is_root=False), - Ns(conversation_id="b", is_root=True), + Ns(conversation_id="skip", is_root=True, turns=[], branches=[]), + Ns(conversation_id="a", is_root=True, turns=[Ns(timestamp_ms=0), Ns(timestamp_ms=100)], branches=[]), + Ns(conversation_id="child", is_root=False, turns=[], branches=[]), + Ns(conversation_id="b", is_root=True, turns=[Ns(timestamp_ms=0), Ns(timestamp_ms=100)], branches=[]), ], ) @@ -35,7 +36,14 @@ def reference(**kwargs): ] ) - result = export_snapshot(metadata(), {"skip": "s", "a": "A", "b": "B"}, lanes=2, seed=42, trajectory_type=reference) + result = export_snapshot( + metadata(), + {"skip": "s", "a": "A", "b": "B"}, + lanes=2, + seed=42, + trajectory_type=reference, + trace_format="agentic_mooncake", + ) assert result == { "seed": 42, "initial_selections": [{"source_play_id": "A", "t_star_ms": 12.5}, {"source_play_id": "B", "t_star_ms": 40.5}], @@ -43,6 +51,47 @@ def reference(**kwargs): } +def test_skipped_timestamp_less_root_does_not_require_an_origin(): + source = metadata() + source.conversations[0].turns = [Ns(timestamp_ms=None)] + + def reference(**kwargs): + sampler = kwargs["dataset_sampler"] + assert sampler.next_conversation_id() == "skip" + assert sampler.next_conversation_id() == "a" + return Ns(trajectories=[Ns(conversation_id="a", snapshot=Ns(t_star_ms=50))]) + + evidence = [] + result = export_snapshot( + source, + {"skip": "S", "a": "A", "b": "B"}, + lanes=1, + seed=42, + trajectory_type=reference, + trace_format="weka", + clock_evidence=evidence, + ) + assert result["corpus_cursor"] == 2 + assert [entry["reference_root_id"] for entry in evidence] == ["a"] + + +@pytest.mark.parametrize( + "children,error", [(["missing"], "missing reference child"), (["a"], "cyclic"), (["b"], "sampleable root")] +) +def test_unmappable_child_metadata_fails_closed(children, error): + source = metadata() + source.conversations[1].branches = [Ns(child_conversation_ids=children)] + with pytest.raises(ValueError, match=error): + export_snapshot( + source, + {"skip": "S", "a": "A", "b": "B"}, + lanes=1, + seed=42, + trajectory_type=lambda **_: Ns(trajectories=[Ns(conversation_id="a", snapshot=Ns(t_star_ms=50))]), + trace_format="weka", + ) + + def test_sequential_cursor_wraps_without_resetting_consumed_draws(): cursor = SequentialCursor(["a", "b"]) assert [cursor.next_conversation_id() for _ in range(5)] == ["a", "b", "a", "b", "a"] @@ -52,13 +101,13 @@ def test_sequential_cursor_wraps_without_resetting_consumed_draws(): @pytest.mark.parametrize("strategy", ["random", "shuffle"]) def test_export_rejects_unimplemented_recycle_samplers(strategy): with pytest.raises(ValueError, match="sequential"): - export_snapshot(metadata(strategy), {}, lanes=2, seed=42, trajectory_type=None) + export_snapshot(metadata(strategy), {}, lanes=2, seed=42, trajectory_type=None, trace_format="weka") @pytest.mark.parametrize("mapping", [{}, {"skip": "S", "a": "A", "b": "A"}, {"skip": "S", "a": "A", "b": ""}]) def test_export_requires_complete_unambiguous_normalized_play_mapping(mapping): with pytest.raises(ValueError, match="source.play|distinct"): - export_snapshot(metadata(), mapping, lanes=2, seed=42, trajectory_type=None) + export_snapshot(metadata(), mapping, lanes=2, seed=42, trajectory_type=None, trace_format="weka") @pytest.mark.parametrize("trajectories", [[], [Ns(conversation_id="a", snapshot=None)]]) @@ -70,4 +119,164 @@ def test_export_fails_closed_for_missing_lanes_or_timestamp_less_source(trajecto lanes=1, seed=42, trajectory_type=lambda **_: Ns(trajectories=trajectories), + trace_format="weka", + ) + + +def _export_cut(times, cut, *, trace_format="weka", child_times=()): + parent = Ns(conversation_id="root", is_root=True, turns=[Ns(timestamp_ms=t) for t in times], branches=[]) + conversations = [parent] + if child_times: + parent.branches = [Ns(child_conversation_ids=["child"], start_timestamp_ms=1)] + conversations.append( + Ns(conversation_id="child", is_root=False, turns=[Ns(timestamp_ms=t) for t in child_times], branches=[]) + ) + source = Ns(sampling_strategy="sequential", conversations=conversations) + + def reference(**kwargs): + kwargs["dataset_sampler"].next_conversation_id() + return Ns(trajectories=[Ns(conversation_id="root", snapshot=Ns(t_star_ms=cut))]) + + clock_evidence = [] + result = export_snapshot( + source, + {"root": "PLAY"}, + lanes=1, + seed=42, + trajectory_type=reference, + trace_format=trace_format, + clock_evidence=clock_evidence, + ) + return result, clock_evidence + + +@pytest.mark.parametrize("trace_format,expected", [("weka", 594.0), ("agentic_mooncake", 10594.0)]) +def test_export_explicitly_converts_weka_but_preserves_mooncake_clock(trace_format, expected): + result, evidence = _export_cut([10000.0, 11000.0], 10594.0, trace_format=trace_format) + assert result["initial_selections"][0]["t_star_ms"] == expected + assert evidence[0]["source_origin_ms"] == 10000.0 + assert evidence[0]["reference_t_star_ms"] == 10594.0 + assert evidence[0]["imported_t_star_ms"] == expected + + +@pytest.mark.parametrize("parent,child", [([10000, 30000], [15000]), ([20000, 30000], [10000])]) +def test_origin_uses_earliest_request_in_whole_tree_not_marker_or_root_only(parent, child): + result, _ = _export_cut(parent, 21000, child_times=child) + assert result["initial_selections"][0]["t_star_ms"] == 11000 + + +@pytest.mark.parametrize("timestamp", [None, float("nan"), float("inf"), -1, True, "10000"]) +def test_incomplete_or_invalid_source_clock_fails_closed(timestamp): + with pytest.raises(ValueError, match="finite nonnegative timestamp"): + _export_cut([timestamp, 30000], 21000) + + +def test_unknown_clock_format_fails_closed(): + with pytest.raises(ValueError, match="explicitly select"): + _export_cut([10000, 30000], 21000, trace_format="auto") + + +@pytest.mark.parametrize("last_seconds,absolute_cut,expected_history", [(11, 10594.0, 1), (40, 21000.0, 2)]) +def test_exported_cuts_reach_real_weka_loader_without_rejection_or_silent_frontier_shift( + tmp_path, last_seconds, absolute_cut, expected_history +): + from aisimulate import _runtime + + # Original fixture: in the long case, passing 21s straight to the normalized + # [0,30]s graph is legal but wrongly makes the 30s source turn historical. + times = [10, last_seconds] if last_seconds == 11 else [10, 20, 30, last_seconds] + path = tmp_path / "nonzero-origin.json" + path.write_text( + json.dumps( + { + "id": "nonzero-origin", + "models": ["model"], + "block_size": 4, + "hash_id_scope": "local", + "requests": [ + { + "t": t, + "type": "s", + "model": "model", + "in": 4 * (i + 1), + "out": 1, + "api_time": 0.1, + "hash_ids": list(range(i + 1)), + } + for i, t in enumerate(times) + ], + } ) + ) + payload = { + "spec": { + "version": 1, + "topology": {"kind": "aggregated", "workers": {"initial_workers": 1}}, + "engine": { + "tensor_parallel_size": 1, + "rank": { + "backend": "vllm", + "block_size": 4, + "num_gpu_blocks": 128, + "timing_model": {"type": "fixed", "prefill_ms": 1, "decode_ms": 1}, + }, + }, + "requests": [], + }, + "traffic": { + "source_type": "trace", + "load_type": "trace_timestamps", + "trace_format": "weka", + "trace_path": str(path), + "trace_block_size": 4, + "execution_model": "model", + "agentic_lanes": 1, + "agentic_snapshot": {"seed": 42}, + }, + } + + def run(): + return json.loads(_runtime.run_replay_json(json.dumps(payload)))["agentic_snapshots"][0] + + play_id = run()["source_play_id"] + options, _ = _export_cut([t * 1000 for t in times], absolute_cut) + options["initial_selections"][0]["source_play_id"] = play_id + payload["traffic"]["agentic_snapshot"] = options + actual = run() + assert actual["t_star_ms"] == absolute_cut - 10000 + assert sum(request["historical"] for request in actual["requests"]) == expected_history + options["initial_selections"][0]["t_star_ms"] = absolute_cut + if last_seconds == 11: + with pytest.raises(RuntimeError, match="bounds"): + run() + else: + assert sum(request["historical"] for request in run()["requests"]) == 3 + + +def test_export_persists_sidecar_and_refuses_to_overwrite(tmp_path, capsys): + output = tmp_path / "snapshot.json" + result, evidence = _export_cut([10000, 11000], 10594) + provenance = {"reference_revision": "pinned", "clock_conversions": evidence} + write_export(result, provenance, output=output, provenance_out=None) + assert json.loads(output.read_text()) == result + sidecar = output.with_name("snapshot.json.provenance.json") + saved = json.loads(sidecar.read_text()) + assert saved["clock_conversions"] == evidence + assert len(saved["snapshot_options_sha256"]) == 64 + before = (output.read_bytes(), sidecar.read_bytes()) + with pytest.raises(FileExistsError): + write_export(result, provenance, output=output, provenance_out=None) + assert (output.read_bytes(), sidecar.read_bytes()) == before + with pytest.raises(FileExistsError): + write_export(result, provenance, output=None, provenance_out=sidecar) + assert capsys.readouterr().out == "" + with pytest.raises(ValueError, match="requires --provenance-out"): + write_export(result, provenance, output=None, provenance_out=None) + + +def test_stdout_remains_importable_options_with_durable_sidecar(tmp_path, capsys): + result, evidence = _export_cut([10000, 11000], 10594) + sidecar = tmp_path / "provenance.json" + write_export(result, {"clock_conversions": evidence}, output=None, provenance_out=sidecar) + assert json.loads(capsys.readouterr().out) == result + assert json.loads(sidecar.read_text())["clock_conversions"] == evidence From fa5ed6578ec7c0978279303e91cbf78424bcd904 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Mon, 19 Oct 2026 17:02:44 -0700 Subject: [PATCH 11/13] docs: explain AgentX live warmup handoff idle Signed-off-by: hongkuanz --- docs/replay/agentic/warmup.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/replay/agentic/warmup.md b/docs/replay/agentic/warmup.md index f4e9a9ec4..919875832 100644 --- a/docs/replay/agentic/warmup.md +++ b/docs/replay/agentic/warmup.md @@ -200,6 +200,18 @@ The policy follows the accelerated request-count warmup contract in cursor. Restore original output lengths and dependency delays. A failed measurement reset commits neither this new frontier nor the cursor update. +Live warmup compresses runtime waits, not the source timeline. At handoff, +surviving timestamped streams retain their offsets from the original sampled +cut; the profile subtracts one global minimum, not a separate minimum per +lane. Consequently, source time skipped by accelerated warmup can reappear as +initial idle in other lanes (subject to the configured per-tree/global idle +caps). For example, a lane that advances far into a long play can wait while +a freshly recycled lane starts at zero; this is not a scheduler stall. This +matches the pinned reference's `_handoff_replay_offset_ms` and +`_profiling_phase_t0_offset_ms` with its default spread phase starts. The +reference's optional `--burst-phase-starts` mode instead re-anchors each lane; +that alternative is not the policy modeled here. + The formal request and trajectory cohorts exclude all preparation-only plays. Partially warmed active trees retain their cache/request identity but start their formal trajectory observation at the barrier. Formal lifecycle ordinals From 1b79f20a1b9dd28a1521f61ff5dfb4784ce32fc8 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Tue, 20 Oct 2026 14:03:55 -0700 Subject: [PATCH 12/13] fix: scale live AgentX primer leads and document SDK migration Signed-off-by: hongkuanz --- crates/core/src/replay/loadgen/driver.rs | 1 + crates/core/src/replay/loadgen/phase.rs | 39 +++++++++++++++++++- docs/aic-backward-compatibility/migration.md | 24 ++++++++++++ docs/replay/agentic/warmup.md | 5 ++- docs/sweeper/sdk.md | 18 ++++++--- 5 files changed, 78 insertions(+), 9 deletions(-) diff --git a/crates/core/src/replay/loadgen/driver.rs b/crates/core/src/replay/loadgen/driver.rs index 1274b3a61..afb707b07 100644 --- a/crates/core/src/replay/loadgen/driver.rs +++ b/crates/core/src/replay/loadgen/driver.rs @@ -1461,6 +1461,7 @@ impl WorkloadDriver { include_replay_hashes, policy, primer_idle_cap_ms, + speedup, )? }; let mut driver = Self::new_agentic_snapshots( diff --git a/crates/core/src/replay/loadgen/phase.rs b/crates/core/src/replay/loadgen/phase.rs index 01767fe19..2f8e02645 100644 --- a/crates/core/src/replay/loadgen/phase.rs +++ b/crates/core/src/replay/loadgen/phase.rs @@ -149,6 +149,7 @@ impl AgenticPreparation { include_replay_hashes, AgenticWarmupPolicy::RepeatedPrefixes, 10_000.0, + 1.0, ) } @@ -158,7 +159,12 @@ impl AgenticPreparation { include_replay_hashes: bool, policy: AgenticWarmupPolicy, primer_idle_cap_ms: f64, + speedup: f64, ) -> Result { + ensure!( + speedup.is_finite() && speedup > 0.0, + "warmup speedup must be finite and greater than zero" + ); ensure!( engine_block_size > 0, "warmup engine block size must be positive" @@ -288,7 +294,7 @@ impl AgenticPreparation { .iter() .map(|&index| { let start = lane.snapshot.context.recorded_start_ms(index); - let lead = (lane.snapshot.evidence.t_star_ms - start) + let lead = ((lane.snapshot.evidence.t_star_ms - start) / speedup) .max(0.0) .min(primer_idle_cap_ms); max_lead_ms = max_lead_ms.max(lead); @@ -1043,6 +1049,37 @@ mod tests { ); } + #[test] + fn live_primers_scale_source_leads_before_runtime_cap() { + // The fixture's two historical requests start at 0/20ms, with a 50ms + // cut: source leads are 50/30ms. Divide first, cap in runtime units, + // then align the largest lead to zero (no engine work is scaled). + for (speedup, cap_ms, second_ms) in [ + (2.0, 10_000.0, 10.0), + (0.5, 10_000.0, 40.0), + (2.0, 20.0, 5.0), + ] { + let mut driver = WorkloadDriver::new_agentic_warmup_with_policy( + prepared(1, 50.0), + 64, + true, + speedup, + AgenticWarmupPolicy::LiveTrajectoryV1, + cap_ms, + ) + .unwrap(); + let first = driver.pop_ready(0.0, 100); + assert_eq!(first.len(), 1); + assert_eq!(first[0].authored_request_id.as_deref(), Some("a-history")); + assert_eq!(driver.next_ready_time_ms(), Some(second_ms)); + assert!(driver.pop_ready(second_ms - 1.0, 100).is_empty()); + let second = driver.pop_ready(second_ms, 100); + assert_eq!(second.len(), 1); + assert_eq!(second[0].authored_request_id.as_deref(), Some("b-history")); + assert_eq!(second[0].request.arrival_timestamp_ms, Some(second_ms)); + } + } + #[test] fn live_primers_are_concurrent_and_globally_aligned() { let mut driver = live_driver(2, 50.0); diff --git a/docs/aic-backward-compatibility/migration.md b/docs/aic-backward-compatibility/migration.md index be1cef87a..3a5da8ff7 100644 --- a/docs/aic-backward-compatibility/migration.md +++ b/docs/aic-backward-compatibility/migration.md @@ -607,6 +607,30 @@ explains the retained workload routing and feature mathematics. the four legacy non-KV fields cannot coexist. See the [resource contract](../perf-model/memory.md). +### Rust AgentX snapshot options + +The snapshot-import extension changes the public Rust +`aisimulate_core::replay::loadgen::AgenticSnapshotOptions` source API. Replace +seed-only literals such as `AgenticSnapshotOptions { seed }` with +`AgenticSnapshotOptions::seeded(seed)`; the constructor leaves imported +selections and the corpus cursor unset. Code importing exact selections must +provide those two optional fields together. + +The options now own imported selections and are no longer `Copy`. Move the +options into their consumer, borrow them where the API permits, or explicitly +`.clone()` when the same selection must be consumed more than once. Floating-point +cuts also mean the type retains `PartialEq` but no longer implements `Eq`; update +generic bounds accordingly. These changes break Rust source compatibility, not +legacy seed-only JSON: +`{"seed": 42}` still deserializes and serializes without optional fields. + +Before the next stable `aisimulate-core` publication, the release owner must +record this source API break in the release notes, select a compatible release +version, and coordinate downstream caller migration and exact-artifact tests. +This migration note neither changes package versions nor certifies downstream +compatibility; nightly availability and the stable-release checklist remain +separate steps under the [release contract](../ci/release.md#nightly-ci-and-fpe-qualification). + ### Release and downstream contracts Wheel and crate versions must match. Namespace changes do not preserve removed diff --git a/docs/replay/agentic/warmup.md b/docs/replay/agentic/warmup.md index 919875832..cca8c19ba 100644 --- a/docs/replay/agentic/warmup.md +++ b/docs/replay/agentic/warmup.md @@ -185,8 +185,9 @@ unchanged when the policy is omitted. The policy follows the accelerated request-count warmup contract in [AIPerf `754356e9`](https://github.com/SemiAnalysisAI/aiperf/blob/754356e9a39acc6cc6afb242d123bb57c3fb6f75/src/aiperf/timing/strategies/agentic_replay.py): -1. Prime each continuing conversation's last historical full input. Primer - leads are globally aligned and capped by the configured global idle cap. +1. Prime each continuing conversation's last historical full input. Convert + source-clock primer leads to runtime milliseconds by dividing by `speedup`, + then apply the configured global idle cap and align them globally. Independent streams on the same lane may prime concurrently. No live request starts until every primer succeeds and settles. 2. Advance the existing dependency graph, including dispatch-triggered spawns, diff --git a/docs/sweeper/sdk.md b/docs/sweeper/sdk.md index 49995fac4..85e25eeeb 100644 --- a/docs/sweeper/sdk.md +++ b/docs/sweeper/sdk.md @@ -213,9 +213,10 @@ Every `Workload` field: | `arrival_speedup_ratio` | `float` | `1.0` | Scales the trace's inter-arrival times (open-loop trace only). `>1` speeds arrivals up. | | `agentic_lanes` | `int \| None` | `None` | Positive number of client play lanes for `weka`, `agentic_mooncake`, or agentic `dynamo` timestamp replay; does not cap concurrent child requests. Requires `source_type: trace`, `trace_path`, and `load_type: trace_timestamps`, with no `replay_concurrency`. | | `agentic_snapshot` | `AgenticSnapshotOptions \| None` | `None` (unset) | Optional object with required unsigned 64-bit `seed` (`0` through `2^64 - 1`) and the optional import fields below. Requires positive `agentic_lanes`, `source_type: trace`, `trace_path`, and `load_type: trace_timestamps`, with no `replay_concurrency`; supported formats are `weka`, `agentic_mooncake`, and agentic `dynamo`. Unset preserves turn-zero execution. | -| `agentic_snapshot.initial_selections` | `list[AgenticSnapshotSelection] \| None` | `None` (unset) | Exact imported `{source_play_id, t_star_ms}` cuts in lane order, with exactly one selection per `agentic_lanes`. Requires `corpus_cursor`; source IDs must exist in the input corpus and finite source-millisecond cuts must lie within the selected play's request-start bounds. Cuts are not re-sampled. | +| `agentic_snapshot.initial_selections` | `list[AgenticSnapshotSelection] \| None` | `None` (unset) | Exact imported `{source_play_id, t_star_ms}` cuts in lane order, with exactly one selection per `agentic_lanes`. Requires `corpus_cursor`; source IDs must exist in the input corpus and finite cuts must lie within the selected play's request-start bounds in the loaded graph's clock (Weka: play-relative milliseconds). Cuts are not re-sampled. | | `agentic_snapshot.corpus_cursor` | `int \| None` | `None` (unset) | Unsigned 64-bit next sequential source draw, including skipped draws. Requires `initial_selections` and must be at least their count; it may exceed the lane count. | | `agentic_warmup` | `bool` | `False` | Requires `agentic_snapshot` and positive lanes when enabled. Runs primers and ten warmup requests per lane before measured replay, preserving native cache state; see [warmup behavior](../replay/workloads.md#agentic-warmup). Failed preparation is retained as candidate evidence and excluded from ranking. | +| `agentic_warmup_policy` | `Literal["repeated_prefixes", "live_trajectory_v1"]` | `"repeated_prefixes"` | The default preserves the saved frontier and is omitted from serialized workloads. `live_trajectory_v1` advances the dependency graph with ten additional one-token admissions per lane; it requires `agentic_warmup: true`, `agentic_profile`, and a runner advertising `supports_agentic_live_warmup`. See [warmup policies](../replay/agentic/warmup.md#live-trajectory-policy). | | `agentic_profile` | `AgenticProfileOptions \| None` | `None` (unset) | Optional object; `{}` enables continuous lane replenishment with the defaults below. Requires `agentic_snapshot` and its agentic trace/lane controls; cannot be combined with `max_sim_time_ms`. Unset preserves finite replay. See [continuous profile](../replay/workloads.md#agentic-profile) for the complete configuration, results, and limits. | | `agentic_profile.duration_seconds` | `float` | `3600.0` when enabled | Positive finite admission duration, starting at the preparation barrier or simulation start without warmup. Stops new workload requests and replacement plays at the deadline. | | `agentic_profile.response_grace_seconds` | `float` | `30.0` when enabled | Nonnegative finite response window after admission closes; remaining client requests are then canceled. | @@ -236,12 +237,17 @@ shared synthetic knobs carried by `ReplaySpec.workload`. Profile controls are concrete workload values, not search dimensions. The built-in Engine runner supports profiles on offline aggregated or P/D vLLM/SGLang replay with HBM-only KV cache. vLLM also supports local or shared [G2 host offload](../replay/engine/kv-cache.md) on a static single -aggregated worker or 1P1D, with attention DP1 on every role. Speculative decoding and G3 -remain unsupported. Warmup is optional; enabling it retains the saved snapshot frontier. +aggregated worker or 1P1D, with attention DP1 on every role. G3 remains unsupported; +speculative decoding follows the narrower [engine SD constraints](../replay/engine/speculation.md). +Warmup is optional: `repeated_prefixes` retains the saved snapshot frontier, +while `live_trajectory_v1` advances it. A completed lane takes a new turn-zero play from the shared corpus cursor, wrapping as necessary -until the admission deadline. Snapshot sampling and warmup still differ from the AgentX reference; -see the [profile limits](../replay/workloads.md#agentic-profile). Injected runners must advertise -`supports_agentic_profile`; unsupported runners fail validation. +until the admission deadline. Seed-only sampling differs from the AgentX reference; import exact +cuts for matched comparisons. The live warmup policy follows the pinned reference's accelerated +request-count contract, not byte-exact HTTP prompt parity; see the +[profile limits](../replay/workloads.md#agentic-profile). Injected runners must advertise +`supports_agentic_profile`, plus `supports_agentic_live_warmup` for `live_trajectory_v1`; +unsupported runners fail validation. ## Validation (`Workload._validate_workload`) From 3fc7094d7011367cca1ce1ae97b9adde695b905f Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Tue, 20 Oct 2026 14:04:23 -0700 Subject: [PATCH 13/13] docs: distinguish AISimulate primer speedup from pinned reference Signed-off-by: hongkuanz --- docs/replay/agentic/warmup.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/replay/agentic/warmup.md b/docs/replay/agentic/warmup.md index cca8c19ba..5c7870d1c 100644 --- a/docs/replay/agentic/warmup.md +++ b/docs/replay/agentic/warmup.md @@ -185,6 +185,10 @@ unchanged when the policy is omitted. The policy follows the accelerated request-count warmup contract in [AIPerf `754356e9`](https://github.com/SemiAnalysisAI/aiperf/blob/754356e9a39acc6cc6afb242d123bb57c3fb6f75/src/aiperf/timing/strategies/agentic_replay.py): +AISimulate additionally scales source-clock primer leads using its trace-arrival +`speedup`; the reference has no such scaling knob. `speedup: 1` preserves the +reference's primer timing. + 1. Prime each continuing conversation's last historical full input. Convert source-clock primer leads to runtime milliseconds by dividing by `speedup`, then apply the configured global idle cap and align them globally.