Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -1,203 +1,178 @@
 """
-Test Llama-2-7b-instruct model integration with vGPU.
-Configure PyTorch to use vGPU as device for text generation.
+Custom PyTorch device implementation that routes operations through our virtual GPU.
 """
-import logging
-import os
-import time
-from contextlib import contextmanager
-from typing import Any, Optional
-
 import torch
-from transformers import pipeline
+from torch.library import Library, impl
+from typing import Optional, Union, Tuple
+import numpy as np
 from virtual_vram import VirtualVRAM
-from http_storage import HTTPGPUStorage
-from torch_vgpu import VGPUDevice, to_vgpu
 
-def setup_vgpu():
-    """Setup vGPU device"""
+# Initialize custom backend
+def init_vgpu_backend():
     try:
-        # Create and register vGPU device
-        vgpu = VGPUDevice()
-        device = vgpu.device()
+        # First rename the backend
+        torch.utils.rename_privateuse1_backend("vgpu")
         
-        # Set as default device for tensor operations
-        torch.set_default_device(device)
+        # Then generate all the necessary methods
+        torch.utils.generate_methods_for_privateuse1_backend(
+            for_tensor=True,
+            for_module=True,
+            for_packed_sequence=True,
+            for_storage=True
+        )
         
-        return device
-
+        # Register our custom library
+        lib = Library("vgpu", "DEF")
+        lib.define("custom_op(Tensor self) -> Tensor")
+        
+        @impl("vgpu", "custom_op", "Tensor")
+        def custom_op_impl(tensor):
+            return tensor.clone()
+            
+        return True
     except Exception as e:
-        logging.error(f"vGPU setup failed: {str(e)}")
-        raise
+        print(f"Backend initialization warning: {e}")
+        return False
 
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
+# Initialize the backend
+VGPU_BACKEND_INITIALIZED = init_vgpu_backend()
 
-@contextmanager
-def gpu_context():
-    """Context manager for vGPU resources"""
-    storage = None
-    try:
-        storage = HTTPGPUStorage()
-        yield storage
-    finally:
-        if storage:
-            storage.close()
-            logger.info("vGPU resources cleaned up")
-
-def get_model_size(model):
-    """Calculate model size in parameters and memory footprint"""
-    param_size = 0
-    for param in model.parameters():
-        param_size += param.nelement() * param.element_size()
-    buffer_size = 0
-    for buffer in model.buffers():
-        buffer_size += buffer.nelement() * buffer.element_size()
-    return param_size + buffer_size
+class VGPUStorage(torch.Storage):
+    """Custom storage class that uses our virtual VRAM"""
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.vram = kwargs.get('vram')
+        if not self.vram:
+            from virtual_vram import VirtualVRAM
+            self.vram = VirtualVRAM()
+        self.tensor_id = kwargs.get('tensor_id', f"tensor_{id(self)}")
+    
+    def _new_shared(self, size):
+        return VGPUStorage(size, vram=self.vram)
 
-def prepare_prompt(instruction: str) -> str:
-    """Prepare a prompt for Llama-2 using its chat format."""
-    # Format: <s>[INST] instruction [/INST] assistant response </s>[INST] ...
-    return f"<s>[INST] {instruction} [/INST]"
+class VGPUTensor:
+    """Tensor implementation that uses vGPU for computations"""
+    @staticmethod
+    def __new__(cls, elem):
+        return torch.Tensor._make_subclass(cls, elem, elem.requires_grad)
 
-def test_ai_integration_http():
-    """Test GPT OSS model on vGPU with text generation"""
-    logger.info("Starting vGPU text generation test")
-    
-    status = {
-        'pipeline_loaded': False,
-        'model_on_vgpu': False,
-        'generation_complete': False,
-        'cleanup_success': False
-    }
+class VGPUDevice:
+    """
+    Custom PyTorch device implementation that routes operations through vGPU.
+    Usage:
+        vgpu = VGPUDevice()
+        with vgpu.mode():
+            tensor = torch.randn(2, 3)  # Will be on vGPU
+    """
+    _VGPU_INSTANCES = {}  # Class-level dict to track instances
     
-    with gpu_context() as storage:
+    def __init__(self, vram: Optional[VirtualVRAM] = None):
+        self.vram = vram or VirtualVRAM()
+        self.tensor_cores = None  # Will be initialized when needed
+        self.device_name = "vgpu"  # Both internal and user-facing name
+        self._register_device()
+
+    def _register_device(self):
+        """Register vGPU device using PyTorch's device system"""
         try:
-            # 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)
-            
-            # Initialize vGPU device
-            device = setup_vgpu()
-            logger.info(f"vGPU initialized with device {device}")
-            
-            # Load model using pipeline
-            model_id = "openai/gpt-oss-20b"
-            logger.info(f"Loading {model_id}")
-            
-            try:
-                # Disable transformers logging temporarily
-                transformers_logger = logging.getLogger("transformers")
-                original_level = transformers_logger.level
-                transformers_logger.setLevel(logging.ERROR)
+            if not VGPU_BACKEND_INITIALIZED:
+                raise RuntimeError("VGPU backend not properly initialized")
                 
-                try:
-                    # Create pipeline
-                    # Create pipeline with vGPU device
-                    pipe = pipeline(
-                        "text-generation",
-                        model=model_id,
-                        torch_dtype=torch.float32,  # Use full precision
-                        device=device,  # Use our vGPU device
-                        use_safetensors=True,
-                        trust_remote_code=True
-                    )
-                    status['pipeline_loaded'] = True
-                    
-                    # Move pipeline model to vGPU
-                    pipe.model = to_vgpu(pipe.model, vram=vram)
-                    pipe.model.eval()
-                    status['model_on_vgpu'] = True
-                    
-                    # Log model details
-                    logger.info(f"Pipeline created with model: {model_id}")
-                    
-                    # Log model size
-                    model_size = get_model_size(pipe.model)
-                    logger.info(f"Model loaded: {model_size/1e9:.2f} GB in parameters")
-                    logger.info(f"Model architecture: {pipe.model.__class__.__name__}")
-                    
-                    # Verify model location
-                    with torch.device(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")
-                
-                finally:
-                    # Restore original logging level
-                    transformers_logger.setLevel(original_level)
-                    
-            except Exception as e:
-                logger.error(f"Model loading failed: {str(e)}")
-                raise
-            except Exception as e:
-                logger.error(f"Model transfer to vGPU failed: {str(e)}")
-                raise
+            # Create device with explicit index
+            self._device = torch.device("vgpu")
             
-            # Run text generation
-            logger.info("Running text generation...")
-            start = time.time()
-            peak_mem = initial_mem
+            # Store this instance for reuse
+            VGPUDevice._VGPU_INSTANCES[self.device_name] = self
             
-            try:
-                # Prepare input prompt
-                prompt = "Explain how virtual GPUs work in simple terms."
+            # Define custom operations for the device
+            class VGPUAllocator:
+                def __init__(self, vram, device):
+                    self.vram = vram
+                    self.device = device
                 
-                with torch.no_grad():
-                    # Generate text
-                    outputs = pipe(
-                        prompt,
-                        max_new_tokens=256,
-                        temperature=0.7,
-                        top_p=0.95,
-                        top_k=40,
-                        num_beams=1,
-                        do_sample=True,
-                        return_full_text=True
-                    )
-                    
-                    if hasattr(storage, 'get_used_memory'):
-                        peak_mem = max(peak_mem, storage.get_used_memory())
-                    
-                    inference_time = time.time() - start
-                    status['generation_complete'] = True
-                    
-                    # Log performance metrics
-                    logger.info(f"\nGeneration stats:")
-                    logger.info(f"- Time: {inference_time:.4f}s")
-                    logger.info(f"- Memory peak: {(peak_mem - initial_mem)/1e9:.2f} GB")
-                    logger.info(f"- Generated text: {outputs[0]['generated_text']}")
-                        
-            except Exception as e:
-                logger.error(f"Text generation failed: {str(e)}")
-                raise
+                def __call__(self, size, dtype=None, device=None):
+                    # Create tensor on CPU first
+                    cpu_tensor = torch.empty(size, dtype=dtype, device='cpu')
+                    # Move to vGPU storage
+                    return to_vgpu(cpu_tensor, self.vram)
             
-        except Exception as e:
-            logger.error(f"Test failed: {str(e)}")
-            raise
-        finally:
-            # Cleanup and status report
-            try:
-                if 'pipe' in locals():
-                    del pipe
-                if 'outputs' in locals():
-                    del outputs
-                torch.cuda.empty_cache() if hasattr(torch, 'cuda') else None
-                status['cleanup_success'] = True
-            except Exception as e:
-                logger.error(f"Cleanup error: {str(e)}")
+            # Set up allocator
+            self._allocator = VGPUAllocator(self.vram, self._device)
             
-            logger.info("\nTest Summary:")
-            for key, value in status.items():
-                logger.info(f"- {key}: {'✓' if value else '✗'}")
-            
-            final_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0
-            if final_mem > initial_mem:
-                logger.warning(f"Memory leak detected: {(final_mem - initial_mem)/1e6:.2f} MB")
+        except Exception as e:
+            raise RuntimeError(f"Failed to register vGPU device: {str(e)}")
+    
+    @property
+    def type(self):
+        return self.internal_name
+        
+    def __str__(self):
+        return f"{self.internal_name}:0"
+        
+    def __repr__(self):
+        return f"vgpu(device='{self.internal_name}:0')"
+        
+    def device(self):
+        """Get the PyTorch device object that maps to our vGPU"""
+        return self._device  # Return the already created device object
+    
+    def mode(self):
+        """Get a context manager for vGPU operations"""
+        return torch.device(self._device)
+    
+    def _init_tensor_cores(self):
+        if self.tensor_cores is None:
+            from tensor_core import TensorCoreArray
+            self.tensor_cores = TensorCoreArray()
             
-if __name__ == "__main__":
-    test_ai_integration_http()
+    def _to_vram(self, tensor: torch.Tensor) -> str:
+        """Store tensor data in virtual VRAM"""
+        tensor_id = f"tensor_{id(tensor)}"
+        data = tensor.detach().cpu().numpy()
+        self.vram.storage.store_tensor(tensor_id, data)
+        return tensor_id
+        
+    def _from_vram(self, tensor_id: str) -> torch.Tensor:
+        """Retrieve tensor data from virtual VRAM"""
+        data = self.vram.storage.load_tensor(tensor_id)
+        return torch.from_numpy(data)
+
+    def matmul(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """Matrix multiplication using tensor cores"""
+        self._init_tensor_cores()
+        
+        # Store inputs in VRAM
+        a_id = self._to_vram(a)
+        b_id = self._to_vram(b)
+        
+        # Perform matmul using tensor cores
+        result = self.tensor_cores.matmul(
+            self.vram.storage.load_tensor(a_id),
+            self.vram.storage.load_tensor(b_id)
+        )
+        
+        # Create new tensor with result
+        return torch.from_numpy(result)
 
+def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.Tensor:
+    """Move a tensor to vGPU device"""
+    if not isinstance(tensor, torch.Tensor):
+        tensor = torch.tensor(tensor)
+    
+    # Get or create vGPU device
+    if not VGPUDevice._VGPU_INSTANCES:
+        device = VGPUDevice(vram)
+    else:
+        device = next(iter(VGPUDevice._VGPU_INSTANCES.values()))
+        if vram is not None:
+            device.vram = vram
+    
+    # Move data to vRAM
+    tensor_id = device._to_vram(tensor)
+    result = device._from_vram(tensor_id)
+    result.requires_grad = tensor.requires_grad
+    
+    # Set the device using the internal name
+    result.data = result.data.to(device._device)
+    return result