Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -17,10 +17,6 @@ from torch_vgpu import VGPUDevice, to_vgpu
 def setup_vgpu():
     """Setup vGPU device"""
     try:
-        # Register vGPU device type
-        if not hasattr(torch, 'vgpu'):
-            torch.register_privateuseone_backend()
-        
         # Create and register vGPU device
         vgpu = VGPUDevice()
         device = vgpu.device()
@@ -100,18 +96,21 @@ def test_ai_integration_http():
                 transformers_logger.setLevel(logging.ERROR)
                 
                 try:
-                    # Create pipeline and manually move to vGPU
+                    # Create pipeline
+                    # Create pipeline with vGPU device
                     pipe = pipeline(
                         "text-generation",
                         model=model_id,
-                        torch_dtype=torch.float32,
-                        device_map=None  # Don't auto-place on devices
+                        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 model to vGPU
-                    pipe.model = pipe.model.to(device)
+                    # Move pipeline model to vGPU
                     pipe.model = to_vgpu(pipe.model, vram=vram)
+                    pipe.model.eval()
                     status['model_on_vgpu'] = True
                     
                     # Log model details
@@ -121,6 +120,11 @@ def test_ai_integration_http():
                     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
@@ -129,6 +133,9 @@ def test_ai_integration_http():
             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
             
             # Run text generation
             logger.info("Running text generation...")
@@ -136,18 +143,20 @@ def test_ai_integration_http():
             peak_mem = initial_mem
             
             try:
-                # Prepare input text
-                text = "Explain how virtual GPUs work in simple terms."
+                # Prepare input prompt
+                prompt = "Explain how virtual GPUs work in simple terms."
                 
                 with torch.no_grad():
                     # Generate text
                     outputs = pipe(
-                        text,
+                        prompt,
                         max_new_tokens=256,
                         temperature=0.7,
                         top_p=0.95,
                         top_k=40,
-                        do_sample=True
+                        num_beams=1,
+                        do_sample=True,
+                        return_full_text=True
                     )
                     
                     if hasattr(storage, 'get_used_memory'):
@@ -160,7 +169,6 @@ 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]['generated_text'])}")
                     logger.info(f"- Generated text: {outputs[0]['generated_text']}")
                         
             except Exception as e: