dynamo_mocker/scheduler/
protocol.rs1use 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}