dynamo-mocker 1.6.2

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

mod artifacts;
mod entrypoints;
pub(crate) mod offline;
mod online;
mod router_shared;
mod validate;

use std::collections::VecDeque;
use std::sync::Arc;

use crate::common::protocols::{DirectRequest, MockerConfig};
use dynamo_kv_router::PrefillLoadEstimator;

/// Backward-compatible Dynamo Mocker name for [`aisimulate_core::ReplayReport`].
pub use aisimulate_core::ReplayReport as TraceSimulationReport;
pub(crate) use aisimulate_core::replay::TraceCollector;
pub use aisimulate_core::replay::{
    CanonicalReplayCoverage, CanonicalReplayRecord, LifecycleOperation, OfflineRuntimeEvidence,
    PerRequestRecord, ReplayCaptureOptions, ReplayDeterminism, ReplayTerminalStatus, SlaThresholds,
    TraceDistributionStats, TraceGoodputStats, TraceInterTokenLatencyStats, TraceLatencyStats,
    TraceRequestCounts, TraceThroughputStats,
};
#[cfg(any(test, feature = "test-support"))]
#[doc(hidden)]
pub use artifacts::native_g1_parent_chain_artifact;
pub use artifacts::{
    ReplayTimedKvEvent, ReplayTimedOutputSignal, ReplayTimedRequest, ReplayWorkerArtifacts,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ReplayRouterMode {
    RoundRobin,
    KvRouter,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ReplayArgsMode {
    Aggregated,
    Disagg,
}

pub type ReplayPrefillLoadEstimator = Arc<dyn PrefillLoadEstimator>;

#[derive(Clone, Debug)]
pub struct OfflineDisaggReplayConfig {
    pub prefill_args: MockerConfig,
    pub decode_args: MockerConfig,
    pub num_prefill_workers: usize,
    pub num_decode_workers: usize,
}

impl OfflineDisaggReplayConfig {
    pub fn normalized(self) -> anyhow::Result<Self> {
        Ok(Self {
            prefill_args: self.prefill_args.normalized()?,
            decode_args: self.decode_args.normalized()?,
            num_prefill_workers: self.num_prefill_workers,
            num_decode_workers: self.num_decode_workers,
        })
    }
}

pub use aisimulate_core::replay::TrafficStats;
pub use aisimulate_core::replay::{
    ReplayScalingDecision, ReplayScalingPolicy, ReplayScalingSnapshot,
    ReplaySchedulerIntervalMetrics, ReplaySchedulerMetricsSnapshot, ReplayTelemetryObserver,
    ReplayTelemetrySampleKind, ReplayTelemetrySnapshot, ReplayTrafficMetricsSnapshot,
};

/// Optional policy-neutral telemetry sink for one offline replay.
///
/// Telemetry is deliberately configured independently from
/// [`ReplayScalingPolicy`]: observing a replay must not change the planner
/// callback contract or its decision cadence.
pub struct ReplayTelemetryOptions {
    pub sample_interval_ms: f64,
    pub observer: Box<dyn ReplayTelemetryObserver>,
}

/// Optional replay observers installed for one invocation.
///
/// Scaling and telemetry share an invocation without coupling their callback
/// contracts or sampling cadence.
#[derive(Default)]
pub struct ReplayRuntimeObservers {
    pub scaling_policy: Option<Box<dyn ReplayScalingPolicy>>,
    pub telemetry: Option<ReplayTelemetryOptions>,
}
pub use entrypoints::{
    ReplayKvEventVisibility, generate_trace_worker_artifacts_offline,
    generate_trace_worker_artifacts_offline_with_kv_event_visibility,
    simulate_agentic_trace_live_workload_with_router_mode_and_options,
    simulate_agentic_trace_workload_disagg_with_router_mode,
    simulate_agentic_trace_workload_disagg_with_router_mode_and_runtime_observers,
    simulate_agentic_trace_workload_disagg_with_router_mode_and_telemetry,
    simulate_agentic_trace_workload_with_router_mode,
    simulate_agentic_trace_workload_with_router_mode_and_runtime_observers,
    simulate_agentic_trace_workload_with_router_mode_and_telemetry, simulate_concurrency_file,
    simulate_concurrency_file_disagg_with_router_mode,
    simulate_concurrency_file_disagg_with_router_mode_and_format,
    simulate_concurrency_file_disagg_with_router_mode_and_format_and_runtime_observers,
    simulate_concurrency_file_with_router_mode,
    simulate_concurrency_file_with_router_mode_and_format,
    simulate_concurrency_file_with_router_mode_and_format_and_runtime_observers,
    simulate_concurrency_live_file, simulate_concurrency_live_file_with_router_mode,
    simulate_concurrency_live_file_with_router_mode_and_format,
    simulate_concurrency_live_file_with_router_mode_and_format_and_options,
    simulate_concurrency_live_requests, simulate_concurrency_live_requests_with_router_mode,
    simulate_concurrency_live_requests_with_router_mode_and_options,
    simulate_concurrency_live_workload, simulate_concurrency_live_workload_with_router_mode,
    simulate_concurrency_live_workload_with_router_mode_and_options, simulate_concurrency_requests,
    simulate_concurrency_requests_disagg_with_router_mode,
    simulate_concurrency_requests_disagg_with_router_mode_and_runtime_observers,
    simulate_concurrency_requests_with_router_mode,
    simulate_concurrency_requests_with_router_mode_and_runtime_observers,
    simulate_concurrency_workload, simulate_concurrency_workload_disagg_with_router_mode,
    simulate_concurrency_workload_disagg_with_router_mode_and_options,
    simulate_concurrency_workload_disagg_with_router_mode_and_options_and_runtime_observers,
    simulate_concurrency_workload_with_router_mode,
    simulate_concurrency_workload_with_router_mode_and_options,
    simulate_concurrency_workload_with_router_mode_and_options_and_runtime_observers,
    simulate_loaded_trace_disagg_with_router_mode_and_capture_options,
    simulate_loaded_trace_disagg_with_router_mode_and_options,
    simulate_loaded_trace_disagg_with_router_mode_and_options_and_runtime_observers,
    simulate_loaded_trace_live_with_router_mode,
    simulate_loaded_trace_live_with_router_mode_and_options,
    simulate_loaded_trace_with_router_mode_and_capture_options,
    simulate_loaded_trace_with_router_mode_and_options,
    simulate_loaded_trace_with_router_mode_and_options_and_runtime_observers, simulate_trace_file,
    simulate_trace_file_disagg_with_router_mode,
    simulate_trace_file_disagg_with_router_mode_and_format,
    simulate_trace_file_disagg_with_router_mode_and_format_and_runtime_observers,
    simulate_trace_file_with_router_mode, simulate_trace_file_with_router_mode_and_format,
    simulate_trace_file_with_router_mode_and_format_and_runtime_observers,
    simulate_trace_live_file, simulate_trace_live_file_with_router_mode,
    simulate_trace_live_file_with_router_mode_and_format,
    simulate_trace_live_file_with_router_mode_and_format_and_options, simulate_trace_live_requests,
    simulate_trace_live_requests_with_router_mode,
    simulate_trace_live_requests_with_router_mode_and_options, simulate_trace_live_workload,
    simulate_trace_live_workload_with_router_mode,
    simulate_trace_live_workload_with_router_mode_and_options, simulate_trace_requests,
    simulate_trace_requests_disagg_with_router_mode,
    simulate_trace_requests_disagg_with_router_mode_and_runtime_observers,
    simulate_trace_requests_with_router_mode,
    simulate_trace_requests_with_router_mode_and_runtime_observers, simulate_trace_workload,
    simulate_trace_workload_disagg_with_router_mode,
    simulate_trace_workload_disagg_with_router_mode_and_options_and_runtime_observers,
    simulate_trace_workload_with_router_mode,
    simulate_trace_workload_with_router_mode_and_options_and_runtime_observers,
};
#[doc(hidden)]
pub use offline::run_offline_handoff_conformance;
pub use validate::validate_replay_args_mode;

thread_local! {
    static KV_EVENT_LAG_MS: std::cell::Cell<f64> = const { std::cell::Cell::new(0.0) };
}

/// Run `replay` with a simulated KV event lag for offline KV-router replay.
///
/// Delays the KV cache events (blocks stored and removed) the router's indexer observes by
/// `lag_ms` of simulated time. Prefill and request completions stay immediate, as a live router
/// observes them in-band on the response path. Applies to offline KV-router replays constructed
/// on this thread inside `replay`; `0.0` keeps synchronous updates.
pub fn with_kv_event_lag_ms<T>(lag_ms: f64, replay: impl FnOnce() -> T) -> anyhow::Result<T> {
    anyhow::ensure!(
        lag_ms.is_finite() && lag_ms >= 0.0,
        "kv_event_lag_ms must be finite and non-negative, got {lag_ms}"
    );
    struct Restore(f64);
    impl Drop for Restore {
        fn drop(&mut self) {
            KV_EVENT_LAG_MS.with(|lag| lag.set(self.0));
        }
    }
    let _restore = Restore(KV_EVENT_LAG_MS.with(|lag| lag.replace(lag_ms)));
    Ok(replay())
}

pub(crate) fn kv_event_lag_ms() -> f64 {
    KV_EVENT_LAG_MS.with(std::cell::Cell::get)
}

pub(crate) fn normalize_trace_requests(
    mut requests: Vec<DirectRequest>,
    arrival_speedup_ratio: f64,
) -> anyhow::Result<VecDeque<DirectRequest>> {
    if !arrival_speedup_ratio.is_finite() || arrival_speedup_ratio <= 0.0 {
        anyhow::bail!(
            "arrival_speedup_ratio must be a finite positive number, got {arrival_speedup_ratio}"
        );
    }

    requests.sort_by(|left, right| {
        let left_ts = left
            .arrival_timestamp_ms
            .expect("trace replay requests must have an arrival timestamp");
        let right_ts = right
            .arrival_timestamp_ms
            .expect("trace replay requests must have an arrival timestamp");
        left_ts.total_cmp(&right_ts)
    });

    let first_arrival_ms = requests
        .first()
        .and_then(|request| request.arrival_timestamp_ms)
        .ok_or_else(|| anyhow::anyhow!("trace replay requires at least one timestamped request"))?;

    Ok(VecDeque::from(
        requests
            .into_iter()
            .map(|mut request| {
                let arrival_timestamp_ms = request
                    .arrival_timestamp_ms
                    .expect("trace replay requests must have an arrival timestamp")
                    - first_arrival_ms;
                let arrival_timestamp_ms = arrival_timestamp_ms / arrival_speedup_ratio;
                request.arrival_timestamp_ms = Some(arrival_timestamp_ms);
                request
            })
            .collect::<Vec<_>>(),
    ))
}

pub(crate) fn effective_agentic_lanes(
    requested_lanes: Option<usize>,
    play_count: usize,
) -> Option<usize> {
    requested_lanes.map(|lane_count| lane_count.min(play_count))
}

#[cfg(test)]
mod tests {
    use super::*;
    use uuid::Uuid;

    #[test]
    fn test_normalize_trace_requests_applies_arrival_speedup_ratio() {
        let requests = vec![
            DirectRequest {
                tokens: vec![1; 4],
                max_output_tokens: 1,
                output_token_ids: None,
                uuid: Some(Uuid::from_u128(1)),
                dp_rank: 0,
                arrival_timestamp_ms: Some(100.0),
                ..Default::default()
            },
            DirectRequest {
                tokens: vec![2; 4],
                max_output_tokens: 1,
                output_token_ids: None,
                uuid: Some(Uuid::from_u128(2)),
                dp_rank: 0,
                arrival_timestamp_ms: Some(200.0),
                ..Default::default()
            },
        ];

        let normalized = normalize_trace_requests(requests, 10.0).unwrap();
        let arrivals = normalized
            .into_iter()
            .map(|request| request.arrival_timestamp_ms.unwrap())
            .collect::<Vec<_>>();

        assert_eq!(arrivals, vec![0.0, 10.0]);
    }

    #[test]
    fn test_effective_agentic_lanes_are_bounded_by_play_count() {
        assert_eq!(effective_agentic_lanes(None, 3), None);
        assert_eq!(effective_agentic_lanes(Some(2), 3), Some(2));
        assert_eq!(effective_agentic_lanes(Some(usize::MAX), 3), Some(3));
    }
}