From 8f91a0633654fa82bb4d3008faed2c3b211e39f4 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sat, 25 Jul 2026 13:00:20 +0200 Subject: [PATCH] Recover the model runtime from panics --- src/engine/metal.rs | 30 ++++-- src/runtime.rs | 222 +++++++++++++++++++++++++++----------------- 2 files changed, 159 insertions(+), 93 deletions(-) diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 5b0ac2d..992a4bd 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -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( diff --git a/src/runtime.rs b/src/runtime.rs index 42970de..341d27a 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -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, @@ -121,97 +123,24 @@ fn run(commands: Receiver, metrics: Arc) { 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 loaded - .as_ref() - .is_none_or(|(settings, _)| settings != &command.engine) - { + 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 _ = command.events.send(GenerationEvent::Loading); - 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, - tokens_per_second, - }); - }; - let result = match command.checkpoint { - CheckpointTarget::Local(path) => generator.generate( - &path, - &command.messages, - &command.turn, - &command.cancel, - &mut emit, - &mut progress, - ), - 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( - 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(); + let _ = events.send(GenerationEvent::Finished(Err(error))); } } Err(mpsc::RecvTimeoutError::Timeout) => { @@ -225,6 +154,116 @@ fn run(commands: Receiver, metrics: Arc) { } } +#[allow(clippy::too_many_arguments)] +fn run_command( + command: Command, + loaded: &mut Option<(EngineSettings, Generator)>, + metrics: &Arc, + 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) + { + if loaded.take().is_some() { + metrics.unloaded(); + } + let _ = command.events.send(GenerationEvent::Loading); + 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)) = 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, + tokens_per_second, + }); + }; + let result = match command.checkpoint { + CheckpointTarget::Local(path) => generator.generate( + &path, + &command.messages, + &command.turn, + &command.cancel, + &mut emit, + &mut progress, + ), + 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( + 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(); + } +} + +fn catch_runtime_panic(operation: impl FnOnce() -> T) -> Result { + // 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)] mod tests { use super::*; @@ -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)); + } }