Add MIG config support when MIG-backed vGPU type

This change adds MIG configuration support via nvidia-mig-parted when a
MIG-backed vGPU type is selected before the vGPU configuration takes
place.

Specifically:

- Add CLI flags for MIG configuration options.
- Include the nvidia-mig-parted binary in the container image.
- Parse vGPU type to detect MIG requirements and configure MIG before
  vGPUs via nvidia-mig-parted.
- nvidia-mig-parted requires NVML. When NVML is not available, skip
  MIG configuration and proceed to vGPU configuration. This keeps
  the vGPU Device Manager backwards compatible with components that do
  not make NVML available to it (e.g. via the env var LD_PRELOAD).

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

Author: mresvanis

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

@@ -26,6 +26,7 @@ import (
 	"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 +39,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 +49,13 @@ const (
 	vGPUConfigStateLabel = "nvidia.com/vgpu.config.state"
 	pluginStateLabel     = "nvidia.com/gpu.deploy.sandbox-device-plugin"
 	validatorStateLabel  = "nvidia.com/gpu.deploy.sandbox-validator"
+
+	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 +65,24 @@ var (
 	configFileFlag        string
 	defaultVGPUConfigFlag 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 +171,70 @@ func main() {
 			Destination: &defaultVGPUConfigFlag,
 			EnvVars:     []string{"DEFAULT_VGPU_CONFIG"},
 		},
+		&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 +383,10 @@ func updateConfig(clientset *kubernetes.Clientset, selectedConfig string) error
 		return fmt.Errorf("unable to shutdown gpu operands: %v", err)
 	}
 
+	if err := handleMIGConfiguration(clientset, selectedConfig); err != nil {
+		return fmt.Errorf("unable to handle MIG configuration: %v", err)
+	}
+
 	log.Info("Applying the selected vGPU device configuration to the node")
 	err = applyConfig(selectedConfig)
 	if err != nil {
@@ -504,3 +595,101 @@ func setNodeLabelValue(clientset *kubernetes.Clientset, label, value string) err
 
 	return nil
 }
+
+func handleMIGConfiguration(clientset kubernetes.Interface, selectedConfig string) error {
+	if err := isNVMLAvailable(); err != nil {
+		log.Infof("Skipping MIG configuration due to NVML error: %v, proceeding with vGPU configuration", err)
+		return nil
+	}
+
+	migConfig, err := determineMIGConfig(selectedConfig)
+	if err != nil {
+		return err
+	}
+
+	log.Infof("Selected MIG configuration: %s", migConfig)
+	return updateMIGConfig(clientset.(*kubernetes.Clientset), migConfig)
+}
+
+func determineMIGConfig(selectedConfig string) (string, error) {
+	vgpuType, err := types.ParseVGPUType(selectedConfig)
+	if err != nil {
+		return "", fmt.Errorf("unable to parse vGPU type: %s", err)
+	}
+
+	if vgpuType.G > 0 {
+		migConfig, err := convertToMIGConfigFormat(selectedConfig)
+		if err != nil {
+			return "", fmt.Errorf("unable to convert vGPU type config to MIG config: %s", err)
+		}
+		return migConfig, nil
+	}
+
+	return migConfigDisabled, nil
+}
+
+// convertToMIGConfigFormat converts a vGPU type string to the MIG config format.
+// Examples: "A100-1-5C" -> "all-1g.5gb", "A100-1-5CME" -> "all-1g.5gb.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
+}
+
+func updateMIGConfig(clientset *kubernetes.Clientset, migConfigValue string) error {
+	gpuClients, err := parseGPUCLientsFile(gpuClientsFileFlag)
+	if err != nil {
+		return fmt.Errorf("error parsing host's GPU clients file: %s", err)
+	}
+
+	opts := &reconfigureMIGOptions{
+		NodeName:                   nodeNameFlag,
+		MIGPartedConfigFile:        migPartedConfigFileFlag,
+		SelectedMIGConfig:          migConfigValue,
+		WithReboot:                 withRebootFlag,
+		WithShutdownHostGPUClients: withShutdownHostGPUClientsFlag,
+		HostRootMount:              hostRootMountFlag,
+		HostNvidiaDir:              hostNvidiaDirFlag,
+		HostMIGManagerStateFile:    hostMigManagerStateFileFlag,
+		HostGPUClientServices:      gpuClients.SystemdServices,
+		HostKubeletService:         hostKubeletSystemdServiceFlag,
+	}
+
+	return reconfigureMIG(clientset, opts)
+}
+
+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
+}

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

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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
+ *
+ *     http://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.
+ */
+
+package main
+
+import (
+	log "github.com/sirupsen/logrus"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+func isNVMLAvailable() error {
+	nvmlLib := nvml.New()
+
+	if ret := nvmlLib.Init(); ret != nvml.SUCCESS {
+		return ret
+	}
+
+	if ret := nvmlLib.Shutdown(); ret != nvml.SUCCESS {
+		log.Warnf("error shutting down NVML: %v", ret)
+	}
+
+	return nil
+}