Add fast path optimization to skip driver reinstall when configuration digest matches

Signed-off-by: Karthik Vetrivel <kvetrivel@nvidia.com>

Author: karthikvetrivel

cmd/driver-manager/main.go

@@ -42,6 +42,7 @@ import (
 const (
 	driverRoot          = "/run/nvidia/driver"
 	driverPIDFile       = "/run/nvidia/nvidia-driver.pid"
+	driverConfigStateFile = "/run/nvidia/nvidia-driver.state"
 	operatorNamespace   = "gpu-operator"
 	pausedStr           = "paused-for-driver-upgrade"
 	defaultDrainTimeout = time.Second * 0
@@ -77,6 +78,8 @@ type config struct {
 	gpuDirectRDMAEnabled       bool
 	useHostMofed               bool
 	kubeconfig                 string
+	driverVersion              string
+	forceReinstall             bool
 }
 
 // ComponentState tracks the deployment state of GPU operator components
@@ -208,6 +211,20 @@ func main() {
 			EnvVars:     []string{"KUBECONFIG"},
 			Value:       "",
 		},
+		&cli.StringFlag{
+			Name:        "driver-version",
+			Usage:       "Desired NVIDIA driver version",
+			Destination: &cfg.driverVersion,
+			EnvVars:     []string{"DRIVER_VERSION"},
+			Value:       "",
+		},
+		&cli.BoolFlag{
+			Name:        "force-reinstall",
+			Usage:       "Force driver reinstall regardless of current state",
+			Destination: &cfg.forceReinstall,
+			EnvVars:     []string{"FORCE_REINSTALL"},
+			Value:       false,
+		},
 	}
 
 	app.Commands = []*cli.Command{
@@ -288,6 +305,26 @@ func (dm *DriverManager) uninstallDriver() error {
 		return fmt.Errorf("failed to evict GPU operator components: %w", err)
 	}
 
+	if dm.shouldSkipUninstall() {
+		dm.log.Info("Fast path activated: desired driver version and configuration already present")
+
+		// Clean up stale artifacts from previous container before rescheduling operands
+		dm.log.Info("Cleaning up stale mounts and state files...")
+
+		// Unmount stale rootfs from previous container
+		if err := dm.unmountRootfs(); err != nil {
+			return fmt.Errorf("failed to unmount stale rootfs: %w", err)
+		}
+
+		// Remove stale PID file from previous container
+		dm.removePIDFile()
+
+		if err := dm.rescheduleGPUOperatorComponents(); err != nil {
+			dm.log.Warnf("Failed to reschedule GPU operator components: %v", err)
+		}
+		return nil
+	}
+
 	drainOpts := kube.DrainOptions{
 		Force:              dm.config.drainUseForce,
 		DeleteEmptyDirData: dm.config.drainDeleteEmptyDirData,
@@ -629,6 +666,61 @@ func (dm *DriverManager) isDriverLoaded() bool {
 	return err == nil
 }
 
+// readStoredDigest reads the driver configuration digest from the state file
+func readStoredDigest() (string, error) {
+	data, err := os.ReadFile(driverConfigStateFile)
+	if err != nil {
+		return "", err
+	}
+	return strings.TrimSpace(string(data)), nil
+}
+
+// getCurrentDigest returns the digest from the environment variable
+func getCurrentDigest() string {
+	return os.Getenv("DRIVER_CONFIG_DIGEST")
+}
+
+// shouldUpdateDriverConfig checks if the driver configuration needs to be updated
+func (dm *DriverManager) shouldUpdateDriverConfig() bool {
+	if !dm.isDriverLoaded() {
+		return true
+	}
+
+	currentDigest := getCurrentDigest()
+	if currentDigest == "" {
+		dm.log.Warn("DRIVER_CONFIG_DIGEST env var not set, assuming config changed")
+		return true
+	}
+
+	storedDigest, err := readStoredDigest()
+	if err != nil {
+		if os.IsNotExist(err) {
+			dm.log.Info("No previous driver configuration found")
+		} else {
+			dm.log.Warnf("Failed to read driver config state file: %v", err)
+		}
+		return true
+	}
+
+	return currentDigest != storedDigest
+}
+
+func (dm *DriverManager) shouldSkipUninstall() bool {
+	if dm.config.forceReinstall {
+		dm.log.Info("Force reinstall is enabled, proceeding with driver uninstall")
+		return false
+	}
+
+	if !dm.shouldUpdateDriverConfig() {
+		dm.log.Info("Driver is loaded with matching config, enabling fast path")
+		return true
+	}
+
+	// Driver not loaded or config changed - proceed with cleanup
+	dm.log.Info("Proceeding with cleanup operations")
+	return false
+}
+
 func (dm *DriverManager) isNouveauLoaded() bool {
 	_, err := os.Stat("/sys/module/nouveau/refcnt")
 	return err == nil
@@ -639,6 +731,12 @@ func (dm *DriverManager) unloadNouveau() error {
 	return unix.DeleteModule("nouveau", 0)
 }
 
+func (dm *DriverManager) removePIDFile() {
+	if err := os.Remove(driverPIDFile); err != nil && !os.IsNotExist(err) {
+		dm.log.Warnf("Failed to remove PID file %s: %v", driverPIDFile, err)
+	}
+}
+
 func (dm *DriverManager) cleanupDriver() error {
 	dm.log.Info("Cleaning up NVIDIA driver")
 
@@ -652,12 +750,7 @@ func (dm *DriverManager) cleanupDriver() error {
 		return fmt.Errorf("failed to unmount rootfs: %w", err)
 	}
 
-	// Remove PID file
-	if _, err := os.Stat(driverPIDFile); err == nil {
-		if err := os.Remove(driverPIDFile); err != nil {
-			dm.log.Warnf("Failed to remove PID file %s: %v", driverPIDFile, err)
-		}
-	}
+	dm.removePIDFile()
 
 	return nil
 }