mesh_llm_node/
serving.rs

1use anyhow::Result;
2use serde::Serialize;
3use std::future::Future;
4use std::pin::Pin;
5use std::time::Duration;
6
7use crate::models::ModelCapabilities;
8
9pub type ServingFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T>> + Send + 'a>>;
10
11#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize)]
12#[serde(rename_all = "snake_case")]
13pub enum DevicePolicy {
14    #[default]
15    Auto,
16    Cpu,
17    Gpu {
18        device_ids: Vec<String>,
19    },
20}
21
22#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize)]
23pub struct LoadModelRequest {
24    pub model_ref: String,
25    pub device_policy: DevicePolicy,
26    #[serde(default)]
27    #[serde(skip_serializing_if = "String::is_empty")]
28    pub profile: String,
29}
30
31#[derive(Clone, Debug, Eq, PartialEq)]
32pub struct UnloadModelRequest {
33    pub target: UnloadTarget,
34    pub options: UnloadOptions,
35}
36
37impl Default for UnloadModelRequest {
38    fn default() -> Self {
39        Self {
40            target: UnloadTarget::Model(String::new()),
41            options: UnloadOptions::default(),
42        }
43    }
44}
45
46#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
47#[serde(rename_all = "snake_case")]
48pub enum UnloadTarget {
49    Model(String),
50    Instance(String),
51}
52
53impl UnloadTarget {
54    pub fn as_runtime_target(&self) -> &str {
55        match self {
56            Self::Model(value) | Self::Instance(value) => value,
57        }
58    }
59}
60
61#[derive(Clone, Debug, Eq, PartialEq)]
62pub struct UnloadOptions {
63    pub drain_timeout: Duration,
64    pub force: bool,
65}
66
67impl Default for UnloadOptions {
68    fn default() -> Self {
69        Self {
70            drain_timeout: Duration::from_secs(30),
71            force: false,
72        }
73    }
74}
75
76#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize)]
77#[serde(rename_all = "snake_case")]
78pub enum ServingModelState {
79    Loading,
80    #[default]
81    Ready,
82    Failed,
83    Unloading,
84    Stopped,
85    Unknown(String),
86}
87
88#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
89pub struct ServedModel {
90    pub model_ref: String,
91    #[serde(default)]
92    #[serde(skip_serializing_if = "String::is_empty")]
93    pub profile: String,
94    pub model_id: String,
95    pub instance_id: Option<String>,
96    pub state: ServingModelState,
97    pub backend: Option<String>,
98    pub capabilities: ModelCapabilities,
99    pub context_length: Option<u32>,
100    pub error: Option<String>,
101}
102
103#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize)]
104pub struct ServingStatus {
105    pub enabled: bool,
106    pub models: Vec<ServedModel>,
107}
108
109#[derive(Clone, Debug, Eq, PartialEq)]
110pub enum ServingError {
111    ModelNotFound {
112        model_ref: String,
113    },
114    DownloadRequired {
115        model_ref: String,
116    },
117    LoadFailed {
118        model_ref: String,
119        message: String,
120    },
121    UnloadFailed {
122        target: UnloadTarget,
123        message: String,
124    },
125    UnsupportedDevicePolicy {
126        policy: DevicePolicy,
127    },
128    RuntimeUnavailable {
129        message: String,
130    },
131}
132
133impl std::fmt::Display for ServingError {
134    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
135        match self {
136            Self::ModelNotFound { model_ref } => write!(formatter, "model not found: {model_ref}"),
137            Self::DownloadRequired { model_ref } => {
138                write!(
139                    formatter,
140                    "model must be downloaded before serving: {model_ref}"
141                )
142            }
143            Self::LoadFailed { model_ref, message } => {
144                write!(formatter, "failed to load {model_ref}: {message}")
145            }
146            Self::UnloadFailed { target, message } => {
147                write!(formatter, "failed to unload {target}: {message}")
148            }
149            Self::UnsupportedDevicePolicy { policy } => {
150                write!(formatter, "unsupported device policy: {policy:?}")
151            }
152            Self::RuntimeUnavailable { message } => {
153                write!(formatter, "runtime unavailable: {message}")
154            }
155        }
156    }
157}
158
159impl std::fmt::Display for UnloadTarget {
160    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
161        match self {
162            Self::Model(value) => write!(formatter, "model {value}"),
163            Self::Instance(value) => write!(formatter, "instance {value}"),
164        }
165    }
166}
167
168impl std::error::Error for ServingError {}
169
170pub trait ServingController: Send + Sync {
171    fn load<'a>(&'a self, request: LoadModelRequest) -> ServingFuture<'a, ServedModel>;
172
173    fn unload<'a>(&'a self, request: UnloadModelRequest) -> ServingFuture<'a, ()>;
174
175    fn served_models<'a>(&'a self) -> ServingFuture<'a, Vec<ServedModel>>;
176
177    fn status<'a>(&'a self) -> ServingFuture<'a, ServingStatus>;
178
179    fn set_device_policy<'a>(&'a self, policy: DevicePolicy) -> ServingFuture<'a, ()>;
180}
mesh_llm_node/serving.rs

mesh_llm_node/
serving.rs
