Skip to main content

dynamo_mocker/scheduler/
protocol.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Stable Dynamo command/lifecycle protocol for the generalized engine.
5
6use uuid::Uuid;
7
8use crate::common::handoff::{HandoffId, HandoffTransferTiming};
9use crate::common::protocols::DirectRequest;
10
11pub enum SchedulerCommand {
12    Submit(DirectRequest),
13    CancelRequest {
14        request_id: Uuid,
15    },
16    SubmitHandoffPrefill {
17        handoff_id: HandoffId,
18        request: DirectRequest,
19    },
20    ReleaseSource {
21        handoff_id: HandoffId,
22    },
23    CancelSource {
24        handoff_id: HandoffId,
25    },
26    ReserveDestination {
27        handoff_id: HandoffId,
28        request: DirectRequest,
29    },
30    ActivateDestination {
31        handoff_id: HandoffId,
32    },
33    CancelDestination {
34        handoff_id: HandoffId,
35    },
36}
37
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub enum SchedulerCommandResult {
40    Submitted(Uuid),
41    DestinationAccepted { request_id: Uuid },
42    Applied,
43    Noop,
44}
45
46#[derive(Clone, Copy, Debug, PartialEq)]
47pub enum SchedulerLifecycleEvent {
48    SourceHeld {
49        handoff_id: HandoffId,
50        request_id: Uuid,
51        transfer_timing: HandoffTransferTiming,
52    },
53    DestinationReserved {
54        handoff_id: HandoffId,
55        request_id: Uuid,
56        transferable_prompt_tokens: usize,
57    },
58}
59
60impl SchedulerLifecycleEvent {
61    pub fn handoff_id(&self) -> HandoffId {
62        match *self {
63            Self::SourceHeld { handoff_id, .. } | Self::DestinationReserved { handoff_id, .. } => {
64                handoff_id
65            }
66        }
67    }
68}
69
70#[derive(Debug)]
71pub struct SchedulerCommandEffects {
72    pub result: SchedulerCommandResult,
73    pub lifecycle_events: Vec<SchedulerLifecycleEvent>,
74    pub kv_events: Vec<dynamo_kv_router::protocols::RouterEvent>,
75}