Upload test_ai_integration.py

test_ai_integration.py

@@ -230,39 +230,39 @@ def test_ai_integration():
         if shared_vram is None:
             shared_vram = VirtualVRAM()
         shared_vram.storage = shared_storage
+        
+        for i in range(num_chips):
+            # Configure each chip with shared WebSocket storage
+            chip = Chip(chip_id=i, vram_size_gb=None, storage=shared_storage)
+            chips.append(chip)
             
-            for i in range(num_chips):
-                # Configure each chip with shared WebSocket storage
-                chip = Chip(chip_id=i, vram_size_gb=None, storage=shared_storage)
-                chips.append(chip)
-                
-                # Connect chips in a ring topology
-                if i > 0:
-                    chip.connect_chip(chips[i-1], optical_link)
-                
-                # Initialize AI accelerator with shared resources
-                ai_accelerator = chip.ai_accelerator
-                ai_accelerator.vram = shared_vram
-                ai_accelerator.storage = shared_storage  # Ensure storage is set
-                ai_accelerators.append(ai_accelerator)
-                
-                # Verify WebSocket connection before loading model
-                if not shared_storage.wait_for_connection():
-                    raise RuntimeError(f"Lost WebSocket connection during chip {i} initialization")
-                
-                # Load model weights from WebSocket storage (no CPU transfer)
-                try:
-                    ai_accelerator.load_model(model_id, None, None)  # Model already in WebSocket storage
-                except Exception as e:
-                    print(f"Warning: Failed to load model on chip {i}: {e}")
-                    continue
-                
-                # Track total processing units
-                total_sms += chip.num_sms
-                total_cores += chip.num_sms * chip.cores_per_sm
+            # Connect chips in a ring topology
+            if i > 0:
+                chip.connect_chip(chips[i-1], optical_link)
+            
+            # Initialize AI accelerator with shared resources
+            ai_accelerator = chip.ai_accelerator
+            ai_accelerator.vram = shared_vram
+            ai_accelerator.storage = shared_storage  # Ensure storage is set
+            ai_accelerators.append(ai_accelerator)
+            
+            # Verify WebSocket connection before loading model
+            if not shared_storage.wait_for_connection():
+                raise RuntimeError(f"Lost WebSocket connection during chip {i} initialization")
+            
+            # Load model weights from WebSocket storage (no CPU transfer)
+            try:
+                ai_accelerator.load_model(model_id, None, None)  # Model already in WebSocket storage
+            except Exception as e:
+                print(f"Warning: Failed to load model on chip {i}: {e}")
+                continue
+            
+            # Track total processing units
+            total_sms += chip.num_sms
+            total_cores += chip.num_sms * chip.cores_per_sm
             
             # Store chip configuration in WebSocket storage
-            storage.store_state(f"chips/{i}/config", "state", {
+            shared_storage.store_state(f"chips/{i}/config", "state", {
                 "num_sms": chip.num_sms,
                 "cores_per_sm": chip.cores_per_sm,
                 "total_cores": chip.num_sms * chip.cores_per_sm,