Update utils/hardware/device_manager.py

utils/hardware/device_manager.py

@@ -3,7 +3,12 @@
 Handles device detection, optimization, and hardware compatibility
 """
 
+# CRITICAL: Set OMP_NUM_THREADS before ANY other imports to prevent libgomp error
 import os
+if 'OMP_NUM_THREADS' not in os.environ:
+    os.environ['OMP_NUM_THREADS'] = '4'
+    os.environ['MKL_NUM_THREADS'] = '4'
+
 import sys
 import platform
 import subprocess
@@ -171,6 +176,14 @@ def fix_cuda_compatibility(self):
     def setup_optimal_threading(self):
         """Configure optimal threading for the system"""
         try:
+            # Skip if already configured (to avoid overwriting the early setting)
+            current_omp = os.environ.get('OMP_NUM_THREADS')
+            if current_omp and current_omp.isdigit() and int(current_omp) > 0:
+                logger.info(f"Threading already configured: OMP_NUM_THREADS={current_omp}")
+                # Just ensure PyTorch uses the same settings
+                torch.set_num_threads(int(current_omp))
+                return
+            
             # Get physical CPU count
             physical_cores = psutil.cpu_count(logical=False)
             if physical_cores is None:
@@ -201,8 +214,10 @@ def setup_optimal_threading(self):
         except Exception as e:
             logger.warning(f"Error setting up threading: {e}")
             # Set safe defaults
-            os.environ['OMP_NUM_THREADS'] = '4'
-            os.environ['MKL_NUM_THREADS'] = '4'
+            if 'OMP_NUM_THREADS' not in os.environ:
+                os.environ['OMP_NUM_THREADS'] = '4'
+            if 'MKL_NUM_THREADS' not in os.environ:
+                os.environ['MKL_NUM_THREADS'] = '4'
     
     def get_system_diagnostics(self) -> Dict[str, Any]:
         """Get comprehensive system diagnostics"""
@@ -333,17 +348,11 @@ def get_system_diagnostics() -> Dict[str, Any]:
 
 # Initialize and configure on module import
 if __name__ != "__main__":
-    # When imported, automatically set up threading to avoid the libgomp error
+    # When imported, automatically set up the device manager
     try:
-        # Ensure OMP_NUM_THREADS is set before any OpenMP operations
-        if 'OMP_NUM_THREADS' not in os.environ:
-            # Set a safe default immediately
-            os.environ['OMP_NUM_THREADS'] = '4'
-        
-        # Get the manager instance and configure threading properly
+        # Get the manager instance (threading is already configured at top of file)
         manager = get_device_manager()
+        # Only run setup_optimal_threading if needed (it will check internally)
         manager.setup_optimal_threading()
     except Exception as e:
-        logger.warning(f"Error during device manager initialization: {e}")
-        # Ensure we have safe defaults even if initialization fails
-        os.environ['OMP_NUM_THREADS'] = '4'
+        logger.warning(f"Error during device manager initialization: {e}")