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,
@@ -26,11 +26,11 @@ from torch_vgpu import VGPUDevice, to_vgpu
 class VGPUMode(TorchFunctionMode):
     """Custom device mode for vGPU operations"""
     
-    def __init__(self, vram):
+    def __init__(self, vram, device_name="vgpu"):
         self.vram = vram
+        self.device_name = device_name
         self.device = VGPUDevice(vram)
     
-    @torch.override
     def __torch_function__(
         self,
         func: Any,
@@ -41,15 +41,16 @@ class VGPUMode(TorchFunctionMode):
         """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
+        # Handle tensor creation and device placement
+        if func is torch.tensor or 'device' in kwargs:
+            kwargs['device'] = f"{self.device_name}:0"
         
-        # Convert any tensor inputs to vGPU
+        # Handle tensor operations
         new_args = []
         for arg in args:
-            if isinstance(arg, torch.Tensor) and not hasattr(arg, 'device_type'):
-                arg = to_vgpu(arg, self.vram)
+            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)
@@ -61,17 +62,18 @@ class VGPUMode(TorchFunctionMode):
         pass
 
 def register_vgpu_device():
-    """Register vGPU as a custom device type"""
+    """Register vGPU as a custom device type using privateuse1 backend"""
     try:
-        # Initialize vGPU device type if not already registered
-        if not hasattr(torch._C, "_vgpu_device"):
-            torch.backends.register_custom_device("vgpu", VGPUDevice)
-            
-        # Create and enable vGPU mode
+        device_name = "vgpu"
+        
+        # Register device using privateuse1 backend
+        torch._C._dispatch._rename_privateuse1_backend(device_name)
+        
         def init_vgpu_mode(vram):
-            mode = VGPUMode(vram)
+            # Create device mode with the registered device name
+            mode = VGPUMode(vram, device_name)
             torch.set_mode(mode)
-            return mode
+            return mode, torch.device(f"{device_name}:0")
             
         return init_vgpu_mode
 
@@ -149,8 +151,8 @@ def test_ai_integration_http():
             
             # 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")
+            vgpu_mode, vgpu_device = init_vgpu_mode(vram)
+            logger.info(f"vGPU mode initialized with device {vgpu_device}")
             
             # Load Florence model and processor
             model_name = "microsoft/florence-2-large"
@@ -183,12 +185,13 @@ def test_ai_integration_http():
                 status['model_on_vgpu'] = True
                 
                 # Verify model location and device mode
-                for param in model.parameters():
-                    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
-                logger.info(f"Model memory usage: {(current_mem - initial_mem)/1e9:.2f} GB")
+                with vgpu_mode:
+                    for param in model.parameters():
+                        if not str(param.device).startswith('vgpu'):
+                            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
+                    logger.info(f"Model memory usage: {(current_mem - initial_mem)/1e9:.2f} GB")
             except Exception as e:
                 logger.error(f"Model transfer to vGPU failed: {str(e)}")
                 raise