diff --git a/overlays/mig/internal/mig_test.go b/overlays/mig/internal/mig_test.go new file mode 100644 index 00000000..1af93a97 --- /dev/null +++ b/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 +` diff --git a/overlays/mig/internal/nodes.conf b/overlays/mig/internal/nodes.conf new file mode 100644 index 00000000..493c2d31 --- /dev/null +++ b/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" \ No newline at end of file diff --git a/overlays/mig/rootfs/etc/systemd/system/multi-user.target.wants/nvidia-mig.service b/overlays/mig/rootfs/etc/systemd/system/multi-user.target.wants/nvidia-mig.service new file mode 120000 index 00000000..fc15c38c --- /dev/null +++ b/overlays/mig/rootfs/etc/systemd/system/multi-user.target.wants/nvidia-mig.service @@ -0,0 +1 @@ +../nvidia-mig.service \ No newline at end of file diff --git a/overlays/mig/rootfs/etc/systemd/system/nvidia-mig.service.ww b/overlays/mig/rootfs/etc/systemd/system/nvidia-mig.service.ww new file mode 100644 index 00000000..c632f232 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/overlays/mig/rootfs/usr/local/sbin/_config_mig b/overlays/mig/rootfs/usr/local/sbin/_config_mig new file mode 100644 index 00000000..b0d48526 --- /dev/null +++ b/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 " >&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 \ No newline at end of file diff --git a/overlays/mig/rootfs/usr/local/sbin/mig2gres.sh b/overlays/mig/rootfs/usr/local/sbin/mig2gres.sh new file mode 100644 index 00000000..a877b152 --- /dev/null +++ b/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 \ No newline at end of file