From 12ea1e7febd868d931cc3c38c618704f4337e755 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 10:27:54 -0700 Subject: [PATCH 1/5] 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 2/5] 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 17356c1328ea5939ec60be83ff590974585060ce Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 11:05:45 -0700 Subject: [PATCH 3/5] 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 93517a7ab99fa93f9b06c251c6fd96194bd6715c Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Wed, 14 Oct 2026 13:26:35 -0700 Subject: [PATCH 4/5] 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 c218f118b180a872e3d4ce6c2f9bf4c83b5260b0 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Mon, 19 Oct 2026 17:02:24 -0700 Subject: [PATCH 5/5] 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