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

Recover the model runtime from panics

Browse Source
This commit is contained in:
Georg Bauer
2026-07-25 13:00:20 +02:00
parent 90400701fa
commit 8f91a06336
2 changed files with 159 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
src/engine/metal.rs
View File

@@ -1651,7 +1651,10 @@ fn encode_batch_layer(
"batch raw prefill attention",
)?;
} else if pos == 0 && ratio == 4 && compressed_rows > shape.indexer_top_k as u32 {
let indexer = state.indexer.as_ref().unwrap();
let indexer = state
.indexer
.as_ref()
.ok_or("ratio-4 layer is missing indexer state")?;
call(
unsafe {
ds4_gpu_indexer_scores_prefill_tensor(
@@ -1683,7 +1686,10 @@ fn encode_batch_layer(
},
"batch indexer top-k",
)?;
let compression = state.compression.as_ref().unwrap();
let compression = state
.compression
.as_ref()
.ok_or("compressed layer is missing attention state")?;
call(
unsafe {
ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
@@ -1712,7 +1718,10 @@ fn encode_batch_layer(
"batch indexed prefill attention",
)?;
} else if pos == 0 {
let compression = state.compression.as_ref().unwrap();
let compression = state
.compression
.as_ref()
.ok_or("compressed layer is missing attention state")?;
call(
unsafe {
ds4_gpu_attention_prefill_static_mixed_heads_tensor(
@@ -1759,7 +1768,10 @@ fn encode_batch_layer(
"batch resumed raw attention",
)?;
} else if ratio == 4 && compressed_rows > shape.indexer_top_k as u32 {
let indexer = state.indexer.as_ref().unwrap();
let indexer = state
.indexer
.as_ref()
.ok_or("ratio-4 layer is missing indexer state")?;
call(
unsafe {
ds4_gpu_indexer_scores_decode_batch_tensor(
@@ -1792,7 +1804,10 @@ fn encode_batch_layer(
},
"batch resumed indexer top-k",
)?;
let compression = state.compression.as_ref().unwrap();
let compression = state
.compression
.as_ref()
.ok_or("compressed layer is missing attention state")?;
call(
unsafe {
ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
@@ -1821,7 +1836,10 @@ fn encode_batch_layer(
"batch resumed indexed attention",
)?;
} else {
let compression = state.compression.as_ref().unwrap();
let compression = state
.compression
.as_ref()
.ok_or("compressed layer is missing attention state")?;
call(
unsafe {
ds4_gpu_attention_decode_mixed_batch_heads_tensor(

80
src/runtime.rs
View File

@@ -8,6 +8,8 @@ use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
use std::time::{Duration, Instant};
const RUNTIME_PANIC_ERROR: &str = "The model runtime hit an internal error and was reset.";
#[derive(Clone)]
pub(crate) struct GenerationService {
commands: Sender<Command>,
@@ -121,8 +123,48 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
Ok(command) => {
let request_started = Instant::now();
let source = command.checkpoint.source();
let events = command.events.clone();
metrics.request_started(source);
idle_timeout = command.idle_timeout;
if let Err(error) = catch_runtime_panic(|| {
run_command(
command,
&mut loaded,
&metrics,
request_started,
source,
&mut last_used,
&mut idle_timeout,
);
}) {
metrics.request_failed(request_started.elapsed());
if loaded.take().is_some() {
metrics.unloaded();
}
let _ = events.send(GenerationEvent::Finished(Err(error)));
}
}
Err(mpsc::RecvTimeoutError::Timeout) => {
if loaded.is_some() && last_used.elapsed() >= idle_timeout {
loaded = None;
metrics.unloaded();
}
}
Err(mpsc::RecvTimeoutError::Disconnected) => break,
}
}
}
#[allow(clippy::too_many_arguments)]
fn run_command(
command: Command,
loaded: &mut Option<(EngineSettings, Generator)>,
metrics: &Arc<Metrics>,
request_started: Instant,
source: WorkSource,
last_used: &mut Instant,
idle_timeout: &mut Duration,
) {
*idle_timeout = command.idle_timeout;
if loaded
.as_ref()
.is_none_or(|(settings, _)| settings != &command.engine)
@@ -133,7 +175,7 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
let _ = command.events.send(GenerationEvent::Loading);
metrics.loading();
let load_started = Instant::now();
loaded = match Generator::open(&command.engine, Arc::clone(&metrics)) {
*loaded = match Generator::open(&command.engine, Arc::clone(metrics)) {
Ok(generator) => {
let summary = generator.summary();
metrics.loaded(
@@ -152,7 +194,7 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
}
};
}
if let Some((_, generator)) = &mut loaded {
if let Some((_, generator)) = loaded {
let mut prefill_started = None::<(Instant, u32)>;
let mut emit = |reasoning, content| {
let _ = command
@@ -188,15 +230,16 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
&mut emit,
&mut progress,
),
CheckpointTarget::Transient(directory)
| CheckpointTarget::OneShot(directory) => generator.generate_transient(
CheckpointTarget::Transient(directory) | CheckpointTarget::OneShot(directory) => {
generator.generate_transient(
&directory,
&command.messages,
&command.turn,
&command.cancel,
&mut emit,
&mut progress,
),
)
}
};
match &result {
Ok(output) => metrics.request_finished(
@@ -211,18 +254,14 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
Err(_) => metrics.request_failed(request_started.elapsed()),
}
let _ = command.events.send(GenerationEvent::Finished(result));
last_used = Instant::now();
}
}
Err(mpsc::RecvTimeoutError::Timeout) => {
if loaded.is_some() && last_used.elapsed() >= idle_timeout {
loaded = None;
metrics.unloaded();
}
}
Err(mpsc::RecvTimeoutError::Disconnected) => break,
*last_used = Instant::now();
}
}
fn catch_runtime_panic<T>(operation: impl FnOnce() -> T) -> Result<T, String> {
// Generator owns GPU handles and is not unwind-safe; callers discard it on error.
std::panic::catch_unwind(std::panic::AssertUnwindSafe(operation))
.map_err(|_| RUNTIME_PANIC_ERROR.to_owned())
}
#[cfg(test)]
@@ -241,4 +280,13 @@ mod tests {
assert!(cancel.load(Ordering::Relaxed));
}
#[test]
fn runtime_panic_recovery_accepts_the_next_operation() {
assert_eq!(
catch_runtime_panic(|| panic!("injected runtime panic")),
Err(RUNTIME_PANIC_ERROR.to_owned())
);
assert_eq!(catch_runtime_panic(|| 42), Ok(42));
}
}