Add MIG config support when MIG-backed vGPU type

Signed-off-by: Michail Resvanis <mresvani@redhat.com>

Author: mresvanis

cmd/nvidia-k8s-vgpu-dm/main.go

@@ -20,12 +20,14 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"strings"
 
 	log "github.com/sirupsen/logrus"
 	cli "github.com/urfave/cli/v2"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/tools/cache"
 	"k8s.io/client-go/tools/clientcmd"
+	"sigs.k8s.io/yaml"
 
 	"context"
 	"sync"
@@ -38,6 +40,7 @@ import (
 	"k8s.io/apimachinery/pkg/util/wait"
 
 	"github.com/NVIDIA/vgpu-device-manager/internal/info"
+	"github.com/NVIDIA/vgpu-device-manager/pkg/types"
 )
 
 const (
@@ -47,6 +50,14 @@ const (
 	vGPUConfigStateLabel = "nvidia.com/vgpu.config.state"
 	pluginStateLabel     = "nvidia.com/gpu.deploy.sandbox-device-plugin"
 	validatorStateLabel  = "nvidia.com/gpu.deploy.sandbox-validator"
+
+	defaultReconfigureMIGScript      = "/usr/bin/reconfigure-mig.sh"
+	defaultHostRootMount             = "/host"
+	defaultHostNvidiaDir             = "/usr/local/nvidia"
+	defaultHostMigManagerStateFile   = "/etc/systemd/system/nvidia-mig-manager.service.d/override.conf"
+	defaultHostKubeletSystemdService = "kubelet.service"
+
+	migConfigDisabled = "all-disabled"
 )
 
 var (
@@ -56,10 +67,25 @@ var (
 	configFileFlag        string
 	defaultVGPUConfigFlag string
 
+	reconfigureMIGScriptFlag       string
+	migPartedConfigFileFlag        string
+	hostRootMountFlag              string
+	hostNvidiaDirFlag              string
+	hostMigManagerStateFileFlag    string
+	hostKubeletSystemdServiceFlag  string
+	gpuClientsFileFlag             string
+	withRebootFlag                 bool
+	withShutdownHostGPUClientsFlag bool
+
 	pluginDeployed    string
 	validatorDeployed string
 )
 
+type GPUClients struct {
+	Version         string   `json:"version"          yaml:"version"`
+	SystemdServices []string `json:"systemd-services" yaml:"systemd-services"`
+}
+
 // SyncableVGPUConfig is used to synchronize on changes to a configuration value.
 // That is, callers of Get() will block until a call to Set() is made.
 // Multiple calls to Set() do not queue, meaning that only calls to Get() made
@@ -148,6 +174,78 @@ func main() {
 			Destination: &defaultVGPUConfigFlag,
 			EnvVars:     []string{"DEFAULT_VGPU_CONFIG"},
 		},
+		&cli.StringFlag{
+			Name:        "reconfigure-mig-script",
+			Aliases:     []string{"s"},
+			Value:       defaultReconfigureMIGScript,
+			Usage:       "script to run to do the actual MIG reconfiguration",
+			Destination: &reconfigureMIGScriptFlag,
+			EnvVars:     []string{"RECONFIGURE_MIG_SCRIPT"},
+		},
+		&cli.StringFlag{
+			Name:        "mig-parted-config-file",
+			Aliases:     []string{"mc"},
+			Value:       "",
+			Usage:       "the path to the mig-parted configuration file",
+			Destination: &migPartedConfigFileFlag,
+			EnvVars:     []string{"MIG_PARTED_CONFIG_FILE"},
+		},
+		&cli.StringFlag{
+			Name:        "host-root-mount",
+			Aliases:     []string{"m"},
+			Value:       defaultHostRootMount,
+			Usage:       "container path where host root directory is mounted",
+			Destination: &hostRootMountFlag,
+			EnvVars:     []string{"HOST_ROOT_MOUNT"},
+		},
+		&cli.StringFlag{
+			Name:        "host-nvidia-dir",
+			Aliases:     []string{"i"},
+			Value:       defaultHostNvidiaDir,
+			Usage:       "host path of the directory where NVIDIA managed software directory is typically located",
+			Destination: &hostNvidiaDirFlag,
+			EnvVars:     []string{"HOST_NVIDIA_DIR"},
+		},
+		&cli.StringFlag{
+			Name:        "host-mig-manager-state-file",
+			Aliases:     []string{"o"},
+			Value:       defaultHostMigManagerStateFile,
+			Usage:       "host path where the host's systemd mig-manager state file is located",
+			Destination: &hostMigManagerStateFileFlag,
+			EnvVars:     []string{"HOST_MIG_MANAGER_STATE_FILE"},
+		},
+		&cli.StringFlag{
+			Name:        "host-kubelet-systemd-service",
+			Aliases:     []string{"k"},
+			Value:       defaultHostKubeletSystemdService,
+			Usage:       "name of the host's 'kubelet' systemd service which may need to be shutdown/restarted across a MIG mode reconfiguration",
+			Destination: &hostKubeletSystemdServiceFlag,
+			EnvVars:     []string{"HOST_KUBELET_SYSTEMD_SERVICE"},
+		},
+		&cli.StringFlag{
+			Name:        "gpu-clients-file",
+			Aliases:     []string{"g"},
+			Value:       "",
+			Usage:       "the path to the file listing the GPU clients that need to be shutdown across a MIG configuration",
+			Destination: &gpuClientsFileFlag,
+			EnvVars:     []string{"GPU_CLIENTS_FILE"},
+		},
+		&cli.BoolFlag{
+			Name:        "with-reboot",
+			Aliases:     []string{"r"},
+			Value:       false,
+			Usage:       "reboot the node if changing the MIG mode fails for any reason",
+			Destination: &withRebootFlag,
+			EnvVars:     []string{"WITH_REBOOT"},
+		},
+		&cli.BoolFlag{
+			Name:        "with-shutdown-host-gpu-clients",
+			Aliases:     []string{"w"},
+			Value:       false,
+			Usage:       "shutdown/restart any required host GPU clients across a MIG configuration",
+			Destination: &withShutdownHostGPUClientsFlag,
+			EnvVars:     []string{"WITH_SHUTDOWN_HOST_GPU_CLIENTS"},
+		},
 	}
 
 	log.Infof("version: %s", c.Version)
@@ -296,6 +394,30 @@ func updateConfig(clientset *kubernetes.Clientset, selectedConfig string) error
 		return fmt.Errorf("unable to shutdown gpu operands: %v", err)
 	}
 
+	vgpuType, err := types.ParseVGPUType(selectedConfig)
+	if err != nil {
+		return fmt.Errorf("unable to parse vGPU type: %s", err)
+	}
+	if vgpuType.G > 0 {
+		log.Info("Running reconfigure MIG script")
+
+		selectedMigConfig, err := convertToMIGConfigFormat(selectedConfig)
+		if err != nil {
+			return fmt.Errorf("unable to convert vGPU type config to MIG config: %s", err)
+		}
+
+		err = runReconfigureMIGScript(selectedMigConfig)
+		if err != nil {
+			return fmt.Errorf("unable to run reconfigure MIG script: %s", err)
+		}
+	} else {
+		log.Info("Disabling MIG if enabled")
+		err = runReconfigureMIGScript(migConfigDisabled)
+		if err != nil {
+			return fmt.Errorf("unable to run reconfigure MIG script: %s", err)
+		}
+	}
+
 	log.Info("Applying the selected vGPU device configuration to the node")
 	err = applyConfig(selectedConfig)
 	if err != nil {
@@ -504,3 +626,78 @@ func setNodeLabelValue(clientset *kubernetes.Clientset, label, value string) err
 
 	return nil
 }
+
+func parseGPUCLientsFile(file string) (*GPUClients, error) {
+	var err error
+	var yamlBytes []byte
+
+	if file == "" {
+		return &GPUClients{}, nil
+	}
+
+	yamlBytes, err = os.ReadFile(file)
+	if err != nil {
+		return nil, fmt.Errorf("read error: %v", err)
+	}
+
+	var clients GPUClients
+	err = yaml.Unmarshal(yamlBytes, &clients)
+	if err != nil {
+		return nil, fmt.Errorf("unmarshal error: %v", err)
+	}
+
+	return &clients, nil
+}
+
+func runReconfigureMIGScript(migConfigValue string) error {
+	gpuClients, err := parseGPUCLientsFile(gpuClientsFileFlag)
+	if err != nil {
+		return fmt.Errorf("error parsing host's GPU clients file: %s", err)
+	}
+
+	args := []string{
+		"-n", nodeNameFlag,
+		"-f", migPartedConfigFileFlag,
+		"-c", migConfigValue,
+		"-m", hostRootMountFlag,
+		"-i", hostNvidiaDirFlag,
+		"-o", hostMigManagerStateFileFlag,
+		"-g", strings.Join(gpuClients.SystemdServices, ","),
+		"-k", hostKubeletSystemdServiceFlag,
+	}
+	if withRebootFlag {
+		args = append(args, "-r")
+	}
+	if withShutdownHostGPUClientsFlag {
+		args = append(args, "-w")
+	}
+
+	stdLogger := log.StandardLogger()
+
+	cmd := exec.Command(reconfigureMIGScriptFlag, args...)
+	cmd.Stdout = stdLogger.Writer()
+	cmd.Stderr = stdLogger.WriterLevel(log.ErrorLevel)
+	return cmd.Run()
+}
+
+// convertToMIGConfigFormat converts a vGPU type string to the MIG config
+// equivalent, e.g. "all-{g}g.{gb}gb[.me]".
+func convertToMIGConfigFormat(s string) (string, error) {
+	vgpu, err := types.ParseVGPUType(s)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse vGPU type: %v", err)
+	}
+
+	// Base format: all-{g}g.{gb}gb
+	result := fmt.Sprintf("all-%dg.%dgb", vgpu.G, vgpu.GB)
+
+	// Add .me suffix if media extension attribute is present
+	for _, attr := range vgpu.Attr {
+		if attr == types.AttributeMediaExtensions {
+			result += ".me"
+			break
+		}
+	}
+
+	return result, nil
+}

deployments/container/Dockerfile.ubi9

@@ -38,13 +38,21 @@ ARG VERSION="N/A"
 ARG GIT_COMMIT="unknown"
 RUN make PREFIX=/artifacts cmds
 
+RUN cp ./deployments/container/reconfigure-mig.sh /artifacts/reconfigure-mig.sh
+
+ARG TARGETARCH
+RUN OS_ARCH=${TARGETARCH/x86_64/amd64} && OS_ARCH=${OS_ARCH/aarch64/arm64} && curl -o /artifacts/kubectl -L "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/${OS_ARCH}/kubectl"
+RUN chmod a+x /artifacts/kubectl
 
 FROM nvcr.io/nvidia/distroless/go:v3.1.8
 
 ENV NVIDIA_VISIBLE_DEVICES=void
 
 COPY --from=build /artifacts/nvidia-vgpu-dm /usr/bin/nvidia-vgpu-dm
 COPY --from=build /artifacts/nvidia-k8s-vgpu-dm /usr/bin/nvidia-k8s-vgpu-dm
+COPY --from=build /artifacts/reconfigure-mig.sh /usr/bin/reconfigure-mig.sh
+COPY --from=build /artifacts/kubectl /usr/bin/kubectl
+COPY --from=nvcr.io/nvidia/cloud-native/k8s-mig-manager:v0.12.1-ubi9  /usr/bin/nvidia-mig-parted /usr/bin/nvidia-mig-parted
 
 LABEL version="${VERSION}"
 LABEL release="N/A"