dynamo-kv-router 1.3.0

KV Router - Radix tree for LLM KV cache routing
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

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

use axum::Json;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};

use crate::scheduling::KvSchedulerError;
use crate::sequences::SequenceError;

#[derive(Debug, thiserror::Error)]
pub enum SelectionError {
    #[error("{0}")]
    BadRequest(String),
    #[error("{0}")]
    NotReady(String),
    #[error("{0}")]
    NotFound(String),
    #[error("{0}")]
    Conflict(String),
    #[error("{0}")]
    Internal(String),
    #[error(transparent)]
    Scheduler(#[from] KvSchedulerError),
    #[error(transparent)]
    Sequence(#[from] SequenceError),
}

impl SelectionError {
    fn status(&self) -> StatusCode {
        match self {
            Self::BadRequest(_) => StatusCode::BAD_REQUEST,
            Self::NotReady(_) => StatusCode::SERVICE_UNAVAILABLE,
            Self::NotFound(_) => StatusCode::NOT_FOUND,
            Self::Conflict(_) => StatusCode::CONFLICT,
            Self::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR,
            Self::Scheduler(error) => scheduler_error_status(error),
            Self::Sequence(error) => sequence_error_status(error),
        }
    }

    /// HTTP-style status code for this error, for callers that consume the
    /// service in-process without an HTTP layer.
    pub fn status_code(&self) -> u16 {
        self.status().as_u16()
    }

    /// Stable, machine-readable category for this error.
    pub fn kind(&self) -> &'static str {
        match self {
            Self::BadRequest(_) => "bad_request",
            Self::NotReady(_) => "not_ready",
            Self::NotFound(_) => "not_found",
            Self::Conflict(_) => "conflict",
            Self::Internal(_) => "internal",
            Self::Scheduler(_) => "scheduler",
            Self::Sequence(_) => "sequence",
        }
    }
}

fn scheduler_error_status(error: &KvSchedulerError) -> StatusCode {
    match error {
        KvSchedulerError::NoEndpoints
        | KvSchedulerError::SubscriberShutdown
        | KvSchedulerError::InitFailed(_) => StatusCode::SERVICE_UNAVAILABLE,
        KvSchedulerError::AllEligibleWorkersOverloaded
        | KvSchedulerError::PinnedWorkerOverloaded { .. } => StatusCode::TOO_MANY_REQUESTS,
        KvSchedulerError::QueueRejected(_) => StatusCode::SERVICE_UNAVAILABLE,
        KvSchedulerError::PinnedWorkerNotAllowed { .. } => StatusCode::BAD_REQUEST,
        KvSchedulerError::BookingFailed(_) => StatusCode::CONFLICT,
    }
}

fn sequence_error_status(error: &SequenceError) -> StatusCode {
    match error {
        SequenceError::WorkerNotFound { .. } | SequenceError::RequestNotFound { .. } => {
            StatusCode::NOT_FOUND
        }
        SequenceError::DuplicateRequest { .. } => StatusCode::CONFLICT,
        SequenceError::ReplicaSyncPublishFailed(_) => StatusCode::INTERNAL_SERVER_ERROR,
    }
}

impl IntoResponse for SelectionError {
    fn into_response(self) -> Response {
        if let Self::Scheduler(KvSchedulerError::QueueRejected(rejection)) = &self {
            return (
                self.status(),
                Json(serde_json::json!({
                    "error": self.to_string(),
                    "details": rejection,
                })),
            )
                .into_response();
        }

        (
            self.status(),
            Json(serde_json::json!({"error": self.to_string()})),
        )
            .into_response()
    }
}