gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(ml): recover from mutex poisoning in Rayon quantization with unwrap_or_else

Browse Source 
Replace all 8 `.lock().unwrap()` calls in the Rayon parallel closure of
`quantize_varmap_parallel` with `.lock().unwrap_or_else(|e| e.into_inner())`
so that mutex poisoning (caused by a panicking Rayon thread) does not
cascade and crash all other worker threads — the guard is recovered from
the poisoned state and processing continues safely.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-21 23:26:40 +01:00
parent 36c1a61fc6
commit 322c43db44
1 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
ml/src/tft/varmap_quantization.rs
View File

@@ -341,12 +341,12 @@ pub fn quantize_varmap_parallel(
// Validate and quantize
match validate_tensor_for_quantization(tensor, name) {
Ok(_) => {
let mut quantizer_guard = quantizer.lock().unwrap();
let mut quantizer_guard = quantizer.lock().unwrap_or_else(|e| e.into_inner());
match quantizer_guard.quantize_tensor(tensor, name) {
Ok(quantized) => Some(quantized),
Err(e) => {
warn!("Failed to quantize {}: {}", name, e);
let mut skip_counters = skipped_counter.lock().unwrap();
let mut skip_counters = skipped_counter.lock().unwrap_or_else(|e| e.into_inner());
skip_counters.3 += 1; // error count
None
},
@@ -354,7 +354,7 @@ pub fn quantize_varmap_parallel(
},
Err(e) => {
warn!("Skipping invalid tensor {}: {}", name, e);
let mut skip_counters = skipped_counter.lock().unwrap();
let mut skip_counters = skipped_counter.lock().unwrap_or_else(|e| e.into_inner());
skip_counters.3 += 1; // error count
None
},
@@ -362,26 +362,26 @@ pub fn quantize_varmap_parallel(
},
TensorCategory::Bias => {
debug!("Skipping bias tensor: {}", name);
let mut skip_counters = skipped_counter.lock().unwrap();
let mut skip_counters = skipped_counter.lock().unwrap_or_else(|e| e.into_inner());
skip_counters.0 += 1; // bias count
None
},
TensorCategory::LayerNorm => {
debug!("Skipping LayerNorm tensor: {}", name);
let mut skip_counters = skipped_counter.lock().unwrap();
let mut skip_counters = skipped_counter.lock().unwrap_or_else(|e| e.into_inner());
skip_counters.1 += 1; // layernorm count
None
},
TensorCategory::Small => {
debug!("Skipping small tensor: {}", name);
let mut skip_counters = skipped_counter.lock().unwrap();
let mut skip_counters = skipped_counter.lock().unwrap_or_else(|e| e.into_inner());
skip_counters.2 += 1; // small count
None
},
};
// Update progress counter
let mut counter = progress_counter.lock().unwrap();
let mut counter = progress_counter.lock().unwrap_or_else(|e| e.into_inner());
*counter += 1;
let current_progress = *counter;
drop(counter);
@@ -419,7 +419,7 @@ pub fn quantize_varmap_parallel(
let elapsed = start_time.elapsed().as_secs_f32();
let rate = total_tensors as f32 / elapsed;
let skip_counters = skipped_counter.lock().unwrap();
let skip_counters = skipped_counter.lock().unwrap_or_else(|e| e.into_inner());
let (bias_count, layernorm_count, small_count, error_count) = *skip_counters;
info!(