Add MIG overlay support with configuration and service files

overlays/mig/internal/mig_test.go

@@ -0,0 +1,112 @@
+package mig
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/warewulf/warewulf/internal/app/wwctl/overlay/show"
+	"github.com/warewulf/warewulf/internal/pkg/testenv"
+	"github.com/warewulf/warewulf/internal/pkg/wwlog"
+)
+
+func Test_migOverlay(t *testing.T) {
+	env := testenv.New(t)
+	defer env.RemoveAll()
+	env.ImportFile("etc/warewulf/nodes.conf", "nodes.conf")
+	env.ImportFile("var/lib/warewulf/overlays/mig/rootfs/etc/systemd/system/nvidia-mig.service.ww", "../rootfs/etc/systemd/system/nvidia-mig.service.ww")
+
+	tests := []struct {
+		name string
+		args []string
+		log  string
+	}{
+		{
+			name: "mig:nvidia-mig.service (default)",
+			args: []string{"--render", "node1", "mig", "etc/systemd/system/nvidia-mig.service.ww"},
+			log:  migServiceDefault,
+		},
+		{
+			name: "mig:nvidia-mig.service (per-gpu)",
+			args: []string{"--render", "node2", "mig", "etc/systemd/system/nvidia-mig.service.ww"},
+			log:  migServicePerGpu,
+		},
+		{
+			name: "mig:nvidia-mig.service (custom)",
+			args: []string{"--render", "node3", "mig", "etc/systemd/system/nvidia-mig.service.ww"},
+			log:  migServiceCustom,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cmd := show.GetCommand()
+			cmd.SetArgs(tt.args)
+			stdout := bytes.NewBufferString("")
+			stderr := bytes.NewBufferString("")
+			logbuf := bytes.NewBufferString("")
+			cmd.SetOut(stdout)
+			cmd.SetErr(stderr)
+			wwlog.SetLogWriter(logbuf)
+			err := cmd.Execute()
+			assert.NoError(t, err)
+			assert.Empty(t, stdout.String())
+			assert.Empty(t, stderr.String())
+			assert.Equal(t, tt.log, logbuf.String())
+		})
+	}
+}
+
+const migServiceDefault string = `backupFile: true
+writeFile: true
+Filename: etc/systemd/system/nvidia-mig.service
+[Unit]
+DefaultDependencies=no
+Description=Configure NVIDIA MIG (Multi-Instance GPU) partitions at boot time
+Before=nvidia-persistenced.service nvidia-dcgm.service dcgm_exporter.service
+
+[Service]
+Type=oneshot
+RemainAfterExit=yes
+ExecStart=/usr/local/sbin/_config_mig "2,2"
+ExecStop=sh -c 'nvidia-smi mig -dci ; nvidia-smi mig -dgi ; nvidia-smi -mig 0'
+
+[Install]
+WantedBy=multi-user.target
+`
+
+const migServicePerGpu string = `backupFile: true
+writeFile: true
+Filename: etc/systemd/system/nvidia-mig.service
+[Unit]
+DefaultDependencies=no
+Description=Configure NVIDIA MIG (Multi-Instance GPU) partitions at boot time
+Before=nvidia-persistenced.service nvidia-dcgm.service dcgm_exporter.service
+
+[Service]
+Type=oneshot
+RemainAfterExit=yes
+ExecStart=/usr/local/sbin/_config_mig "14,14,14,15 7,7,7,8"
+ExecStop=sh -c 'nvidia-smi mig -dci ; nvidia-smi mig -dgi ; nvidia-smi -mig 0'
+
+[Install]
+WantedBy=multi-user.target
+`
+
+const migServiceCustom string = `backupFile: true
+writeFile: true
+Filename: etc/systemd/system/nvidia-mig.service
+[Unit]
+DefaultDependencies=no
+Description=Configure NVIDIA MIG (Multi-Instance GPU) partitions at boot time
+Before=nvidia-persistenced.service nvidia-dcgm.service dcgm_exporter.service
+
+[Service]
+Type=oneshot
+RemainAfterExit=yes
+ExecStart=/usr/local/sbin/_config_mig "1,1,1,1"
+ExecStop=sh -c 'nvidia-smi mig -dci ; nvidia-smi mig -dgi ; nvidia-smi -mig 0'
+
+[Install]
+WantedBy=multi-user.target
+`

overlays/mig/internal/nodes.conf

@@ -0,0 +1,8 @@
+nodes:
+  node1: {}
+  node2:
+    tags:
+      gpuMigType: "14,14,14,15 7,7,7,8"
+  node3:
+    tags:
+      gpuMigType: "1,1,1,1"

overlays/mig/rootfs/etc/systemd/system/multi-user.target.wants/nvidia-mig.service

@@ -0,0 +1 @@
+../nvidia-mig.service

overlays/mig/rootfs/etc/systemd/system/nvidia-mig.service.ww

@@ -0,0 +1,13 @@
+[Unit]
+DefaultDependencies=no
+Description=Configure NVIDIA MIG (Multi-Instance GPU) partitions at boot time
+Before=nvidia-persistenced.service nvidia-dcgm.service dcgm_exporter.service
+
+[Service]
+Type=oneshot
+RemainAfterExit=yes
+ExecStart=/usr/local/sbin/_config_mig "{{ .Tags.gpuMigType | default "2,2" }}"
+ExecStop=sh -c 'nvidia-smi mig -dci ; nvidia-smi mig -dgi ; nvidia-smi -mig 0'
+
+[Install]
+WantedBy=multi-user.target

overlays/mig/rootfs/usr/local/sbin/_config_mig

@@ -0,0 +1,37 @@
+#!/usr/bin/bash
+set -euo pipefail
+
+# This script configures NVIDIA MIG (Multi-Instance GPU) partitions based on the provided
+# gpuMigType string argument. It supports both a single configuration for all GPUs or
+# individual configurations per GPU.
+# Example inputs:
+#   "14,14,14,15"               # single config for all GPUs
+#   "14,14,14,15 7,7,7,8"  # per-GPU configs for GPU0 and GPU1
+
+# See https://docs.nvidia.com/datacenter/tesla/mig-user-guide/supported-mig-profiles.html for details on supported MIG profiles
+
+# Get the gpuMigType string from the first argument
+MIG_STRING="${1}"
+if [[ -z "$MIG_STRING" ]]; then
+  echo "Usage: $0 <gpuMigType string>" >&2
+  exit 1
+fi
+
+# Check if nvidia-smi command exists
+if ! command -v nvidia-smi &> /dev/null; then
+  echo "Error: nvidia-smi command not found" >&2
+  exit 1
+fi
+
+# enable MIG mode
+nvidia-smi -mig 1
+
+# per-GPU space-separated list (one token per GPU), otherwise single-string common case
+if [[ "$MIG_STRING" == *" "* ]]; then
+  read -a PER_GPU_MIG_CONFIGS <<< "$MIG_STRING"
+  for gpu_index in "${!PER_GPU_MIG_CONFIGS[@]}"; do
+    nvidia-smi mig -cgi "${PER_GPU_MIG_CONFIGS[$gpu_index]}" -i "$gpu_index" -C
+  done
+else
+  nvidia-smi mig -cgi "$MIG_STRING" -C
+fi

overlays/mig/rootfs/usr/local/sbin/mig2gres.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/bash
+
+set -euo pipefail
+
+# Generates NVIDIA MIG device entries for Slurm's gres.conf.
+# Queries MIG GPU instances and maps them to /dev/nvidia-caps/ device files.
+#
+# Process: nvidia-smi -> extract GPU/GI/profile -> lookup minor in 
+# /proc/driver/nvidia-caps/mig-minors -> map to device file -> output gres format
+
+
+if ! command -v nvidia-smi &> /dev/null; then
+  echo "Error: nvidia-smi command not found" >&2
+  exit 1
+fi
+
+hostname=$(hostname -s)
+
+nvidia-smi mig -lgi | grep -E '^\|' | grep -v '===' | grep -v 'GPU.*Name' | while read line; do
+  gpu=$(echo "$line" | awk '{print $2}')
+  name=$(echo "$line" | awk '{print $4}')
+  gi=$(echo "$line" | awk '{print $6}')  # Instance ID is column 6
+
+  [[ "$gpu" =~ ^[0-9]+$ ]] || continue
+
+  pattern="gpu${gpu}/gi${gi}/access"
+  minor=$(grep "^${pattern} " /proc/driver/nvidia-caps/mig-minors | awk '{print $2}')
+
+  echo "NodeName=${hostname} Name=gpu Type=${name} File=/dev/nvidia-caps/nvidia-cap${minor}"
+done