NPU (wip); docker

NN/models/model_interfaces.py

@@ -3,20 +3,64 @@ Model Interfaces Module
 
 Defines abstract base classes and concrete implementations for various model types
 to ensure consistent interaction within the trading system.
+Includes NPU acceleration support for Strix Halo processors.
 """
 
 import logging
-from typing import Dict, Any, Optional, List
+import os
+from typing import Dict, Any, Optional, List, Union
 from abc import ABC, abstractmethod
 import numpy as np
 
+# Try to import NPU acceleration utilities
+try:
+    from utils.npu_acceleration import NPUAcceleratedModel, is_npu_available
+    from utils.npu_detector import get_npu_info
+    HAS_NPU_SUPPORT = True
+except ImportError:
+    HAS_NPU_SUPPORT = False
+    NPUAcceleratedModel = None
+
 logger = logging.getLogger(__name__)
 
 class ModelInterface(ABC):
-    """Base interface for all models"""
+    """Base interface for all models with NPU acceleration support"""
     
-    def __init__(self, name: str):
+    def __init__(self, name: str, enable_npu: bool = True):
         self.name = name
+        self.enable_npu = enable_npu and HAS_NPU_SUPPORT
+        self.npu_model = None
+        self.npu_available = False
+        
+        # Initialize NPU acceleration if available
+        if self.enable_npu:
+            self._setup_npu_acceleration()
+    
+    def _setup_npu_acceleration(self):
+        """Setup NPU acceleration for this model"""
+        try:
+            if HAS_NPU_SUPPORT and is_npu_available():
+                self.npu_available = True
+                logger.info(f"NPU acceleration available for model: {self.name}")
+            else:
+                logger.info(f"NPU acceleration not available for model: {self.name}")
+        except Exception as e:
+            logger.warning(f"Failed to setup NPU acceleration: {e}")
+            self.npu_available = False
+    
+    def get_acceleration_info(self) -> Dict[str, Any]:
+        """Get acceleration information"""
+        info = {
+            'model_name': self.name,
+            'npu_support_available': HAS_NPU_SUPPORT,
+            'npu_enabled': self.enable_npu,
+            'npu_available': self.npu_available
+        }
+        
+        if HAS_NPU_SUPPORT:
+            info.update(get_npu_info())
+        
+        return info
     
     @abstractmethod
     def predict(self, data):
@@ -29,15 +73,39 @@ class ModelInterface(ABC):
         pass
 
 class CNNModelInterface(ModelInterface):
-    """Interface for CNN models"""
+    """Interface for CNN models with NPU acceleration support"""
     
-    def __init__(self, model, name: str):
-        super().__init__(name)
+    def __init__(self, model, name: str, enable_npu: bool = True, input_shape: tuple = None):
+        super().__init__(name, enable_npu)
         self.model = model
+        self.input_shape = input_shape
+        
+        # Setup NPU acceleration for CNN model
+        if self.enable_npu and self.npu_available and input_shape:
+            self._setup_cnn_npu_acceleration()
+    
+    def _setup_cnn_npu_acceleration(self):
+        """Setup NPU acceleration for CNN model"""
+        try:
+            if HAS_NPU_SUPPORT and NPUAcceleratedModel:
+                self.npu_model = NPUAcceleratedModel(
+                    pytorch_model=self.model,
+                    model_name=f"{self.name}_cnn",
+                    input_shape=self.input_shape
+                )
+                logger.info(f"CNN NPU acceleration setup for: {self.name}")
+        except Exception as e:
+            logger.warning(f"Failed to setup CNN NPU acceleration: {e}")
+            self.npu_model = None
     
     def predict(self, data):
-        """Make CNN prediction"""
+        """Make CNN prediction with NPU acceleration if available"""
         try:
+            # Use NPU acceleration if available
+            if self.npu_model and self.npu_available:
+                return self.npu_model.predict(data)
+            
+            # Fallback to original model
             if hasattr(self.model, 'predict'):
                 return self.model.predict(data)
             return None
@@ -47,18 +115,48 @@ class CNNModelInterface(ModelInterface):
     
     def get_memory_usage(self) -> float:
         """Estimate CNN memory usage"""
-        return 50.0  # MB
+        base_memory = 50.0  # MB
+        
+        # Add NPU memory overhead if using NPU acceleration
+        if self.npu_model:
+            base_memory += 25.0  # Additional NPU memory
+        
+        return base_memory
 
 class RLAgentInterface(ModelInterface):
-    """Interface for RL agents"""
+    """Interface for RL agents with NPU acceleration support"""
     
-    def __init__(self, model, name: str):
-        super().__init__(name)
+    def __init__(self, model, name: str, enable_npu: bool = True, input_shape: tuple = None):
+        super().__init__(name, enable_npu)
         self.model = model
+        self.input_shape = input_shape
+        
+        # Setup NPU acceleration for RL model
+        if self.enable_npu and self.npu_available and input_shape:
+            self._setup_rl_npu_acceleration()
+    
+    def _setup_rl_npu_acceleration(self):
+        """Setup NPU acceleration for RL model"""
+        try:
+            if HAS_NPU_SUPPORT and NPUAcceleratedModel:
+                self.npu_model = NPUAcceleratedModel(
+                    pytorch_model=self.model,
+                    model_name=f"{self.name}_rl",
+                    input_shape=self.input_shape
+                )
+                logger.info(f"RL NPU acceleration setup for: {self.name}")
+        except Exception as e:
+            logger.warning(f"Failed to setup RL NPU acceleration: {e}")
+            self.npu_model = None
     
     def predict(self, data):
-        """Make RL prediction"""
+        """Make RL prediction with NPU acceleration if available"""
         try:
+            # Use NPU acceleration if available
+            if self.npu_model and self.npu_available:
+                return self.npu_model.predict(data)
+            
+            # Fallback to original model
             if hasattr(self.model, 'act'):
                 return self.model.act(data)
             elif hasattr(self.model, 'predict'):
@@ -70,7 +168,13 @@ class RLAgentInterface(ModelInterface):
     
     def get_memory_usage(self) -> float:
         """Estimate RL memory usage"""
-        return 25.0  # MB
+        base_memory = 25.0  # MB
+        
+        # Add NPU memory overhead if using NPU acceleration
+        if self.npu_model:
+            base_memory += 15.0  # Additional NPU memory
+        
+        return base_memory
 
 class ExtremaTrainerInterface(ModelInterface):
     """Interface for ExtremaTrainer models, providing context features"""