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
4use anyhow::{Result, bail};
5use rustc_hash::FxHashMap;
6use serde::{Deserialize, Serialize};
7use uuid::Uuid;
8
9use super::protocols::{EngineType, KvTransferTimingMode};
10
11/// Stable identifier for one prefill-to-decode handoff attempt.
12#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
13#[serde(transparent)]
14pub struct HandoffId(Uuid);
15
16impl HandoffId {
17    pub fn new() -> Self {
18        Self(Uuid::new_v4())
19    }
20}
21
22impl Default for HandoffId {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl From<Uuid> for HandoffId {
29    fn from(value: Uuid) -> Self {
30        Self(value)
31    }
32}
33
34impl From<HandoffId> for Uuid {
35    fn from(value: HandoffId) -> Self {
36        value.0
37    }
38}
39
40#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
41pub enum HandoffOrder {
42    SourceFirst,
43    DestinationFirst,
44}
45
46#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
47pub struct HandoffTransferTiming {
48    pub mode: KvTransferTimingMode,
49    pub full_prompt_tokens: usize,
50    pub kv_bytes_per_token: Option<usize>,
51    pub bandwidth_gb_s: Option<f64>,
52}
53
54impl HandoffTransferTiming {
55    pub fn delay_ms(self, destination_missing_tokens: usize) -> Option<f64> {
56        let tokens = match self.mode {
57            KvTransferTimingMode::FullPrompt => self.full_prompt_tokens,
58            KvTransferTimingMode::DestinationMissing => destination_missing_tokens,
59        };
60        let (Some(bytes_per_token), Some(bandwidth_gb_s)) =
61            (self.kv_bytes_per_token, self.bandwidth_gb_s)
62        else {
63            return None;
64        };
65        if bandwidth_gb_s <= 0.0 {
66            return None;
67        }
68        Some(tokens as f64 * bytes_per_token as f64 / (bandwidth_gb_s * 1e9) * 1000.0)
69    }
70
71    pub fn full_prompt_delay_ms(self) -> Option<f64> {
72        let full_prompt = Self {
73            mode: KvTransferTimingMode::FullPrompt,
74            ..self
75        };
76        full_prompt.delay_ms(0)
77    }
78}
79
80#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
81pub enum HandoffFact {
82    SourceHeld {
83        handoff_id: HandoffId,
84        transfer_timing: HandoffTransferTiming,
85    },
86    DestinationReserved {
87        handoff_id: HandoffId,
88        transferable_prompt_tokens: usize,
89    },
90    TransferCompleted {
91        handoff_id: HandoffId,
92    },
93    Failed {
94        handoff_id: HandoffId,
95    },
96    TimedOut {
97        handoff_id: HandoffId,
98    },
99    Canceled {
100        handoff_id: HandoffId,
101    },
102}
103
104impl HandoffFact {
105    fn handoff_id(&self) -> HandoffId {
106        match *self {
107            Self::SourceHeld { handoff_id, .. }
108            | Self::DestinationReserved { handoff_id, .. }
109            | Self::TransferCompleted { handoff_id }
110            | Self::Failed { handoff_id }
111            | Self::TimedOut { handoff_id }
112            | Self::Canceled { handoff_id } => handoff_id,
113        }
114    }
115}
116
117#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
118pub enum HandoffAction {
119    SubmitPrefill {
120        handoff_id: HandoffId,
121    },
122    ReserveDestination {
123        handoff_id: HandoffId,
124    },
125    StartTransfer {
126        handoff_id: HandoffId,
127        delay_ms: f64,
128    },
129    ActivateDestination {
130        handoff_id: HandoffId,
131    },
132    ReleaseSource {
133        handoff_id: HandoffId,
134    },
135    CancelSource {
136        handoff_id: HandoffId,
137    },
138    CancelDestination {
139        handoff_id: HandoffId,
140    },
141    Complete {
142        handoff_id: HandoffId,
143    },
144}
145
146#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
147pub struct HandoffActionId(u64);
148
149#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
150pub struct IssuedHandoffAction {
151    pub id: HandoffActionId,
152    pub action: HandoffAction,
153}
154
155#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
156pub enum HandoffActionOutcome {
157    Submitted,
158    Accepted,
159    Scheduled,
160    Applied,
161    Noop,
162    Failed(String),
163}
164
165#[derive(Clone, Copy, Debug, Eq, PartialEq)]
166enum CoordinatorMode {
167    Active,
168    CleaningUp,
169    Complete,
170}
171
172#[derive(Clone, Copy, Debug, Eq, PartialEq)]
173pub enum HandoffCompletion {
174    Success,
175    Canceled,
176}
177
178#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
179pub enum NormalizedHandoffEvent {
180    SourceHeld,
181    DestinationAccepted,
182    DestinationReserved,
183    DestinationActivated,
184    SourceReleased,
185    Completed,
186}
187
188/// Surface-independent summary used to compare offline replay with the live
189/// handoff driver. This is public only so cross-crate conformance tests can
190/// exercise both implementations.
191#[doc(hidden)]
192#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
193pub struct NormalizedHandoffConformance {
194    pub engine_type: EngineType,
195    pub order: HandoffOrder,
196    pub lifecycle: Vec<NormalizedHandoffEvent>,
197    pub source_output_tokens: usize,
198    pub destination_output_tokens: usize,
199    pub completed_requests: usize,
200    pub destination_stored: NormalizedStoredTiming,
201    pub source_drained: bool,
202    pub destination_drained: bool,
203    pub driver_drained: bool,
204}
205
206#[doc(hidden)]
207#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
208pub struct NormalizedStoredTiming {
209    pub before_activation: usize,
210    pub on_activation: usize,
211    pub repeated_activation_hashes_after_activation: usize,
212}
213
214impl NormalizedHandoffConformance {
215    /// Validate the deterministic one-request fixture shared by offline and
216    /// live conformance tests.
217    #[doc(hidden)]
218    pub fn validate(&self) -> Result<()> {
219        let expected_order = match self.engine_type {
220            EngineType::Vllm => HandoffOrder::SourceFirst,
221            EngineType::Sglang => HandoffOrder::DestinationFirst,
222            EngineType::Trtllm => bail!("TRT-LLM does not support destination handoff"),
223        };
224        if self.order != expected_order {
225            bail!(
226                "normalized handoff order mismatch: expected {expected_order:?}, got {:?}",
227                self.order
228            );
229        }
230        if self.lifecycle != expected_normalized_handoff(self.order) {
231            bail!(
232                "normalized handoff lifecycle mismatch: expected {:?}, got {:?}",
233                expected_normalized_handoff(self.order),
234                self.lifecycle
235            );
236        }
237        if self.source_output_tokens != 1 {
238            bail!(
239                "normalized source output count mismatch: expected 1, got {}",
240                self.source_output_tokens
241            );
242        }
243        if self.destination_output_tokens != 2 {
244            bail!(
245                "normalized destination output count mismatch: expected 2, got {}",
246                self.destination_output_tokens
247            );
248        }
249        if self.completed_requests != 1 {
250            bail!(
251                "normalized completion count mismatch: expected 1, got {}",
252                self.completed_requests
253            );
254        }
255        if self.destination_stored.before_activation != 0 {
256            bail!(
257                "destination published {} KV blocks before activation",
258                self.destination_stored.before_activation
259            );
260        }
261        if self.destination_stored.on_activation == 0 {
262            bail!("destination activation published no KV blocks");
263        }
264        if self
265            .destination_stored
266            .repeated_activation_hashes_after_activation
267            != 0
268        {
269            bail!(
270                "destination republished {} activation KV blocks",
271                self.destination_stored
272                    .repeated_activation_hashes_after_activation
273            );
274        }
275        if !self.source_drained || !self.destination_drained || !self.driver_drained {
276            bail!(
277                "handoff did not drain: source={}, destination={}, driver={}",
278                self.source_drained,
279                self.destination_drained,
280                self.driver_drained
281            );
282        }
283        Ok(())
284    }
285}
286
287pub fn expected_normalized_handoff(order: HandoffOrder) -> &'static [NormalizedHandoffEvent] {
288    use NormalizedHandoffEvent::*;
289    match order {
290        HandoffOrder::SourceFirst => &[
291            SourceHeld,
292            DestinationAccepted,
293            DestinationReserved,
294            DestinationActivated,
295            SourceReleased,
296            Completed,
297        ],
298        HandoffOrder::DestinationFirst => &[
299            DestinationAccepted,
300            DestinationReserved,
301            SourceHeld,
302            DestinationActivated,
303            SourceReleased,
304            Completed,
305        ],
306    }
307}
308
309#[derive(Default)]
310struct ActionJournal {
311    started: bool,
312    next_id: u64,
313    issued: FxHashMap<HandoffActionId, HandoffAction>,
314    outcomes: FxHashMap<HandoffActionId, HandoffActionOutcome>,
315}
316
317#[derive(Default)]
318struct SourceProgress {
319    submit_issued: bool,
320    submitted: bool,
321    held: bool,
322    transfer_timing: Option<HandoffTransferTiming>,
323    release_issued: bool,
324    cancel_issued: bool,
325    cleanup_done: bool,
326}
327
328#[derive(Default)]
329struct DestinationProgress {
330    reserve_issued: bool,
331    accepted: bool,
332    reserved: bool,
333    transferable_prompt_tokens: Option<usize>,
334    activation_issued: bool,
335    activation_applied: bool,
336    cancel_issued: bool,
337    cleanup_done: bool,
338}
339
340#[derive(Default)]
341struct TransferProgress {
342    issued: bool,
343    scheduled: bool,
344    completed: bool,
345}
346
347/// Pure state machine for one prefill-to-decode ownership handoff.
348///
349/// Drivers execute returned actions and feed action outcomes and asynchronous
350/// facts back into this core. The core owns ordering only; it owns no engine or
351/// transport resources.
352pub struct HandoffCoordinatorCore {
353    handoff_id: HandoffId,
354    order: HandoffOrder,
355    mode: CoordinatorMode,
356    actions: ActionJournal,
357    source: SourceProgress,
358    destination: DestinationProgress,
359    transfer: TransferProgress,
360    completion: Option<HandoffCompletion>,
361}
362
363impl HandoffCoordinatorCore {
364    pub fn new(handoff_id: HandoffId, order: HandoffOrder) -> Self {
365        Self {
366            handoff_id,
367            order,
368            mode: CoordinatorMode::Active,
369            actions: ActionJournal::default(),
370            source: SourceProgress::default(),
371            destination: DestinationProgress::default(),
372            transfer: TransferProgress::default(),
373            completion: None,
374        }
375    }
376
377    pub fn start(&mut self) -> Result<Vec<IssuedHandoffAction>> {
378        if self.actions.started {
379            return Ok(Vec::new());
380        }
381        self.actions.started = true;
382        let action = match self.order {
383            HandoffOrder::SourceFirst => self.issue_submit_prefill(),
384            HandoffOrder::DestinationFirst => self.issue_reserve_destination(),
385        };
386        Ok(vec![action])
387    }
388
389    pub fn on_fact(&mut self, fact: HandoffFact) -> Result<Vec<IssuedHandoffAction>> {
390        self.validate_handoff(fact.handoff_id())?;
391        if self.mode != CoordinatorMode::Active {
392            return Ok(Vec::new());
393        }
394
395        match fact {
396            HandoffFact::SourceHeld {
397                transfer_timing, ..
398            } => {
399                if self.source.held {
400                    return Ok(Vec::new());
401                }
402                if !self.source.submitted {
403                    bail!("source held before prefill submission was acknowledged");
404                }
405                validate_transfer_timing(transfer_timing)?;
406                self.source.held = true;
407                self.source.transfer_timing = Some(transfer_timing);
408                self.advance_active()
409            }
410            HandoffFact::DestinationReserved {
411                transferable_prompt_tokens,
412                ..
413            } => {
414                if self.destination.reserved {
415                    return Ok(Vec::new());
416                }
417                if !self.destination.accepted {
418                    bail!("destination reserved before ownership was accepted");
419                }
420                self.destination.reserved = true;
421                self.destination.transferable_prompt_tokens = Some(transferable_prompt_tokens);
422                self.advance_active()
423            }
424            HandoffFact::TransferCompleted { .. } => {
425                if self.transfer.completed {
426                    return Ok(Vec::new());
427                }
428                if !self.transfer.scheduled {
429                    bail!("transfer completed before it was scheduled");
430                }
431                self.transfer.completed = true;
432                self.advance_active()
433            }
434            HandoffFact::Failed { .. }
435            | HandoffFact::TimedOut { .. }
436            | HandoffFact::Canceled { .. } => self.begin_cleanup(),
437        }
438    }
439
440    pub fn on_action_outcome(
441        &mut self,
442        action_id: HandoffActionId,
443        outcome: HandoffActionOutcome,
444    ) -> Result<Vec<IssuedHandoffAction>> {
445        if self.mode == CoordinatorMode::Complete {
446            return Ok(Vec::new());
447        }
448        let Some(action) = self.actions.issued.get(&action_id).copied() else {
449            bail!("unknown handoff action {action_id:?}");
450        };
451        if let Some(previous) = self.actions.outcomes.get(&action_id) {
452            if previous != &outcome {
453                bail!("conflicting outcome for handoff action {action_id:?}");
454            }
455            return Ok(Vec::new());
456        }
457        self.actions.outcomes.insert(action_id, outcome.clone());
458
459        if let HandoffActionOutcome::Failed(_) = outcome {
460            if matches!(
461                action,
462                HandoffAction::CancelSource { .. } | HandoffAction::CancelDestination { .. }
463            ) {
464                bail!("handoff cleanup action {action_id:?} failed");
465            }
466            return self.begin_cleanup();
467        }
468
469        match action {
470            HandoffAction::SubmitPrefill { .. } => {
471                require_outcome(&outcome, &[HandoffActionOutcome::Submitted])?;
472                self.source.submitted = true;
473            }
474            HandoffAction::ReserveDestination { .. } => {
475                require_outcome(&outcome, &[HandoffActionOutcome::Accepted])?;
476                self.destination.accepted = true;
477            }
478            HandoffAction::StartTransfer { .. } => {
479                require_outcome(&outcome, &[HandoffActionOutcome::Scheduled])?;
480                self.transfer.scheduled = true;
481            }
482            HandoffAction::ActivateDestination { .. } => {
483                require_outcome(&outcome, &[HandoffActionOutcome::Applied])?;
484                self.destination.activation_applied = true;
485            }
486            HandoffAction::ReleaseSource { .. } => {
487                require_outcome(
488                    &outcome,
489                    &[HandoffActionOutcome::Applied, HandoffActionOutcome::Noop],
490                )?;
491                self.source.cleanup_done = true;
492            }
493            HandoffAction::CancelSource { .. } => {
494                require_outcome(
495                    &outcome,
496                    &[HandoffActionOutcome::Applied, HandoffActionOutcome::Noop],
497                )?;
498                self.source.cleanup_done = true;
499            }
500            HandoffAction::CancelDestination { .. } => {
501                require_outcome(
502                    &outcome,
503                    &[HandoffActionOutcome::Applied, HandoffActionOutcome::Noop],
504                )?;
505                self.destination.cleanup_done = true;
506            }
507            HandoffAction::Complete { .. } => return Ok(Vec::new()),
508        }
509
510        match self.mode {
511            CoordinatorMode::Active => self.advance_active(),
512            CoordinatorMode::CleaningUp => self.advance_cleanup(),
513            CoordinatorMode::Complete => Ok(Vec::new()),
514        }
515    }
516
517    pub fn is_complete(&self) -> bool {
518        self.mode == CoordinatorMode::Complete
519    }
520
521    pub fn completion(&self) -> Option<HandoffCompletion> {
522        self.completion
523    }
524
525    fn advance_active(&mut self) -> Result<Vec<IssuedHandoffAction>> {
526        if self.order == HandoffOrder::SourceFirst
527            && self.source.held
528            && !self.destination.reserve_issued
529        {
530            return Ok(vec![self.issue_reserve_destination()]);
531        }
532        if self.order == HandoffOrder::DestinationFirst
533            && self.destination.reserved
534            && !self.source.submit_issued
535        {
536            return Ok(vec![self.issue_submit_prefill()]);
537        }
538        if self.source.held && self.destination.reserved && !self.transfer.issued {
539            self.transfer.issued = true;
540            let transfer_timing = self
541                .source
542                .transfer_timing
543                .expect("held source must retain transfer timing");
544            let transferable_prompt_tokens = self
545                .destination
546                .transferable_prompt_tokens
547                .expect("reserved destination must report its transferable footprint");
548            return Ok(vec![
549                self.issue(HandoffAction::StartTransfer {
550                    handoff_id: self.handoff_id,
551                    delay_ms: transfer_timing
552                        .delay_ms(transferable_prompt_tokens)
553                        .unwrap_or_default(),
554                }),
555            ]);
556        }
557        if self.transfer.completed && !self.destination.activation_issued {
558            self.destination.activation_issued = true;
559            return Ok(vec![self.issue(HandoffAction::ActivateDestination {
560                handoff_id: self.handoff_id,
561            })]);
562        }
563        if self.destination.activation_applied && !self.source.release_issued {
564            self.source.release_issued = true;
565            return Ok(vec![self.issue(HandoffAction::ReleaseSource {
566                handoff_id: self.handoff_id,
567            })]);
568        }
569        if self.source.cleanup_done {
570            return Ok(vec![self.complete()]);
571        }
572        Ok(Vec::new())
573    }
574
575    fn begin_cleanup(&mut self) -> Result<Vec<IssuedHandoffAction>> {
576        if self.mode == CoordinatorMode::Complete {
577            return Ok(Vec::new());
578        }
579        self.mode = CoordinatorMode::CleaningUp;
580        self.advance_cleanup()
581    }
582
583    fn advance_cleanup(&mut self) -> Result<Vec<IssuedHandoffAction>> {
584        let mut actions = Vec::new();
585        if self.source.submit_issued && !self.source.cancel_issued && !self.source.cleanup_done {
586            self.source.cancel_issued = true;
587            actions.push(self.issue(HandoffAction::CancelSource {
588                handoff_id: self.handoff_id,
589            }));
590        }
591        if self.destination.reserve_issued
592            && !self.destination.cancel_issued
593            && !self.destination.cleanup_done
594        {
595            self.destination.cancel_issued = true;
596            actions.push(self.issue(HandoffAction::CancelDestination {
597                handoff_id: self.handoff_id,
598            }));
599        }
600        if actions.is_empty()
601            && (!self.source.submit_issued || self.source.cleanup_done)
602            && (!self.destination.reserve_issued || self.destination.cleanup_done)
603        {
604            actions.push(self.complete());
605        }
606        Ok(actions)
607    }
608
609    fn issue_submit_prefill(&mut self) -> IssuedHandoffAction {
610        self.source.submit_issued = true;
611        self.issue(HandoffAction::SubmitPrefill {
612            handoff_id: self.handoff_id,
613        })
614    }
615
616    fn issue_reserve_destination(&mut self) -> IssuedHandoffAction {
617        self.destination.reserve_issued = true;
618        self.issue(HandoffAction::ReserveDestination {
619            handoff_id: self.handoff_id,
620        })
621    }
622
623    fn complete(&mut self) -> IssuedHandoffAction {
624        self.completion = Some(if self.mode == CoordinatorMode::CleaningUp {
625            HandoffCompletion::Canceled
626        } else {
627            HandoffCompletion::Success
628        });
629        self.mode = CoordinatorMode::Complete;
630        let action = self.issue(HandoffAction::Complete {
631            handoff_id: self.handoff_id,
632        });
633        self.actions.issued = FxHashMap::default();
634        self.actions.outcomes = FxHashMap::default();
635        action
636    }
637
638    fn issue(&mut self, action: HandoffAction) -> IssuedHandoffAction {
639        let id = HandoffActionId(self.actions.next_id);
640        self.actions.next_id = self
641            .actions
642            .next_id
643            .checked_add(1)
644            .expect("handoff action ID overflow");
645        let previous = self.actions.issued.insert(id, action);
646        debug_assert!(previous.is_none());
647        IssuedHandoffAction { id, action }
648    }
649
650    fn validate_handoff(&self, handoff_id: HandoffId) -> Result<()> {
651        if handoff_id != self.handoff_id {
652            bail!("fact belongs to a different handoff");
653        }
654        Ok(())
655    }
656}
657
658pub fn validate_transfer_delay_ms(transfer_delay_ms: Option<f64>) -> Result<()> {
659    let Some(delay_ms) = transfer_delay_ms else {
660        return Ok(());
661    };
662    if !delay_ms.is_finite() || delay_ms < 0.0 {
663        bail!("invalid handoff transfer delay {delay_ms}");
664    }
665    Ok(())
666}
667
668pub fn validate_transfer_timing(transfer_timing: HandoffTransferTiming) -> Result<()> {
669    if let Some(bandwidth_gb_s) = transfer_timing.bandwidth_gb_s
670        && (!bandwidth_gb_s.is_finite() || bandwidth_gb_s <= 0.0)
671    {
672        bail!("invalid handoff transfer bandwidth {bandwidth_gb_s}");
673    }
674    validate_transfer_delay_ms(transfer_timing.full_prompt_delay_ms())
675}
676
677fn require_outcome(outcome: &HandoffActionOutcome, allowed: &[HandoffActionOutcome]) -> Result<()> {
678    if allowed.contains(outcome) {
679        return Ok(());
680    }
681    bail!("invalid handoff action outcome {outcome:?}")
682}
683
684#[cfg(test)]
685#[path = "handoff_tests.rs"]
686mod coordinator_tests;
dynamo_mocker/common/handoff.rs

dynamo_mocker/common/
handoff.rs
