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,39 @@
 """
 
 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,21 +72,93 @@ 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' based mock if mock mode is enabled and node_name is available
+        2. Naive mock with GPU-0, GPU-1, etc. if mock mode is enabled
+        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):
                 uuid = f"GPU-{index}"
                 self.mapping[uuid] = index
                 self.reverse_mapping[index] = uuid
             self.device_count = self.mock_gpu_count
-            logger.info("GPU Translator initialized in mock mode with %d mock GPUs", self.mock_gpu_count)
+            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()
@@ -90,7 +172,9 @@ def _populate_mapping(self):
                 )
                 self.mapping[uuid] = index
             pynvml.nvmlShutdown()
-            logger.info("GPU Translator initialized with %d real GPUs", self.device_count)
+            logger.info(
+                "GPU Translator initialized with %d real GPUs", self.device_count
+            )
 
         except pynvml.NVMLError as error:
             logger.error("Failed to initialize pynvml: %s", error)

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,16 @@ 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
@@ -575,45 +585,44 @@ def set_env_vars(env_vars: Dict[str, Any]):
     parser.add_argument(
         "--mock-gpus",
         action="store_true",
-        help="Enable mock GPU mode for CPU-only testing environments"
+        help="Enable mock GPU mode for CPU-only testing environments",
     )
     parser.add_argument(
         "--mock-gpu-count",
         type=int,
         default=8,
-        help="Number of mock GPUs to create in mock mode (default: 8)"
+        help="Number of mock GPUs to create in mock mode (default: 8)",
     )
     parser.add_argument(
         "--host",
         type=str,
         default="0.0.0.0",
-        help="Host to bind the server to (default: 0.0.0.0)"
+        help="Host to bind the server to (default: 0.0.0.0)",
     )
     parser.add_argument(
         "--port",
         type=int,
         default=8001,
-        help="Port to bind the server to (default: 8001)"
+        help="Port to bind the server to (default: 8001)",
     )
     parser.add_argument(
         "--log-level",
         type=str,
         default="info",
         choices=["critical", "error", "warning", "info", "debug"],
-        help="Logging level (default: info)"
+        help="Logging level (default: info)",
     )
 
     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(
-        app,
-        host=args.host,
-        port=args.port,
-        log_level=args.log_level
-    )
+    uvicorn.run(app, host=args.host, port=args.port, log_level=args.log_level)

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"

inference_server/launcher/tests/test_launcher.py

@@ -17,6 +17,7 @@
 Run as:
 python -m pytest tests/test_launcher.py -v
 """
+
 import signal
 import sys
 from unittest.mock import MagicMock, patch

pkg/controller/dual-pods/inference-server.go

@@ -462,7 +462,7 @@ func (item infSvrItem) process(urCtx context.Context, ctl *controller, nodeDat *
 
 	cfg, iscHash, err := ctl.configInferenceServer(isc, serverDat.GPUIDs)
 	if err != nil {
-		return fmt.Errorf("parse inference server config: %w", err), true
+		return fmt.Errorf("failed to configure inference server config: %w", err), true
 	}
 	logger.V(5).Info("Nominal hash of InferenceServerConfig", "hash", iscHash)
 

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