Update test_ai_integration_http.py

test_ai_integration_http.py

@@ -1,4 +1,3 @@
-
 import logging
 import os
 import time
@@ -96,21 +95,22 @@ def test_ai_integration_http():
                 transformers_logger.setLevel(logging.ERROR)
                 
                 try:
-                    # Create pipeline
-                    # Create pipeline with vGPU device
+                    # Create pipeline with direct vGPU device mapping
                     pipe = pipeline(
                         "text-generation",
                         model=model_id,
-                        torch_dtype=torch.float32,  # Use full precision,
+                        torch_dtype=torch.float32,  # Use full precision
+                        device=device,  # Load directly to vGPU
                         use_safetensors=True,
-                        trust_remote_code=True
+                        trust_remote_code=True,
+                        model_kwargs={
+                            "device_map": device,  # Ensure all model parts go to vGPU
+                            "vram": vram  # Pass our vRAM manager
+                        }
                     )
                     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
+                    pipe.model.eval()
                     
                     # Log model details
                     logger.info(f"Pipeline created with model: {model_id}")