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,
@@ -23,67 +23,19 @@ from virtual_vram import VirtualVRAM
 from http_storage import HTTPGPUStorage
 from torch_vgpu import VGPUDevice, to_vgpu
 
-class VGPUMode(TorchFunctionMode):
-    """Custom device mode for vGPU operations"""
-    
-    def __init__(self, vram, device_name="vgpu"):
-        self.vram = vram
-        self.device_name = device_name
-        self.device = VGPUDevice(vram)
-    
-    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 tensor creation and device placement
-        if func is torch.tensor or 'device' in kwargs:
-            kwargs['device'] = f"{self.device_name}:0"
-        
-        # Handle tensor operations
-        new_args = []
-        for arg in args:
-            if isinstance(arg, torch.Tensor):
-                if not hasattr(arg, 'device') or not str(arg.device).startswith(self.device_name):
-                    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 using privateuse1 backend"""
+def setup_vgpu():
+    """Setup vGPU device and mode"""
     try:
-        device_name = "vgpu"
-        
-        # Register device using privateuse1 backend
-        torch._C._dispatch._rename_privateuse1_backend(device_name)
+        # Create device and get its mode
+        device = VGPUDevice()
+        mode = device.mode()
         
-        def init_vgpu_mode(vram):
-            # Create device mode with the registered device name
-            mode = VGPUMode(vram, device_name)
-            torch.set_mode(mode)
-            return mode, torch.device(f"{device_name}:0")
-            
-        return init_vgpu_mode
+        return mode, device
 
     except Exception as e:
-        logging.error(f"vGPU device registration failed: {str(e)}")
+        logging.error(f"vGPU setup failed: {str(e)}")
         raise
 
-# Register vGPU device
-register_vgpu_device()
-
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
@@ -149,10 +101,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 mode and register device
-            init_vgpu_mode = register_vgpu_device()
-            vgpu_mode, vgpu_device = init_vgpu_mode(vram)
-            logger.info(f"vGPU mode initialized with device {vgpu_device}")
+            # Initialize vGPU device and mode
+            vgpu_mode, vgpu_device = setup_vgpu()
+            logger.info(f"vGPU initialized with device {vgpu_device}")
             
             # Load Florence model and processor
             model_name = "microsoft/florence-2-large"