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

@@ -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,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 = reconfigureMIG(clientset, selectedMigConfig)
+		if err != nil {
+			return fmt.Errorf("unable to run reconfigure MIG script: %s", err)
+		}
+	} else {
+		log.Info("Disabling MIG if enabled")
+		err = reconfigureMIG(clientset, 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 +615,69 @@ 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 reconfigureMIG(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 := &MIGReconfigOptions{
+		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)
+}
+
+// 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
+}