feat: import reference AgentX snapshot selections

.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

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

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

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 {

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(

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)

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)

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] {

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};

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)

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<Vec<AgenticSnapshotSelection>>,
+    /// Next draw from the shared sequential corpus sampler, including skipped draws.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub corpus_cursor: Option<u64>,
+}
+
+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<String, usize>,
     node_to_play: Vec<usize>,
     topological_rank: Vec<usize>,
@@ -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,16 +258,44 @@ 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,
             topological_rank,
             recorded_intervals,
             token_stride,
         });
-        let plays = (0..lanes)
-            .map(|lane| context.prepare_play(lane, 0, None))
-            .collect::<Result<Vec<_>>>()?;
+        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::<FxHashMap<_, _>>();
+            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::<Result<Vec<_>>>()?
+        } else {
+            (0..lanes)
+                .map(|lane| context.prepare_play(lane, 0, None))
+                .collect::<Result<Vec<_>>>()?
+        };
         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