From e36698ef1494ca2f5233f37f1678b164d467f333 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Mon, 23 Feb 2026 10:06:40 +0100
Subject: [PATCH] fix(ml): GPU OOM detection with automatic CPU fallback in
 inference engine

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 ml/src/inference.rs | 40 ++++++++++++++++++++++++++++++----------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/ml/src/inference.rs b/ml/src/inference.rs
index 2a6ab690a..21b681ceb 100644
--- a/ml/src/inference.rs
+++ b/ml/src/inference.rs
@@ -544,17 +544,37 @@ impl RealMLInferenceEngine {
         let device = match self.config.device_preference.as_str() {
             "cuda" | "gpu" => match Device::new_cuda(0) {
                 Ok(cuda_device) => {
-                    info!("✅ Using CUDA device for model: {}", model_id);
-                    cuda_device
-                },
+                    match crate::memory_optimization::auto_batch_size::detect_gpu_memory() {
+                        Ok((_, free_mb, _)) if free_mb > 500.0 => {
+                            info!(
+                                "Using CUDA device for model: {} (free VRAM: {:.0}MB)",
+                                model_id, free_mb
+                            );
+                            cuda_device
+                        }
+                        Ok((_, free_mb, _)) => {
+                            warn!(
+                                "GPU VRAM too low ({:.0}MB free), falling back to CPU for model: {}",
+                                free_mb, model_id
+                            );
+                            Device::Cpu
+                        }
+                        Err(e) => {
+                            warn!(
+                                "Cannot detect GPU memory ({}), falling back to CPU for model: {}",
+                                e, model_id
+                            );
+                            Device::Cpu
+                        }
+                    }
+                }
                 Err(e) => {
-                    return Err(MLSafetyError::from(RealInferenceError::GpuRequired {
-                        reason: format!(
-                            "GPU acceleration required for production model {}: {}",
-                            model_id, e
-                        ),
-                    }));
-                },
+                    warn!(
+                        "CUDA not available ({}), falling back to CPU for model: {}",
+                        e, model_id
+                    );
+                    Device::Cpu
+                }
             },
             _ => {
                 info!("Using CPU device for model: {}", model_id);