Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -12,7 +12,7 @@ from typing import Any, Optional
 import torch
 from torch import nn
 import torch.nn.functional as F
-from torch.overrides import TorchFunctionMode
+from torch.utils._python_dispatch import TorchFunctionMode
 from PIL import Image
 from transformers import (
     AutoTokenizer,
@@ -24,13 +24,16 @@ from http_storage import HTTPGPUStorage
 from torch_vgpu import VGPUDevice, to_vgpu
 
 def setup_vgpu():
-    """Setup vGPU device and mode"""
+    """Setup vGPU device"""
     try:
-        # Create device and get its mode
-        device = VGPUDevice()
-        mode = device.mode()
+        # Create and register vGPU device
+        vgpu = VGPUDevice()
+        device = vgpu.device()
         
-        return mode, device
+        # Set as default device for tensor operations
+        torch.set_default_device(device)
+        
+        return device
 
     except Exception as e:
         logging.error(f"vGPU setup failed: {str(e)}")
@@ -101,9 +104,9 @@ def test_ai_integration_http():
             initial_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0
             vram = VirtualVRAM(size_gb=None, storage=storage)
             
-            # Initialize vGPU device and mode
-            vgpu_mode, vgpu_device = setup_vgpu()
-            logger.info(f"vGPU initialized with device {vgpu_device}")
+            # Initialize vGPU device
+            device = setup_vgpu()
+            logger.info(f"vGPU initialized with device {device}")
             
             # Load Florence model and processor
             model_name = "microsoft/florence-2-large"
@@ -135,10 +138,10 @@ def test_ai_integration_http():
                 model.eval()
                 status['model_on_vgpu'] = True
                 
-                # Verify model location and device mode
-                with vgpu_mode:
+                # Verify model location
+                with torch.device(device):
                     for param in model.parameters():
-                        if not str(param.device).startswith('vgpu'):
+                        if param.device != device:
                             raise RuntimeError(f"Model parameter not on vGPU device. Found device: {param.device}")
                     
                     current_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0