fix(neuron,shutdown): time-bound unloads, fast-exit past tokio drain

Two failure modes from the 2026-05-26 beast incident:

1. `unload_all_models` looped through models calling `unload_model`,
   logging individual failures at warn. The cumulative effect was a
   single warn line for the failed unload then "shutdown complete" —
   no signal that the model was actually still loaded. Now each unload
   is bounded by a 20s timeout, failures escalate to error, and a
   summary "leaving N model(s) loaded" line fires when anything is
   stuck so the operator knows the OS will reclaim VRAM after exit.

2. Returning `Ok(())` from `main` after the unload sweep dropped the
   tokio runtime, which then waited indefinitely on a CUDA-stuck
   spawn_blocking thread (the journal's "Stack trace of thread
   2951308" — spinning on `cuCtxGetCurrent`). systemd's TimeoutStopSec
   fired 2 minutes later, SIGABRT, core dump. Replacing the return
   with `std::process::exit(0)` skips the runtime drain and hands the
   OS a clean exit code; stuck threads get reaped with the process.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

crates/neuron/src/harness/candle.rs

@@ -29,9 +29,13 @@ use serde_json::json;
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+#[cfg(feature = "cuda")]
+use std::time::Duration;
 use std::time::{SystemTime, UNIX_EPOCH};
 use tokenizers::Tokenizer;
 use tokio::sync::{Mutex, RwLock, mpsc};
+use tracing::Instrument;
 
 /// In-process candle harness. Owns the loaded model registry.
 pub struct CandleHarness {
@@ -71,6 +75,18 @@ impl LoadedHandle {
             LoadedHandle::Tp(m) => m.devices.clone(),
         }
     }
+
+    /// True if an earlier inference left the device context in an
+    /// unrecoverable state. Surfaced in `/models` so cortex (and an
+    /// operator running `curl beast:13131/models`) can see at a glance
+    /// that the model needs unload+reload.
+    pub fn is_poisoned(&self) -> bool {
+        match self {
+            LoadedHandle::Single(m) => m.poisoned.load(Ordering::Acquire),
+            #[cfg(feature = "cuda")]
+            LoadedHandle::Tp(m) => m.poisoned.load(Ordering::Acquire),
+        }
+    }
 }
 
 /// A loaded model with its tokenizer, device placement, and architecture-
@@ -83,6 +99,15 @@ pub struct LoadedModel {
     pub device: Device,
     pub quant: Option<String>,
     pub devices: Vec<u32>,
+    /// Set to `true` after any forward / kv-cache call fails. A CUDA
+    /// driver error (OOM, illegal address) leaves the device's context
+    /// in an unrecoverable state — subsequent kernels can hang, return
+    /// garbage, or hit another illegal address. The harness refuses
+    /// further inference against a poisoned model and reports a clear
+    /// error so an operator knows to unload+reload to recover. See
+    /// the 2026-05-26 beast incident where a 14k-token prefill OOM
+    /// silently turned every subsequent request into a stuck wait.
+    pub poisoned: AtomicBool,
 }
 
 /// Tensor-parallel loaded model. Holds the leader's rank-0 shard
@@ -102,6 +127,16 @@ pub struct TpLoadedModel {
     /// story.
     pub pool: tokio::sync::Mutex<super::tp::WorkerPool>,
     pub leader_model: Arc<tokio::sync::Mutex<super::tp::TpLeaderModel>>,
+    /// Candle device for rank 0. Mirrors what `leader_model.device()`
+    /// would return, but stored separately so the request path can
+    /// query VRAM without locking the leader (which would contend with
+    /// the in-flight forward).
+    pub leader_device: Device,
+    /// Same poisoning gate as [`LoadedModel::poisoned`]. A TP forward
+    /// failure (CUDA OOM on any rank, NCCL desync, illegal address) is
+    /// terminal: the leader's and workers' CUDA contexts cannot be
+    /// reliably reset without restarting the worker subprocesses.
+    pub poisoned: AtomicBool,
 }
 
 /// Architecture-specific weights. Each variant covers one (family,
@@ -351,6 +386,108 @@ fn resolve_hf_cache(explicit: Option<PathBuf>) -> Option<PathBuf> {
     None
 }
 
+/// Build the InferenceError reported to a client when their request
+/// hits a model that's been marked poisoned by an earlier driver
+/// failure. The message names the model and the recovery procedure so
+/// the operator doesn't have to chase the original failure to know
+/// what to do.
+fn poisoned_error(model_id: &str) -> InferenceError {
+    InferenceError::Other(anyhow::anyhow!(
+        "model '{model_id}' is in a poisoned state \
+         (an earlier inference hit a CUDA driver error and the device \
+          context cannot be safely reused); unload and reload the model \
+          to recover"
+    ))
+}
+
+/// Free/total VRAM on the candle `Device` in MiB. Returns `(0, 0)` if
+/// the query fails or the device is the CPU fallback so logging never
+/// crashes the request path. Mirrors the existing helper in
+/// `tp_qwen3_5.rs`; kept separate to avoid coupling the inference path
+/// to the TP-specific module.
+#[cfg(feature = "cuda")]
+fn device_vram_mb(device: &Device) -> (u64, u64) {
+    use candle_core::cuda::cudarc::driver::result;
+    use candle_core::cuda_backend::WrapErr;
+    let Device::Cuda(dev) = device else {
+        return (0, 0);
+    };
+    let Ok(()) = dev.cuda_stream().context().bind_to_thread().w() else {
+        return (0, 0);
+    };
+    match result::mem_get_info() {
+        Ok((free, total)) => (
+            (free / (1024 * 1024)) as u64,
+            (total / (1024 * 1024)) as u64,
+        ),
+        Err(_) => (0, 0),
+    }
+}
+
+#[cfg(not(feature = "cuda"))]
+#[allow(dead_code)]
+fn device_vram_mb(_device: &Device) -> (u64, u64) {
+    (0, 0)
+}
+
+/// A short hex tag used to group every log line emitted on behalf of
+/// one chat-completion request. Six hex digits is unique enough across
+/// a 4-hour journal window (24 bits ≈ 16M values, while a busy neuron
+/// sees ~10³ requests/hour) and fits cleanly inside `req_id=…` in the
+/// fmt subscriber's span-prefix output.
+fn new_req_id() -> String {
+    format!("{:06x}", unix_subsec_nanos() & 0xFFFFFF)
+}
+
+/// Threshold above which `pool.lock().await` blocking is interesting
+/// enough to warn about. Healthy concurrent requests serialise behind
+/// the pool in single-digit ms — anything past 2 seconds is either a
+/// huge in-flight prompt or, more often, a stuck request holding the
+/// lock against a poisoned CUDA context. See the 2026-05-26 4-hour
+/// silence on beast where dozens of requests piled up invisibly here.
+#[cfg(feature = "cuda")]
+const POOL_LOCK_WARN_THRESHOLD: Duration = Duration::from_secs(2);
+
+/// Acquire the TP pool lock, emitting a warn-level breadcrumb if the
+/// wait exceeds [`POOL_LOCK_WARN_THRESHOLD`]. Wrapped in a helper so
+/// the warn happens at the call site — the request whose lock-wait is
+/// slow is the one that knows its prompt_len and other context.
+#[cfg(feature = "cuda")]
+async fn acquire_pool_lock(
+    pool: &tokio::sync::Mutex<super::tp::WorkerPool>,
+    model_id: &str,
+) -> tokio::sync::MutexGuard<'_, super::tp::WorkerPool> {
+    let start = std::time::Instant::now();
+    // Tick once at the threshold so a stuck request shows up in
+    // journalctl even while it's still waiting. Without this the wait
+    // looks like silence in the log right up until the lock is freed.
+    tokio::pin! {
+        let lock = pool.lock();
+    }
+    loop {
+        tokio::select! {
+            guard = &mut lock => {
+                let elapsed = start.elapsed();
+                if elapsed >= POOL_LOCK_WARN_THRESHOLD {
+                    tracing::warn!(
+                        model = %model_id,
+                        waited_ms = elapsed.as_millis(),
+                        "TP chat_completion: pool lock acquired after long wait"
+                    );
+                }
+                return guard;
+            }
+            _ = tokio::time::sleep(POOL_LOCK_WARN_THRESHOLD) => {
+                tracing::warn!(
+                    model = %model_id,
+                    waited_ms = start.elapsed().as_millis(),
+                    "TP chat_completion: still waiting on pool lock"
+                );
+            }
+        }
+    }
+}
+
 /// Apply the repetition penalty (if any) to the prediction logits and
 /// then sample. Centralises the prefill / generation-loop call sites
 /// so they share identical sampling behaviour.
@@ -746,76 +883,146 @@ impl CandleHarness {
             }
         };
 
-        let prompt = format_qwen3_prompt(&request.messages);
-
-        let encoding = loaded
-            .tokenizer
-            .encode(prompt.as_str(), true)
-            .map_err(|e| InferenceError::Other(anyhow::anyhow!("tokenize: {e}")))?;
-        let prompt_tokens: Vec<u32> = encoding.get_ids().to_vec();
-        let prompt_len = prompt_tokens.len();
-
-        let temperature = request.temperature.unwrap_or(0.7);
-        let top_p = request.top_p;
-        let max_new = request.max_tokens.unwrap_or(512) as usize;
-        let seed = unix_subsec_nanos();
-
-        let eos_id = loaded
-            .tokenizer
-            .token_to_id("<|im_end|>")
-            .or_else(|| loaded.tokenizer.token_to_id("<|endoftext|>"));
-
-        let arch_arc = Arc::clone(&loaded.arch);
-        let device = loaded.device.clone();
+        // Span every line of this request with a short req_id +
+        // model so `grep req_id=…` over the journal can reconstruct
+        // one request even when dozens overlap. Add a terminal log
+        // line on both success and failure — the single-GPU path
+        // used to log nothing on either side, so a failing request
+        // looked exactly like an idle neuron.
+        let req_id = new_req_id();
         let model_id = request.model.clone();
+        let span = tracing::info_span!("chat", req_id = %req_id, model = %model_id);
+        let req_start = std::time::Instant::now();
 
-        let (generated_ids, finish_reason) =
-            tokio::task::spawn_blocking(move || -> Result<(Vec<u32>, String)> {
-                let mut guard = arch_arc.blocking_lock();
-                run_inference(
-                    &mut guard,
-                    &device,
-                    &prompt_tokens,
-                    max_new,
-                    temperature,
-                    top_p,
-                    seed,
-                    eos_id,
-                )
-            })
-            .await
-            .map_err(|e| InferenceError::Other(anyhow::anyhow!("inference task panicked: {e}")))?
-            .map_err(InferenceError::Other)?;
+        // Refuse the request up front if a prior inference poisoned
+        // the device context — otherwise we hand the doomed forward
+        // off to spawn_blocking and stall waiting for CUDA to fail.
+        if loaded.poisoned.load(Ordering::Acquire) {
+            let _g = span.enter();
+            tracing::warn!("chat_completion: refusing request, model poisoned");
+            return Err(poisoned_error(&model_id));
+        }
 
-        let completion_text = loaded
-            .tokenizer
-            .decode(&generated_ids, true)
-            .map_err(|e| InferenceError::Other(anyhow::anyhow!("detokenize: {e}")))?;
+        let result = async {
+            let prompt = format_qwen3_prompt(&request.messages);
 
-        let usage = Usage {
-            prompt_tokens: prompt_len as u64,
-            completion_tokens: generated_ids.len() as u64,
-            total_tokens: (prompt_len + generated_ids.len()) as u64,
-        };
+            let encoding = loaded
+                .tokenizer
+                .encode(prompt.as_str(), true)
+                .map_err(|e| InferenceError::Other(anyhow::anyhow!("tokenize: {e}")))?;
+            let prompt_tokens: Vec<u32> = encoding.get_ids().to_vec();
+            let prompt_len = prompt_tokens.len();
 
-        Ok(ChatCompletionResponse {
-            id: format!("chatcmpl-{:x}", unix_subsec_nanos()),
-            object: "chat.completion".into(),
-            created: unix_now_secs(),
-            model: model_id,
-            choices: vec![ChatCompletionChoice {
-                index: 0,
-                message: ChatMessage {
-                    role: "assistant".into(),
-                    content: MessageContent::Text(completion_text),
+            let temperature = request.temperature.unwrap_or(0.7);
+            let top_p = request.top_p;
+            let max_new = request.max_tokens.unwrap_or(512) as usize;
+            let seed = unix_subsec_nanos();
+
+            let eos_id = loaded
+                .tokenizer
+                .token_to_id("<|im_end|>")
+                .or_else(|| loaded.tokenizer.token_to_id("<|endoftext|>"));
+
+            let (vram_free_mb, vram_total_mb) = device_vram_mb(&loaded.device);
+            tracing::info!(
+                prompt_len,
+                max_new,
+                temperature,
+                ?top_p,
+                ?eos_id,
+                vram_free_mb,
+                vram_total_mb,
+                "chat_completion: starting"
+            );
+
+            let arch_arc = Arc::clone(&loaded.arch);
+            let device = loaded.device.clone();
+
+            let inference_result =
+                tokio::task::spawn_blocking(move || -> Result<(Vec<u32>, String)> {
+                    let mut guard = arch_arc.blocking_lock();
+                    run_inference(
+                        &mut guard,
+                        &device,
+                        &prompt_tokens,
+                        max_new,
+                        temperature,
+                        top_p,
+                        seed,
+                        eos_id,
+                    )
+                })
+                .await;
+
+            // Any failure inside the spawn_blocking touched CUDA via
+            // candle's forward / cache code, so we treat it as a
+            // device-poisoning event. The terminal log at the bottom
+            // of the wrapper reports the error; this flag stops the
+            // NEXT request from going down the same path.
+            let (generated_ids, finish_reason) = match inference_result {
+                Ok(Ok(v)) => v,
+                Ok(Err(e)) => {
+                    loaded.poisoned.store(true, Ordering::Release);
+                    return Err(InferenceError::Other(e));
+                }
+                Err(e) => {
+                    loaded.poisoned.store(true, Ordering::Release);
+                    return Err(InferenceError::Other(anyhow::anyhow!(
+                        "inference task panicked: {e}"
+                    )));
+                }
+            };
+
+            let completion_text = loaded
+                .tokenizer
+                .decode(&generated_ids, true)
+                .map_err(|e| InferenceError::Other(anyhow::anyhow!("detokenize: {e}")))?;
+
+            let usage = Usage {
+                prompt_tokens: prompt_len as u64,
+                completion_tokens: generated_ids.len() as u64,
+                total_tokens: (prompt_len + generated_ids.len()) as u64,
+            };
+
+            tracing::info!(
+                prompt_tokens = prompt_len,
+                completion_tokens = generated_ids.len(),
+                finish_reason = %finish_reason,
+                total_ms = req_start.elapsed().as_millis(),
+                "chat_completion: done"
+            );
+
+            Ok::<_, InferenceError>(ChatCompletionResponse {
+                id: format!("chatcmpl-{:x}", unix_subsec_nanos()),
+                object: "chat.completion".into(),
+                created: unix_now_secs(),
+                model: request.model.clone(),
+                choices: vec![ChatCompletionChoice {
+                    index: 0,
+                    message: ChatMessage {
+                        role: "assistant".into(),
+                        content: MessageContent::Text(completion_text),
+                        extra: serde_json::Value::Object(Default::default()),
+                    },
+                    finish_reason: Some(finish_reason),
                     extra: serde_json::Value::Object(Default::default()),
-                },
-                finish_reason: Some(finish_reason),
+                }],
+                usage: Some(usage),
                 extra: serde_json::Value::Object(Default::default()),
-            }],
-            usage: Some(usage),
-            extra: serde_json::Value::Object(Default::default()),
-        })
+            })
+        }
+        .instrument(span.clone())
+        .await;
+
+        if let Err(ref e) = result {
+            let _g = span.enter();
+            tracing::error!(
+                error = %format!("{e:#}"),
+                total_ms = req_start.elapsed().as_millis(),
+                "chat_completion: failed"
+            );
+        }
+        result
     }
 
     /// Run a streaming chat completion against a loaded model.
@@ -897,15 +1104,48 @@ impl CandleHarness {
             usage: None,
             extra: serde_json::Value::Object(Default::default()),
         };
+        // Refuse if the model is already poisoned. No point opening
+        // an SSE stream just to send the role chunk and then bail.
+        if loaded.poisoned.load(Ordering::Acquire) {
+            return Err(poisoned_error(&model_id));
+        }
+
         // If sending the role chunk fails the receiver is already gone;
         // bail before kicking off the heavy blocking work.
         tx.send(role_chunk)
             .await
             .map_err(|_| InferenceError::Other(anyhow::anyhow!("client disconnected")))?;
 
+        // Span context — spawn_blocking detaches from the async
+        // executor so we capture the span explicitly and re-enter it
+        // inside the closure to keep the req_id on every emitted line.
+        let req_id = new_req_id();
+        let span = tracing::info_span!("chat_stream", req_id = %req_id, model = %model_id);
+        let prompt_len = prompt_tokens.len();
+        let req_start = std::time::Instant::now();
+        // Cloned `Arc<LoadedModel>` so the spawned task can mark the
+        // model poisoned if its forward fails.
+        let loaded_for_task = Arc::clone(&loaded);
+        let span_for_starting = span.clone();
+        let span_for_task = span.clone();
+        {
+            let _g = span_for_starting.enter();
+            let (vram_free_mb, vram_total_mb) = device_vram_mb(&loaded.device);
+            tracing::info!(
+                prompt_len,
+                max_new,
+                temperature,
+                ?top_p,
+                ?eos_id,
+                vram_free_mb,
+                vram_total_mb,
+                "chat_completion (stream): starting"
+            );
+        }
         tokio::task::spawn_blocking(move || {
+            let _g = span_for_task.enter();
             let mut guard = arch_arc.blocking_lock();
-            if let Err(e) = run_inference_streaming(
+            match run_inference_streaming(
                 &mut guard,
                 &device,
                 &tokenizer,
@@ -920,7 +1160,20 @@ impl CandleHarness {
                 &model_id,
                 &tx,
             ) {
-                tracing::warn!(model = %model_id, error = %e, "streaming inference failed");
+                Ok(()) => tracing::info!(
+                    prompt_tokens = prompt_len,
+                    total_ms = req_start.elapsed().as_millis(),
+                    "chat_completion (stream): done"
+                ),
+                Err(e) => {
+                    loaded_for_task.poisoned.store(true, Ordering::Release);
+                    tracing::error!(
+                        error = %format!("{e:#}"),
+                        prompt_tokens = prompt_len,
+                        total_ms = req_start.elapsed().as_millis(),
+                        "chat_completion (stream): failed, model marked poisoned"
+                    );
+                }
             }
         });
 
@@ -949,7 +1202,11 @@ impl Harness for CandleHarness {
             .map(|h| ModelInfo {
                 id: h.model_id().into(),
                 harness: "candle".into(),
-                status: "loaded".into(),
+                status: if h.is_poisoned() {
+                    "poisoned".into()
+                } else {
+                    "loaded".into()
+                },
                 devices: h.devices(),
                 vram_used_mb: None,
             })
@@ -1007,6 +1264,7 @@ impl Harness for CandleHarness {
             device,
             quant: spec.quant.clone(),
             devices,
+            poisoned: AtomicBool::new(false),
         });
 
         let mut models = self.models.write().await;
@@ -1155,6 +1413,8 @@ impl CandleHarness {
             devices: devices.clone(),
             pool: TMutex::new(pool),
             leader_model,
+            leader_device: leader_device.clone(),
+            poisoned: AtomicBool::new(false),
         });
 
         let mut models = self.models.write().await;
@@ -1187,13 +1447,47 @@ impl CandleHarness {
         tp: Arc<TpLoadedModel>,
         request: ChatCompletionRequest,
     ) -> Result<ChatCompletionResponse, InferenceError> {
-        let handle = tokio::spawn(chat_completion_tp_inner(tp, request));
-        match handle.await {
-            Ok(result) => result,
+        // Tag every line of this request with a short req_id so a
+        // grep over journalctl reconstructs one request even when
+        // dozens are queued and interleaved. The span prefix is added
+        // by the fmt subscriber to every event emitted within the
+        // instrumented future, including events from `WorkerPool::*`
+        // since those run on the leader's task.
+        let req_id = new_req_id();
+        let model_id = request.model.clone();
+        let span = tracing::info_span!("tp_chat", req_id = %req_id, model = %model_id);
+        let req_start = std::time::Instant::now();
+
+        if tp.poisoned.load(Ordering::Acquire) {
+            let _g = span.enter();
+            tracing::warn!("TP chat_completion: refusing request, model poisoned");
+            return Err(poisoned_error(&model_id));
+        }
+
+        let tp_for_marker = Arc::clone(&tp);
+        let handle = tokio::spawn(chat_completion_tp_inner(tp, request).instrument(span.clone()));
+        let result = match handle.await {
+            Ok(r) => r,
             Err(join_err) => Err(InferenceError::Other(anyhow::anyhow!(
                 "TP inference task panicked or was cancelled: {join_err}"
             ))),
+        };
+        if let Err(ref e) = result {
+            // Mark poisoned: a failure inside the spawned task either
+            // hit a CUDA/NCCL driver error directly or surfaced as a
+            // task panic. Both cases leave the worker subprocesses in
+            // an unknown state — refuse subsequent requests until an
+            // operator unload+reloads. This is the gate that turned
+            // the 2026-05-26 silent-hang into a clean 5xx.
+            tp_for_marker.poisoned.store(true, Ordering::Release);
+            let _g = span.enter();
+            tracing::error!(
+                error = %format!("{e:#}"),
+                total_ms = req_start.elapsed().as_millis(),
+                "TP chat_completion: failed, model marked poisoned"
+            );
         }
+        result
     }
 
     /// Streaming counterpart to `chat_completion_tp`. Same per-step
@@ -1213,6 +1507,10 @@ impl CandleHarness {
         tp: Arc<TpLoadedModel>,
         request: ChatCompletionRequest,
     ) -> Result<mpsc::Receiver<ChatCompletionChunk>, InferenceError> {
+        if tp.poisoned.load(Ordering::Acquire) {
+            return Err(poisoned_error(&request.model));
+        }
+
         let prompt = format_qwen3_prompt(&request.messages);
         let encoding = tp
             .tokenizer
@@ -1263,143 +1561,193 @@ impl CandleHarness {
         // The orchestration task. Holds the pool lock for the lifetime
         // of this inference; concurrent requests against the same TP
         // model serialise behind it.
+        //
+        // Tagged with the same req_id span as the non-streaming path
+        // so the journal can be reconstructed regardless of which API
+        // surface the client hit.
+        let req_id = new_req_id();
+        let span = tracing::info_span!(
+            "tp_chat_stream",
+            req_id = %req_id,
+            model = %model_id
+        );
+        let req_start = std::time::Instant::now();
+        let (vram_free_mb, vram_total_mb) = device_vram_mb(&tp.leader_device);
+        tracing::info!(
+            parent: &span,
+            prompt_len,
+            max_new,
+            temperature,
+            ?top_p,
+            ?eos_id,
+            vram_free_mb,
+            vram_total_mb,
+            "TP chat_completion (stream): starting"
+        );
         let tp_for_task = Arc::clone(&tp);
-        tokio::spawn(async move {
-            let mut pool = tp_for_task.pool.lock().await;
-            let leader_arc = tp_for_task.leader_model.clone();
+        tokio::spawn(
+            async move {
+                let mut failure: Option<String> = None;
+                let mut pool = acquire_pool_lock(&tp_for_task.pool, &model_id).await;
+                let leader_arc = tp_for_task.leader_model.clone();
 
-            if let Err(e) = pool.clear_kv_cache(&model_id, leader_arc.clone()).await {
-                tracing::warn!(model = %model_id, error = %e, "TP stream: clear_kv_cache failed");
-                return;
-            }
+                let mut all_tokens: Vec<u32> = Vec::new();
+                let mut decoded_prefix = String::new();
+                let mut finish_reason = "length".to_string();
 
-            let mut logits_processor = {
-                let sampling = if temperature <= 0.0 {
-                    Sampling::ArgMax
-                } else {
-                    match top_p {
-                        Some(p) => Sampling::TopP { p, temperature },
-                        None => Sampling::All { temperature },
+                'work: {
+                    if let Err(e) = pool.clear_kv_cache(&model_id, leader_arc.clone()).await {
+                        failure = Some(format!("clear_kv_cache: {e:#}"));
+                        break 'work;
                     }
-                };
-                LogitsProcessor::from_sampling(seed, sampling)
-            };
 
-            let mut all_tokens: Vec<u32> = Vec::new();
-            let mut decoded_prefix = String::new();
-            let mut finish_reason = "length".to_string();
+                    let mut logits_processor = {
+                        let sampling = if temperature <= 0.0 {
+                            Sampling::ArgMax
+                        } else {
+                            match top_p {
+                                Some(p) => Sampling::TopP { p, temperature },
+                                None => Sampling::All { temperature },
+                            }
+                        };
+                        LogitsProcessor::from_sampling(seed, sampling)
+                    };
 
-            // Prefill — every rank embeds the prompt, offset = 0.
-            let logits = match pool
-                .generate_step(&model_id, leader_arc.clone(), prompt_tokens.clone(), 0)
-                .await
-            {
-                Ok(l) => l,
-                Err(e) => {
-                    tracing::warn!(model = %model_id, error = %e, "TP stream: prefill failed");
-                    return;
-                }
-            };
-            let mut next_token = match sample_with_penalty(
-                &logits,
-                &all_tokens,
-                &mut logits_processor,
-            ) {
-                Ok(t) => t,
-                Err(e) => {
-                    tracing::warn!(model = %model_id, error = %e, "TP stream: prefill sample failed");
-                    return;
-                }
-            };
-
-            if Some(next_token) == eos_id {
-                finish_reason = "stop".into();
-            } else {
-                all_tokens.push(next_token);
-                if !emit_chunk(
-                    &all_tokens,
-                    &mut decoded_prefix,
-                    &tokenizer,
-                    &tx,
-                    &id,
-                    created,
-                    &model_id,
-                )
-                .await
-                {
-                    return;
-                }
-
-                for index in 0..max_new.saturating_sub(1) {
+                    // Prefill — every rank embeds the prompt, offset = 0.
                     let logits = match pool
-                        .generate_step(
-                            &model_id,
-                            leader_arc.clone(),
-                            vec![next_token],
-                            prompt_len + index,
-                        )
+                        .generate_step(&model_id, leader_arc.clone(), prompt_tokens.clone(), 0)
                         .await
                     {
                         Ok(l) => l,
                         Err(e) => {
-                            tracing::warn!(
-                                model = %model_id,
-                                error = %e,
-                                "TP stream: decode step failed"
-                            );
-                            return;
+                            failure = Some(format!("prefill: {e:#}"));
+                            break 'work;
                         }
                     };
-                    next_token =
+                    let mut next_token =
                         match sample_with_penalty(&logits, &all_tokens, &mut logits_processor) {
                             Ok(t) => t,
                             Err(e) => {
-                                tracing::warn!(
-                                    model = %model_id,
-                                    error = %e,
-                                    "TP stream: decode sample failed"
-                                );
-                                return;
+                                failure = Some(format!("prefill sample: {e:#}"));
+                                break 'work;
                             }
                         };
+
                     if Some(next_token) == eos_id {
                         finish_reason = "stop".into();
-                        break;
-                    }
-                    all_tokens.push(next_token);
-                    if !emit_chunk(
-                        &all_tokens,
-                        &mut decoded_prefix,
-                        &tokenizer,
-                        &tx,
-                        &id,
-                        created,
-                        &model_id,
-                    )
-                    .await
-                    {
-                        return;
+                    } else {
+                        all_tokens.push(next_token);
+                        if !emit_chunk(
+                            &all_tokens,
+                            &mut decoded_prefix,
+                            &tokenizer,
+                            &tx,
+                            &id,
+                            created,
+                            &model_id,
+                        )
+                        .await
+                        {
+                            // Client gone — treat as normal stream end,
+                            // not a failure. No log spam.
+                            break 'work;
+                        }
+
+                        for index in 0..max_new.saturating_sub(1) {
+                            let logits = match pool
+                                .generate_step(
+                                    &model_id,
+                                    leader_arc.clone(),
+                                    vec![next_token],
+                                    prompt_len + index,
+                                )
+                                .await
+                            {
+                                Ok(l) => l,
+                                Err(e) => {
+                                    failure = Some(format!("decode step {index}: {e:#}"));
+                                    break 'work;
+                                }
+                            };
+                            next_token = match sample_with_penalty(
+                                &logits,
+                                &all_tokens,
+                                &mut logits_processor,
+                            ) {
+                                Ok(t) => t,
+                                Err(e) => {
+                                    failure = Some(format!("decode sample {index}: {e:#}"));
+                                    break 'work;
+                                }
+                            };
+                            if Some(next_token) == eos_id {
+                                finish_reason = "stop".into();
+                                break;
+                            }
+                            all_tokens.push(next_token);
+                            if !emit_chunk(
+                                &all_tokens,
+                                &mut decoded_prefix,
+                                &tokenizer,
+                                &tx,
+                                &id,
+                                created,
+                                &model_id,
+                            )
+                            .await
+                            {
+                                break 'work;
+                            }
+                        }
                     }
                 }
-            }
 
-            // Final chunk carrying finish_reason.
-            let final_chunk = ChatCompletionChunk {
-                id: id.clone(),
-                object: "chat.completion.chunk".into(),
-                created,
-                model: model_id.clone(),
-                choices: vec![ChunkChoice {
-                    index: 0,
-                    delta: serde_json::Value::Object(Default::default()),
-                    finish_reason: Some(finish_reason),
-                    extra: serde_json::Value::Object(Default::default()),
-                }],
-                usage: None,
-                extra: serde_json::Value::Object(Default::default()),
-            };
-            let _ = tx.send(final_chunk).await;
-        });
+                // One terminal line per request, success or failure. The
+                // success branch was previously implicit (the SSE final
+                // chunk went out and the spawned task just ended); now
+                // there's always a log line for the operator.
+                if let Some(err) = &failure {
+                    tp_for_task.poisoned.store(true, Ordering::Release);
+                    tracing::error!(
+                        error = %err,
+                        completion_tokens = all_tokens.len(),
+                        total_ms = req_start.elapsed().as_millis(),
+                        "TP chat_completion (stream): failed, model marked poisoned"
+                    );
+                } else {
+                    tracing::info!(
+                        prompt_tokens = prompt_len,
+                        completion_tokens = all_tokens.len(),
+                        finish_reason = %finish_reason,
+                        total_ms = req_start.elapsed().as_millis(),
+                        "TP chat_completion (stream): done"
+                    );
+                }
+
+                // Final chunk carrying finish_reason — only on the success
+                // path. On failure we drop the channel so the client sees
+                // the SSE stream end abruptly (matches pre-change behaviour
+                // when the failed-path early-returned without final chunk).
+                if failure.is_none() {
+                    let final_chunk = ChatCompletionChunk {
+                        id: id.clone(),
+                        object: "chat.completion.chunk".into(),
+                        created,
+                        model: model_id.clone(),
+                        choices: vec![ChunkChoice {
+                            index: 0,
+                            delta: serde_json::Value::Object(Default::default()),
+                            finish_reason: Some(finish_reason),
+                            extra: serde_json::Value::Object(Default::default()),
+                        }],
+                        usage: None,
+                        extra: serde_json::Value::Object(Default::default()),
+                    };
+                    let _ = tx.send(final_chunk).await;
+                }
+            }
+            .instrument(span),
+        );
 
         Ok(rx)
     }
@@ -1441,6 +1789,7 @@ async fn chat_completion_tp_inner(
         .token_to_id("<|im_end|>")
         .or_else(|| tp.tokenizer.token_to_id("<|endoftext|>"));
 
+    let (vram_free_mb, vram_total_mb) = device_vram_mb(&tp.leader_device);
     tracing::info!(
         model = %model_id,
         prompt_len,
@@ -1448,6 +1797,8 @@ async fn chat_completion_tp_inner(
         temperature,
         ?top_p,
         ?eos_id,
+        vram_free_mb,
+        vram_total_mb,
         "TP chat_completion: starting"
     );
 
@@ -1455,13 +1806,10 @@ async fn chat_completion_tp_inner(
     // leader_model's own Mutex is acquired step-by-step inside
     // pool.generate_step (so spawn_blocking can grab it without
     // holding the pool lock across the blocking_lock call).
-    let lock_start = std::time::Instant::now();
-    let mut pool = tp.pool.lock().await;
-    tracing::debug!(
-        model = %model_id,
-        elapsed_ms = lock_start.elapsed().as_millis(),
-        "TP chat_completion: pool lock acquired"
-    );
+    // `acquire_pool_lock` warns periodically while we wait so a
+    // stuck holder doesn't make the queueing requests look like
+    // silence in the journal.
+    let mut pool = acquire_pool_lock(&tp.pool, &model_id).await;
     let leader_arc = tp.leader_model.clone();
 
     // Reset every rank's KV cache so this request doesn't attend

crates/neuron/src/harness/tp/mod.rs

@@ -656,10 +656,32 @@ impl WorkerPool {
         .await
         .context("leader forward task panicked");
         let leader_ok = matches!(leader_result, Ok(Ok(_)));
+        let leader_ms = leader_start.elapsed().as_millis();
+        // Surface the leader's own error at WARN. Previously this was
+        // silently coerced to `leader_ok=false` while only worker
+        // ranks' errors got logged — when both the leader and a worker
+        // fail together (the typical "CUDA context is now poisoned"
+        // pattern after an OOM), the operator could see only the
+        // worker side and had to guess what hit rank 0.
+        if !leader_ok {
+            let detail = match &leader_result {
+                Ok(Err(e)) => format!("{e:#}"),
+                Err(e) => format!("task: {e:#}"),
+                Ok(Ok(_)) => unreachable!("leader_ok=false implies an error path"),
+            };
+            tracing::warn!(
+                model = %model_id,
+                tokens = tokens_len,
+                offset,
+                leader_ms,
+                error = %detail,
+                "WorkerPool::generate_step: leader forward failed"
+            );
+        }
         tracing::debug!(
             model = %model_id,
             tokens = tokens_len,
-            leader_ms = leader_start.elapsed().as_millis(),
+            leader_ms,
             leader_ok,
             "WorkerPool::generate_step: leader forward returned"
         );

crates/neuron/src/main.rs

@@ -211,6 +211,13 @@ async fn daemon(args: Args) -> Result<()> {
     let registry = state.registry.read().await;
     startup::unload_all_models(&registry).await;
     tracing::info!("shutdown complete");
-
-    Ok(())
+    // Fast-exit instead of returning. Returning lets `#[tokio::main]`
+    // drop the runtime, which in turn waits on the blocking thread
+    // pool to drain. After a CUDA driver error (OOM → illegal address)
+    // a spawn_blocking thread can be wedged inside `cuCtxGetCurrent`,
+    // and tokio's drain has no timeout. systemd then SIGABRTs us and
+    // dumps core. Skipping the drain hands the OS a clean exit code;
+    // the OS reaps the stuck threads. See the 2026-05-26 incident
+    // captured under "Stack trace of thread 2951308" in the journal.
+    std::process::exit(0);
 }

crates/neuron/src/startup.rs

@@ -7,9 +7,17 @@
 
 use crate::harness::HarnessRegistry;
 use cortex_core::harness::ModelSpec;
-use std::time::Instant;
+use std::time::{Duration, Instant};
 use tokio::signal;
 
+/// Maximum time we wait on a single `unload_model` call during
+/// shutdown. The TP unload path tries `Arc::try_unwrap`, which fails
+/// fast when an inference is in flight, so a healthy unload returns
+/// in milliseconds. The timeout exists to bound a *future* unload
+/// path that might genuinely block on a stuck worker, so a single
+/// wedged model can't burn the whole systemd TimeoutStopSec window.
+const UNLOAD_TIMEOUT: Duration = Duration::from_secs(20);
+
 /// Load each spec sequentially against the registry, treating
 /// individual failures as warnings rather than fatal errors.
 ///
@@ -79,19 +87,44 @@ pub async fn unload_all_models(registry: &HarnessRegistry) {
     }
 
     tracing::info!(count = listed.len(), "unloading models for shutdown");
+    let mut stuck = 0;
     for model in listed {
         let start = Instant::now();
-        match registry.unload_model(&model.id).await {
-            Ok(()) => tracing::info!(
+        match tokio::time::timeout(UNLOAD_TIMEOUT, registry.unload_model(&model.id)).await {
+            Ok(Ok(())) => tracing::info!(
                 model = %model.id,
                 elapsed_ms = start.elapsed().as_millis() as u64,
                 "unloaded"
             ),
-            Err(e) => tracing::warn!(
-                model = %model.id,
-                error = %e,
-                "unload failed during shutdown"
-            ),
+            // Most common shape today: TP unload bails because an
+            // inference is still mid-flight (the spawned task holds
+            // an `Arc<TpLoadedModel>` clone). Promoted from warn to
+            // error and tagged with the request-state so the operator
+            // can correlate with the chat_completion logs above.
+            Ok(Err(e)) => {
+                stuck += 1;
+                tracing::error!(
+                    model = %model.id,
+                    error = %e,
+                    elapsed_ms = start.elapsed().as_millis() as u64,
+                    "unload failed during shutdown"
+                );
+            }
+            Err(_) => {
+                stuck += 1;
+                tracing::error!(
+                    model = %model.id,
+                    timeout_secs = UNLOAD_TIMEOUT.as_secs(),
+                    "unload timed out during shutdown, continuing"
+                );
+            }
         }
     }
+    if stuck > 0 {
+        tracing::error!(
+            stuck,
+            "shutdown leaving {stuck} model(s) loaded; VRAM will be \
+             reclaimed by the OS on process exit"
+        );
+    }
 }