Configmap-based GPU mock for the launcher

Signed-off-by: Jun Duan <jun.duan.phd@outlook.com>

Author: waltforme

dockerfiles/Dockerfile.launcher.cpu

@@ -17,6 +17,6 @@ WORKDIR /app
 COPY inference_server/launcher/launcher.py inference_server/launcher/gputranslator.py /app/
 
 # Install uvicorn for serving the launcher API and nvidia-ml-py for gputranslator
-RUN pip install --root-user-action=ignore --no-cache-dir uvicorn nvidia-ml-py
+RUN pip install --root-user-action=ignore --no-cache-dir uvicorn nvidia-ml-py kubernetes
 
 ENTRYPOINT ["uvicorn", "--app-dir", "/app", "launcher:app"]

inference_server/launcher/gputranslator.py

@@ -18,29 +18,34 @@
 """
 
 import importlib.metadata
+import json
 import logging
-from typing import Dict
+import os
+from typing import Dict, Optional
 
 import pynvml
+from kubernetes import client, config
 
 logger = logging.getLogger(__name__)
 
 
 # VLLM process manager
 class GpuTranslator:
-    def __init__(self, mock_gpus: bool = False, mock_gpu_count: int = 8):
+    def __init__(self, mock_gpus: bool = False, mock_gpu_count: int = 8, node_name: Optional[str] = None):
         """
         Initialize GPU Translator
 
         Args:
             mock_gpus: If True, skip pynvml and use mock mode for testing
             mock_gpu_count: Number of mock GPUs to create (default: 8)
+            node_name: Kubernetes node name for ConfigMap-based GPU discovery
         """
         self.mapping = {}
         self.reverse_mapping = {}
         self.device_count = 0
         self.mock_mode = mock_gpus
         self.mock_gpu_count = mock_gpu_count
+        self.node_name = node_name or os.getenv("NODE_NAME")
         if not self.mock_mode:
             self._check_library()
         self._populate_mapping()
@@ -62,11 +67,79 @@ def _check_library(self):
                 f"package {package_name} not found. Please install it."
             )
 
+    def _load_gpu_map_from_configmap(self) -> Optional[Dict[str, int]]:
+        """
+        Load GPU mapping from Kubernetes ConfigMap 'gpu-map'.
+        
+        Returns:
+            Dict[str, int]: GPU UUID to index mapping, or None if ConfigMap not available
+        """
+        if not self.node_name:
+            logger.info("No node name provided, skipping ConfigMap GPU discovery")
+            return None
+        
+        try:
+            # Try to load in-cluster config first, fall back to kubeconfig
+            try:
+                config.load_incluster_config()
+            except config.ConfigException:
+                config.load_kube_config()
+            
+            v1 = client.CoreV1Api()
+            
+            # Read the ConfigMap
+            namespace = os.getenv("NAMESPACE", "default")
+            cm = v1.read_namespaced_config_map(name="gpu-map", namespace=namespace)
+            
+            if not cm.data or self.node_name not in cm.data:
+                logger.warning(
+                    "Node '%s' not found in ConfigMap 'gpu-map' in namespace '%s'",
+                    self.node_name,
+                    namespace
+                )
+                return None
+            
+            # Parse the JSON mapping for this node
+            node_gpu_data = cm.data[self.node_name]
+            gpu_mapping = json.loads(node_gpu_data)
+            
+            logger.info(
+                "Loaded GPU mapping from ConfigMap for node '%s': %s",
+                self.node_name,
+                gpu_mapping
+            )
+            return gpu_mapping
+            
+        except Exception as e:
+            logger.warning(
+                "Failed to load GPU mapping from ConfigMap: %s. Falling back to mock mode.",
+                e
+            )
+            return None
+
     def _populate_mapping(self):
         """
         Creates mapping and reverse_mapping for the GPU Translator.
-        In mock mode, pre-populates with mock GPU UUIDs following the pattern GPU-{index}.
+        Priority order:
+        1. ConfigMap 'gpu-map' (if in Kubernetes and node_name available)
+        2. Mock mode (if mock_gpus=True)
+        3. Real GPUs via pynvml
         """
+        # Try ConfigMap first if in mock mode and node_name is available
+        if self.mock_mode and self.node_name:
+            configmap_mapping = self._load_gpu_map_from_configmap()
+            if configmap_mapping:
+                self.mapping = configmap_mapping
+                self.reverse_mapping = {v: k for k, v in self.mapping.items()}
+                self.device_count = len(self.mapping)
+                logger.info(
+                    "GPU Translator initialized from ConfigMap with %d GPUs for node '%s'",
+                    self.device_count,
+                    self.node_name
+                )
+                return
+        
+        # Fall back to hardcoded mock mode
         if self.mock_mode:
             # Pre-populate with mock GPUs following the test pattern: GPU-0, GPU-1, etc.
             for index in range(self.mock_gpu_count):
@@ -77,6 +150,7 @@ def _populate_mapping(self):
             logger.info("GPU Translator initialized in mock mode with %d mock GPUs", self.mock_gpu_count)
             return
 
+        # Use real GPUs via pynvml
         try:
             pynvml.nvmlInit()
             self.device_count = pynvml.nvmlDeviceGetCount()

inference_server/launcher/launcher.py

@@ -81,6 +81,9 @@ def __init__(
             for uuid_str in config.gpu_uuids:
                 index = gpu_translator.uuid_to_index(uuid_str)
                 cuda_indices.append(str(index))
+            logger.info(
+                f"Translated GPU UUIDs {config.gpu_uuids} to indices {cuda_indices}."
+            )
 
             if config.env_vars is None:
                 config.env_vars = {}
@@ -193,9 +196,13 @@ def get_logs(
 
 # Multi-instance vLLM process manager
 class VllmMultiProcessManager:
-    def __init__(self, mock_gpus: bool = False, mock_gpu_count: int = 8):
+    def __init__(self, mock_gpus: bool = False, mock_gpu_count: int = 8, node_name: Optional[str] = None):
         self.instances: Dict[str, VllmInstance] = {}
-        self.gpu_translator = GpuTranslator(mock_gpus=mock_gpus, mock_gpu_count=mock_gpu_count)
+        self.gpu_translator = GpuTranslator(
+            mock_gpus=mock_gpus,
+            mock_gpu_count=mock_gpu_count,
+            node_name=node_name
+        )
 
     def create_instance(
         self, vllm_config: VllmConfig, instance_id: Optional[str] = None
@@ -605,10 +612,14 @@ def set_env_vars(env_vars: Dict[str, Any]):
 
     args = parser.parse_args()
 
+    # Get node name from environment variable
+    node_name = os.getenv("NODE_NAME")
+
     # Reinitialize the global manager with mock mode settings
     vllm_manager = VllmMultiProcessManager(
         mock_gpus=args.mock_gpus,
-        mock_gpu_count=args.mock_gpu_count
+        mock_gpu_count=args.mock_gpu_count,
+        node_name=node_name
     )
 
     uvicorn.run(

inference_server/launcher/requirements.txt

@@ -3,5 +3,6 @@ pydantic
 uvicorn
 uvloop
 nvidia-ml-py
+kubernetes
 # WARNING: vllm must be built from source on a macOS Silicon
 vllm; sys_platform != "darwin" or platform_machine != "arm64"

test/e2e/mkobjs.sh

@@ -1,7 +1,6 @@
 #!/usr/bin/env bash
 
 inst=$(date +%d-%H-%M-%S)
-server_img=$(make echo-var VAR=TEST_SERVER_IMG)
 requester_img=$(make echo-var VAR=TEST_REQUESTER_IMG)
 launcher_img=$(make echo-var VAR=TEST_LAUNCHER_IMG)
 if out=$(kubectl apply -f - 2>&1 <<EOF
@@ -78,6 +77,7 @@ spec:
   maxSleepingInstances: 1
   podTemplate:
     spec:
+      serviceAccount: testlauncher
       containers:
         - name: inference-server
           image: $launcher_img
@@ -92,6 +92,13 @@ spec:
             --host 0.0.0.0 \
             --port 8001 \
             --log-level info
+          env:
+            - name: NODE_NAME
+              valueFrom:
+                fieldRef: { fieldPath: spec.nodeName }
+            - name: NAMESPACE
+              valueFrom:
+                fieldRef: { fieldPath: metadata.namespace }
 ---
 apiVersion: apps/v1
 kind: ReplicaSet

test/e2e/run-launcher-based.sh

@@ -114,6 +114,26 @@ kubectl create rolebinding testreq --role=testreq --serviceaccount=$(kubectl get
 kubectl create clusterrolebinding testreq-view --clusterrole=view --serviceaccount=$(kubectl get sa default -o jsonpath={.metadata.namespace}):testreq
 
 kubectl create sa testreq
+
+kubectl apply -f - <<EOF
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+  name: testlauncher
+rules:
+- apiGroups:
+  - ""
+  resources:
+  - configmaps
+  verbs:
+  - get
+  - list
+  - watch
+EOF
+
+kubectl create rolebinding testlauncher --role=testlauncher --serviceaccount=$(kubectl get sa default -o jsonpath={.metadata.namespace}):testlauncher
+
+kubectl create sa testlauncher
 kubectl create cm gpu-map
 kubectl get nodes -o name | sed 's%^node/%%' | while read node; do
     kubectl label node $node nvidia.com/gpu.present=true nvidia.com/gpu.product=NVIDIA-L40S nvidia.com/gpu.count=2 --overwrite=true