Upload 2 files

http_storage.py

@@ -48,6 +48,29 @@ class LocalStorage:
             
         self.lock = threading.Lock()
         self._closing = False
+        self._connected = True
+        
+        # Initialize monitoring
+        self.resource_monitor = {
+            'vram_used': 0,
+            'active_tensors': 0,
+            'loaded_models': set(),
+            'last_updated': time.time()
+        }
+        
+        # Initialize registries
+        self.model_registry = {}
+        
+        self.initialized = True
+        
+    def is_connected(self) -> bool:
+        """Check if storage is connected (always True for local storage)"""
+        return self._connected and not self._closing and self.ping()
+        
+    def close(self):
+        """Close storage connection"""
+        self._closing = True
+        self._connected = False
         
         # Initialize resource monitoring
         self.resource_monitor = {

test_ai_integration_http.py

@@ -172,9 +172,16 @@ def test_ai_integration_http():
                 
                 if success:
                     print(f"Model '{model_id}' loaded successfully to local storage")
-                    # Verify model is loaded in both accelerator and storage
+                    
+                    # Verify model is loaded in storage
+                    if not ai_accelerator_for_loading.storage.is_model_loaded(model_id):
+                        raise RuntimeError(f"Model {model_id} not found in storage")
+                        
+                    # Verify model is registered in accelerator
                     if not ai_accelerator_for_loading.has_model(model_id):
-                        raise RuntimeError(f"Model {model_id} not found in accelerator registry after loading")
+                        raise RuntimeError(f"Model {model_id} not registered in accelerator")
+                        
+                    print(f"Model {model_id} verified in both storage and accelerator")
                 else:
                     raise RuntimeError("Failed to load model in local storage")
                     
@@ -262,31 +269,14 @@ def test_ai_integration_http():
     ai_accelerators = []
 
     try:
-        # Try to reuse existing connection with verification
-        shared_storage = None
-        max_connection_attempts = 3
-        
-        for attempt in range(max_connection_attempts):
-            try:
-                if (components['storage'] and 
-                    components['storage'].is_connected()):
-                    shared_storage = components['storage']
-                    logging.info("Successfully reused existing HTTP connection")
-                    break
-                else:
-                    logging.warning("Existing connection unavailable, creating new HTTP connection...")
-                    with storage_manager() as new_storage:
-                        if new_storage and new_storage.is_connected():
-                            components['storage'] = new_storage
-                            shared_storage = new_storage
-                            logging.info("Successfully established new HTTP connection")
-                            break
-            except Exception as e:
-                logging.error(f"HTTP connection attempt {attempt + 1} failed: {e}")
-                if attempt < max_connection_attempts - 1:
-                    time.sleep(2)
-                    continue
-                raise RuntimeError(f"Failed to establish HTTP connection after {max_connection_attempts} attempts")
+        # Reuse existing storage instance
+        shared_storage = components['storage']
+        if not shared_storage or not shared_storage.is_connected():
+            logging.warning("Storage not available, creating new instance...")
+            with storage_manager() as new_storage:
+                components['storage'] = new_storage
+                shared_storage = new_storage
+                logging.info("Successfully initialized local storage")
         
         # Initialize high-performance chip array with HTTP storage
         total_sms = 0