Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -1,178 +1,203 @@
 """
-Custom PyTorch device implementation that routes operations through our virtual GPU.
+Test Llama-2-7b-instruct model integration with vGPU.
+Configure PyTorch to use vGPU as device for text generation.
 """
+import logging
+import os
+import time
+from contextlib import contextmanager
+from typing import Any, Optional
+
 import torch
-from torch.library import Library, impl
-from typing import Optional, Union, Tuple
-import numpy as np
+from transformers import pipeline
 from virtual_vram import VirtualVRAM
+from http_storage import HTTPGPUStorage
+from torch_vgpu import VGPUDevice, to_vgpu
 
-# Initialize custom backend
-def init_vgpu_backend():
+def setup_vgpu():
+    """Setup vGPU device"""
     try:
-        # First rename the backend
-        torch.utils.rename_privateuse1_backend("vgpu")
-        
-        # 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
-        )
+        # Create and register vGPU device
+        vgpu = VGPUDevice()
+        device = vgpu.device()
         
-        # Register our custom library
-        lib = Library("vgpu", "DEF")
-        lib.define("custom_op(Tensor self) -> Tensor")
+        # Set as default device for tensor operations
+        torch.set_default_device(device)
         
-        @impl("vgpu", "custom_op", "Tensor")
-        def custom_op_impl(tensor):
-            return tensor.clone()
-            
-        return True
+        return device
+
     except Exception as e:
-        print(f"Backend initialization warning: {e}")
-        return False
+        logging.error(f"vGPU setup failed: {str(e)}")
+        raise
 
-# Initialize the backend
-VGPU_BACKEND_INITIALIZED = init_vgpu_backend()
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
 
-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)
+@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")
 
-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 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 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
-    
-    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 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]"
 
-    def _register_device(self):
-        """Register vGPU device using PyTorch's device system"""
+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
+    }
+    
+    with gpu_context() as storage:
         try:
-            if not VGPU_BACKEND_INITIALIZED:
-                raise RuntimeError("VGPU backend not properly initialized")
-                
-            # Create device with explicit index
-            self._device = torch.device("vgpu")
+            # 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)
             
-            # Store this instance for reuse
-            VGPUDevice._VGPU_INSTANCES[self.device_name] = self
+            # Initialize vGPU device
+            device = setup_vgpu()
+            logger.info(f"vGPU initialized with device {device}")
             
-            # Define custom operations for the device
-            class VGPUAllocator:
-                def __init__(self, vram, device):
-                    self.vram = vram
-                    self.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)
+                
+                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")
                 
-                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)
+                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
             
-            # Set up allocator
-            self._allocator = VGPUAllocator(self.vram, self._device)
+            # Run text generation
+            logger.info("Running text generation...")
+            start = time.time()
+            peak_mem = initial_mem
+            
+            try:
+                # Prepare input prompt
+                prompt = "Explain how virtual GPUs work in simple terms."
+                
+                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
             
         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()
+            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)}")
             
-    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)
+            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")
+            
+if __name__ == "__main__":
+    test_ai_integration_http()
 
-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