[model-gateway] add GPU allocator and model pool infrastructure for parallel E2E tests (#16460)

2026-01-04 23:30:32 -08:00
parent 399ca037b1
commit 7f6a678f8f
5 changed files with 957 additions and 0 deletions
@@ -0,0 +1,31 @@
+"""Infrastructure for parallel GPU test execution."""
+
+from .gpu_allocator import (
+    GPUAllocator,
+    GPUInfo,
+    GPUSlot,
+    get_gpu_memory_usage,
+    get_open_port,
+    get_physical_device_indices,
+    nvml_context,
+    wait_for_gpu_memory_to_clear,
+)
+from .model_pool import ModelInstance, ModelPool
+from .model_specs import MODEL_SPECS
+
+__all__ = [
+    # GPU allocation
+    "GPUAllocator",
+    "GPUInfo",
+    "GPUSlot",
+    # GPU utilities
+    "nvml_context",
+    "get_open_port",
+    "get_physical_device_indices",
+    "get_gpu_memory_usage",
+    "wait_for_gpu_memory_to_clear",
+    # Model management
+    "ModelInstance",
+    "ModelPool",
+    "MODEL_SPECS",
+]