Update ai.py

ai.py

@@ -49,6 +49,75 @@ class AIAccelerator:
         )
         self.tensor_cores_initialized = False
         
+        # Initialize model and tensor tracking
+        self.model_registry: Dict[str, Dict[str, Any]] = {}  # Track loaded models
+        self.tensor_registry: Dict[str, Dict[str, Any]] = {}  # Track tensor metadata
+        self.resource_monitor = {
+            'vram_used': 0,
+            'active_tensors': 0,
+            'loaded_models': set()
+        }
+        
+    def _serialize_model_config(self, config: Any) -> dict:
+        """Convert model config to a serializable format."""
+        if hasattr(config, '__dict__'):
+            # Convert object attributes to dict
+            config_dict = {}
+            for key, value in config.__dict__.items():
+                if isinstance(value, (int, float, str, bool, type(None))):
+                    config_dict[key] = value
+                elif isinstance(value, (list, tuple)):
+                    config_dict[key] = [self._serialize_model_config(item) for item in value]
+                elif isinstance(value, dict):
+                    config_dict[key] = {k: self._serialize_model_config(v) for k, v in value.items()}
+                elif hasattr(value, '__dict__'):
+                    config_dict[key] = self._serialize_model_config(value)
+                else:
+                    config_dict[key] = str(value)  # Fallback to string representation
+            return config_dict
+        elif isinstance(config, (list, tuple)):
+            return [self._serialize_model_config(item) for item in config]
+        elif isinstance(config, dict):
+            return {k: self._serialize_model_config(v) for k, v in config.items()}
+        elif isinstance(config, (int, float, str, bool, type(None))):
+            return config
+        else:
+            return str(config)  # Fallback to string representation
+            
+    def store_model_state(self, model_name: str, model_info: Dict[str, Any]) -> bool:
+        """Store model state in WebSocket storage with proper serialization."""
+        try:
+            # Convert any non-serializable parts of model_info
+            serializable_info = self._serialize_model_config(model_info)
+            
+            # Store in model registry
+            self.model_registry[model_name] = serializable_info
+            
+            # Save to storage
+            if self.storage:
+                # Store model info
+                info_success = self.storage.store_state(
+                    "models",
+                    f"{model_name}/info",
+                    serializable_info
+                )
+                
+                # Store model state
+                state_success = self.storage.store_state(
+                    "models",
+                    f"{model_name}/state",
+                    {"loaded": True, "timestamp": time.time()}
+                )
+                
+                if info_success and state_success:
+                    self.resource_monitor['loaded_models'].add(model_name)
+                    return True
+                    
+            return False
+        except Exception as e:
+            print(f"Error storing model state: {str(e)}")
+            return False
+        
     def initialize_tensor_cores(self):
         """Initialize tensor cores and verify they're ready for computation"""
         if self.tensor_cores_initialized:
@@ -499,41 +568,22 @@ class AIAccelerator:
         total_ops = total_params * batch_size * ops_per_param
         return (total_ops / inference_time) / 1e12  # Convert to TFLOPS
     
-    def load_model(self, model_id: str, model: Any = None, processor: Any = None, model_config: Dict[str, Any] = None):
+    def load_model(self, model_id: str, model: Any, processor: Any):
         """Loads a model directly into WebSocket storage without CPU intermediary."""
         try:
-            # Use provided config or create default
-            config = model_config or {}
-            
-            if model is not None:
-                # Extract model metadata if model is provided
-                model_info = {
-                    "architecture": model.__class__.__name__,
-                    "processor": processor.__class__.__name__ if processor else "None",
-                    "config": model.config.to_dict() if hasattr(model, "config") else {},
-                    "model_config": {
-                        k: str(v) if not isinstance(v, (bool, int, float, str, list, dict)) else v 
-                        for k, v in config.items()
-                    }  # Ensure config is JSON serializable
-                }
-            else:
-                # Use provided config for zero-copy mode
-                model_info = {
-                    "architecture": "ZeroCopy",
-                    "processor": "None",
-                    "config": {},
-                    "model_config": {
-                        k: str(v) if not isinstance(v, (bool, int, float, str, list, dict)) else v 
-                        for k, v in config.items()
-                    }  # Ensure config is JSON serializable
-                }
+            # Extract model metadata
+            model_info = {
+                "architecture": model.__class__.__name__,
+                "processor": processor.__class__.__name__,
+                "config": model.config.to_dict() if hasattr(model, "config") else {}
+            }
             
             # Store model state in WebSocket storage
             self.storage.store_state(f"models/{model_id}", "info", model_info)
             
-            # Map weight tensors directly to WebSocket storage if model is provided
-            model_weights = {}
-            if model is not None and hasattr(model, "state_dict"):
+            # Map weight tensors directly to WebSocket storage
+            if hasattr(model, "state_dict"):
+                model_weights = {}
                 
                 for name, param in model.state_dict().items():
                     tensor_id = f"{model_id}/weights/{name}"
@@ -554,37 +604,12 @@ class AIAccelerator:
                 self.storage.store_state(f"models/{model_id}", "state", model)
                 self.model_registry[model_id] = tensor_id
             
-            # Store processor if provided
-            if processor is not None:
-                self.tokenizer_registry[model_id] = processor
-            
+            self.tokenizer_registry[model_id] = processor
             self.model_loaded = True
             print(f"Model '{model_id}' loaded into WebSocket storage")
-            
-            # Additional setup for zero-copy mode
-            if model_config and model_config.get("zero_copy"):
-                # Register empty tensors for zero-copy mode
-                tensor_id = f"{model_id}/zero_copy"
-                self.model_registry[model_id] = {
-                    "mode": "zero_copy",
-                    "config": model_config,
-                    "tensor_id": tensor_id
-                }
-                
-            return True
-            
         except Exception as e:
             print(f"Error loading model into WebSocket storage: {str(e)}")
-            if "is not JSON serializable" in str(e):
-                print("Attempting to serialize config differently...")
-                try:
-                    # Try again with string conversion for non-serializable types
-                    if model_config:
-                        model_config = {k: str(v) for k, v in model_config.items()}
-                        return self.load_model(model_id, model, processor, model_config)
-                except Exception as e2:
-                    print(f"Second attempt also failed: {e2}")
-            return False
+            raise
 
     def has_model(self, model_id: str) -> bool:
         """Checks if a model is loaded in the accelerator's registry."""