dynamo_mocker/common/
handoff.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Dynamo compatibility types for the shared Replay handoff coordinator.
5//!
6//! Dynamo retains its UUID transport DTOs at this boundary. Ordering and
7//! cleanup are owned by `aisimulate_core::replay`; this module only converts between
8//! the public Dynamo surface and Replay's runtime-neutral value types.
9
10use anyhow::Result;
11use serde::{Deserialize, Serialize};
12use uuid::Uuid;
13
14use super::protocols::{EngineType, KvTransferTimingMode};
15
16pub use aisimulate_core::replay::{
17    HandoffActionId, HandoffActionOutcome, HandoffCompletion, HandoffOrder, NormalizedHandoffEvent,
18    NormalizedStoredTiming, expected_normalized_handoff, validate_transfer_delay_ms,
19};
20
21/// Stable identifier for one prefill-to-decode handoff attempt.
22#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
23#[serde(transparent)]
24pub struct HandoffId(Uuid);
25
26impl HandoffId {
27    pub fn new() -> Self {
28        Self(Uuid::new_v4())
29    }
30}
31
32impl Default for HandoffId {
33    fn default() -> Self {
34        Self::new()
35    }
36}
37
38impl From<Uuid> for HandoffId {
39    fn from(value: Uuid) -> Self {
40        Self(value)
41    }
42}
43
44impl From<HandoffId> for Uuid {
45    fn from(value: HandoffId) -> Self {
46        value.0
47    }
48}
49
50impl From<HandoffId> for aisimulate_core::replay::HandoffId {
51    fn from(value: HandoffId) -> Self {
52        Self::from(value.0)
53    }
54}
55
56#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
57pub struct HandoffTransferTiming {
58    pub mode: KvTransferTimingMode,
59    pub full_prompt_tokens: usize,
60    pub kv_bytes_per_token: Option<usize>,
61    pub bandwidth_gb_s: Option<f64>,
62}
63
64impl HandoffTransferTiming {
65    pub fn delay_ms(self, destination_missing_tokens: usize) -> Option<f64> {
66        replay_timing(self).delay_ms(destination_missing_tokens)
67    }
68
69    pub fn full_prompt_delay_ms(self) -> Option<f64> {
70        replay_timing(self).full_prompt_delay_ms()
71    }
72}
73
74#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
75pub enum HandoffFact {
76    SourceHeld {
77        handoff_id: HandoffId,
78        transfer_timing: HandoffTransferTiming,
79    },
80    DestinationReserved {
81        handoff_id: HandoffId,
82        transferable_prompt_tokens: usize,
83    },
84    TransferCompleted {
85        handoff_id: HandoffId,
86    },
87    Failed {
88        handoff_id: HandoffId,
89    },
90    TimedOut {
91        handoff_id: HandoffId,
92    },
93    Canceled {
94        handoff_id: HandoffId,
95    },
96}
97
98#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
99pub enum HandoffAction {
100    SubmitPrefill {
101        handoff_id: HandoffId,
102    },
103    ReserveDestination {
104        handoff_id: HandoffId,
105    },
106    StartTransfer {
107        handoff_id: HandoffId,
108        delay_ms: f64,
109    },
110    ActivateDestination {
111        handoff_id: HandoffId,
112    },
113    ReleaseSource {
114        handoff_id: HandoffId,
115    },
116    CancelSource {
117        handoff_id: HandoffId,
118    },
119    CancelDestination {
120        handoff_id: HandoffId,
121    },
122    Complete {
123        handoff_id: HandoffId,
124    },
125}
126
127#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
128pub struct IssuedHandoffAction {
129    pub id: HandoffActionId,
130    pub action: HandoffAction,
131}
132
133/// Compatibility summary used by Dynamo's live/offline conformance tests.
134#[doc(hidden)]
135#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
136pub struct NormalizedHandoffConformance {
137    pub engine_type: EngineType,
138    pub order: HandoffOrder,
139    pub lifecycle: Vec<NormalizedHandoffEvent>,
140    pub source_output_tokens: usize,
141    pub destination_output_tokens: usize,
142    pub completed_requests: usize,
143    pub destination_stored: NormalizedStoredTiming,
144    pub source_drained: bool,
145    pub destination_drained: bool,
146    pub driver_drained: bool,
147}
148
149impl NormalizedHandoffConformance {
150    #[doc(hidden)]
151    pub fn validate(&self) -> Result<()> {
152        aisimulate_core::replay::NormalizedHandoffConformance {
153            engine_type: match self.engine_type {
154                EngineType::Vllm => aisimulate_core::engine::Backend::Vllm,
155                EngineType::Sglang => aisimulate_core::engine::Backend::Sglang,
156                EngineType::Trtllm => aisimulate_core::engine::Backend::Trtllm,
157            },
158            order: self.order,
159            lifecycle: self.lifecycle.clone(),
160            source_output_tokens: self.source_output_tokens,
161            destination_output_tokens: self.destination_output_tokens,
162            completed_requests: self.completed_requests,
163            destination_stored: self.destination_stored.clone(),
164            source_drained: self.source_drained,
165            destination_drained: self.destination_drained,
166            driver_drained: self.driver_drained,
167        }
168        .validate()
169    }
170}
171
172impl From<aisimulate_core::replay::NormalizedHandoffConformance> for NormalizedHandoffConformance {
173    fn from(value: aisimulate_core::replay::NormalizedHandoffConformance) -> Self {
174        Self {
175            engine_type: match value.engine_type {
176                aisimulate_core::engine::Backend::Vllm => EngineType::Vllm,
177                aisimulate_core::engine::Backend::Sglang => EngineType::Sglang,
178                aisimulate_core::engine::Backend::Trtllm => EngineType::Trtllm,
179            },
180            order: value.order,
181            lifecycle: value.lifecycle,
182            source_output_tokens: value.source_output_tokens,
183            destination_output_tokens: value.destination_output_tokens,
184            completed_requests: value.completed_requests,
185            destination_stored: value.destination_stored,
186            source_drained: value.source_drained,
187            destination_drained: value.destination_drained,
188            driver_drained: value.driver_drained,
189        }
190    }
191}
192
193/// Thin UUID-compatibility wrapper around Replay's single handoff state machine.
194pub struct HandoffCoordinatorCore {
195    inner: aisimulate_core::replay::HandoffCoordinatorCore,
196}
197
198impl HandoffCoordinatorCore {
199    pub fn new(handoff_id: HandoffId, order: HandoffOrder) -> Self {
200        Self {
201            inner: aisimulate_core::replay::HandoffCoordinatorCore::new(handoff_id.into(), order),
202        }
203    }
204
205    pub fn start(&mut self) -> Result<Vec<IssuedHandoffAction>> {
206        Ok(self
207            .inner
208            .start()?
209            .into_iter()
210            .map(convert_action)
211            .collect())
212    }
213
214    pub fn on_fact(&mut self, fact: HandoffFact) -> Result<Vec<IssuedHandoffAction>> {
215        Ok(self
216            .inner
217            .on_fact(convert_fact(fact))?
218            .into_iter()
219            .map(convert_action)
220            .collect())
221    }
222
223    pub fn on_action_outcome(
224        &mut self,
225        action_id: HandoffActionId,
226        outcome: HandoffActionOutcome,
227    ) -> Result<Vec<IssuedHandoffAction>> {
228        Ok(self
229            .inner
230            .on_action_outcome(action_id, outcome)?
231            .into_iter()
232            .map(convert_action)
233            .collect())
234    }
235
236    pub fn is_complete(&self) -> bool {
237        self.inner.is_complete()
238    }
239
240    pub fn completion(&self) -> Option<HandoffCompletion> {
241        self.inner.completion()
242    }
243}
244
245pub fn validate_transfer_timing(transfer_timing: HandoffTransferTiming) -> Result<()> {
246    aisimulate_core::replay::validate_transfer_timing(replay_timing(transfer_timing))
247}
248
249fn replay_timing(timing: HandoffTransferTiming) -> aisimulate_core::replay::HandoffTransferTiming {
250    aisimulate_core::replay::HandoffTransferTiming {
251        mode: match timing.mode {
252            KvTransferTimingMode::FullPrompt => {
253                aisimulate_core::engine::TransferTimingMode::FullPrompt
254            }
255            KvTransferTimingMode::DestinationMissing => {
256                aisimulate_core::engine::TransferTimingMode::DestinationMissing
257            }
258        },
259        full_prompt_tokens: timing.full_prompt_tokens,
260        kv_bytes_per_token: timing.kv_bytes_per_token,
261        bandwidth_gb_s: timing.bandwidth_gb_s,
262    }
263}
264
265fn convert_fact(fact: HandoffFact) -> aisimulate_core::replay::HandoffFact {
266    match fact {
267        HandoffFact::SourceHeld {
268            handoff_id,
269            transfer_timing,
270        } => aisimulate_core::replay::HandoffFact::SourceHeld {
271            handoff_id: handoff_id.into(),
272            transfer_timing: replay_timing(transfer_timing),
273        },
274        HandoffFact::DestinationReserved {
275            handoff_id,
276            transferable_prompt_tokens,
277        } => aisimulate_core::replay::HandoffFact::DestinationReserved {
278            handoff_id: handoff_id.into(),
279            transferable_prompt_tokens,
280        },
281        HandoffFact::TransferCompleted { handoff_id } => {
282            aisimulate_core::replay::HandoffFact::TransferCompleted {
283                handoff_id: handoff_id.into(),
284            }
285        }
286        HandoffFact::Failed { handoff_id } => aisimulate_core::replay::HandoffFact::Failed {
287            handoff_id: handoff_id.into(),
288        },
289        HandoffFact::TimedOut { handoff_id } => aisimulate_core::replay::HandoffFact::TimedOut {
290            handoff_id: handoff_id.into(),
291        },
292        HandoffFact::Canceled { handoff_id } => aisimulate_core::replay::HandoffFact::Canceled {
293            handoff_id: handoff_id.into(),
294        },
295    }
296}
297
298fn convert_action(action: aisimulate_core::replay::IssuedHandoffAction) -> IssuedHandoffAction {
299    let aisimulate_core::replay::IssuedHandoffAction { id, action } = action;
300    let action = match action {
301        aisimulate_core::replay::HandoffAction::SubmitPrefill { handoff_id } => {
302            HandoffAction::SubmitPrefill {
303                handoff_id: HandoffId::from(handoff_id.get()),
304            }
305        }
306        aisimulate_core::replay::HandoffAction::ReserveDestination { handoff_id } => {
307            HandoffAction::ReserveDestination {
308                handoff_id: HandoffId::from(handoff_id.get()),
309            }
310        }
311        aisimulate_core::replay::HandoffAction::StartTransfer {
312            handoff_id,
313            delay_ms,
314        } => HandoffAction::StartTransfer {
315            handoff_id: HandoffId::from(handoff_id.get()),
316            delay_ms,
317        },
318        aisimulate_core::replay::HandoffAction::ActivateDestination { handoff_id } => {
319            HandoffAction::ActivateDestination {
320                handoff_id: HandoffId::from(handoff_id.get()),
321            }
322        }
323        aisimulate_core::replay::HandoffAction::ReleaseSource { handoff_id } => {
324            HandoffAction::ReleaseSource {
325                handoff_id: HandoffId::from(handoff_id.get()),
326            }
327        }
328        aisimulate_core::replay::HandoffAction::CancelSource { handoff_id } => {
329            HandoffAction::CancelSource {
330                handoff_id: HandoffId::from(handoff_id.get()),
331            }
332        }
333        aisimulate_core::replay::HandoffAction::CancelDestination { handoff_id } => {
334            HandoffAction::CancelDestination {
335                handoff_id: HandoffId::from(handoff_id.get()),
336            }
337        }
338        aisimulate_core::replay::HandoffAction::Complete { handoff_id } => {
339            HandoffAction::Complete {
340                handoff_id: HandoffId::from(handoff_id.get()),
341            }
342        }
343    };
344    IssuedHandoffAction { id, action }
345}
346
347#[cfg(test)]
348#[path = "handoff_tests.rs"]
349mod coordinator_tests;
dynamo_mocker/common/handoff.rs

dynamo_mocker/common/
handoff.rs
