1mod artifacts;
5mod collector;
6mod entrypoints;
7pub(crate) mod offline;
8mod online;
9mod router_shared;
10mod validate;
11
12use std::collections::VecDeque;
13use std::sync::Arc;
14
15use crate::common::protocols::{DirectRequest, MockEngineArgs};
16use dynamo_kv_router::PrefillLoadEstimator;
17
18#[cfg(any(test, feature = "test-support"))]
19#[doc(hidden)]
20pub use artifacts::native_g1_parent_chain_artifact;
21pub use artifacts::{
22 ReplayTimedKvEvent, ReplayTimedOutputSignal, ReplayTimedRequest, ReplayWorkerArtifacts,
23};
24pub(crate) use collector::TraceCollector;
25#[cfg(test)]
26pub(crate) use collector::TraceRequestStatsSnapshot;
27pub use collector::{
28 PerRequestRecord, ReplayTerminalStatus, SlaThresholds, TraceDistributionStats,
29 TraceGoodputStats, TraceInterTokenLatencyStats, TraceLatencyStats, TraceRequestCounts,
30 TraceSimulationReport, TraceThroughputStats,
31};
32#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33pub enum ReplayRouterMode {
34 RoundRobin,
35 KvRouter,
36}
37
38#[derive(Clone, Copy, Debug, PartialEq, Eq)]
39pub enum ReplayArgsMode {
40 Aggregated,
41 Disagg,
42}
43
44pub type ReplayPrefillLoadEstimator = Arc<dyn PrefillLoadEstimator>;
45
46#[derive(Clone, Debug)]
47pub struct OfflineDisaggReplayConfig {
48 pub prefill_args: MockEngineArgs,
49 pub decode_args: MockEngineArgs,
50 pub num_prefill_workers: usize,
51 pub num_decode_workers: usize,
52}
53
54impl OfflineDisaggReplayConfig {
55 pub fn normalized(self) -> anyhow::Result<Self> {
56 Ok(Self {
57 prefill_args: self.prefill_args.normalized()?,
58 decode_args: self.decode_args.normalized()?,
59 num_prefill_workers: self.num_prefill_workers,
60 num_decode_workers: self.num_decode_workers,
61 })
62 }
63}
64
65pub use entrypoints::{
66 ReplayKvEventVisibility, generate_trace_worker_artifacts_offline,
67 generate_trace_worker_artifacts_offline_with_kv_event_visibility,
68 simulate_agentic_trace_live_workload_with_router_mode_and_options,
69 simulate_agentic_trace_workload_with_router_mode, simulate_concurrency_file,
70 simulate_concurrency_file_disagg_with_router_mode,
71 simulate_concurrency_file_disagg_with_router_mode_and_format,
72 simulate_concurrency_file_disagg_with_router_mode_and_format_and_scaling_policy,
73 simulate_concurrency_file_with_router_mode,
74 simulate_concurrency_file_with_router_mode_and_format,
75 simulate_concurrency_file_with_router_mode_and_format_and_scaling_policy,
76 simulate_concurrency_live_file, simulate_concurrency_live_file_with_router_mode,
77 simulate_concurrency_live_file_with_router_mode_and_format,
78 simulate_concurrency_live_file_with_router_mode_and_format_and_options,
79 simulate_concurrency_live_requests, simulate_concurrency_live_requests_with_router_mode,
80 simulate_concurrency_live_requests_with_router_mode_and_options,
81 simulate_concurrency_live_workload, simulate_concurrency_live_workload_with_router_mode,
82 simulate_concurrency_live_workload_with_router_mode_and_options, simulate_concurrency_requests,
83 simulate_concurrency_requests_disagg_with_router_mode,
84 simulate_concurrency_requests_disagg_with_router_mode_and_scaling_policy,
85 simulate_concurrency_requests_with_router_mode,
86 simulate_concurrency_requests_with_router_mode_and_scaling_policy,
87 simulate_concurrency_workload, simulate_concurrency_workload_disagg_with_router_mode,
88 simulate_concurrency_workload_disagg_with_router_mode_and_options,
89 simulate_concurrency_workload_disagg_with_router_mode_and_options_and_scaling_policy,
90 simulate_concurrency_workload_with_router_mode,
91 simulate_concurrency_workload_with_router_mode_and_options,
92 simulate_concurrency_workload_with_router_mode_and_options_and_scaling_policy,
93 simulate_loaded_trace_disagg_with_router_mode_and_options,
94 simulate_loaded_trace_disagg_with_router_mode_and_options_and_scaling_policy,
95 simulate_loaded_trace_live_with_router_mode,
96 simulate_loaded_trace_live_with_router_mode_and_options,
97 simulate_loaded_trace_with_router_mode_and_options,
98 simulate_loaded_trace_with_router_mode_and_options_and_scaling_policy, simulate_trace_file,
99 simulate_trace_file_disagg_with_router_mode,
100 simulate_trace_file_disagg_with_router_mode_and_format,
101 simulate_trace_file_disagg_with_router_mode_and_format_and_scaling_policy,
102 simulate_trace_file_with_router_mode, simulate_trace_file_with_router_mode_and_format,
103 simulate_trace_file_with_router_mode_and_format_and_scaling_policy, simulate_trace_live_file,
104 simulate_trace_live_file_with_router_mode,
105 simulate_trace_live_file_with_router_mode_and_format,
106 simulate_trace_live_file_with_router_mode_and_format_and_options, simulate_trace_live_requests,
107 simulate_trace_live_requests_with_router_mode,
108 simulate_trace_live_requests_with_router_mode_and_options, simulate_trace_live_workload,
109 simulate_trace_live_workload_with_router_mode,
110 simulate_trace_live_workload_with_router_mode_and_options, simulate_trace_requests,
111 simulate_trace_requests_disagg_with_router_mode,
112 simulate_trace_requests_disagg_with_router_mode_and_scaling_policy,
113 simulate_trace_requests_with_router_mode,
114 simulate_trace_requests_with_router_mode_and_scaling_policy, simulate_trace_workload,
115 simulate_trace_workload_disagg_with_router_mode,
116 simulate_trace_workload_disagg_with_router_mode_and_options_and_scaling_policy,
117 simulate_trace_workload_with_router_mode,
118 simulate_trace_workload_with_router_mode_and_options_and_scaling_policy,
119};
120pub use offline::components::TrafficStats;
121#[doc(hidden)]
122pub use offline::run_offline_handoff_conformance;
123pub use offline::scaling::{ReplayScalingDecision, ReplayScalingPolicy, ReplayScalingSnapshot};
124pub use validate::validate_replay_args_mode;
125
126pub(crate) fn normalize_trace_requests(
127 mut requests: Vec<DirectRequest>,
128 arrival_speedup_ratio: f64,
129) -> anyhow::Result<VecDeque<DirectRequest>> {
130 if !arrival_speedup_ratio.is_finite() || arrival_speedup_ratio <= 0.0 {
131 anyhow::bail!(
132 "arrival_speedup_ratio must be a finite positive number, got {arrival_speedup_ratio}"
133 );
134 }
135
136 requests.sort_by(|left, right| {
137 let left_ts = left
138 .arrival_timestamp_ms
139 .expect("trace replay requests must have an arrival timestamp");
140 let right_ts = right
141 .arrival_timestamp_ms
142 .expect("trace replay requests must have an arrival timestamp");
143 left_ts.total_cmp(&right_ts)
144 });
145
146 let first_arrival_ms = requests
147 .first()
148 .and_then(|request| request.arrival_timestamp_ms)
149 .ok_or_else(|| anyhow::anyhow!("trace replay requires at least one timestamped request"))?;
150
151 Ok(VecDeque::from(
152 requests
153 .into_iter()
154 .map(|mut request| {
155 let arrival_timestamp_ms = request
156 .arrival_timestamp_ms
157 .expect("trace replay requests must have an arrival timestamp")
158 - first_arrival_ms;
159 let arrival_timestamp_ms = arrival_timestamp_ms / arrival_speedup_ratio;
160 request.arrival_timestamp_ms = Some(arrival_timestamp_ms);
161 request
162 })
163 .collect::<Vec<_>>(),
164 ))
165}
166
167#[cfg(test)]
168mod tests {
169 use super::*;
170 use uuid::Uuid;
171
172 #[test]
173 fn test_replay_itl_uses_per_token_gaps() {
174 fn assert_ddsketch_relative_error(actual: f64, expected: f64) {
175 assert!((actual - expected).abs() <= expected.abs() * 0.001 + f64::EPSILON);
176 }
177
178 let mut collector = TraceCollector::default();
179 let uuid = Uuid::from_u128(11);
180
181 collector.on_arrival(uuid, 0.0, 4, 4);
182 collector.on_admit(uuid, 0.0, 0);
183 collector.on_token(uuid, 10.0);
184 collector.on_token(uuid, 11.0);
185 collector.on_token(uuid, 12.0);
186 collector.on_token(uuid, 110.0);
187 collector.on_terminal(uuid, 110.0, ReplayTerminalStatus::Completed);
188
189 let report = collector.finish();
190
191 assert!((report.latency.tpot.mean_ms - (100.0 / 3.0)).abs() < 1e-9);
192 assert!((report.latency.itl.distribution.mean_ms - (100.0 / 3.0)).abs() < 1e-9);
193 assert_ddsketch_relative_error(report.latency.itl.distribution.median_ms, 1.0);
194 assert_ddsketch_relative_error(report.latency.itl.distribution.p75_ms, 98.0);
195 assert_ddsketch_relative_error(report.latency.itl.distribution.p90_ms, 98.0);
196 assert_ddsketch_relative_error(report.latency.itl.distribution.p95_ms, 98.0);
197 assert_eq!(report.latency.itl.max_ms, 98.0);
198 assert_eq!(report.latency.ttst.min_ms, 1.0);
199 assert_eq!(report.latency.ttst.max_ms, 1.0);
200 assert_eq!(
201 report.latency.output_token_throughput_per_user.min_ms,
202 1000.0 / 98.0
203 );
204 assert_eq!(
205 report.latency.output_token_throughput_per_user.max_ms,
206 1000.0
207 );
208 }
209
210 #[test]
211 fn test_normalize_trace_requests_applies_arrival_speedup_ratio() {
212 let requests = vec![
213 DirectRequest {
214 tokens: vec![1; 4],
215 max_output_tokens: 1,
216 output_token_ids: None,
217 uuid: Some(Uuid::from_u128(1)),
218 dp_rank: 0,
219 arrival_timestamp_ms: Some(100.0),
220 ..Default::default()
221 },
222 DirectRequest {
223 tokens: vec![2; 4],
224 max_output_tokens: 1,
225 output_token_ids: None,
226 uuid: Some(Uuid::from_u128(2)),
227 dp_rank: 0,
228 arrival_timestamp_ms: Some(200.0),
229 ..Default::default()
230 },
231 ];
232
233 let normalized = normalize_trace_requests(requests, 10.0).unwrap();
234 let arrivals = normalized
235 .into_iter()
236 .map(|request| request.arrival_timestamp_ms.unwrap())
237 .collect::<Vec<_>>();
238
239 assert_eq!(arrivals, vec![0.0, 10.0]);
240 }
241}