dynamo_mocker/common/
utils.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::time::{Duration, Instant};
5
6use crate::common::handoff::HandoffTransferTiming;
7use crate::common::protocols::{KvTransferTimingMode, MockEngineArgs, WorkerType};
8
9pub fn prefill_handoff_transfer_timing(
10    num_input_tokens: usize,
11    kv_transfer_bandwidth: Option<f64>,
12    kv_bytes_per_token: Option<usize>,
13    mode: KvTransferTimingMode,
14) -> HandoffTransferTiming {
15    HandoffTransferTiming {
16        mode,
17        full_prompt_tokens: num_input_tokens,
18        kv_bytes_per_token,
19        bandwidth_gb_s: kv_transfer_bandwidth,
20    }
21}
22
23/// Compute the modeled handoff delay after a prefill worker emits its terminal token.
24///
25/// NOTE: this intentionally does not model the internal prefill TTFT itself accurately, and the
26/// exact prefill/decode boundary is backend dependent. For now we only care about decode-visible
27/// TTFT, which is what the client observes, so modeling the delay as prefill-to-decode handoff is
28/// good enough.
29pub fn compute_prefill_handoff_delay_ms(
30    worker_type: WorkerType,
31    completed: bool,
32    num_input_tokens: usize,
33    kv_transfer_bandwidth: Option<f64>,
34    kv_bytes_per_token: Option<usize>,
35) -> Option<f64> {
36    if worker_type != WorkerType::Prefill || !completed {
37        return None;
38    }
39    let timing = prefill_handoff_transfer_timing(
40        num_input_tokens,
41        kv_transfer_bandwidth,
42        kv_bytes_per_token,
43        KvTransferTimingMode::FullPrompt,
44    );
45    match timing.full_prompt_delay_ms() {
46        Some(delay_ms) => {
47            tracing::debug!(
48                num_input_tokens,
49                bandwidth_gb_s = kv_transfer_bandwidth,
50                delay_ms = format!("{delay_ms:.2}"),
51                "KV handoff delay for prefill completion"
52            );
53            Some(delay_ms)
54        }
55        None => None,
56    }
57}
58
59/// Compute the KV transfer delay duration for a given number of input tokens.
60///
61/// Returns `None` if KV transfer simulation is disabled (bandwidth is 0 or not configured).
62pub fn compute_kv_transfer_delay(
63    args: &MockEngineArgs,
64    num_input_tokens: usize,
65) -> Option<Duration> {
66    compute_prefill_handoff_delay_ms(
67        args.worker_type,
68        true,
69        num_input_tokens,
70        args.kv_transfer_bandwidth,
71        args.kv_bytes_per_token,
72    )
73    .map(|delay_ms| Duration::from_secs_f64(delay_ms / 1000.0))
74}
75
76/// Sleep for the specified duration using timerfd on Linux for precision.
77pub async fn sleep_precise(duration: Duration) {
78    sleep_until_precise(Instant::now() + duration).await;
79}
80
81/// Sleep until the specified deadline using timerfd on Linux for precision.
82///
83/// Unlike `sleep_precise`, this accounts for time already elapsed since the
84/// deadline's reference point, making it suitable for simulation loops where
85/// computation time should be subtracted from the sleep.
86pub async fn sleep_until_precise(deadline: Instant) {
87    // Scheduler work may consume the modeled delay, especially at high speedup ratios. Avoid
88    // allocating and registering a timerfd when there is no remaining time to sleep. Preserve
89    // the scheduler loop's cooperative yield so other tasks on the runtime can make progress.
90    if deadline <= Instant::now() {
91        tokio::task::yield_now().await;
92        return;
93    }
94
95    #[cfg(target_os = "linux")]
96    {
97        if let Ok(delay) = tokio_timerfd::Delay::new(deadline) {
98            let _ = delay.await;
99        } else {
100            tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await;
101        }
102    }
103    #[cfg(not(target_os = "linux"))]
104    {
105        tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await;
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use std::task::Poll;
113
114    #[tokio::test(flavor = "current_thread")]
115    async fn test_expired_precise_sleep_yields_to_runtime() {
116        let sleep = sleep_until_precise(Instant::now());
117        tokio::pin!(sleep);
118
119        let first_poll = futures::poll!(sleep.as_mut());
120
121        assert!(matches!(first_poll, Poll::Pending));
122        sleep.await;
123    }
124
125    #[test]
126    fn test_prefill_handoff_delay_only_applies_to_completed_prefill() {
127        let delay_ms = compute_prefill_handoff_delay_ms(
128            WorkerType::Prefill,
129            true,
130            128,
131            Some(1.0),
132            Some(1_000_000),
133        )
134        .expect("prefill completion should produce a handoff delay");
135        assert!((delay_ms - 128.0).abs() < 1e-9);
136
137        assert!(
138            compute_prefill_handoff_delay_ms(
139                WorkerType::Prefill,
140                false,
141                128,
142                Some(1.0),
143                Some(1_000_000),
144            )
145            .is_none()
146        );
147        assert!(
148            compute_prefill_handoff_delay_ms(
149                WorkerType::Decode,
150                true,
151                128,
152                Some(1.0),
153                Some(1_000_000),
154            )
155            .is_none()
156        );
157    }
158}
dynamo_mocker/common/utils.rs

dynamo_mocker/common/
utils.rs
