Update test_ai_integration.py

test_ai_integration.py

@@ -2,6 +2,7 @@
 Test AI integration with WebSocket-based storage and zero CPU memory usage.
 All operations are performed through WebSocket storage with direct tensor core access.
 """
+import asyncio
 from gpu_arch import Chip
 from ai import AIAccelerator
 from virtual_vram import VirtualVRAM
@@ -10,9 +11,16 @@ import numpy as np
 from websocket_storage import WebSocketGPUStorage
 import time
 import os
+import platform
 import contextlib
-import resource
 import atexit
+import logging
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
 
 # Increase system file descriptor limit
 def increase_file_limit():
@@ -23,16 +31,34 @@ def increase_file_limit():
     except Exception as e:
         print(f"Warning: Could not increase file descriptor limit: {e}")
 
-# WebSocket connection manager
+# WebSocket connection manager with retry
 @contextlib.contextmanager
-def websocket_manager():
-    storage = WebSocketGPUStorage()
+def websocket_manager(max_retries=3, retry_delay=2):
+    storage = None
+    for attempt in range(max_retries):
+        try:
+            storage = WebSocketGPUStorage()
+            if storage.wait_for_connection(timeout=10.0):
+                logging.info("Successfully connected to GPU storage server")
+                break
+            else:
+                logging.warning(f"Connection attempt {attempt + 1} failed, retrying...")
+                if storage:
+                    storage.close()
+                time.sleep(retry_delay)
+        except Exception as e:
+            logging.error(f"Connection attempt {attempt + 1} failed with error: {e}")
+            if storage:
+                storage.close()
+            if attempt == max_retries - 1:
+                raise RuntimeError(f"Could not connect to GPU storage server after {max_retries} attempts")
+            time.sleep(retry_delay)
+    
     try:
-        if not storage.wait_for_connection():
-            raise RuntimeError("Could not connect to GPU storage server")
         yield storage
     finally:
-        storage.close()  # Ensure connection is closed
+        if storage:
+            storage.close()  # Ensure connection is closed
         
 # Cleanup handler
 def cleanup_resources():
@@ -288,6 +314,4 @@ def test_ai_integration():
         print(f"Matrix operations test failed: {e}")
         return
 
-    print("\n--- All AI Integration Tests Completed ---")
-
-
+    print("\n--- All AI Integration Tests Completed ---")