Add Dynamic MIG tests to Lambda CI with functional nvmm replacement

Enable the existing DynMIG test suite on Lambda by creating a
functional nvmm replacement and wiring up GPU-type-aware filtering.

Lambda nvmm (tests/bats/lib/lambda/nvmm):
- Drop-in replacement for the original nvmm (which requires GPU
  Operator's nvidia-mig-manager pod). Runs nvidia-smi commands on
  the host via a privileged ephemeral pod with host /usr, /dev,
  and /proc/driver/nvidia mounted. Same interface as the original.
- NVMM_PATH env var selects which nvmm to use. Defaults to the
  original (tests/bats/lib). Lambda CI sets it to the lambda
  version (tests/bats/lib/lambda).

Dynamic MIG on Lambda:
- test_gpu_dynmig.bats tagged 'dynmig' for auto-filtering
- e2e-test.sh: skip dynmig on non-MIG GPUs (V100, A10); run on
  A100, H100, GH200, B200
- test_gpu_dynmig.bats added to tests-gpu-single Makefile target
- cleanup-from-previous-run.sh: make nvmm MIG cleanup non-fatal

Also removes SKIP_CLEANUP by handling MIG pre-cleanup on the host
via SSH and using the lambda nvmm for in-test MIG operations.

Signed-off-by: Davanum Srinivas <davanum@gmail.com>

Author: dims

hack/ci/lambda/e2e-test.sh

@@ -84,8 +84,19 @@ fi
 case "${LAMBDA_GPU_TYPE}" in
   *v100*|*a10) FILTER="${FILTER},!gpu-busgrind" ;;
 esac
+# Dynamic MIG requires MIG-capable GPUs (A100/H100/GH200/B200).
+case "${LAMBDA_GPU_TYPE}" in
+  *a100*|*h100*|*gh200*|*b200*) ;;
+  *) FILTER="${FILTER},!dynmig" ;;
+esac
 echo "Test filter: ${FILTER}"
 
+# --- Pre-cleanup: MIG teardown on host ---
+# Run MIG cleanup directly on the host where nvidia-smi is available.
+# The BATS Docker container uses the lambda nvmm stub (privileged pod)
+# for nvmm calls within tests, but pre-cleanup is faster via SSH.
+lambda_remote sh -c 'nvidia-smi mig -dci 2>/dev/null; nvidia-smi mig -dgi 2>/dev/null; nvidia-smi -mig 0 2>/dev/null; echo "MIG cleanup done"' || true
+
 # --- Run BATS tests ---
 # Tests local artifacts: local chart + local image built from the PR.
 # This is a real presubmit -- it validates the repo's code, chart, and specs.
@@ -103,10 +114,11 @@ export KUBECONFIG=\$HOME/.kube/config
 export CI=true
 export TEST_NVIDIA_DRIVER_ROOT=/
 export TEST_CHART_LOCAL=true
-export SKIP_CLEANUP=true
 export DISABLE_COMPUTE_DOMAINS=true
 export TEST_FILTER_TAGS='${FILTER}'
 export GIT_COMMIT_SHORT=${GIT_COMMIT_SHORT}
+# Use lambda nvmm (privileged pod, no GPU Operator).
+export NVMM_PATH=/cwd/tests/bats/lib/lambda
 
 make -f tests/bats/Makefile tests-gpu-single GIT_COMMIT_SHORT=${GIT_COMMIT_SHORT}
 EOF

tests/bats/Makefile

@@ -105,7 +105,8 @@ DOCKER_ENVS = \
 	--env TEST_NVIDIA_DRIVER_ROOT=$(TEST_NVIDIA_DRIVER_ROOT) \
 	--env TEST_EXPECTED_IMAGE_SPEC_SUBSTRING=$(TEST_EXPECTED_IMAGE_SPEC_SUBSTRING) \
 	--env SKIP_CLEANUP=$(SKIP_CLEANUP) \
-	--env DISABLE_COMPUTE_DOMAINS=$(DISABLE_COMPUTE_DOMAINS)
+	--env DISABLE_COMPUTE_DOMAINS=$(DISABLE_COMPUTE_DOMAINS) \
+	--env NVMM_PATH=$(NVMM_PATH)
 
 DOCKER_UID := $(shell id -u)
 DOCKER_GID := $(shell id -g)
@@ -169,13 +170,14 @@ runner-image:
 # cmdline args).
 .PHONY: tests-gpu tests-gpu-single tests-cd
 
-# Single-GPU-safe subset. Suitable for CI environments with one GPU and no
-# GPU Operator (e.g., Lambda Cloud). Excludes test_basics.bats (expects GPU
-# Operator + clean state), MIG, stress, and upgrade tests.
+# Lambda CI subset. Excludes test_basics.bats (expects GPU Operator + clean
+# state), static MIG, stress, and upgrade tests. DynMIG tests are included
+# but tagged 'dynmig' — auto-skipped on non-MIG GPUs via TEST_FILTER_TAGS.
 tests-gpu-single: runner-image
 	$(call RUN_BATS, \
 		tests/bats/test_gpu_basic.bats \
-		tests/bats/test_gpu_cuda_workloads.bats)
+		tests/bats/test_gpu_cuda_workloads.bats \
+		tests/bats/test_gpu_dynmig.bats)
 
 # Run a subset covering mainly the GPU plugin
 tests-gpu: runner-image

tests/bats/cleanup-from-previous-run.sh

@@ -104,7 +104,8 @@ set -e
 bash tests/bats/clean-state-dirs-all-nodes.sh
 
 # Remove any stray MIG devices and disable MIG mode on all nodes.
-nvmm all sh -c 'nvidia-smi mig -dci; nvidia-smi mig -dgi; nvidia-smi -mig 0'
+# Non-fatal: MIG may not be supported on all GPU types (V100, A10).
+nvmm all sh -c 'nvidia-smi mig -dci; nvidia-smi mig -dgi; nvidia-smi -mig 0' || echo "nvmm MIG cleanup skipped (non-fatal)"
 
 set +x
 echo "cleanup: done"

tests/bats/helpers.sh

@@ -18,9 +18,11 @@
 export TEST_HELM_RELEASE_NAME="dra-driver-nvidia-gpu-batssuite"
 
 
-# Extend PATH, for example for the `nvmm` utility.
+# Extend PATH for the `nvmm` utility. NVMM_PATH allows overriding the
+# directory containing `nvmm` — set it to `tests/bats/lib/lambda` on
+# Lambda CI (no GPU Operator) so the privileged-pod stub is used instead.
 SELF_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
-export PATH="${SELF_DIR}/lib:${PATH}"
+export PATH="${NVMM_PATH:-${SELF_DIR}/lib}:${PATH}"
 
 
 _common_setup() {

tests/bats/lib/lambda/nvmm

@@ -0,0 +1,76 @@
+#!/usr/bin/env bash
+
+# Copyright The Kubernetes Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# nvmm replacement for Lambda CI (single-node, no GPU Operator).
+#
+# The original nvmm execs into GPU Operator's nvidia-mig-manager pod.
+# This replacement runs commands on the host via a privileged pod that
+# mounts the host's /usr (for nvidia-smi) and /dev (for GPU devices).
+#
+# Usage (same as nvmm):
+#   nvmm <node-hint> [command...]
+#   nvmm all sh -c 'nvidia-smi mig -dci; nvidia-smi mig -dgi; nvidia-smi -mig 0'
+
+set -o errexit
+set -o nounset
+
+if [ -z "${1:-}" ]; then
+	echo "Usage: nvmm <node-hint> [command...]"
+	exit 1
+fi
+
+# Discard node-hint — Lambda is single-node.
+shift
+
+# If no command, nothing to do.
+if [ $# -eq 0 ]; then
+	exit 0
+fi
+
+# Build the command string. All remaining args are the command to run.
+CMD="$*"
+
+# Run via a privileged ephemeral pod with host's /usr and /dev mounted.
+# This gives us access to nvidia-smi and GPU devices.
+POD_NAME="nvmm-lambda-$(date +%s)"
+kubectl run "${POD_NAME}" \
+	--rm \
+	--attach \
+	--wait \
+	--restart=Never \
+	--quiet \
+	--image=ubuntu:24.04 \
+	--overrides='{
+		"spec": {
+			"hostPID": true,
+			"containers": [{
+				"name": "nvmm",
+				"image": "ubuntu:24.04",
+				"securityContext": {"privileged": true},
+				"command": ["/bin/sh", "-c", "export PATH=$PATH:/host/usr/bin:/host/usr/local/bin && '"$(echo "${CMD}" | sed "s/'/'\\\\''/g")"'"],
+				"volumeMounts": [
+					{"mountPath": "/dev", "name": "dev"},
+					{"mountPath": "/host/usr", "name": "host-usr", "readOnly": true},
+					{"mountPath": "/proc/driver/nvidia", "name": "proc-nvidia"}
+				]
+			}],
+			"volumes": [
+				{"name": "dev", "hostPath": {"path": "/dev"}},
+				{"name": "host-usr", "hostPath": {"path": "/usr"}},
+				{"name": "proc-nvidia", "hostPath": {"path": "/proc/driver/nvidia"}}
+			]
+		}
+	}' 2>/dev/null

tests/bats/test_gpu_dynmig.bats

@@ -57,7 +57,7 @@ confirm_mig_mode_disabled_all_nodes() {
 }
 
 
-# bats test_tags=fastfeedback
+# bats test_tags=fastfeedback,dynmig
 @test "DynMIG: inspect device attributes in resource slice (gpu)" {
   local reference=(
     "architecture"
@@ -80,7 +80,7 @@ confirm_mig_mode_disabled_all_nodes() {
 }
 
 
-# bats test_tags=fastfeedback
+# bats test_tags=fastfeedback,dynmig
 @test "DynMIG: inspect device attributes in resource slice (mig)" {
   local reference=(
     "architecture"
@@ -104,7 +104,7 @@ confirm_mig_mode_disabled_all_nodes() {
 }
 
 
-# bats test_tags=fastfeedback
+# bats test_tags=fastfeedback,dynmig
 @test "DynMIG: 1 pod, 1 MIG" {
   confirm_mig_mode_disabled_all_nodes
   kubectl apply -f tests/bats/specs/gpu-simple-mig.yaml
@@ -127,7 +127,7 @@ confirm_mig_mode_disabled_all_nodes() {
 }
 
 
-# bats test_tags=fastfeedback
+# bats test_tags=fastfeedback,dynmig
 @test "DynMIG: 1 pod, 2 containers (1 MIG each)" {
   confirm_mig_mode_disabled_all_nodes
 
@@ -159,7 +159,7 @@ confirm_mig_mode_disabled_all_nodes() {
 }
 
 
-# bats test_tags=fastfeedback
+# bats test_tags=fastfeedback,dynmig
 @test "DynMIG: 1 pod, 1 MIG + TimeSlicing config" {
   local _iargs=("--set" "logVerbosity=6" "--set" "featureGates.DynamicMIG=true" "--set" "featureGates.TimeSlicingSettings=true")
   iupgrade_wait "${TEST_CHART_REPO}" "${TEST_CHART_VERSION}" _iargs