Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -9,12 +9,7 @@ from contextlib import contextmanager
 from typing import Any, Optional
 
 import torch
-import torch.nn.functional as F
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    TextStreamer
-)
+from transformers import pipeline
 from virtual_vram import VirtualVRAM
 from http_storage import HTTPGPUStorage
 from torch_vgpu import VGPUDevice, to_vgpu
@@ -70,12 +65,11 @@ def prepare_prompt(instruction: str) -> str:
     return f"<s>[INST] {instruction} [/INST]"
 
 def test_ai_integration_http():
-    """Test Llama-2-7b-instruct model on vGPU with text generation"""
+    """Test GPT OSS model on vGPU with text generation"""
     logger.info("Starting vGPU text generation test")
     
     status = {
-        'model_loaded': False,
-        'tokenizer_loaded': False,
+        'pipeline_loaded': False,
         'model_on_vgpu': False,
         'generation_complete': False,
         'cleanup_success': False
@@ -91,9 +85,9 @@ def test_ai_integration_http():
             device = setup_vgpu()
             logger.info(f"vGPU initialized with device {device}")
             
-            # Load Llama model and tokenizer
-            model_name = "meta-llama/Llama-2-7b-chat-hf"
-            logger.info(f"Loading {model_name}")
+            # Load model using pipeline
+            model_id = "openai/gpt-oss-20b"
+            logger.info(f"Loading {model_id}")
             
             try:
                 # Disable transformers logging temporarily
@@ -102,32 +96,26 @@ def test_ai_integration_http():
                 transformers_logger.setLevel(logging.ERROR)
                 
                 try:
-                    # Load tokenizer first
-                    tokenizer = AutoTokenizer.from_pretrained(
-                        model_name,
-                        trust_remote_code=True,
-                        use_fast=True
+                    # Create pipeline
+                    pipe = pipeline(
+                        "text-generation",
+                        model=model_id,
+                        torch_dtype="auto",
+                        device=device  # Use our vGPU device
                     )
-                    status['tokenizer_loaded'] = True
+                    status['pipeline_loaded'] = True
                     
-                    # Load model with full precision
-                    model = AutoModelForCausalLM.from_pretrained(
-                        model_name,
-                        trust_remote_code=True,
-                        torch_dtype=torch.float32,  # Use full precision
-                        device_map=None,  # Don't auto-map devices
-                        use_safetensors=True
-                    )
-                    status['model_loaded'] = True
+                    # Move pipeline to vGPU
+                    pipe.model = to_vgpu(pipe.model, vram=vram)
+                    status['model_on_vgpu'] = True
                     
                     # Log model details
-                    logger.info(f"Tokenizer type: {type(tokenizer).__name__}")
-                    logger.info(f"Model type: {type(model).__name__}")
+                    logger.info(f"Pipeline created with model: {model_id}")
                     
-                    # Log model architecture
-                    model_size = get_model_size(model)
+                    # 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: {model.__class__.__name__}")
+                    logger.info(f"Model architecture: {pipe.model.__class__.__name__}")
                 
                 finally:
                     # Restore original logging level
@@ -137,52 +125,25 @@ def test_ai_integration_http():
                 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
-                with torch.device(device):
-                    for param in model.parameters():
-                        if param.device != device:
-                            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
-            
             # Run text generation
             logger.info("Running text generation...")
             start = time.time()
             peak_mem = initial_mem
             
             try:
-                # Prepare input prompt
-                instruction = "Explain how virtual GPUs work in simple terms."
-                prompt = prepare_prompt(instruction)
-                
-                # Tokenize input
-                inputs = tokenizer(prompt, return_tensors="pt")
-                inputs = {k: to_vgpu(v, vram=vram) for k, v in inputs.items()}
-                
-                # Set up streamer for token-by-token output
-                streamer = TextStreamer(tokenizer)
+                # Prepare messages
+                messages = [
+                    {"role": "user", "content": "Explain how virtual GPUs work in simple terms."}
+                ]
                 
                 with torch.no_grad():
                     # Generate text
-                    outputs = model.generate(
-                        **inputs,
-                        max_length=512,
+                    outputs = pipe(
+                        messages,
+                        max_new_tokens=256,
                         temperature=0.7,
                         top_p=0.95,
-                        top_k=40,
-                        num_beams=1,
-                        streamer=streamer,
-                        pad_token_id=tokenizer.pad_token_id
+                        top_k=40
                     )
                     
                     if hasattr(storage, 'get_used_memory'):
@@ -195,8 +156,8 @@ def test_ai_integration_http():
                     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"- Output length: {len(outputs[0])}")
-                    logger.info(f"- Output device: {outputs.device}")
+                    logger.info(f"- Output length: {len(outputs[0]['generated_text'])}")
+                    logger.info(f"- Generated text: {outputs[0]['generated_text']}")
                         
             except Exception as e:
                 logger.error(f"Text generation failed: {str(e)}")