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 diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 6d262e270..b40b59e74 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 33f3f2ce1..6e9323433 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -1553,7 +1553,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 e5e5a5a78..4e7199753 100644 --- a/crates/core/src/replay/agg.rs +++ b/crates/core/src/replay/agg.rs @@ -1752,7 +1752,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( @@ -1912,7 +1912,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 85b690156..ef719259c 100644 --- a/crates/core/src/replay/disagg_tests.rs +++ b/crates/core/src/replay/disagg_tests.rs @@ -3119,7 +3119,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 bbcfa28f7..bbe684b97 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), @@ -4968,7 +4968,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 240eff4da..0d5ad4af3 100644 --- a/crates/core/src/replay/loadgen/phase.rs +++ b/crates/core/src/replay/loadgen/phase.rs @@ -681,7 +681,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) @@ -982,7 +982,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..2f0040ebc 100644 --- a/crates/core/src/replay/loadgen/snapshot.rs +++ b/crates/core/src/replay/loadgen/snapshot.rs @@ -24,10 +24,33 @@ 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, + /// Loaded graph clock: per-play-relative for Weka, source clock for Mooncake. + 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 +103,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 +153,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 +258,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 +266,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 +325,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 377a23bf7..e0cf8b3e0 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/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/crates/tests/public-api/src/lib.rs b/crates/tests/public-api/src/lib.rs index d5bdd992f..535e13847 100644 --- a/crates/tests/public-api/src/lib.rs +++ b/crates/tests/public-api/src/lib.rs @@ -611,7 +611,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 a5f169c55..ed5c9efaa 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, 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). | | `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. | @@ -275,7 +277,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. | @@ -333,6 +335,85 @@ 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 +``` + +`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 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. + +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 --trace-format weka \ + --lanes 15 --seed 42 --output snapshot-options.json +``` + +`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. +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. It fails on unsupported samplers, timestamp-less trajectories, +or incomplete lane selection. Neither AIPerf nor NumPy becomes an AISimulate +runtime dependency. + ### Warmup @@ -403,7 +484,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/THIRD_PARTY_NOTICES.md b/python/aisimulate/THIRD_PARTY_NOTICES.md index 6d262e270..b40b59e74 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/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/python/aisimulate/src/aisimulate/sweeper/replay.py b/python/aisimulate/src/aisimulate/sweeper/replay.py index 29c14c89a..860de0824 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 @@ -539,15 +539,14 @@ 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") + 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: + 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..74372a582 --- /dev/null +++ b/scripts/prediction_regression/export_agentx_snapshots.py @@ -0,0 +1,217 @@ +#!/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 math +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 _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") + 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") + 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": 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) + 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) + 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() + 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) + clock_evidence = [] + result = export_snapshot( + metadata, + json.loads(mapping_bytes), + lanes=args.lanes, + seed=args.seed, + trajectory_type=module.TrajectorySource, + trace_format=args.trace_format, + clock_evidence=clock_evidence, + ) + 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, + ) + + +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..b13b5774d --- /dev/null +++ b/tests/test_agentx_snapshot_export.py @@ -0,0 +1,282 @@ +# 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, write_export + + +def metadata(strategy="sequential"): + return Ns( + sampling_strategy=strategy, + conversations=[ + 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=[]), + ], + ) + + +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, + 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}], + "corpus_cursor": 3, + } + + +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"] + 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, 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, trace_format="weka") + + +@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), + 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 diff --git a/tests/test_cli_config.py b/tests/test_cli_config.py index fbe0e2d99..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 @@ -1658,6 +1670,77 @@ 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("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 316545b6d..1395dd740 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -2221,9 +2221,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): @@ -2231,26 +2260,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