Switch the local inference provider MLX backend to use the safemlx crate (#10304)

This commit is contained in:
Jasper
2026-07-07 14:01:15 -07:00
committed by GitHub
parent f96f62d985
commit 1bced6616b
58 changed files with 1835 additions and 646 deletions
+40 -14
View File
@@ -53,7 +53,7 @@ use std::collections::{HashMap, HashSet};
use std::io::IsTerminal;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;
use std::time::{Duration, Instant};
use tokio;
use tokio_util::sync::CancellationToken;
use tracing::warn;
@@ -1769,25 +1769,54 @@ fn print_run_stats(
usage: Option<&ProviderUsage>,
) {
let elapsed = run_started.elapsed();
let stats = usage.and_then(|usage| usage.stats.as_ref());
let generation_elapsed = stats
.and_then(|stats| stats.elapsed_ms)
.map(Duration::from_millis);
let output_tokens = usage
.and_then(|usage| usage.usage.output_tokens)
.and_then(|tokens| usize::try_from(tokens).ok())
.or_else(|| usage.and_then(|usage| usage.stats.as_ref()?.output_tokens));
.or_else(|| stats.and_then(|stats| stats.output_tokens));
let tokens_per_second = output_tokens.map(|tokens| {
if elapsed.as_secs_f64() > 0.0 {
tokens as f64 / elapsed.as_secs_f64()
let rate_elapsed = generation_elapsed.unwrap_or(elapsed);
if rate_elapsed.as_secs_f64() > 0.0 {
tokens as f64 / rate_elapsed.as_secs_f64()
} else {
0.0
}
});
let model_load_ms = stats.and_then(|stats| stats.model_load_ms);
let generation_time_to_first_token_ms = stats.and_then(|stats| stats.time_to_first_token_ms);
eprintln!("\nStats:");
match first_token_at {
Some(first) => eprintln!(
" Time to first token: {:.2}s",
first.duration_since(run_started).as_secs_f64()
),
None => eprintln!(" Time to first token: unavailable"),
if let Some(ms) = model_load_ms {
eprintln!(" Model load: {:.2}s", ms as f64 / 1000.0);
}
if model_load_ms.is_some() {
match generation_time_to_first_token_ms {
Some(ms) => eprintln!(
" Generation time to first token: {:.2}s",
ms as f64 / 1000.0
),
None => eprintln!(" Generation time to first token: unavailable"),
}
match first_token_at {
Some(first) => eprintln!(
" End-to-end time to first token: {:.2}s",
first.duration_since(run_started).as_secs_f64()
),
None => eprintln!(" End-to-end time to first token: unavailable"),
}
} else if let Some(ms) = generation_time_to_first_token_ms {
eprintln!(" Time to first token: {:.2}s", ms as f64 / 1000.0);
} else {
match first_token_at {
Some(first) => eprintln!(
" Time to first token: {:.2}s",
first.duration_since(run_started).as_secs_f64()
),
None => eprintln!(" Time to first token: unavailable"),
}
}
match tokens_per_second {
Some(rate) => eprintln!(" Tokens/sec: {:.2}", rate),
@@ -1797,10 +1826,7 @@ fn print_run_stats(
eprintln!(" Output tokens: {tokens}");
}
if let Some(draft) = usage
.and_then(|usage| usage.stats.as_ref())
.and_then(|stats| stats.draft.as_ref())
{
if let Some(draft) = stats.and_then(|stats| stats.draft.as_ref()) {
eprintln!(" Draft accept rate: {:.1}%", draft.accept_rate * 100.0);
eprintln!(
" Draft tokens: {} accepted: {} target verified: {} rounds: {}",
+4 -2
View File
@@ -248,7 +248,8 @@ pub fn render_message(message: &Message, debug: bool) {
}
MessageContent::SystemNotification(notification) => {
match notification.notification_type {
SystemNotificationType::ThinkingMessage => {
SystemNotificationType::ThinkingMessage
| SystemNotificationType::ProgressMessage => {
show_thinking();
set_thinking_message(&notification.msg);
}
@@ -330,7 +331,8 @@ pub fn render_message_streaming(
}
MessageContent::SystemNotification(notification) => {
match notification.notification_type {
SystemNotificationType::ThinkingMessage => {
SystemNotificationType::ThinkingMessage
| SystemNotificationType::ProgressMessage => {
show_thinking();
set_thinking_message(&notification.msg);
}