Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -1,231 +1,231 @@
-"""
-Test Florence-2-Large model integration with vGPU.
-Configure PyTorch to use vGPU as device and run image inference.
-"""
-import logging
-import os
-import time
-from contextlib import contextmanager
-from io import BytesIO
-
-import torch
-from torch import nn
-import torch.nn.functional as F
-from PIL import Image
-from transformers import (
-    AutoTokenizer,
-    Florence2ForConditionalGeneration,
-    Florence2Processor
-)
-from virtual_vram import VirtualVRAM
-from http_storage import HTTPGPUStorage
-from torch_vgpu import VGPUDevice, to_vgpu
-
-# Register vGPU device type
-def register_vgpu_device():
-    """Register vGPU as a custom device type"""
-    try:
-        if hasattr(torch.backends, 'register_custom_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")
-    except Exception as e:
-        logger.error(f"vGPU device registration failed: {str(e)}")
-        raise
-
-# Register vGPU device
-register_vgpu_device()
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
-
-@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
-
-def load_image(image_name):
-    """Load and preprocess image from sample_task folder"""
-    try:
-        image_path = os.path.join("sample_task", image_name)
-        if not os.path.exists(image_path):
-            raise FileNotFoundError(f"Image not found: {image_path}")
-            
-        image = Image.open(image_path)
-        # Convert to RGB if needed
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        logger.info(f"Loaded image from {image_path}: size={image.size}")
-        return image
-    except Exception as e:
-        logger.error(f"Image loading failed: {str(e)}")
-        raise
-
-def test_ai_integration_http():
-    """Test Florence-2-Large model on vGPU with image inference"""
-    logger.info("Starting vGPU image inference test")
-    
-    status = {
-        'model_loaded': False,
-        'processor_loaded': False,
-        'model_on_vgpu': False,
-        'image_processed': False,
-        'inference_complete': False,
-        'cleanup_success': False
-    }
-    
-    with gpu_context() as storage:
-        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)
-            device = VGPUDevice(vram=vram)
-            logger.info("vGPU device initialized with HTTP storage backend")
-            
-            # Load Florence model and processor
-            model_name = "microsoft/florence-2-large"
-            logger.info(f"Loading {model_name}")
-            
-            try:
-                processor = Florence2Processor.from_pretrained(
-                    model_name,
-                    trust_remote_code=True
-                )
-                model = Florence2ForConditionalGeneration.from_pretrained(
-                    model_name,
-                    trust_remote_code=True
-                )
-                status['processor_loaded'] = True
-                status['model_loaded'] = True
-                
-                # Log model architecture
-                model_size = get_model_size(model)
-                logger.info(f"Model loaded: {model_size/1e9:.2f} GB in parameters")
-                logger.info(f"Model architecture: {model.__class__.__name__}")
-            except Exception as e:
-                logger.error(f"Model loading failed: {str(e)}")
-                raise
-            
-            # Move model to vGPU with verification
-            try:
-                model = to_vgpu(model, vram=vram)
-                model.eval()
-                status['model_on_vgpu'] = True
-                
-                # Verify model location
-                for param in model.parameters():
-                    if not hasattr(param, 'device') or param.device != device:
-                        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")
-            except Exception as e:
-                logger.error(f"Model transfer to vGPU failed: {str(e)}")
-                raise
-            
-            # Prepare image input from sample_task folder
-            try:
-                # Load image from sample_task directory
-                image_name = "sample1.jpg"  # Replace with your image name
-                image = load_image(image_name)
-                
-                # Process image with Florence processor
-                inputs = processor(images=image, return_tensors="pt")
-                if not inputs or 'pixel_values' not in inputs:
-                    raise ValueError("Invalid processor output")
-                
-                # Move inputs to vGPU
-                inputs = {k: to_vgpu(v, vram=vram) for k, v in inputs.items()}
-                status['image_processed'] = True
-                logger.info(f"Image processed: shape={inputs['pixel_values'].shape}")
-            except Exception as e:
-                logger.error(f"Image preparation failed: {str(e)}")
-                raise
-            
-            # Run image inference with monitoring
-            logger.info("Running image inference...")
-            start = time.time()
-            peak_mem = initial_mem
-            
-            try:
-                with torch.no_grad():
-                    # Get image embeddings
-                    outputs = model(**inputs)
-                    image_features = outputs.last_hidden_state[:, 0]  # Take [CLS] token features
-                    
-                    # Normalize features
-                    image_features = F.normalize(image_features, dim=-1)
-                    
-                    if hasattr(storage, 'get_used_memory'):
-                        peak_mem = max(peak_mem, storage.get_used_memory())
-                    
-                    inference_time = time.time() - start
-                    status['inference_complete'] = True
-                    
-                    # Log performance metrics
-                    logger.info(f"Inference stats:")
-                    logger.info(f"- Time: {inference_time:.4f}s")
-                    logger.info(f"- Memory peak: {(peak_mem - initial_mem)/1e9:.2f} GB")
-                    logger.info(f"- Image features shape: {image_features.shape}")
-                    logger.info(f"- Feature norm: {torch.norm(image_features).item():.4f}")
-                    logger.info(f"- Output device: {image_features.device}")
-                    
-                    # Optionally compute confidence scores
-                    if hasattr(outputs, 'logits'):
-                        logits = outputs.logits
-                        probs = F.softmax(logits, dim=-1)
-                        confidence = torch.max(probs).item()
-                        logger.info(f"- Confidence: {confidence:.4f}")
-                        
-            except Exception as e:
-                logger.error(f"Image inference failed: {str(e)}")
-                raise
-            
-        except Exception as e:
-            logger.error(f"Test failed: {str(e)}")
-            raise
-        finally:
-            # Cleanup and status report
-            try:
-                del model
-                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)}")
-            
-            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()
-
+"""
+Test Florence-2-Large model integration with vGPU.
+Configure PyTorch to use vGPU as device and run image inference.
+"""
+import logging
+import os
+import time
+from contextlib import contextmanager
+from io import BytesIO
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from PIL import Image
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoProcessor
+)
+from virtual_vram import VirtualVRAM
+from http_storage import HTTPGPUStorage
+from torch_vgpu import VGPUDevice, to_vgpu
+
+# Register vGPU device type
+def register_vgpu_device():
+    """Register vGPU as a custom device type"""
+    try:
+        if hasattr(torch.backends, 'register_custom_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")
+    except Exception as e:
+        logger.error(f"vGPU device registration failed: {str(e)}")
+        raise
+
+# Register vGPU device
+register_vgpu_device()
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+@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
+
+def load_image(image_name):
+    """Load and preprocess image from sample_task folder"""
+    try:
+        image_path = os.path.join("sample_task", image_name)
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image not found: {image_path}")
+            
+        image = Image.open(image_path)
+        # Convert to RGB if needed
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        logger.info(f"Loaded image from {image_path}: size={image.size}")
+        return image
+    except Exception as e:
+        logger.error(f"Image loading failed: {str(e)}")
+        raise
+
+def test_ai_integration_http():
+    """Test Florence-2-Large model on vGPU with image inference"""
+    logger.info("Starting vGPU image inference test")
+    
+    status = {
+        'model_loaded': False,
+        'processor_loaded': False,
+        'model_on_vgpu': False,
+        'image_processed': False,
+        'inference_complete': False,
+        'cleanup_success': False
+    }
+    
+    with gpu_context() as storage:
+        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)
+            device = VGPUDevice(vram=vram)
+            logger.info("vGPU device initialized with HTTP storage backend")
+            
+            # Load Florence model and processor
+            model_name = "microsoft/florence-2-large"
+            logger.info(f"Loading {model_name}")
+            
+            try:
+                processor = AutoProcessor.from_pretrained(
+                    model_name,
+                    trust_remote_code=True
+                )
+                model = AutoModelForCausalLM.from_pretrained(
+                    model_name,
+                    trust_remote_code=True
+                )
+                status['processor_loaded'] = True
+                status['model_loaded'] = True
+                
+                # Log model architecture
+                model_size = get_model_size(model)
+                logger.info(f"Model loaded: {model_size/1e9:.2f} GB in parameters")
+                logger.info(f"Model architecture: {model.__class__.__name__}")
+            except Exception as e:
+                logger.error(f"Model loading failed: {str(e)}")
+                raise
+            
+            # Move model to vGPU with verification
+            try:
+                model = to_vgpu(model, vram=vram)
+                model.eval()
+                status['model_on_vgpu'] = True
+                
+                # Verify model location
+                for param in model.parameters():
+                    if not hasattr(param, 'device') or param.device != device:
+                        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")
+            except Exception as e:
+                logger.error(f"Model transfer to vGPU failed: {str(e)}")
+                raise
+            
+            # Prepare image input from sample_task folder
+            try:
+                # Load image from sample_task directory
+                image_name = "sample1.jpg"  # Replace with your image name
+                image = load_image(image_name)
+                
+                # Process image with Florence processor
+                inputs = processor(images=image, return_tensors="pt")
+                if not inputs or 'pixel_values' not in inputs:
+                    raise ValueError("Invalid processor output")
+                
+                # Move inputs to vGPU
+                inputs = {k: to_vgpu(v, vram=vram) for k, v in inputs.items()}
+                status['image_processed'] = True
+                logger.info(f"Image processed: shape={inputs['pixel_values'].shape}")
+            except Exception as e:
+                logger.error(f"Image preparation failed: {str(e)}")
+                raise
+            
+            # Run image inference with monitoring
+            logger.info("Running image inference...")
+            start = time.time()
+            peak_mem = initial_mem
+            
+            try:
+                with torch.no_grad():
+                    # Get image embeddings
+                    outputs = model(**inputs)
+                    image_features = outputs.last_hidden_state[:, 0]  # Take [CLS] token features
+                    
+                    # Normalize features
+                    image_features = F.normalize(image_features, dim=-1)
+                    
+                    if hasattr(storage, 'get_used_memory'):
+                        peak_mem = max(peak_mem, storage.get_used_memory())
+                    
+                    inference_time = time.time() - start
+                    status['inference_complete'] = True
+                    
+                    # Log performance metrics
+                    logger.info(f"Inference stats:")
+                    logger.info(f"- Time: {inference_time:.4f}s")
+                    logger.info(f"- Memory peak: {(peak_mem - initial_mem)/1e9:.2f} GB")
+                    logger.info(f"- Image features shape: {image_features.shape}")
+                    logger.info(f"- Feature norm: {torch.norm(image_features).item():.4f}")
+                    logger.info(f"- Output device: {image_features.device}")
+                    
+                    # Optionally compute confidence scores
+                    if hasattr(outputs, 'logits'):
+                        logits = outputs.logits
+                        probs = F.softmax(logits, dim=-1)
+                        confidence = torch.max(probs).item()
+                        logger.info(f"- Confidence: {confidence:.4f}")
+                        
+            except Exception as e:
+                logger.error(f"Image inference failed: {str(e)}")
+                raise
+            
+        except Exception as e:
+            logger.error(f"Test failed: {str(e)}")
+            raise
+        finally:
+            # Cleanup and status report
+            try:
+                del model
+                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)}")
+            
+            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()
+