fix 1. todo: squash later

Signed-off-by: Léiyì Zhang <leiyiz@google.com>

Author: leiyiz

cmd/gpu-kubelet-plugin/main.go

@@ -73,11 +73,11 @@ func main() {
 		fmt.Fprintf(os.Stderr, "Error masking NVIDIA driver params: %v\n", err)
 	}
 
-	if len(os.Args) > 1 && os.Args[1] == "prestart-init" {
+	if len(os.Args) > 1 && os.Args[1] == "prestart" {
 		ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGTERM)
 		defer cancel()
 
-		if err := runPrestartInit(ctx); err != nil {
+		if err := runPrestart(ctx); err != nil {
 			fmt.Fprintf(os.Stderr, "Error: %v\n", err)
 			os.Exit(1)
 		}

cmd/gpu-kubelet-plugin/prestart.go

@@ -29,11 +29,9 @@ import (
 
 // Main intent: help users to self-troubleshoot when the GPU driver is not set up
 // properly before installing this DRA driver. In that case, the log of the init
-// container running this script is meant to yield an actionable error message.
+// container running the prestart code is meant to yield an actionable error message.
 // For now, rely on k8s to implement a high-level retry with back-off.
-func runPrestartInit(ctx context.Context) error {
-	// Design goal: long-running init container that retries at constant frequency,
-	// and leaves only upon success (with code 0).
+func runPrestart(ctx context.Context) error {
 	waitS := 10 * time.Second
 	attempt := 0
 
@@ -45,7 +43,7 @@ func runPrestartInit(ctx context.Context) error {
 	}
 
 	driverRootParent := "/driver-root-parent"
-	// Remove trailing slash (if existing) and get last path element.
+	// filepath.Base removes trailing slash (if existing) and get last path element.
 	driverRootPath := filepath.Join(driverRootParent, filepath.Base(nvidiaDriverRoot))
 
 	// Ensure the /driver-root-parent directory exists
@@ -58,7 +56,7 @@ func runPrestartInit(ctx context.Context) error {
 	// once the driver becomes mounted (e.g., once GPU operator provides the driver
 	// on the host at /run/nvidia/driver).
 	fmt.Printf("create symlink: /driver-root -> %s\n", driverRootPath)
-	_ = os.Remove("/driver-root")
+
 	if err := os.Symlink(driverRootPath, "/driver-root"); err != nil {
 		klog.Warningf("Failed to create symlink: %v", err)
 	}
@@ -96,6 +94,10 @@ func validateAndExitOnSuccess(ctx context.Context, nvidiaDriverRoot string, atte
 	// Search specific set of directories (not recursively: not required, and
 	// /driver-root may be a big tree). Limit to first result (multiple results
 	// are a bit of a pathological state, but continue with validation logic).
+
+	// original script does not follow symlink for nvpath but since symlinkm
+	// can also execute so reuse findFirstFile to avoid new func that's largely
+	// duplicative.
 	nvPath := findFirstFile(
 		"nvidia-smi",
 		"/driver-root/opt/bin",
@@ -105,7 +107,6 @@ func validateAndExitOnSuccess(ctx context.Context, nvidiaDriverRoot string, atte
 		"/driver-root/sbin",
 	)
 
-	// Follow symlinks (-L), because `libnvidia-ml.so.1` is typically a link.
 	nvLibPath := findFirstFile(
 		"libnvidia-ml.so.1",
 		"/driver-root/usr/lib64",
@@ -145,6 +146,7 @@ func validateAndExitOnSuccess(ctx context.Context, nvidiaDriverRoot string, atte
 		// hang).
 		fmt.Printf("invoke: env -i LD_PRELOAD=%s %s\n", nvLibPath, nvPath)
 
+		// override default env to just LD_PRELOAD
 		cmd := exec.CommandContext(ctx, nvPath)
 		cmd.Env = []string{"LD_PRELOAD=" + nvLibPath}
 		cmd.Stdout = os.Stdout
@@ -162,6 +164,7 @@ func validateAndExitOnSuccess(ctx context.Context, nvidiaDriverRoot string, atte
 		if exitErr, ok := err.(*exec.ExitError); ok {
 			fmt.Printf("exit code: %d\n", exitErr.ExitCode())
 		} else {
+			// nvidia-smi fails to start. e.g. permission denied etc.
 			fmt.Printf("execution failed: %v, exit code: -1\n", err)
 		}
 	}
@@ -209,6 +212,9 @@ func validateAndExitOnSuccess(ctx context.Context, nvidiaDriverRoot string, atte
 	return false
 }
 
+// findFirstFile finds the first occurrence of filename in the provided
+// directories not recursively.
+// It follows symlinks (similar to find -L).
 func findFirstFile(filename string, dirs ...string) string {
 	for _, dir := range dirs {
 		path := filepath.Join(dir, filename)

cmd/gpu-kubelet-plugin/vfio-device.go

@@ -260,7 +260,7 @@ func (vm *VfioPciManager) acquireUnbindLock(gpu string) error {
 		klog.Infof("[retry %d/%d] Attempting to acquire unbindLock for %s", attempt, lockRetries, gpu)
 
 		// Try to write 1 to acquire the lock
-		err := os.WriteFile(unbindLockFile, []byte("1\n"), 0200)
+		err := os.WriteFile(unbindLockFile, []byte("1\n"), 0644)
 		if err != nil {
 			klog.Warningf("failed to write to unbindLock file %s: %v", unbindLockFile, err)
 		}
@@ -300,8 +300,7 @@ func (vm *VfioPciManager) unbindFromDriver(pciAddress string) error {
 		}
 	}
 
-	unbindFile := filepath.Join(existingDriver, "unbind")
-	if err := os.WriteFile(unbindFile, []byte(pciAddress+"\n"), 0200); err != nil {
+	if err := os.WriteFile(filepath.Join(existingDriver, "unbind"), []byte(pciAddress+"\n"), 0644); err != nil {
 		klog.Errorf("Attempting to unbind %s from its driver failed; err: %v", pciAddress, err)
 		return err
 	}
@@ -318,7 +317,7 @@ func (vm *VfioPciManager) bindToDriver(pciAddress, driver string) error {
 		return fmt.Errorf("driver_override file not found: %v", err)
 	}
 
-	if err := os.WriteFile(driverOverrideFile, []byte(driver+"\n"), 0200); err != nil {
+	if err := os.WriteFile(driverOverrideFile, []byte(driver+"\n"), 0644); err != nil {
 		klog.Errorf("failed to write '%s' to %s", driver, driverOverrideFile)
 		return fmt.Errorf("failed to write to driver_override: %v", err)
 	}
@@ -328,10 +327,10 @@ func (vm *VfioPciManager) bindToDriver(pciAddress, driver string) error {
 		return fmt.Errorf("bind file not found: %v", err)
 	}
 
-	if err := os.WriteFile(bindFile, []byte(pciAddress+"\n"), 0200); err != nil {
-		klog.Errorf("Attempting to bind %s to %s driver failed; err: %v", pciAddress, driver, err)
+	if err := os.WriteFile(bindFile, []byte(pciAddress+"\n"), 0644); err != nil {
+		klog.Errorf("failed to write %s to %s; err: %v", pciAddress, bindFile, err)
 		// attempt to revert driver_override
-		_ = os.WriteFile(driverOverrideFile, []byte("\n"), 0200)
+		_ = os.WriteFile(driverOverrideFile, []byte("\n"), 0644)
 		return fmt.Errorf("failed to write to bind file: %v", err)
 	}
 	return nil

deployments/helm/nvidia-dra-driver-gpu/templates/kubeletplugin.yaml

@@ -54,7 +54,7 @@ spec:
         image: {{ include "nvidia-dra-driver-gpu.fullimage" . }}
         securityContext:
           privileged: true
-        command: [/usr/bin/gpu-kubelet-plugin, prestart-init]
+        command: [/usr/bin/gpu-kubelet-plugin, prestart]
         env:
         - name: NVIDIA_DRIVER_ROOT
           value: "{{ .Values.nvidiaDriverRoot }}"