defer format! until invalid comparison

mistralrs-core/src/sampler.rs

@@ -1,6 +1,7 @@
 #![allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
 
 use std::{
+    cmp::Ordering,
     collections::{HashMap, HashSet},
     iter::zip,
     sync::{Arc, Mutex},
@@ -470,9 +471,15 @@ impl Sampler {
     ) -> Result<Logprobs> {
         let mut argsort_indices = (0..probs.len()).collect::<Vec<_>>();
         // Sort by descending probability.
-        argsort_indices
-            .sort_unstable_by(|&i, &j| probs[j].partial_cmp(&probs[i]).expect("No ordering."));
 
+        argsort_indices.sort_unstable_by(|&i, &j| match (probs[i].is_nan(), probs[j].is_nan()) {
+            (true, true) => Ordering::Equal,
+            (true, false) => Ordering::Greater,
+            (false, true) => Ordering::Less,
+            _ => probs[j].partial_cmp(&probs[i]).unwrap_or_else(|| {
+                panic!("Incomparable log probs at indices i={}, j={}. Cannot compare probs[i]={} & probs[j]={}", i, j, probs[i], probs[j])
+            })
+        });
         if top_k > 0 {
             // Clamp smaller probabilities to zero.
             for (index, val) in argsort_indices.iter().enumerate() {

mistralrs-core/src/utils/mod.rs

@@ -147,13 +147,14 @@ macro_rules! handle_pipeline_forward_error {
                             usage: group.get_usage(),
                         };
 
-                        seq.responder()
+                        if let Err(e) = seq.responder()
                             .send(Response::ModelError(
                                 e.to_string(),
                                 partial_completion_response
                             ))
-                            .await
-                            .unwrap();
+                            .await {
+                                error!("{} - Failed to send partial chat response after error: {:?}", $stage, &e);
+                            };
                     } else {
                         let partial_completion_response = CompletionResponse {
                             id: seq.id().to_string(),
@@ -165,13 +166,14 @@ macro_rules! handle_pipeline_forward_error {
                             usage: group.get_usage(),
                         };
 
-                        seq.responder()
+                        if let Err(e) = seq.responder()
                             .send(Response::CompletionModelError(
                                 e.to_string(),
                                 partial_completion_response
                             ))
-                            .await
-                            .unwrap();
+                            .await {
+                                error!("{} - Failed to send partial completion response after error: {:?}", $stage, &e);
+                            };
                     }
                 }
                 for seq in $seq_slice.iter_mut() {
@@ -184,7 +186,9 @@ macro_rules! handle_pipeline_forward_error {
                 // - The sequence is gone
                 // - We should reset the state then, including draft.
                 p.set_none_cache($seq_slice, true, true, false);
-                $prefix_cacher.evict_all_to_cpu().unwrap();
+                if let Err(e) = $prefix_cacher.evict_all_to_cpu() {
+                    error!("{} - Failed to evict prefix caches from CPU: {:?}", $stage, &e);
+                }
 
                 continue $label;
             }