Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -17,6 +17,10 @@ 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()
@@ -96,16 +100,17 @@ def test_ai_integration_http():
                 transformers_logger.setLevel(logging.ERROR)
                 
                 try:
-                    # Create pipeline
+                    # Create pipeline and manually move to vGPU
                     pipe = pipeline(
                         "text-generation",
                         model=model_id,
-                        torch_dtype="auto",
-                        device=device  # Use our vGPU device
+                        torch_dtype=torch.float32,
+                        device_map=None  # Don't auto-place on devices
                     )
                     status['pipeline_loaded'] = True
                     
-                    # Move pipeline to vGPU
+                    # Move model to vGPU
+                    pipe.model = pipe.model.to(device)
                     pipe.model = to_vgpu(pipe.model, vram=vram)
                     status['model_on_vgpu'] = True
                     
@@ -131,19 +136,18 @@ def test_ai_integration_http():
             peak_mem = initial_mem
             
             try:
-                # Prepare messages
-                messages = [
-                    {"role": "user", "content": "Explain how virtual GPUs work in simple terms."}
-                ]
+                # Prepare input text
+                text = "Explain how virtual GPUs work in simple terms."
                 
                 with torch.no_grad():
                     # Generate text
                     outputs = pipe(
-                        messages,
+                        text,
                         max_new_tokens=256,
                         temperature=0.7,
                         top_p=0.95,
-                        top_k=40
+                        top_k=40,
+                        do_sample=True
                     )
                     
                     if hasattr(storage, 'get_used_memory'):
@@ -169,8 +173,10 @@ def test_ai_integration_http():
         finally:
             # Cleanup and status report
             try:
-                del model
-                del outputs
+                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: