Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -7,10 +7,12 @@ import os
 import time
 from contextlib import contextmanager
 from io import BytesIO
+from typing import Any, Optional
 
 import torch
 from torch import nn
 import torch.nn.functional as F
+from torch.utils._python_dispatch import TorchFunctionMode
 from PIL import Image
 from transformers import (
     AutoTokenizer,
@@ -21,19 +23,60 @@ from virtual_vram import VirtualVRAM
 from http_storage import HTTPGPUStorage
 from torch_vgpu import VGPUDevice, to_vgpu
 
-# Register vGPU device type
+class VGPUMode(TorchFunctionMode):
+    """Custom device mode for vGPU operations"""
+    
+    def __init__(self, vram):
+        self.vram = vram
+        self.device = VGPUDevice(vram)
+    
+    @torch.override
+    def __torch_function__(
+        self,
+        func: Any,
+        types: Any,
+        args: Any = (),
+        kwargs: Optional[dict] = None
+    ):
+        """Override torch functions to handle vGPU device operations"""
+        kwargs = kwargs or {}
+        
+        # Handle device placement
+        if 'device' in kwargs and kwargs['device'] == 'vgpu':
+            kwargs['device'] = self.device
+        
+        # Convert any tensor inputs to vGPU
+        new_args = []
+        for arg in args:
+            if isinstance(arg, torch.Tensor) and not hasattr(arg, 'device_type'):
+                arg = to_vgpu(arg, self.vram)
+            new_args.append(arg)
+        
+        return func(*new_args, **kwargs)
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
 def register_vgpu_device():
     """Register vGPU as a custom device type"""
     try:
-        if hasattr(torch.backends, 'register_custom_device'):
+        # Initialize vGPU device type if not already registered
+        if not hasattr(torch._C, "_vgpu_device"):
             torch.backends.register_custom_device("vgpu", VGPUDevice)
-        else:
-            # Fallback: Add device type to torch._C
-            if not hasattr(torch._C, "_vgpu_device"):
-                torch._C._vgpu_device = VGPUDevice
-            logger.info("Using fallback vGPU device registration")
+            
+        # Create and enable vGPU mode
+        def init_vgpu_mode(vram):
+            mode = VGPUMode(vram)
+            torch.set_mode(mode)
+            return mode
+            
+        return init_vgpu_mode
+
     except Exception as e:
-        logger.error(f"vGPU device registration failed: {str(e)}")
+        logging.error(f"vGPU device registration failed: {str(e)}")
         raise
 
 # Register vGPU device
@@ -103,8 +146,11 @@ def test_ai_integration_http():
             # Initialize vRAM with monitoring
             initial_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0
             vram = VirtualVRAM(size_gb=None, storage=storage)
-            device = VGPUDevice(vram=vram)
-            logger.info("vGPU device initialized with HTTP storage backend")
+            
+            # Initialize vGPU mode and register device
+            init_vgpu_mode = register_vgpu_device()
+            vgpu_mode = init_vgpu_mode(vram)
+            logger.info("vGPU mode initialized with HTTP storage backend")
             
             # Load Florence model and processor
             model_name = "microsoft/florence-2-large"
@@ -136,9 +182,9 @@ def test_ai_integration_http():
                 model.eval()
                 status['model_on_vgpu'] = True
                 
-                # Verify model location
+                # Verify model location and device mode
                 for param in model.parameters():
-                    if not hasattr(param, 'device') or param.device != device:
+                    if not hasattr(param, 'device') or not isinstance(param.device, VGPUDevice):
                         raise RuntimeError("Model not properly moved to vGPU")
                 
                 current_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0