hugo/DS4Server
1
0
Fork 0
Code Issues 8 Pull Requests Actions Packages Projects Releases Wiki Activity

Implement batched server parity and observability

Browse Source
This commit is contained in:
Georg Bauer
2026-07-25 08:32:35 +02:00
parent d554b77b9d
commit 2fbe605b5d
15 changed files with 6478 additions and 266 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
src/app.rs
View File

@@ -5,6 +5,7 @@ pub(crate) use view::app_theme;
use crate::database::{AppPreferences, Database, ProjectWithSessions, StoredMessage};
#[cfg(target_os = "macos")]
use crate::engine::ChatTurn;
use crate::metrics::{Metrics, MetricsSnapshot};
use crate::model::{self, DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice};
#[cfg(target_os = "macos")]
use crate::runtime::{ActiveGeneration, CheckpointTarget, GenerationEvent, GenerationService};
@@ -16,6 +17,7 @@ use crate::settings::{
use iced::widget::{markdown, scrollable};
use iced::{Size, Subscription, Task, keyboard, window};
use rfd::AsyncFileDialog;
use std::collections::VecDeque;
use std::fs;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
@@ -25,6 +27,7 @@ use std::thread;
use std::time::{Duration, Instant};
const APP_ID: &str = "DS4Server.rfc1437.de";
const METRICS_SAMPLE_INTERVAL: Duration = Duration::from_millis(200);
#[derive(Clone)]
struct PreferenceDraft {
@@ -414,6 +417,13 @@ pub(crate) struct App {
pub(super) context_used: u32,
pub(super) context_limit: u32,
pub(super) tokens_per_second: Option<f32>,
pub(super) detail_tab: DetailTab,
metrics: Arc<Metrics>,
pub(super) metrics_snapshot: MetricsSnapshot,
pub(super) metrics_history: VecDeque<MetricsPoint>,
last_http_requests: u64,
last_kv_read_bytes: u64,
last_kv_write_bytes: u64,
#[cfg(target_os = "macos")]
generation_service: Option<GenerationService>,
#[cfg(target_os = "macos")]
@@ -425,6 +435,22 @@ pub(crate) struct App {
error: Option<String>,
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub(super) enum DetailTab {
#[default]
Chat,
Stats,
}
#[derive(Clone, Copy, Debug, Default)]
pub(super) struct MetricsPoint {
pub(super) decode_tokens_per_second: f32,
pub(super) prefill_tokens_per_second: f32,
pub(super) http_requests_per_second: f32,
pub(super) kv_read_bytes_per_second: f32,
pub(super) kv_write_bytes_per_second: f32,
}
#[derive(Clone, Debug)]
pub(super) struct ChatMessage {
id: i32,
@@ -571,6 +597,9 @@ pub(crate) enum Message {
CreateSession,
SelectSession(i32, i32),
DeleteSession(i32),
ShowChat,
ShowStats,
MetricsTick,
}
impl App {
@@ -584,9 +613,12 @@ impl App {
Err(error) => return Self::failed(error, main_window),
};
let context_limit = preferences.context_tokens.max(0) as u32;
let metrics =
Arc::new(Metrics::new(&application_support_path().join("kv-cache")));
let metrics_snapshot = metrics.snapshot();
#[cfg(target_os = "macos")]
let (runtime_preferences, generation_service, endpoint, service_error) =
spawn_services(&preferences);
spawn_services(&preferences, Arc::clone(&metrics));
Self {
main_window,
model_manager_window: None,
@@ -613,6 +645,13 @@ impl App {
context_used: 0,
context_limit,
tokens_per_second: None,
detail_tab: DetailTab::Chat,
metrics,
metrics_snapshot,
metrics_history: VecDeque::with_capacity(120),
last_http_requests: 0,
last_kv_read_bytes: 0,
last_kv_write_bytes: 0,
#[cfg(target_os = "macos")]
generation_service,
#[cfg(target_os = "macos")]
@@ -644,9 +683,11 @@ impl App {
let preference_draft = PreferenceDraft::from_saved(&preferences)
.expect("default preferences must use a supported model");
let context_limit = preferences.context_tokens.max(0) as u32;
let metrics = Arc::new(Metrics::new(&application_support_path().join("kv-cache")));
let metrics_snapshot = metrics.snapshot();
#[cfg(target_os = "macos")]
let (runtime_preferences, generation_service, endpoint, service_error) =
spawn_services(&preferences);
spawn_services(&preferences, Arc::clone(&metrics));
#[cfg(target_os = "macos")]
let startup_error = service_error;
#[cfg(not(target_os = "macos"))]
@@ -677,6 +718,13 @@ impl App {
context_used: 0,
context_limit,
tokens_per_second: None,
detail_tab: DetailTab::Chat,
metrics,
metrics_snapshot,
metrics_history: VecDeque::with_capacity(120),
last_http_requests: 0,
last_kv_read_bytes: 0,
last_kv_write_bytes: 0,
#[cfg(target_os = "macos")]
generation_service,
#[cfg(target_os = "macos")]
@@ -740,6 +788,9 @@ impl App {
self.error = None;
}
}
Message::ShowChat => self.detail_tab = DetailTab::Chat,
Message::ShowStats => self.detail_tab = DetailTab::Stats,
Message::MetricsTick => self.sample_metrics(),
Message::PreferenceModelChanged(model) => {
self.preference_draft.model = model;
if !model.supports_dspark() {
@@ -1153,6 +1204,8 @@ impl App {
window::close_requests().map(Message::WindowClosed),
window::close_events().map(Message::WindowClosed),
];
subscriptions
.push(iced::time::every(METRICS_SAMPLE_INTERVAL).map(|_| Message::MetricsTick));
#[cfg(target_os = "macos")]
subscriptions.push(iced::time::every(Duration::from_millis(50)).map(|_| {
match crate::native_menu::next_event() {
@@ -1327,6 +1380,7 @@ impl App {
models_path(),
application_support_path().join("kv-cache").join("http"),
endpoint_port,
Arc::clone(&self.metrics),
) {
Ok(endpoint) => Some(endpoint),
Err(error) => {
@@ -1569,6 +1623,7 @@ impl App {
.iter()
.map(|message| ChatTurn {
user: message.user,
skip_previous_eos: false,
reasoning: message.reasoning.clone(),
reasoning_complete: message.reasoning_complete,
content: message.content.clone(),
@@ -1577,6 +1632,7 @@ impl App {
#[cfg(target_os = "macos")]
messages.push(ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: prompt.clone(),
@@ -1729,11 +1785,53 @@ impl App {
#[cfg(not(target_os = "macos"))]
false
}
fn sample_metrics(&mut self) {
let snapshot = self.metrics.snapshot();
let http_requests_per_second = snapshot
.http_requests
.saturating_sub(self.last_http_requests) as f32
/ METRICS_SAMPLE_INTERVAL.as_secs_f32();
let kv_read_bytes_per_second = snapshot
.kv_read_bytes
.saturating_sub(self.last_kv_read_bytes) as f32
/ METRICS_SAMPLE_INTERVAL.as_secs_f32();
let kv_write_bytes_per_second = snapshot
.kv_write_bytes
.saturating_sub(self.last_kv_write_bytes)
as f32
/ METRICS_SAMPLE_INTERVAL.as_secs_f32();
self.last_http_requests = snapshot.http_requests;
self.last_kv_read_bytes = snapshot.kv_read_bytes;
self.last_kv_write_bytes = snapshot.kv_write_bytes;
self.metrics_history.push_back(MetricsPoint {
decode_tokens_per_second: if snapshot.phase == crate::metrics::RuntimePhase::Generating
{
snapshot.decode_tokens_per_second
} else {
0.0
},
prefill_tokens_per_second: if snapshot.phase == crate::metrics::RuntimePhase::Prefilling
{
snapshot.prefill_tokens_per_second
} else {
0.0
},
http_requests_per_second,
kv_read_bytes_per_second,
kv_write_bytes_per_second,
});
if self.metrics_history.len() > 120 {
self.metrics_history.pop_front();
}
self.metrics_snapshot = snapshot;
}
}
#[cfg(target_os = "macos")]
fn spawn_services(
preferences: &AppPreferences,
metrics: Arc<Metrics>,
) -> (
Arc<RwLock<AppPreferences>>,
Option<GenerationService>,
@@ -1741,7 +1839,7 @@ fn spawn_services(
Option<String>,
) {
let runtime_preferences = Arc::new(RwLock::new(preferences.clone()));
let generation = match GenerationService::spawn() {
let generation = match GenerationService::spawn(Arc::clone(&metrics)) {
Ok(generation) => generation,
Err(error) => return (runtime_preferences, None, None, Some(error)),
};
@@ -1751,6 +1849,7 @@ fn spawn_services(
models_path(),
application_support_path().join("kv-cache").join("http"),
u16::try_from(preferences.endpoint_port).unwrap_or(4000),
metrics,
);
match endpoint {
Ok(endpoint) => (runtime_preferences, Some(generation), Some(endpoint), None),

550
src/app/view.rs
View File

@@ -1,5 +1,6 @@
use super::{
ActiveDownload, App, Message, ModelDownload, ModelOperation, chat_scroll_id, models_path,
ActiveDownload, App, DetailTab, Message, MetricsPoint, ModelDownload, ModelOperation,
chat_scroll_id, models_path,
};
use crate::database::{ProjectWithSessions, Session};
use crate::model::{
@@ -12,6 +13,7 @@ use iced::widget::{
pick_list, progress_bar, row, scrollable, stack, svg, text, text_input,
};
use iced::{Alignment, Background, Border, Color, Element, Length, Theme, window};
use std::collections::VecDeque;
use std::path::Path;
const ICON_FOLDER: &[u8] = include_bytes!("../../assets/icons/folder.svg");
@@ -191,6 +193,43 @@ impl App {
}
fn detail(&self) -> Element<'_, Message> {
let chat_active = self.detail_tab == DetailTab::Chat;
let stats_active = self.detail_tab == DetailTab::Stats;
let tabs = container(
row![
button(text("Chat").size(13))
.width(92)
.padding([7, 18])
.on_press(Message::ShowChat)
.style(move |theme, status| segmented_button_style(theme, status, chat_active)),
button(text("Stats").size(13))
.width(92)
.padding([7, 18])
.on_press(Message::ShowStats)
.style(move |theme, status| segmented_button_style(
theme,
status,
stats_active
)),
]
.spacing(2),
)
.padding(3)
.style(segmented_control_style);
let body = match self.detail_tab {
DetailTab::Chat => self.chat_detail(),
DetailTab::Stats => self.stats_dashboard(),
};
column![
container(tabs).center_x(Length::Fill).padding([12, 0]),
body
]
.width(Length::Fill)
.height(Length::Fill)
.into()
}
fn chat_detail(&self) -> Element<'_, Message> {
let Some(item) = self.selected_project() else {
let open_project_content = row![icon(ICON_FOLDER_PLUS, 17), text("Open project…"),]
.spacing(8)
@@ -396,6 +435,373 @@ impl App {
.into()
}
fn stats_dashboard(&self) -> Element<'_, Message> {
let stats = &self.metrics_snapshot;
let context_fraction = if stats.context_limit == 0 {
0.0
} else {
stats.context_used.min(stats.context_limit) as f32 / stats.context_limit as f32
};
let cache_fraction = if stats.last_prompt_tokens == 0 {
0.0
} else {
stats.last_cached_tokens as f32 / stats.last_prompt_tokens as f32
};
let cache_hit_fraction = if stats.kv_lookups == 0 {
0.0
} else {
stats.kv_hits as f32 / stats.kv_lookups as f32
};
let endpoint = if stats.server_listening {
format!("Listening · 127.0.0.1:{}", stats.server_port)
} else {
"Stopped".to_owned()
};
let phase_color = match stats.phase {
crate::metrics::RuntimePhase::Generating => Color::from_rgb8(84, 170, 255),
crate::metrics::RuntimePhase::Prefilling | crate::metrics::RuntimePhase::Loading => {
Color::from_rgb8(240, 180, 70)
}
crate::metrics::RuntimePhase::Ready => Color::from_rgb8(72, 176, 112),
crate::metrics::RuntimePhase::Failed => Color::from_rgb8(220, 80, 86),
crate::metrics::RuntimePhase::Unloaded => muted_text(),
};
let heading = container(
row![
column![
text("Runtime observability").size(24),
text(format!(
"{} · {} · uptime {}",
stats.model,
stats.source.label(),
format_duration(stats.uptime_seconds as f64)
))
.size(12)
.color(muted_text()),
]
.spacing(5),
Space::with_width(Length::Fill),
container(text(stats.phase.label()).size(12).color(phase_color))
.padding([7, 11])
.style(move |_| status_badge_style(phase_color)),
]
.align_y(Alignment::Center),
)
.padding(16)
.style(overview_style);
let headline = column![
row![
metric_card(
"DECODE",
format!("{:.1} tok/s", stats.decode_tokens_per_second),
format!(
"{} completion tokens total",
format_count(stats.completion_tokens)
),
),
metric_card(
"PREFILL",
format!("{:.1} tok/s", stats.prefill_tokens_per_second),
format!("{} prompt tokens total", format_count(stats.prompt_tokens)),
),
]
.spacing(10),
row![
metric_card(
"CONTEXT",
format!(
"{} / {}",
format_count(u64::from(stats.context_used)),
format_count(u64::from(stats.context_limit))
),
format!("{:.0}% occupied", context_fraction * 100.0),
),
metric_card(
"WORK",
format!(
"{} active · {} queued",
stats.http_active, stats.queue_depth
),
format!("{} runtime requests", format_count(stats.runtime_requests)),
),
]
.spacing(10),
]
.spacing(10);
let throughput = stats_panel(
"MODEL ACTIVITY · LAST 24 SECONDS",
column![
mini_chart(
&self.metrics_history,
|point| point.decode_tokens_per_second,
Color::from_rgb8(84, 170, 255),
),
row![
text("Decode")
.size(12)
.color(Color::from_rgb8(84, 170, 255)),
Space::with_width(Length::Fill),
text(format!("{:.1} tok/s", stats.decode_tokens_per_second))
.size(12)
.color(muted_text()),
],
mini_chart(
&self.metrics_history,
|point| point.prefill_tokens_per_second,
Color::from_rgb8(157, 119, 255),
),
row![
text("Prefill")
.size(12)
.color(Color::from_rgb8(157, 119, 255)),
Space::with_width(Length::Fill),
text(format!("{:.1} tok/s", stats.prefill_tokens_per_second))
.size(12)
.color(muted_text()),
],
]
.spacing(7)
.into(),
);
let requests = stats_panel(
"SERVER REQUEST RATE · LAST 24 SECONDS",
column![
mini_chart(
&self.metrics_history,
|point| point.http_requests_per_second,
Color::from_rgb8(72, 176, 112),
),
row![
text(format!("{} requests", format_count(stats.http_requests))).size(12),
Space::with_width(Length::Fill),
text(format!(
"{} errors · {} streaming",
stats.http_errors, stats.http_streaming_requests
))
.size(12)
.color(muted_text()),
],
]
.spacing(7)
.into(),
);
let latest = self.metrics_history.back().copied().unwrap_or_default();
let kv_io = stats_panel(
"KV CHECKPOINT I/O · LAST 24 SECONDS",
column![
mini_chart(
&self.metrics_history,
|point| point.kv_read_bytes_per_second,
Color::from_rgb8(67, 194, 203),
),
row![
text(if stats.kv_read_active {
"Read · active"
} else {
"Read"
})
.size(12)
.color(Color::from_rgb8(67, 194, 203)),
Space::with_width(Length::Fill),
text(format_rate(latest.kv_read_bytes_per_second))
.size(12)
.color(muted_text()),
],
mini_chart(
&self.metrics_history,
|point| point.kv_write_bytes_per_second,
Color::from_rgb8(240, 180, 70),
),
row![
text(if stats.kv_write_active {
"Write · active"
} else {
"Write"
})
.size(12)
.color(Color::from_rgb8(240, 180, 70)),
Space::with_width(Length::Fill),
text(format_rate(latest.kv_write_bytes_per_second))
.size(12)
.color(muted_text()),
],
]
.spacing(7)
.into(),
);
let model = stats_panel(
"MODEL CORE",
column![
metric_row("State", stats.phase.label()),
metric_row("Loaded model", stats.model),
metric_row("Mapped weights", format_bytes(stats.model_bytes)),
metric_row("Tensors", format_count(stats.tensor_count)),
metric_row("Vocabulary", format_count(stats.vocabulary_size)),
metric_row("Last load", format_milliseconds(stats.model_load_ms)),
metric_row(
"Lifecycle",
format!(
"{} loads · {} unloads",
stats.model_loads, stats.model_unloads
),
),
]
.spacing(9)
.into(),
);
let runtime = stats_panel(
"GENERATION",
column![
metric_row("Last runtime", format_milliseconds(stats.last_runtime_ms)),
metric_row(
"Average runtime",
format_milliseconds(stats.average_runtime_ms)
),
metric_row("Last prompt", format_count(stats.last_prompt_tokens)),
metric_row("Last reused", format_count(stats.last_cached_tokens)),
metric_row(
"Last completion",
format_count(stats.last_completion_tokens)
),
metric_row("Cached tokens total", format_count(stats.cached_tokens)),
metric_row(
"Cache reuse",
format!("{:.0}% of last prompt", cache_fraction * 100.0),
),
metric_row(
"Results",
format!(
"{} completed · {} failed",
stats.completed_requests, stats.failed_requests
),
),
]
.spacing(9)
.into(),
);
let cache = stats_panel(
"KV CACHE",
column![
metric_row(
"Total",
format!(
"{} · {} files",
format_bytes(stats.kv_bytes),
stats.kv_files
)
),
metric_row(
"Local sessions",
format!(
"{} · {} files",
format_bytes(stats.local_kv_bytes),
stats.local_kv_files
),
),
metric_row(
"HTTP transient",
format!(
"{} · {} files",
format_bytes(stats.http_kv_bytes),
stats.http_kv_files
),
),
metric_row("Checkpoint writes", format_count(stats.checkpoint_writes)),
metric_row(
"Exact hits",
format!(
"{} · {} memory / {} disk",
stats.kv_hits, stats.kv_memory_hits, stats.kv_disk_hits
),
),
metric_row(
"Misses",
format!("{} · {} invalid", stats.kv_misses, stats.kv_invalid),
),
metric_row("Lookups", format_count(stats.kv_lookups)),
metric_row(
"Exact hit rate",
format!("{:.1}%", cache_hit_fraction * 100.0)
),
metric_row("Prefix hits", format_count(stats.kv_prefix_hits)),
metric_row(
"Reads",
format!(
"{} · {} · {} errors · last {}",
stats.kv_read_operations,
format_bytes(stats.kv_read_bytes),
stats.kv_read_errors,
format_milliseconds(stats.last_kv_read_ms),
),
),
metric_row(
"Writes",
format!(
"{} · {} · {} errors · last {}",
stats.kv_write_operations,
format_bytes(stats.kv_write_bytes),
stats.kv_write_errors,
format_milliseconds(stats.last_kv_write_ms),
),
),
progress_bar(0.0..=1.0, cache_fraction.min(1.0)).height(4),
]
.spacing(9)
.into(),
);
let server = stats_panel(
"LOCAL SERVER",
column![
metric_row("Endpoint", endpoint),
metric_row("Active", format_count(u64::from(stats.http_active))),
metric_row("Completed", format_count(stats.http_completed)),
metric_row("Chat completions", format_count(stats.http_chat_requests)),
metric_row("Model queries", format_count(stats.http_model_requests)),
metric_row(
"Runtime sources",
format!(
"{} local · {} HTTP",
stats.local_requests, stats.endpoint_generations
),
),
metric_row("Received", format_bytes(stats.http_bytes_received)),
metric_row("Last latency", format_milliseconds(stats.last_http_ms)),
metric_row(
"Average latency",
format_milliseconds(stats.average_http_ms)
),
]
.spacing(9)
.into(),
);
scrollable(
container(
column![
heading,
headline,
throughput,
kv_io,
requests,
row![model, runtime].spacing(10),
row![cache, server].spacing(10),
text("Counters are published by the runtime with relaxed atomics and sampled by the UI every 200 ms.")
.size(11)
.color(muted_text()),
]
.spacing(12),
)
.padding(24)
.max_width(960)
.center_x(Length::Fill),
)
.height(Length::Fill)
.into()
}
fn preferences_panel(&self) -> Element<'_, Message> {
let dspark_toggle: Option<fn(bool) -> Message> = self
.preference_draft
@@ -1218,6 +1624,101 @@ fn action_button<'a>(content: impl Into<Element<'a, Message>>) -> Button<'a, Mes
button(content).padding([8, 14]).style(action_button_style)
}
fn metric_card(label: &'static str, value: String, detail: String) -> Element<'static, Message> {
container(
column![
text(label).size(10).color(muted_text()),
text(value).size(20),
text(detail).size(11).color(muted_text()),
]
.spacing(5),
)
.padding(14)
.width(Length::FillPortion(1))
.style(preference_group_style)
.into()
}
fn metric_row(label: &'static str, value: impl ToString) -> Element<'static, Message> {
row![
text(label).size(12).color(muted_text()),
Space::with_width(Length::Fill),
text(value.to_string()).size(12),
]
.spacing(12)
.align_y(Alignment::Center)
.into()
}
fn stats_panel<'a>(title: &'static str, content: Element<'a, Message>) -> Element<'a, Message> {
container(column![text(title).size(10).color(muted_text()), content].spacing(11))
.padding(14)
.width(Length::FillPortion(1))
.style(preference_group_style)
.into()
}
fn mini_chart(
history: &VecDeque<MetricsPoint>,
value: fn(&MetricsPoint) -> f32,
color: Color,
) -> Element<'static, Message> {
let values = history
.iter()
.rev()
.take(120)
.map(value)
.collect::<Vec<_>>();
let maximum = values.iter().copied().fold(1.0_f32, f32::max);
let mut bars = row![].spacing(1).height(54).align_y(Alignment::End);
for value in values.into_iter().rev() {
let height = if value > 0.0 {
(value / maximum * 52.0).max(2.0)
} else {
1.0
};
bars = bars.push(
container(Space::new(Length::Fill, height))
.width(Length::FillPortion(1))
.style(move |_| chart_bar_style(color)),
);
}
if history.is_empty() {
bars = bars.push(
container(text("Waiting for samples…").size(11).color(muted_text()))
.center_x(Length::Fill)
.center_y(Length::Fill),
);
}
container(bars)
.height(58)
.width(Length::Fill)
.padding([3, 0])
.into()
}
fn format_count(value: u64) -> String {
if value >= 1_000_000 {
format!("{:.1}M", value as f64 / 1_000_000.0)
} else if value >= 1_000 {
format!("{:.1}K", value as f64 / 1_000.0)
} else {
value.to_string()
}
}
fn format_milliseconds(milliseconds: u64) -> String {
if milliseconds >= 1_000 {
format!("{:.2}s", milliseconds as f64 / 1_000.0)
} else {
format!("{milliseconds}ms")
}
}
fn format_rate(bytes_per_second: f32) -> String {
format!("{}/s", format_bytes(bytes_per_second.max(0.0) as u64))
}
fn danger_button<'a>(content: impl Into<Element<'a, Message>>) -> Button<'a, Message> {
button(content).padding([8, 14]).style(danger_button_style)
}
@@ -1250,6 +1751,53 @@ fn danger_button_style(theme: &Theme, status: button::Status) -> button::Style {
style
}
fn segmented_control_style(_: &Theme) -> container::Style {
container::Style {
background: Some(Background::Color(Color::from_rgb8(24, 24, 26))),
border: Border {
color: Color::from_rgb8(57, 57, 60),
width: 1.0,
radius: 18.0.into(),
},
..container::Style::default()
}
}
fn segmented_button_style(_: &Theme, status: button::Status, selected: bool) -> button::Style {
let background = if selected {
Some(Background::Color(Color::from_rgb8(55, 55, 58)))
} else if status == button::Status::Hovered {
Some(Background::Color(Color::from_rgb8(39, 39, 42)))
} else {
None
};
button::Style {
background,
text_color: if selected { Color::WHITE } else { muted_text() },
border: Border {
radius: 14.0.into(),
..Border::default()
},
..button::Style::default()
}
}
fn status_badge_style(color: Color) -> container::Style {
container::Style {
background: Some(Background::Color(color.scale_alpha(0.12))),
border: Border {
color: color.scale_alpha(0.45),
width: 1.0,
radius: 12.0.into(),
},
..container::Style::default()
}
}
fn chart_bar_style(color: Color) -> container::Style {
container::Style::default().background(color.scale_alpha(0.82))
}
fn overview_style(_: &Theme) -> container::Style {
container::Style {
background: Some(Background::Color(Color::from_rgb8(31, 31, 33))),

195
src/engine.rs
View File

@@ -3,12 +3,16 @@ mod gguf;
mod metal;
mod tokenizer;
#[cfg(target_os = "macos")]
use crate::metrics::{KvLookup, Metrics};
use crate::model::ModelChoice;
use crate::settings::TurnSettings;
use crate::settings::{EngineSettings, ReasoningMode};
use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value};
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};
#[cfg(target_os = "macos")]
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(target_os = "macos")]
use std::time::Instant;
@@ -256,8 +260,14 @@ impl Model {
.encode_conversation(system, messages, reasoning)
}
fn render_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec<i32> {
self.tokenizer.encode_continuation(prompt, reasoning)
fn render_continuation(
&self,
prompt: &str,
reasoning: ReasoningMode,
skip_previous_eos: bool,
) -> Vec<i32> {
self.tokenizer
.encode_continuation(prompt, reasoning, skip_previous_eos)
}
pub(crate) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
@@ -295,11 +305,13 @@ impl Model {
pub(crate) struct Generator {
executor: metal::Executor,
checkpoint: Option<PathBuf>,
metrics: Arc<Metrics>,
}
#[derive(Clone)]
pub(crate) struct ChatTurn {
pub(crate) user: bool,
pub(crate) skip_previous_eos: bool,
pub(crate) reasoning: Option<String>,
pub(crate) reasoning_complete: bool,
pub(crate) content: String,
@@ -311,11 +323,13 @@ pub(crate) struct GenerationOutput {
pub(crate) cached_tokens: u32,
pub(crate) completion_tokens: u32,
pub(crate) finish_reason: &'static str,
pub(crate) previous_checkpoint_bytes: Option<u64>,
pub(crate) checkpoint_bytes: u64,
}
#[cfg(target_os = "macos")]
impl Generator {
pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> {
pub(crate) fn open(settings: &EngineSettings, metrics: Arc<Metrics>) -> Result<Self, String> {
if settings.model != ModelChoice::DeepSeekV4Flash {
return Err("local generation currently supports DeepSeek V4 Flash only".into());
}
@@ -330,13 +344,19 @@ impl Generator {
model,
settings.context_tokens.max(1) as u32,
settings.execution.quality,
settings.execution.prefill_chunk,
)?;
Ok(Self {
executor,
checkpoint: None,
metrics,
})
}
pub(crate) fn summary(&self) -> ModelSummary {
self.executor.model().summary()
}
pub(crate) fn generate(
&mut self,
checkpoint: &Path,
@@ -346,12 +366,20 @@ impl Generator {
mut emit: impl FnMut(bool, String),
mut progress: impl FnMut(u32, u32, Option<f32>),
) -> Result<GenerationOutput, String> {
self.select_checkpoint(checkpoint)?;
let (output, prompt_complete) =
let history = messages
.split_last()
.map_or(messages, |(_, history)| history);
self.select_checkpoint(
checkpoint,
conversation_tag(&settings.system_prompt, settings.reasoning_mode, history),
)?;
let (mut output, prompt_complete) =
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
let mut completed = messages.to_vec();
completed.push(output.message.clone());
self.executor.save_checkpoint(
output.previous_checkpoint_bytes =
std::fs::metadata(checkpoint).ok().map(|item| item.len());
self.save_checkpoint(
checkpoint,
if prompt_complete {
conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed)
@@ -359,6 +387,9 @@ impl Generator {
[0; 32]
},
)?;
output.checkpoint_bytes = std::fs::metadata(checkpoint)
.map(|item| item.len())
.unwrap_or(0);
Ok(output)
}
@@ -378,15 +409,17 @@ impl Generator {
conversation_tag(&settings.system_prompt, settings.reasoning_mode, history);
let checkpoint = directory.join(format!("{}.bin", hex_tag(history_tag)));
if self.executor.checkpoint_tag() != history_tag {
self.select_checkpoint(&checkpoint)?;
self.select_checkpoint(&checkpoint, history_tag)?;
if self.executor.checkpoint_tag() != history_tag {
self.executor.reset()?;
self.checkpoint = None;
let _ = std::fs::remove_file(&checkpoint);
}
} else {
self.metrics.kv_lookup(KvLookup::MemoryHit);
}
let (output, prompt_complete) =
let (mut output, prompt_complete) =
self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
let mut completed = messages.to_vec();
completed.push(output.message.clone());
@@ -399,7 +432,10 @@ impl Generator {
)
})?;
let completed_checkpoint = directory.join(format!("{}.bin", hex_tag(completed_tag)));
self.executor.save_checkpoint(
output.previous_checkpoint_bytes = std::fs::metadata(&completed_checkpoint)
.ok()
.map(|item| item.len());
self.save_checkpoint(
&completed_checkpoint,
if prompt_complete {
completed_tag
@@ -407,23 +443,62 @@ impl Generator {
[0; 32]
},
)?;
output.checkpoint_bytes = std::fs::metadata(&completed_checkpoint)
.map(|item| item.len())
.unwrap_or(0);
self.checkpoint = Some(completed_checkpoint);
Ok(output)
}
fn select_checkpoint(&mut self, checkpoint: &Path) -> Result<(), String> {
fn select_checkpoint(
&mut self,
checkpoint: &Path,
expected_tag: [u8; 32],
) -> Result<(), String> {
if self.checkpoint.as_deref() == Some(checkpoint) {
self.metrics
.kv_lookup(if self.executor.checkpoint_tag() == expected_tag {
KvLookup::MemoryHit
} else {
KvLookup::Miss
});
return Ok(());
}
self.executor.reset()?;
if self.executor.load_checkpoint(checkpoint).is_err() {
self.executor.reset()?;
let _ = std::fs::remove_file(checkpoint);
}
self.metrics.kv_read_started();
let started = Instant::now();
let loaded = self
.executor
.load_checkpoint(checkpoint, &mut |bytes| self.metrics.kv_read_bytes(bytes));
self.metrics
.kv_read_finished(started.elapsed(), loaded.is_err());
let lookup = match loaded {
Ok(true) if self.executor.checkpoint_tag() == expected_tag => KvLookup::DiskHit,
Ok(true) | Ok(false) => KvLookup::Miss,
Err(_) => {
self.executor.reset()?;
let _ = std::fs::remove_file(checkpoint);
KvLookup::Invalid
}
};
self.metrics.kv_lookup(lookup);
self.checkpoint = Some(checkpoint.to_owned());
Ok(())
}
fn save_checkpoint(&mut self, checkpoint: &Path, tag: [u8; 32]) -> Result<(), String> {
self.metrics.kv_write_started();
let started = Instant::now();
let result = self
.executor
.save_checkpoint(checkpoint, tag, &mut |bytes| {
self.metrics.kv_write_bytes(bytes);
});
self.metrics
.kv_write_finished(started.elapsed(), result.is_err());
result
}
fn generate_inner(
&mut self,
messages: &[ChatTurn],
@@ -443,11 +518,11 @@ impl Generator {
) =>
{
let mut tokens = self.executor.tokens().to_vec();
tokens.extend(
self.executor
.model()
.render_continuation(&latest.content, settings.reasoning_mode),
);
tokens.extend(self.executor.model().render_continuation(
&latest.content,
settings.reasoning_mode,
latest.skip_previous_eos,
));
tokens
}
_ => self.executor.model().render_conversation(
@@ -467,11 +542,13 @@ impl Generator {
));
}
let reused = self.executor.align_prompt(&tokens)?;
self.metrics.kv_prefix_reused(reused);
progress(self.executor.position(), self.executor.context(), None);
let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct;
let mut generated = ChatTurn {
user: false,
skip_previous_eos: false,
reasoning: reasoning.then(String::new),
reasoning_complete: !reasoning,
content: String::new(),
@@ -479,24 +556,40 @@ impl Generator {
let mut emitted_reasoning = 0;
let mut emitted_content = 0;
let prompt_tokens = tokens.len();
for (index, token) in tokens.into_iter().enumerate().skip(reused) {
if cancelled.load(Ordering::Relaxed) {
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: 0,
finish_reason: "stop",
},
false,
));
}
self.executor.eval(token)?;
if (index + 1).is_multiple_of(16) || index + 1 == prompt_tokens {
let suffix = &tokens[reused..];
let completed = if (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 {
let context = self.executor.context();
self.executor.prefill(suffix, |used| {
progress(used, context, None);
!cancelled.load(Ordering::Relaxed)
})?
} else {
let mut completed = 0;
for &token in suffix {
if cancelled.load(Ordering::Relaxed) {
break;
}
self.executor.eval(token)?;
completed += 1;
progress(self.executor.position(), self.executor.context(), None);
}
completed
};
if completed != suffix.len() {
return Ok((
GenerationOutput {
message: generated,
prompt_tokens: prompt_tokens as u32,
cached_tokens: reused as u32,
completion_tokens: 0,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
}
progress(self.executor.position(), self.executor.context(), Some(0.0));
let generation_started = Instant::now();
let mut generated_tokens = 0_u32;
for _ in 0..settings
@@ -519,6 +612,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
true,
));
@@ -550,6 +645,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
true,
));
@@ -575,6 +672,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
@@ -615,6 +714,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens + 1,
finish_reason: "stop",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
false,
));
@@ -645,6 +746,8 @@ impl Generator {
cached_tokens: reused as u32,
completion_tokens: generated_tokens,
finish_reason: "length",
previous_checkpoint_bytes: None,
checkpoint_bytes: 0,
},
true,
))
@@ -891,6 +994,7 @@ mod sampling_tests {
fn checkpoint_tag_covers_the_canonical_chat_state() {
let messages = [ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
@@ -924,7 +1028,7 @@ mod sampling_tests {
ReasoningMode::Direct,
);
assert_eq!(tokens.len(), 10);
let mut executor = metal::Executor::open(model, 32_768, false).unwrap();
let mut executor = metal::Executor::open(model, 32_768, false, 0).unwrap();
assert_eq!(executor.context(), 32_768);
for &token in &tokens {
executor.eval(token).unwrap();
@@ -947,7 +1051,9 @@ mod sampling_tests {
let next = argmax.0 as i32;
let checkpoint =
std::env::temp_dir().join(format!("ds4-rust-kv-{}.bin", std::process::id()));
executor.save_checkpoint(&checkpoint, [7; 32]).unwrap();
executor
.save_checkpoint(&checkpoint, [7; 32], &mut |_| {})
.unwrap();
executor.eval(next).unwrap();
let continued_logit = executor.logits()[0];
let continued_argmax = executor
@@ -960,8 +1066,8 @@ mod sampling_tests {
drop(executor);
let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).unwrap();
let mut restored = metal::Executor::open(model, 32_768, false).unwrap();
assert!(restored.load_checkpoint(&checkpoint).unwrap());
let mut restored = metal::Executor::open(model, 32_768, false, 0).unwrap();
assert!(restored.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
assert_eq!(restored.position(), tokens.len() as u32);
assert_eq!(restored.tokens(), tokens);
assert_eq!(restored.checkpoint_tag(), [7; 32]);
@@ -1872,6 +1978,7 @@ mod tests {
model.render_prompt("", "Hello", ReasoningMode::Direct),
[0, 128_803, 19_923, 128_804, 128_822]
);
assert_eq!(model.tokenizer.tokenize_rendered("｜DSML｜").len(), 1);
assert!(model.is_stop_token_for_reasoning(128_822, ReasoningMode::Direct));
assert!(!model.is_stop_token_for_reasoning(128_822, ReasoningMode::High));
let think_start = *model
@@ -1884,18 +1991,21 @@ mod tests {
&[
ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
},
ChatTurn {
user: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
},
ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
@@ -1908,27 +2018,34 @@ mod tests {
]
);
assert_eq!(
model.render_continuation("Hello", ReasoningMode::Direct),
model.render_continuation("Hello", ReasoningMode::Direct, false),
[1, 128_803, 19_923, 128_804, 128_822]
);
assert_eq!(
model.render_continuation("Hello", ReasoningMode::Direct, true),
[128_803, 19_923, 128_804, 128_822]
);
assert_eq!(
model.render_conversation(
"",
&[
ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),
},
ChatTurn {
user: false,
skip_previous_eos: false,
reasoning: Some("Hello".into()),
reasoning_complete: true,
content: "Hello".into(),
},
ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: "Hello".into(),

14
src/engine/gguf.rs
View File

@@ -46,6 +46,8 @@ pub(super) struct Tensor {
pub(super) struct Gguf {
path: PathBuf,
map: Mmap,
data_offset: u64,
max_tensor_bytes: u64,
pub(super) metadata: HashMap<String, Value>,
pub(super) tensors: HashMap<String, Tensor>,
}
@@ -141,6 +143,7 @@ impl Gguf {
}
let data_start = align(cursor.position() as u64, alignment)?;
let mut max_tensor_bytes = 0;
for (name, tensor) in &mut tensors {
tensor.offset = data_start
.checked_add(tensor.offset)
@@ -152,11 +155,14 @@ impl Gguf {
if end > map.len() as u64 {
return Err(format!("tensor {name} points outside the GGUF file"));
}
max_tensor_bytes = max_tensor_bytes.max(tensor.bytes);
}
Ok(Self {
path: path.to_owned(),
map,
data_offset: data_start,
max_tensor_bytes,
metadata,
tensors,
})
@@ -174,6 +180,14 @@ impl Gguf {
self.map.as_ptr()
}
pub(super) fn data_offset(&self) -> u64 {
self.data_offset
}
pub(super) fn max_tensor_bytes(&self) -> u64 {
self.max_tensor_bytes
}
pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> {
self.tensors
.get(name)

2122
src/engine/metal.rs
View File

File diff suppressed because it is too large Load Diff

77
src/engine/tokenizer.rs
View File

@@ -21,6 +21,7 @@ pub(super) struct Tokenizer {
sop: i32,
think_start: i32,
think_end: i32,
rendered_specials: Vec<(&'static [u8], i32)>,
}
impl Tokenizer {
@@ -85,6 +86,32 @@ impl Tokenizer {
{
return Err("tokenizer does not provide the required DS4 chat markers".into());
}
let rendered_specials = [
("<｜begin▁of▁sentence｜>".as_bytes(), bos),
("<｜end▁of▁sentence｜>".as_bytes(), eos),
(b"[gMASK]".as_slice(), bos),
(b"<sop>".as_slice(), sop),
(b"<|system|>".as_slice(), system),
("<｜User｜>".as_bytes(), user),
("<｜Assistant｜>".as_bytes(), assistant),
(b"<|user|>".as_slice(), user),
(b"<|assistant|>".as_slice(), assistant),
(b"<|observation|>".as_slice(), observation),
(b"<think>".as_slice(), think_start),
(b"</think>".as_slice(), think_end),
(b"<tool_call>".as_slice(), lookup(b"<tool_call>")),
(b"</tool_call>".as_slice(), lookup(b"</tool_call>")),
(b"<tool_response>".as_slice(), lookup(b"<tool_response>")),
(b"</tool_response>".as_slice(), lookup(b"</tool_response>")),
(b"<arg_key>".as_slice(), lookup(b"<arg_key>")),
(b"</arg_key>".as_slice(), lookup(b"</arg_key>")),
(b"<arg_value>".as_slice(), lookup(b"<arg_value>")),
(b"</arg_value>".as_slice(), lookup(b"</arg_value>")),
("｜DSML｜".as_bytes(), lookup("｜DSML｜".as_bytes())),
]
.into_iter()
.filter(|(_, token)| *token >= 0)
.collect();
Ok(Self {
family,
@@ -100,6 +127,7 @@ impl Tokenizer {
sop,
think_start,
think_end,
rendered_specials,
})
}
@@ -126,6 +154,37 @@ impl Tokenizer {
output
}
pub(super) fn tokenize_rendered(&self, text: &str) -> Vec<i32> {
let bytes = text.as_bytes();
let mut output = Vec::new();
let mut span = 0;
let mut position = 0;
while position < bytes.len() {
let special = self
.rendered_specials
.iter()
.find(|(marker, _)| bytes[position..].starts_with(marker));
if let Some((marker, token)) = special {
self.tokenize_plain(&text[span..position], &mut output);
output.push(*token);
position += marker.len();
span = position;
} else {
position += 1;
}
}
self.tokenize_plain(&text[span..], &mut output);
output
}
fn tokenize_plain(&self, text: &str, output: &mut Vec<i32>) {
if self.family == ModelFamily::Glm {
self.tokenize_glm(text, output);
} else {
self.tokenize_joyai(text, output);
}
}
pub(super) fn encode_chat(
&self,
system_prompt: &str,
@@ -136,6 +195,7 @@ impl Tokenizer {
system_prompt,
&[ChatTurn {
user: true,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: prompt.to_owned(),
@@ -172,7 +232,7 @@ impl Tokenizer {
if self.family == ModelFamily::Glm {
output.push(self.system);
}
output.extend(self.tokenize(system_prompt));
output.extend(self.tokenize_rendered(system_prompt));
}
for message in messages {
output.push(if message.user {
@@ -183,7 +243,7 @@ impl Tokenizer {
if !message.user {
if let Some(reasoning) = &message.reasoning {
output.push(self.think_start);
output.extend(self.tokenize(reasoning));
output.extend(self.tokenize_rendered(reasoning));
if message.reasoning_complete {
output.push(self.think_end);
}
@@ -193,7 +253,7 @@ impl Tokenizer {
output.push(self.think_end);
}
}
output.extend(self.tokenize(&message.content));
output.extend(self.tokenize_rendered(&message.content));
if !message.user && self.family == ModelFamily::DeepSeek {
output.push(self.eos);
}
@@ -209,13 +269,18 @@ impl Tokenizer {
output
}
pub(super) fn encode_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec<i32> {
pub(super) fn encode_continuation(
&self,
prompt: &str,
reasoning: ReasoningMode,
skip_previous_eos: bool,
) -> Vec<i32> {
let mut output = Vec::new();
if self.family == ModelFamily::DeepSeek {
if self.family == ModelFamily::DeepSeek && !skip_previous_eos {
output.push(self.eos);
}
output.push(self.user);
output.extend(self.tokenize(prompt));
output.extend(self.tokenize_rendered(prompt));
output.push(self.assistant);
if reasoning != ReasoningMode::Direct {
output.push(self.think_start);

1
src/main.rs
View File

@@ -1,6 +1,7 @@
mod app;
mod database;
mod engine;
mod metrics;
mod model;
#[cfg(target_os = "macos")]
mod native_edit;

727
src/metrics.rs Normal file
View File

@@ -0,0 +1,727 @@
use crate::model::ModelChoice;
use std::fs;
use std::path::Path;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU16, AtomicU32, AtomicU64, Ordering};
use std::time::{Duration, Instant};
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
#[repr(u8)]
pub(crate) enum RuntimePhase {
#[default]
Unloaded,
Loading,
Prefilling,
Generating,
Ready,
Failed,
}
impl RuntimePhase {
pub(crate) fn label(self) -> &'static str {
match self {
Self::Unloaded => "Unloaded",
Self::Loading => "Loading",
Self::Prefilling => "Prefilling",
Self::Generating => "Generating",
Self::Ready => "Ready",
Self::Failed => "Failed",
}
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
#[repr(u8)]
pub(crate) enum WorkSource {
#[default]
None,
LocalChat,
Http,
}
impl WorkSource {
pub(crate) fn label(self) -> &'static str {
match self {
Self::None => "None",
Self::LocalChat => "Local chat",
Self::Http => "HTTP endpoint",
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum KvLookup {
MemoryHit,
DiskHit,
Miss,
Invalid,
}
#[derive(Clone, Debug, Default)]
pub(crate) struct MetricsSnapshot {
pub(crate) uptime_seconds: u64,
pub(crate) phase: RuntimePhase,
pub(crate) source: WorkSource,
pub(crate) model: &'static str,
pub(crate) queue_depth: u32,
pub(crate) runtime_requests: u64,
pub(crate) local_requests: u64,
pub(crate) endpoint_generations: u64,
pub(crate) completed_requests: u64,
pub(crate) failed_requests: u64,
pub(crate) model_loads: u64,
pub(crate) model_unloads: u64,
pub(crate) model_load_ms: u64,
pub(crate) model_bytes: u64,
pub(crate) tensor_count: u64,
pub(crate) vocabulary_size: u64,
pub(crate) context_used: u32,
pub(crate) context_limit: u32,
pub(crate) decode_tokens_per_second: f32,
pub(crate) prefill_tokens_per_second: f32,
pub(crate) last_runtime_ms: u64,
pub(crate) average_runtime_ms: u64,
pub(crate) last_prompt_tokens: u64,
pub(crate) last_cached_tokens: u64,
pub(crate) last_completion_tokens: u64,
pub(crate) prompt_tokens: u64,
pub(crate) cached_tokens: u64,
pub(crate) completion_tokens: u64,
pub(crate) kv_lookups: u64,
pub(crate) kv_hits: u64,
pub(crate) kv_memory_hits: u64,
pub(crate) kv_disk_hits: u64,
pub(crate) kv_misses: u64,
pub(crate) kv_invalid: u64,
pub(crate) kv_prefix_hits: u64,
pub(crate) checkpoint_writes: u64,
pub(crate) kv_read_active: bool,
pub(crate) kv_write_active: bool,
pub(crate) kv_read_operations: u64,
pub(crate) kv_write_operations: u64,
pub(crate) kv_read_errors: u64,
pub(crate) kv_write_errors: u64,
pub(crate) kv_read_bytes: u64,
pub(crate) kv_write_bytes: u64,
pub(crate) last_kv_read_ms: u64,
pub(crate) last_kv_write_ms: u64,
pub(crate) kv_files: u64,
pub(crate) kv_bytes: u64,
pub(crate) local_kv_files: u64,
pub(crate) local_kv_bytes: u64,
pub(crate) http_kv_files: u64,
pub(crate) http_kv_bytes: u64,
pub(crate) server_listening: bool,
pub(crate) server_port: u16,
pub(crate) http_active: u32,
pub(crate) http_requests: u64,
pub(crate) http_completed: u64,
pub(crate) http_errors: u64,
pub(crate) http_chat_requests: u64,
pub(crate) http_model_requests: u64,
pub(crate) http_streaming_requests: u64,
pub(crate) http_bytes_received: u64,
pub(crate) last_http_ms: u64,
pub(crate) average_http_ms: u64,
}
pub(crate) struct Metrics {
started: Instant,
phase: AtomicU8,
source: AtomicU8,
model: AtomicU8,
queue_depth: AtomicU32,
runtime_requests: AtomicU64,
local_requests: AtomicU64,
endpoint_generations: AtomicU64,
completed_requests: AtomicU64,
failed_requests: AtomicU64,
model_loads: AtomicU64,
model_unloads: AtomicU64,
model_load_ms: AtomicU64,
model_bytes: AtomicU64,
tensor_count: AtomicU64,
vocabulary_size: AtomicU64,
context_used: AtomicU32,
context_limit: AtomicU32,
decode_tps: AtomicU32,
prefill_tps: AtomicU32,
last_runtime_ms: AtomicU64,
total_runtime_ms: AtomicU64,
last_prompt_tokens: AtomicU64,
last_cached_tokens: AtomicU64,
last_completion_tokens: AtomicU64,
prompt_tokens: AtomicU64,
cached_tokens: AtomicU64,
completion_tokens: AtomicU64,
kv_lookups: AtomicU64,
kv_hits: AtomicU64,
kv_memory_hits: AtomicU64,
kv_disk_hits: AtomicU64,
kv_misses: AtomicU64,
kv_invalid: AtomicU64,
kv_prefix_hits: AtomicU64,
checkpoint_writes: AtomicU64,
kv_read_active: AtomicBool,
kv_write_active: AtomicBool,
kv_read_operations: AtomicU64,
kv_write_operations: AtomicU64,
kv_read_errors: AtomicU64,
kv_write_errors: AtomicU64,
kv_read_bytes: AtomicU64,
kv_write_bytes: AtomicU64,
last_kv_read_ms: AtomicU64,
last_kv_write_ms: AtomicU64,
kv_files: AtomicU64,
kv_bytes: AtomicU64,
local_kv_files: AtomicU64,
local_kv_bytes: AtomicU64,
http_kv_files: AtomicU64,
http_kv_bytes: AtomicU64,
server_listening: AtomicBool,
server_port: AtomicU16,
http_active: AtomicU32,
http_requests: AtomicU64,
http_completed: AtomicU64,
http_errors: AtomicU64,
http_chat_requests: AtomicU64,
http_model_requests: AtomicU64,
http_streaming_requests: AtomicU64,
http_bytes_received: AtomicU64,
last_http_ms: AtomicU64,
total_http_ms: AtomicU64,
}
impl Metrics {
pub(crate) fn new(cache_root: &Path) -> Self {
let cache = cache_usage(cache_root);
Self {
started: Instant::now(),
phase: AtomicU8::new(RuntimePhase::Unloaded as u8),
source: AtomicU8::new(WorkSource::None as u8),
model: AtomicU8::new(0),
queue_depth: AtomicU32::new(0),
runtime_requests: AtomicU64::new(0),
local_requests: AtomicU64::new(0),
endpoint_generations: AtomicU64::new(0),
completed_requests: AtomicU64::new(0),
failed_requests: AtomicU64::new(0),
model_loads: AtomicU64::new(0),
model_unloads: AtomicU64::new(0),
model_load_ms: AtomicU64::new(0),
model_bytes: AtomicU64::new(0),
tensor_count: AtomicU64::new(0),
vocabulary_size: AtomicU64::new(0),
context_used: AtomicU32::new(0),
context_limit: AtomicU32::new(0),
decode_tps: AtomicU32::new(0),
prefill_tps: AtomicU32::new(0),
last_runtime_ms: AtomicU64::new(0),
total_runtime_ms: AtomicU64::new(0),
last_prompt_tokens: AtomicU64::new(0),
last_cached_tokens: AtomicU64::new(0),
last_completion_tokens: AtomicU64::new(0),
prompt_tokens: AtomicU64::new(0),
cached_tokens: AtomicU64::new(0),
completion_tokens: AtomicU64::new(0),
kv_lookups: AtomicU64::new(0),
kv_hits: AtomicU64::new(0),
kv_memory_hits: AtomicU64::new(0),
kv_disk_hits: AtomicU64::new(0),
kv_misses: AtomicU64::new(0),
kv_invalid: AtomicU64::new(0),
kv_prefix_hits: AtomicU64::new(0),
checkpoint_writes: AtomicU64::new(0),
kv_read_active: AtomicBool::new(false),
kv_write_active: AtomicBool::new(false),
kv_read_operations: AtomicU64::new(0),
kv_write_operations: AtomicU64::new(0),
kv_read_errors: AtomicU64::new(0),
kv_write_errors: AtomicU64::new(0),
kv_read_bytes: AtomicU64::new(0),
kv_write_bytes: AtomicU64::new(0),
last_kv_read_ms: AtomicU64::new(0),
last_kv_write_ms: AtomicU64::new(0),
kv_files: AtomicU64::new(cache.files),
kv_bytes: AtomicU64::new(cache.bytes),
local_kv_files: AtomicU64::new(cache.local_files),
local_kv_bytes: AtomicU64::new(cache.local_bytes),
http_kv_files: AtomicU64::new(cache.http_files),
http_kv_bytes: AtomicU64::new(cache.http_bytes),
server_listening: AtomicBool::new(false),
server_port: AtomicU16::new(0),
http_active: AtomicU32::new(0),
http_requests: AtomicU64::new(0),
http_completed: AtomicU64::new(0),
http_errors: AtomicU64::new(0),
http_chat_requests: AtomicU64::new(0),
http_model_requests: AtomicU64::new(0),
http_streaming_requests: AtomicU64::new(0),
http_bytes_received: AtomicU64::new(0),
last_http_ms: AtomicU64::new(0),
total_http_ms: AtomicU64::new(0),
}
}
pub(crate) fn request_queued(&self, source: WorkSource) {
self.runtime_requests.fetch_add(1, Ordering::Relaxed);
match source {
WorkSource::LocalChat => self.local_requests.fetch_add(1, Ordering::Relaxed),
WorkSource::Http => self.endpoint_generations.fetch_add(1, Ordering::Relaxed),
WorkSource::None => 0,
};
self.queue_depth.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn request_rejected(&self) {
self.queue_depth.fetch_sub(1, Ordering::Relaxed);
self.failed_requests.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn request_started(&self, source: WorkSource) {
self.queue_depth.fetch_sub(1, Ordering::Relaxed);
self.source.store(source as u8, Ordering::Relaxed);
self.decode_tps.store(0, Ordering::Relaxed);
self.prefill_tps.store(0, Ordering::Relaxed);
}
pub(crate) fn loading(&self) {
self.phase
.store(RuntimePhase::Loading as u8, Ordering::Relaxed);
}
pub(crate) fn loaded(
&self,
model: ModelChoice,
elapsed: Duration,
mapped_bytes: u64,
tensors: usize,
vocabulary_size: usize,
) {
self.model.store(model_code(model), Ordering::Relaxed);
self.model_loads.fetch_add(1, Ordering::Relaxed);
self.model_load_ms
.store(milliseconds(elapsed), Ordering::Relaxed);
self.model_bytes.store(mapped_bytes, Ordering::Relaxed);
self.tensor_count.store(tensors as u64, Ordering::Relaxed);
self.vocabulary_size
.store(vocabulary_size as u64, Ordering::Relaxed);
}
pub(crate) fn prefill_progress(&self, used: u32, limit: u32, speed: f32) {
self.phase
.store(RuntimePhase::Prefilling as u8, Ordering::Relaxed);
self.context_used.store(used, Ordering::Relaxed);
self.context_limit.store(limit, Ordering::Relaxed);
store_f32(&self.prefill_tps, speed);
}
pub(crate) fn generation_progress(&self, used: u32, limit: u32, speed: f32) {
self.phase
.store(RuntimePhase::Generating as u8, Ordering::Relaxed);
self.context_used.store(used, Ordering::Relaxed);
self.context_limit.store(limit, Ordering::Relaxed);
store_f32(&self.decode_tps, speed);
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn request_finished(
&self,
source: WorkSource,
elapsed: Duration,
prompt_tokens: u32,
cached_tokens: u32,
completion_tokens: u32,
previous_checkpoint_bytes: Option<u64>,
checkpoint_bytes: u64,
) {
let elapsed = milliseconds(elapsed);
self.completed_requests.fetch_add(1, Ordering::Relaxed);
self.last_runtime_ms.store(elapsed, Ordering::Relaxed);
self.total_runtime_ms.fetch_add(elapsed, Ordering::Relaxed);
self.last_prompt_tokens
.store(u64::from(prompt_tokens), Ordering::Relaxed);
self.last_cached_tokens
.store(u64::from(cached_tokens), Ordering::Relaxed);
self.last_completion_tokens
.store(u64::from(completion_tokens), Ordering::Relaxed);
self.prompt_tokens
.fetch_add(u64::from(prompt_tokens), Ordering::Relaxed);
self.cached_tokens
.fetch_add(u64::from(cached_tokens), Ordering::Relaxed);
self.completion_tokens
.fetch_add(u64::from(completion_tokens), Ordering::Relaxed);
self.record_checkpoint(source, previous_checkpoint_bytes, checkpoint_bytes);
self.phase
.store(RuntimePhase::Ready as u8, Ordering::Relaxed);
self.source.store(WorkSource::None as u8, Ordering::Relaxed);
}
pub(crate) fn request_failed(&self, elapsed: Duration) {
let elapsed = milliseconds(elapsed);
self.failed_requests.fetch_add(1, Ordering::Relaxed);
self.last_runtime_ms.store(elapsed, Ordering::Relaxed);
self.total_runtime_ms.fetch_add(elapsed, Ordering::Relaxed);
self.phase
.store(RuntimePhase::Failed as u8, Ordering::Relaxed);
self.source.store(WorkSource::None as u8, Ordering::Relaxed);
}
pub(crate) fn unloaded(&self) {
self.phase
.store(RuntimePhase::Unloaded as u8, Ordering::Relaxed);
self.model.store(0, Ordering::Relaxed);
self.model_bytes.store(0, Ordering::Relaxed);
self.tensor_count.store(0, Ordering::Relaxed);
self.vocabulary_size.store(0, Ordering::Relaxed);
self.context_used.store(0, Ordering::Relaxed);
self.decode_tps.store(0, Ordering::Relaxed);
self.prefill_tps.store(0, Ordering::Relaxed);
self.model_unloads.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn server_listening(&self, port: u16) {
self.server_port.store(port, Ordering::Relaxed);
self.server_listening.store(true, Ordering::Relaxed);
}
pub(crate) fn kv_lookup(&self, result: KvLookup) {
self.kv_lookups.fetch_add(1, Ordering::Relaxed);
match result {
KvLookup::MemoryHit => {
self.kv_hits.fetch_add(1, Ordering::Relaxed);
self.kv_memory_hits.fetch_add(1, Ordering::Relaxed);
}
KvLookup::DiskHit => {
self.kv_hits.fetch_add(1, Ordering::Relaxed);
self.kv_disk_hits.fetch_add(1, Ordering::Relaxed);
}
KvLookup::Miss => {
self.kv_misses.fetch_add(1, Ordering::Relaxed);
}
KvLookup::Invalid => {
self.kv_misses.fetch_add(1, Ordering::Relaxed);
self.kv_invalid.fetch_add(1, Ordering::Relaxed);
}
}
}
pub(crate) fn kv_prefix_reused(&self, tokens: usize) {
if tokens > 0 {
self.kv_prefix_hits.fetch_add(1, Ordering::Relaxed);
}
}
pub(crate) fn kv_read_started(&self) {
self.kv_read_active.store(true, Ordering::Relaxed);
self.kv_read_operations.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn kv_read_bytes(&self, bytes: u64) {
self.kv_read_bytes.fetch_add(bytes, Ordering::Relaxed);
}
pub(crate) fn kv_read_finished(&self, elapsed: Duration, failed: bool) {
self.kv_read_active.store(false, Ordering::Relaxed);
self.last_kv_read_ms
.store(milliseconds(elapsed), Ordering::Relaxed);
if failed {
self.kv_read_errors.fetch_add(1, Ordering::Relaxed);
}
}
pub(crate) fn kv_write_started(&self) {
self.kv_write_active.store(true, Ordering::Relaxed);
self.kv_write_operations.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn kv_write_bytes(&self, bytes: u64) {
self.kv_write_bytes.fetch_add(bytes, Ordering::Relaxed);
}
pub(crate) fn kv_write_finished(&self, elapsed: Duration, failed: bool) {
self.kv_write_active.store(false, Ordering::Relaxed);
self.last_kv_write_ms
.store(milliseconds(elapsed), Ordering::Relaxed);
if failed {
self.kv_write_errors.fetch_add(1, Ordering::Relaxed);
}
}
pub(crate) fn server_stopped(&self, port: u16) {
if self.server_port.load(Ordering::Relaxed) == port {
self.server_listening.store(false, Ordering::Relaxed);
self.http_active.store(0, Ordering::Relaxed);
}
}
pub(crate) fn http_started(&self, path: &str, bytes: usize) {
self.http_active.fetch_add(1, Ordering::Relaxed);
self.http_requests.fetch_add(1, Ordering::Relaxed);
self.http_bytes_received
.fetch_add(bytes as u64, Ordering::Relaxed);
match path {
"/v1/chat/completions" => {
self.http_chat_requests.fetch_add(1, Ordering::Relaxed);
}
path if path == "/v1/models" || path.starts_with("/v1/models/") => {
self.http_model_requests.fetch_add(1, Ordering::Relaxed);
}
_ => {}
}
}
pub(crate) fn http_streaming(&self) {
self.http_streaming_requests.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn http_finished(&self, elapsed: Duration, failed: bool) {
self.http_active.fetch_sub(1, Ordering::Relaxed);
self.http_completed.fetch_add(1, Ordering::Relaxed);
if failed {
self.http_errors.fetch_add(1, Ordering::Relaxed);
}
let elapsed = milliseconds(elapsed);
self.last_http_ms.store(elapsed, Ordering::Relaxed);
self.total_http_ms.fetch_add(elapsed, Ordering::Relaxed);
}
pub(crate) fn snapshot(&self) -> MetricsSnapshot {
let completed = self.completed_requests.load(Ordering::Relaxed);
let failed = self.failed_requests.load(Ordering::Relaxed);
let http_completed = self.http_completed.load(Ordering::Relaxed);
MetricsSnapshot {
uptime_seconds: self.started.elapsed().as_secs(),
phase: phase(self.phase.load(Ordering::Relaxed)),
source: source(self.source.load(Ordering::Relaxed)),
model: model_name(self.model.load(Ordering::Relaxed)),
queue_depth: self.queue_depth.load(Ordering::Relaxed),
runtime_requests: self.runtime_requests.load(Ordering::Relaxed),
local_requests: self.local_requests.load(Ordering::Relaxed),
endpoint_generations: self.endpoint_generations.load(Ordering::Relaxed),
completed_requests: completed,
failed_requests: failed,
model_loads: self.model_loads.load(Ordering::Relaxed),
model_unloads: self.model_unloads.load(Ordering::Relaxed),
model_load_ms: self.model_load_ms.load(Ordering::Relaxed),
model_bytes: self.model_bytes.load(Ordering::Relaxed),
tensor_count: self.tensor_count.load(Ordering::Relaxed),
vocabulary_size: self.vocabulary_size.load(Ordering::Relaxed),
context_used: self.context_used.load(Ordering::Relaxed),
context_limit: self.context_limit.load(Ordering::Relaxed),
decode_tokens_per_second: load_f32(&self.decode_tps),
prefill_tokens_per_second: load_f32(&self.prefill_tps),
last_runtime_ms: self.last_runtime_ms.load(Ordering::Relaxed),
average_runtime_ms: average(
self.total_runtime_ms.load(Ordering::Relaxed),
completed + failed,
),
last_prompt_tokens: self.last_prompt_tokens.load(Ordering::Relaxed),
last_cached_tokens: self.last_cached_tokens.load(Ordering::Relaxed),
last_completion_tokens: self.last_completion_tokens.load(Ordering::Relaxed),
prompt_tokens: self.prompt_tokens.load(Ordering::Relaxed),
cached_tokens: self.cached_tokens.load(Ordering::Relaxed),
completion_tokens: self.completion_tokens.load(Ordering::Relaxed),
kv_lookups: self.kv_lookups.load(Ordering::Relaxed),
kv_hits: self.kv_hits.load(Ordering::Relaxed),
kv_memory_hits: self.kv_memory_hits.load(Ordering::Relaxed),
kv_disk_hits: self.kv_disk_hits.load(Ordering::Relaxed),
kv_misses: self.kv_misses.load(Ordering::Relaxed),
kv_invalid: self.kv_invalid.load(Ordering::Relaxed),
kv_prefix_hits: self.kv_prefix_hits.load(Ordering::Relaxed),
checkpoint_writes: self.checkpoint_writes.load(Ordering::Relaxed),
kv_read_active: self.kv_read_active.load(Ordering::Relaxed),
kv_write_active: self.kv_write_active.load(Ordering::Relaxed),
kv_read_operations: self.kv_read_operations.load(Ordering::Relaxed),
kv_write_operations: self.kv_write_operations.load(Ordering::Relaxed),
kv_read_errors: self.kv_read_errors.load(Ordering::Relaxed),
kv_write_errors: self.kv_write_errors.load(Ordering::Relaxed),
kv_read_bytes: self.kv_read_bytes.load(Ordering::Relaxed),
kv_write_bytes: self.kv_write_bytes.load(Ordering::Relaxed),
last_kv_read_ms: self.last_kv_read_ms.load(Ordering::Relaxed),
last_kv_write_ms: self.last_kv_write_ms.load(Ordering::Relaxed),
kv_files: self.kv_files.load(Ordering::Relaxed),
kv_bytes: self.kv_bytes.load(Ordering::Relaxed),
local_kv_files: self.local_kv_files.load(Ordering::Relaxed),
local_kv_bytes: self.local_kv_bytes.load(Ordering::Relaxed),
http_kv_files: self.http_kv_files.load(Ordering::Relaxed),
http_kv_bytes: self.http_kv_bytes.load(Ordering::Relaxed),
server_listening: self.server_listening.load(Ordering::Relaxed),
server_port: self.server_port.load(Ordering::Relaxed),
http_active: self.http_active.load(Ordering::Relaxed),
http_requests: self.http_requests.load(Ordering::Relaxed),
http_completed,
http_errors: self.http_errors.load(Ordering::Relaxed),
http_chat_requests: self.http_chat_requests.load(Ordering::Relaxed),
http_model_requests: self.http_model_requests.load(Ordering::Relaxed),
http_streaming_requests: self.http_streaming_requests.load(Ordering::Relaxed),
http_bytes_received: self.http_bytes_received.load(Ordering::Relaxed),
last_http_ms: self.last_http_ms.load(Ordering::Relaxed),
average_http_ms: average(self.total_http_ms.load(Ordering::Relaxed), http_completed),
}
}
fn record_checkpoint(&self, source: WorkSource, previous_bytes: Option<u64>, bytes: u64) {
if bytes == 0 {
return;
}
self.checkpoint_writes.fetch_add(1, Ordering::Relaxed);
let previous = previous_bytes.unwrap_or(0);
update_bytes(&self.kv_bytes, previous, bytes);
let (files, total) = match source {
WorkSource::LocalChat => (&self.local_kv_files, &self.local_kv_bytes),
WorkSource::Http => (&self.http_kv_files, &self.http_kv_bytes),
WorkSource::None => return,
};
update_bytes(total, previous, bytes);
if previous_bytes.is_none() {
self.kv_files.fetch_add(1, Ordering::Relaxed);
files.fetch_add(1, Ordering::Relaxed);
}
}
}
fn store_f32(target: &AtomicU32, value: f32) {
target.store(value.max(0.0).to_bits(), Ordering::Relaxed);
}
fn load_f32(source: &AtomicU32) -> f32 {
f32::from_bits(source.load(Ordering::Relaxed))
}
fn milliseconds(duration: Duration) -> u64 {
duration.as_millis().min(u128::from(u64::MAX)) as u64
}
fn average(total: u64, count: u64) -> u64 {
total.checked_div(count).unwrap_or(0)
}
fn update_bytes(target: &AtomicU64, previous: u64, current: u64) {
let total = target.load(Ordering::Relaxed);
target.store(total.saturating_sub(previous) + current, Ordering::Relaxed);
}
fn phase(value: u8) -> RuntimePhase {
match value {
1 => RuntimePhase::Loading,
2 => RuntimePhase::Prefilling,
3 => RuntimePhase::Generating,
4 => RuntimePhase::Ready,
5 => RuntimePhase::Failed,
_ => RuntimePhase::Unloaded,
}
}
fn source(value: u8) -> WorkSource {
match value {
1 => WorkSource::LocalChat,
2 => WorkSource::Http,
_ => WorkSource::None,
}
}
fn model_code(model: ModelChoice) -> u8 {
match model {
ModelChoice::DeepSeekV4Flash => 1,
ModelChoice::DeepSeekV4Pro => 2,
ModelChoice::Glm52 => 3,
}
}
fn model_name(value: u8) -> &'static str {
match value {
1 => "DeepSeek V4 Flash",
2 => "DeepSeek V4 Pro",
3 => "GLM 5.2",
_ => "No model loaded",
}
}
#[derive(Default)]
struct CacheUsage {
files: u64,
bytes: u64,
local_files: u64,
local_bytes: u64,
http_files: u64,
http_bytes: u64,
}
fn cache_usage(root: &Path) -> CacheUsage {
let mut usage = CacheUsage::default();
for (path, http) in [(root.to_path_buf(), false), (root.join("http"), true)] {
let Ok(entries) = fs::read_dir(path) else {
continue;
};
for entry in entries.flatten() {
let Ok(metadata) = entry.metadata() else {
continue;
};
if !metadata.is_file() {
continue;
}
usage.files += 1;
usage.bytes += metadata.len();
if http {
usage.http_files += 1;
usage.http_bytes += metadata.len();
} else {
usage.local_files += 1;
usage.local_bytes += metadata.len();
}
}
}
usage
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn snapshots_track_runtime_and_server_counters() {
let metrics = Metrics::new(Path::new("/path/that/does/not/exist"));
metrics.request_queued(WorkSource::LocalChat);
metrics.request_started(WorkSource::LocalChat);
metrics.prefill_progress(100, 1_000, 250.0);
metrics.generation_progress(120, 1_000, 20.0);
metrics.kv_lookup(KvLookup::DiskHit);
metrics.kv_lookup(KvLookup::MemoryHit);
metrics.kv_lookup(KvLookup::Miss);
metrics.kv_lookup(KvLookup::Invalid);
metrics.kv_prefix_reused(80);
metrics.kv_prefix_reused(0);
metrics.kv_read_started();
metrics.kv_read_bytes(2_048);
metrics.kv_read_finished(Duration::from_millis(10), false);
metrics.kv_write_started();
metrics.kv_write_bytes(4_096);
metrics.kv_write_finished(Duration::from_millis(20), false);
metrics.request_finished(
WorkSource::LocalChat,
Duration::from_millis(250),
100,
80,
20,
None,
4_096,
);
metrics.http_started("/v1/models", 0);
metrics.http_finished(Duration::from_millis(5), false);
let snapshot = metrics.snapshot();
assert_eq!(snapshot.phase, RuntimePhase::Ready);
assert_eq!(snapshot.cached_tokens, 80);
assert_eq!(snapshot.kv_lookups, 4);
assert_eq!(snapshot.kv_hits, 2);
assert_eq!(snapshot.kv_memory_hits, 1);
assert_eq!(snapshot.kv_disk_hits, 1);
assert_eq!(snapshot.kv_misses, 2);
assert_eq!(snapshot.kv_invalid, 1);
assert_eq!(snapshot.kv_prefix_hits, 1);
assert_eq!(snapshot.kv_read_bytes, 2_048);
assert_eq!(snapshot.kv_write_bytes, 4_096);
assert_eq!(snapshot.local_kv_bytes, 4_096);
assert_eq!(snapshot.http_model_requests, 1);
}
}

78
src/runtime.rs
View File

@@ -1,4 +1,5 @@
use crate::engine::{ChatTurn, GenerationOutput, Generator};
use crate::metrics::{Metrics, WorkSource};
use crate::settings::{EngineSettings, TurnSettings};
use std::path::PathBuf;
use std::sync::Arc;
@@ -10,6 +11,7 @@ use std::time::{Duration, Instant};
#[derive(Clone)]
pub(crate) struct GenerationService {
commands: Sender<Command>,
metrics: Arc<Metrics>,
}
pub(crate) struct ActiveGeneration {
@@ -47,13 +49,14 @@ struct Command {
}
impl GenerationService {
pub(crate) fn spawn() -> Result<Self, String> {
pub(crate) fn spawn(metrics: Arc<Metrics>) -> Result<Self, String> {
let (commands, receiver) = mpsc::channel::<Command>();
let worker_metrics = Arc::clone(&metrics);
thread::Builder::new()
.name("model-runtime".into())
.spawn(move || run(receiver))
.spawn(move || run(receiver, worker_metrics))
.map_err(|error| format!("Could not start the model runtime: {error}"))?;
Ok(Self { commands })
Ok(Self { commands, metrics })
}
pub(crate) fn generate(
@@ -64,9 +67,15 @@ impl GenerationService {
checkpoint: CheckpointTarget,
idle_timeout: Duration,
) -> Result<ActiveGeneration, String> {
let source = match &checkpoint {
CheckpointTarget::Local(_) => WorkSource::LocalChat,
CheckpointTarget::Transient(_) => WorkSource::Http,
};
let cancel = Arc::new(AtomicBool::new(false));
let (events, receiver) = mpsc::channel();
self.commands
self.metrics.request_queued(source);
if self
.commands
.send(Command {
engine,
turn,
@@ -76,7 +85,11 @@ impl GenerationService {
cancel: Arc::clone(&cancel),
events,
})
.map_err(|_| "The model runtime stopped unexpectedly.".to_owned())?;
.is_err()
{
self.metrics.request_rejected();
return Err("The model runtime stopped unexpectedly.".to_owned());
}
Ok(ActiveGeneration {
events: receiver,
cancel,
@@ -84,34 +97,70 @@ impl GenerationService {
}
}
fn run(commands: Receiver<Command>) {
fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
let mut loaded = None::<(EngineSettings, Generator)>;
let mut last_used = Instant::now();
let mut idle_timeout = Duration::from_secs(15 * 60);
loop {
match commands.recv_timeout(Duration::from_secs(1)) {
Ok(command) => {
let request_started = Instant::now();
let source = match &command.checkpoint {
CheckpointTarget::Local(_) => WorkSource::LocalChat,
CheckpointTarget::Transient(_) => WorkSource::Http,
};
metrics.request_started(source);
idle_timeout = command.idle_timeout;
if loaded
.as_ref()
.is_none_or(|(settings, _)| settings != &command.engine)
{
if loaded.take().is_some() {
metrics.unloaded();
}
let _ = command.events.send(GenerationEvent::Loading);
loaded = match Generator::open(&command.engine) {
Ok(generator) => Some((command.engine.clone(), generator)),
metrics.loading();
let load_started = Instant::now();
loaded = match Generator::open(&command.engine, Arc::clone(&metrics)) {
Ok(generator) => {
let summary = generator.summary();
metrics.loaded(
summary.model,
load_started.elapsed(),
summary.mapped_bytes,
summary.tensor_count,
summary.vocabulary_size,
);
Some((command.engine.clone(), generator))
}
Err(error) => {
metrics.request_failed(request_started.elapsed());
let _ = command.events.send(GenerationEvent::Finished(Err(error)));
None
}
};
}
if let Some((_, generator)) = &mut loaded {
let mut prefill_started = None::<(Instant, u32)>;
let mut emit = |reasoning, content| {
let _ = command
.events
.send(GenerationEvent::Chunk { reasoning, content });
};
let mut progress = |used, limit, tokens_per_second| {
if let Some(speed) = tokens_per_second {
metrics.generation_progress(used, limit, speed);
} else {
let (started, initial) =
prefill_started.get_or_insert_with(|| (Instant::now(), used));
let elapsed = started.elapsed().as_secs_f32();
let speed = if elapsed > 0.0 {
used.saturating_sub(*initial) as f32 / elapsed
} else {
0.0
};
metrics.prefill_progress(used, limit, speed);
}
let _ = command.events.send(GenerationEvent::Context {
used,
limit,
@@ -136,6 +185,18 @@ fn run(commands: Receiver<Command>) {
&mut progress,
),
};
match &result {
Ok(output) => metrics.request_finished(
source,
request_started.elapsed(),
output.prompt_tokens,
output.cached_tokens,
output.completion_tokens,
output.previous_checkpoint_bytes,
output.checkpoint_bytes,
),
Err(_) => metrics.request_failed(request_started.elapsed()),
}
let _ = command.events.send(GenerationEvent::Finished(result));
last_used = Instant::now();
}
@@ -143,6 +204,7 @@ fn run(commands: Receiver<Command>) {
Err(mpsc::RecvTimeoutError::Timeout) => {
if loaded.is_some() && last_used.elapsed() >= idle_timeout {
loaded = None;
metrics.unloaded();
}
}
Err(mpsc::RecvTimeoutError::Disconnected) => break,

2451
src/server.rs
View File

File diff suppressed because it is too large Load Diff