Recover the model runtime from panics

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(

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<Command>,
@@ -121,97 +123,24 @@ 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(|| {
-                if loaded
+                    run_command(
-                    .as_ref()
+                        command,
-                    .is_none_or(|(settings, _)| settings != &command.engine)
+                        &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);
+                    let _ = events.send(GenerationEvent::Finished(Err(error)));
-                    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();
                 }
             }
             Err(mpsc::RecvTimeoutError::Timeout) => {
@@ -225,6 +154,116 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
     }
 }
 
+#[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)
+    {
+        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<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)]
 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));
+    }
 }