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,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::Transient(directory) | CheckpointTarget::OneShot(directory) => {
-                        | CheckpointTarget::OneShot(directory) => generator.generate_transient(
+                generator.generate_transient(
                     &directory,
                     &command.messages,
                     &command.turn,
                     &command.cancel,
                     &mut emit,
                     &mut progress,
-                        ),
+                )
+            }
         };
         match &result {
             Ok(output) => metrics.request_finished(
@@ -211,20 +254,16 @@ 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();
+        *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,
-        }
     }
 }
 
+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));
+    }
 }