aisimulate_core/replay/
spec.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::BTreeSet;
5
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8
9use crate::replay::{ReplayError, ReplayResult, SlaThresholds};
10
11pub const CURRENT_REPLAY_SPEC_VERSION: u32 = 1;
12
13/// Serializable input to one replay execution.
14///
15/// Provider descriptors are data only. A runner resolves them to concrete
16/// placement/scaling implementations before constructing [`crate::replay::Replayer`].
17#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
18pub struct ReplaySpec {
19    #[serde(default = "default_spec_version")]
20    pub version: u32,
21    pub topology: ReplayTopology,
22    #[serde(default)]
23    pub engine: Value,
24    #[serde(default)]
25    pub adapters: ReplayAdapters,
26    /// Soft virtual-time cutoff. Events at the cutoff are processed; replay
27    /// stops before the first event after it and leaves in-flight requests
28    /// non-terminal in the report.
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub max_sim_time_ms: Option<f64>,
31    /// Optional source-side in-flight cap. Requests whose authored arrival is
32    /// ready remain outside the simulated system until an earlier request is
33    /// terminal, matching closed-loop and replay-concurrency workloads.
34    #[serde(default, skip_serializing_if = "Option::is_none")]
35    pub max_in_flight: Option<usize>,
36    /// Whether the report should retain one record per arrived request.
37    ///
38    /// This defaults to `true` to preserve reports produced by version-1
39    /// execution specs written before this control was added.
40    #[serde(
41        default = "default_record_per_request",
42        skip_serializing_if = "is_true"
43    )]
44    pub record_per_request: bool,
45    /// Optional latency targets used to calculate goodput.
46    #[serde(default, skip_serializing_if = "SlaThresholds::is_unset")]
47    pub sla: SlaThresholds,
48    pub requests: Vec<ReplayRequest>,
49}
50
51impl ReplaySpec {
52    pub fn validate(&self) -> ReplayResult<()> {
53        if self.version != CURRENT_REPLAY_SPEC_VERSION {
54            return Err(ReplayError::InvalidSpec(format!(
55                "unsupported version {}; expected {}",
56                self.version, CURRENT_REPLAY_SPEC_VERSION
57            )));
58        }
59        self.topology.validate()?;
60        if let Some(max_sim_time_ms) = self.max_sim_time_ms {
61            validate_time("max_sim_time_ms", max_sim_time_ms)?;
62        }
63        if self.max_in_flight == Some(0) {
64            return Err(ReplayError::InvalidSpec(
65                "max_in_flight must be positive".to_string(),
66            ));
67        }
68        self.sla.validate()?;
69
70        let mut ids = BTreeSet::new();
71        for request in &self.requests {
72            request.validate()?;
73            if !ids.insert(request.id.clone()) {
74                return Err(ReplayError::InvalidSpec(format!(
75                    "duplicate request id {:?}",
76                    request.id
77                )));
78            }
79        }
80        Ok(())
81    }
82}
83
84fn default_spec_version() -> u32 {
85    CURRENT_REPLAY_SPEC_VERSION
86}
87
88fn default_record_per_request() -> bool {
89    true
90}
91
92fn is_true(value: &bool) -> bool {
93    *value
94}
95
96#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
97#[serde(tag = "kind", rename_all = "snake_case")]
98pub enum ReplayTopology {
99    Aggregated {
100        workers: WorkerPoolSpec,
101    },
102    Disaggregated {
103        prefill: WorkerPoolSpec,
104        decode: WorkerPoolSpec,
105        /// Transfer latency used only when the selected engine cannot derive
106        /// one from KV bytes and bandwidth.
107        #[serde(default)]
108        handoff_latency_ms: f64,
109    },
110}
111
112impl ReplayTopology {
113    pub fn aggregated(workers: usize) -> Self {
114        Self::Aggregated {
115            workers: WorkerPoolSpec {
116                initial_workers: workers,
117                ..WorkerPoolSpec::default()
118            },
119        }
120    }
121
122    pub fn validate(&self) -> ReplayResult<()> {
123        match self {
124            Self::Aggregated { workers } => workers.validate("aggregated"),
125            Self::Disaggregated {
126                prefill,
127                decode,
128                handoff_latency_ms,
129            } => {
130                prefill.validate("prefill")?;
131                decode.validate("decode")?;
132                validate_time("handoff_latency_ms", *handoff_latency_ms)
133            }
134        }
135    }
136}
137
138#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
139pub struct WorkerPoolSpec {
140    pub initial_workers: usize,
141    #[serde(default)]
142    pub startup_delay_ms: f64,
143}
144
145impl Default for WorkerPoolSpec {
146    fn default() -> Self {
147        Self {
148            initial_workers: 1,
149            startup_delay_ms: 0.0,
150        }
151    }
152}
153
154impl WorkerPoolSpec {
155    fn validate(&self, name: &str) -> ReplayResult<()> {
156        if self.initial_workers == 0 {
157            return Err(ReplayError::InvalidSpec(format!(
158                "{name} pool must start with at least one worker"
159            )));
160        }
161        validate_time(&format!("{name} startup_delay_ms"), self.startup_delay_ms)
162    }
163}
164
165#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
166pub struct ReplayRequest {
167    pub id: String,
168    pub arrival_time_ms: f64,
169    pub input_tokens: usize,
170    /// Materialized prompt content for KV-aware replay. Length-only engine
171    /// replay may omit it; adapters must never invent tokens when it is needed.
172    #[serde(default, skip_serializing_if = "Option::is_none")]
173    pub input_token_ids: Option<Vec<u32>>,
174    pub output_tokens: usize,
175    /// Optional exact generated-token plan. When present, its length controls
176    /// native generation while `output_tokens` remains the authored maximum.
177    #[serde(default, skip_serializing_if = "Option::is_none")]
178    pub output_token_ids: Option<Vec<u32>>,
179    /// Optional attention-DP rank selected by the workload.
180    #[serde(default, skip_serializing_if = "Option::is_none")]
181    pub dp_rank: Option<u32>,
182    #[serde(default, skip_serializing_if = "Option::is_none")]
183    pub session_id: Option<String>,
184    /// Zero-based turn index within `session_id`, when the workload has one.
185    #[serde(default, skip_serializing_if = "Option::is_none")]
186    pub turn_index: Option<usize>,
187    /// Provider-neutral request metadata resolved by the selected runner.
188    #[serde(default, skip_serializing_if = "Value::is_null")]
189    pub metadata: Value,
190}
191
192impl ReplayRequest {
193    pub fn validate(&self) -> ReplayResult<()> {
194        if self.id.is_empty() {
195            return Err(ReplayError::InvalidSpec(
196                "request id must not be empty".to_string(),
197            ));
198        }
199        validate_time("request arrival_time_ms", self.arrival_time_ms)?;
200        if let Some(input_token_ids) = &self.input_token_ids
201            && input_token_ids.len() != self.input_tokens
202        {
203            return Err(ReplayError::InvalidSpec(format!(
204                "request {:?} declares {} input tokens but materializes {} token IDs",
205                self.id,
206                self.input_tokens,
207                input_token_ids.len()
208            )));
209        }
210        self.routing_metadata()?;
211        Ok(())
212    }
213
214    /// Resolve the small routing subset of provider-neutral metadata while
215    /// leaving the authored metadata value intact for reporting and other
216    /// adapters.
217    pub fn routing_metadata(&self) -> ReplayResult<ReplayRoutingMetadata> {
218        if self.metadata.is_null() {
219            return Ok(ReplayRoutingMetadata::default());
220        }
221        serde_json::from_value(self.metadata.clone()).map_err(|error| {
222            ReplayError::InvalidSpec(format!(
223                "request {:?} has invalid routing metadata: {error}",
224                self.id
225            ))
226        })
227    }
228}
229
230/// Provider-neutral routing controls recognized during ReplaySpec lowering.
231/// Unknown metadata keys remain available in [`ReplayRequest::metadata`].
232#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
233pub struct ReplayRoutingMetadata {
234    #[serde(default)]
235    pub priority: i32,
236    #[serde(default)]
237    pub strict_priority: u32,
238    #[serde(default)]
239    pub policy_class: Option<String>,
240}
241
242#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
243pub struct ReplayAdapters {
244    #[serde(default = "ProviderSpec::round_robin")]
245    pub placement: ProviderSpec,
246    #[serde(default = "ProviderSpec::no_scaling")]
247    pub scaling: ProviderSpec,
248}
249
250impl Default for ReplayAdapters {
251    fn default() -> Self {
252        Self {
253            placement: ProviderSpec::round_robin(),
254            scaling: ProviderSpec::no_scaling(),
255        }
256    }
257}
258
259#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
260pub struct ProviderSpec {
261    pub provider: String,
262    #[serde(default)]
263    pub config: Value,
264}
265
266impl ProviderSpec {
267    pub fn round_robin() -> Self {
268        Self {
269            provider: "round_robin".to_string(),
270            config: Value::Null,
271        }
272    }
273
274    pub fn no_scaling() -> Self {
275        Self {
276            provider: "none".to_string(),
277            config: Value::Null,
278        }
279    }
280}
281
282#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
283#[serde(rename_all = "snake_case")]
284pub enum WorkerStage {
285    Aggregated,
286    Prefill,
287    Decode,
288}
289
290fn validate_time(name: &str, value: f64) -> ReplayResult<()> {
291    if !value.is_finite() || value < 0.0 {
292        return Err(ReplayError::InvalidSpec(format!(
293            "{name} must be finite and non-negative, got {value}"
294        )));
295    }
296    Ok(())
297}
aisimulate_core/replay/spec.rs

aisimulate_core/replay/
spec.rs
