dynamo-mocker 1.4.0

Mock LLM scheduler and KV manager for testing
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Pluggable KV cache block managers.

#[cfg(feature = "kvbm-offload")]
use crate::kvbm_offload::OffloadId;

#[cfg(not(feature = "kvbm-offload"))]
type OffloadId = u64;

/// Result of an atomic G1 capacity acquisition, shared by both G1 backends.
#[cfg_attr(not(feature = "kvbm-offload"), allow(dead_code))]
pub(crate) enum G1Acquire<T> {
    Ready(T),
    CapacityExhausted,
    BlockedOnOffload {
        offload_id: OffloadId,
        deadline_ms: Option<f64>,
    },
    RetryNow {
        capacity_generation: u64,
        released_slots: usize,
    },
}

impl<T> G1Acquire<T> {
    pub(crate) fn map<U>(self, f: impl FnOnce(T) -> U) -> G1Acquire<U> {
        match self {
            Self::Ready(value) => G1Acquire::Ready(f(value)),
            Self::CapacityExhausted => G1Acquire::CapacityExhausted,
            Self::BlockedOnOffload {
                offload_id,
                deadline_ms,
            } => G1Acquire::BlockedOnOffload {
                offload_id,
                deadline_ms,
            },
            Self::RetryNow {
                capacity_generation,
                released_slots,
            } => G1Acquire::RetryNow {
                capacity_generation,
                released_slots,
            },
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq)]
#[doc(hidden)]
pub struct OffloadDependency {
    pub(crate) offload_id: OffloadId,
    pub(crate) deadline_ms: Option<f64>,
}

mod g1_manager;
pub mod kvbm_backend;
pub mod sglang_backend;
mod vllm_backend;
#[cfg(test)]
mod vllm_firewall_tests;

pub(crate) use g1_manager::DestinationReservation;
pub use g1_manager::G1Manager;
pub use kvbm_backend::KvManager;
pub use sglang_backend::SglangKvManager;