contrib: check if nvidia drivers are already installed

Signed-off-by: CrazyMax <[email protected]>

Author: crazy-max

contrib/cdisetup/nvidia/nvidia.go

@@ -25,10 +25,13 @@ import (
 // This is example of experimental on-demand setup of a CDI devices.
 // This code is not currently shipping with BuildKit and will probably change.
 
-const (
-	cdiKind        = "nvidia.com/gpu"
-	defaultVersion = "570.0"
-)
+const cdiKind = "nvidia.com/gpu"
+
+// https://github.com/ollama/ollama/blob/b816ff86c923e0290f58f2275e831fc17c29ba37/discover/gpu_linux.go#L33-L43
+var libcudaGlobs = []string{
+	"/usr/lib/*-linux-gnu/libcuda.so*",
+	"/usr/lib/wsl/drivers/*/libcuda.so*",
+}
 
 func init() {
 	cdidevices.Register(cdiKind, &setup{})
@@ -93,25 +96,15 @@ func (s *setup) Run(ctx context.Context) (err error) {
 	}
 
 	var needsDriver bool
-	if nvidiaSmi, err := exec.LookPath("nvidia-smi"); err == nil && nvidiaSmi != "" {
-		if err := run(ctx, []string{nvidiaSmi, "-L"}, pw, dgst); err != nil {
+	if !hasDriversInstalled() {
+		// Check for nvidia-smi or /proc/driver/nvidia to determine if NVIDIA drivers are installed
+		if nvidiaSmi, err := exec.LookPath("nvidia-smi"); err == nil && nvidiaSmi != "" {
+			if err := run(ctx, []string{nvidiaSmi, "-L"}, pw, dgst); err != nil {
+				needsDriver = true
+			}
+		} else if _, err := os.Stat("/proc/driver/nvidia"); err != nil {
 			needsDriver = true
 		}
-	} else if _, err := os.Stat("/proc/driver/nvidia"); err != nil {
-		needsDriver = true
-	}
-
-	var arch string
-	switch runtime.GOARCH {
-	case "amd64":
-		arch = "x86_64"
-	case "arm64":
-		arch = "sbsa"
-		// for non-sbsa could use https://nvidia.github.io/libnvidia-container/stable/deb
-	}
-
-	if arch == "" {
-		return errors.Errorf("unsupported architecture: %s", runtime.GOARCH)
 	}
 
 	if needsDriver {
@@ -122,21 +115,19 @@ func (s *setup) Run(ctx context.Context) (err error) {
 	}
 
 	var dv string
-	if !hasWSLGPU() {
+	if needsDriver {
+		if hasWSLGPU() {
+			return errors.Errorf("NVIDIA drivers are required for WSL with non PCI-based GPUs")
+		}
 		version, err := readVersion()
-		if err != nil && !needsDriver {
+		if err != nil {
 			return errors.Wrapf(err, "failed to read NVIDIA driver version")
 		}
-		if version == "" {
-			version = defaultVersion
-		}
 		var ok bool
 		dv, _, ok = strings.Cut(version, ".")
 		if !ok {
 			return errors.Errorf("failed to parse NVIDIA driver version %q", version)
 		}
-	} else if needsDriver {
-		return errors.Errorf("NVIDIA drivers are required for WSL with non PCI-based GPUs")
 	}
 
 	if err := run(ctx, []string{"apt-get", "update"}, pw, dgst); err != nil {
@@ -147,100 +138,138 @@ func (s *setup) Run(ctx context.Context) (err error) {
 		return err
 	}
 
+	if needsDriver {
+		if err := installDrivers(ctx, dv, pw, dgst); err != nil {
+			return err
+		}
+	}
+	if err := installContainerToolkit(ctx, pw, dgst); err != nil {
+		return err
+	}
+
+	if err := os.MkdirAll("/etc/cdi", 0700); err != nil {
+		return errors.Wrapf(err, "failed to create /etc/cdi")
+	}
+
+	buf := &bytes.Buffer{}
+
+	cmd := exec.CommandContext(ctx, "nvidia-ctk", "cdi", "generate")
+	cmd.Stdout = buf
+	cmd.Stderr = newStream(pw, 2, dgst)
+	if err := cmd.Run(); err != nil {
+		return errors.Wrapf(err, "failed to generate CDI spec")
+	}
+
+	if len(buf.Bytes()) == 0 {
+		return errors.Errorf("nvidia-ctk output is empty")
+	}
+
+	if err := os.WriteFile("/etc/cdi/nvidia.yaml", buf.Bytes(), 0644); err != nil {
+		return errors.Wrapf(err, "failed to write /etc/cdi/nvidia.yaml")
+	}
+
+	return nil
+}
+
+func run(ctx context.Context, args []string, pw progress.Writer, dgst digest.Digest) error {
+	fmt.Fprintf(newStream(pw, 2, dgst), "> %s\n", strings.Join(args, " "))
+	cmd := exec.CommandContext(ctx, args[0], args[1:]...) //nolint:gosec
+	cmd.Stderr = newStream(pw, 2, dgst)
+	cmd.Stdout = newStream(pw, 1, dgst)
+	return cmd.Run()
+}
+
+func installDrivers(ctx context.Context, dv string, pw progress.Writer, dgst digest.Digest) error {
 	const aptDistro = "ubuntu2404"
-	aptURL := "https://developer.download.nvidia.com/compute/cuda/repos/" + aptDistro + "/" + arch + "/"
 
+	var arch string
+	switch runtime.GOARCH {
+	case "amd64":
+		arch = "x86_64"
+	case "arm64":
+		arch = "sbsa"
+		// for non-sbsa could use https://nvidia.github.io/libnvidia-container/stable/deb
+	}
+	if arch == "" {
+		return errors.Errorf("unsupported architecture: %s", runtime.GOARCH)
+	}
+
+	aptURL := "https://developer.download.nvidia.com/compute/cuda/repos/" + aptDistro + "/" + arch + "/"
 	keyTarget := "/usr/share/keyrings/nvidia-cuda-keyring.gpg"
 
 	if _, err := os.Stat(keyTarget); err != nil {
-		fmt.Fprintf(newStream(pw, 2, dgst), "Downloading NVIDIA GPG key\n")
+		fmt.Fprintf(newStream(pw, 2, dgst), "Downloading NVIDIA cuda GPG key\n")
 
 		req, err := http.NewRequestWithContext(ctx, http.MethodGet, aptURL+"3bf863cc.pub", nil)
 		if err != nil {
-			return errors.Wrapf(err, "failed to create request for NVIDIA GPG key")
+			return errors.Wrapf(err, "failed to create request for NVIDIA cuda GPG key")
 		}
 
 		resp, err := http.DefaultClient.Do(req)
 		if err != nil {
-			return errors.Wrapf(err, "failed to download NVIDIA GPG key")
+			return errors.Wrapf(err, "failed to download NVIDIA cuda GPG key")
 		}
 
 		cmd := exec.CommandContext(ctx, "gpg", "--dearmor", "-o", keyTarget)
 		cmd.Stdin = resp.Body
 		cmd.Stderr = newStream(pw, 2, dgst)
 		if err := cmd.Run(); err != nil {
-			return errors.Wrapf(err, "failed to install NVIDIA GPG key")
+			return errors.Wrapf(err, "failed to install NVIDIA cuda GPG key")
 		}
 		resp.Body.Close()
 	}
 
 	if err := os.WriteFile("/etc/apt/sources.list.d/nvidia-cuda.list", []byte("deb [signed-by="+keyTarget+"] "+aptURL+" /"), 0644); err != nil {
-		return errors.Wrapf(err, "failed to add NVIDIA apt repo")
+		return errors.Wrapf(err, "failed to add NVIDIA cuda apt repo")
 	}
 
 	if err := run(ctx, []string{"apt-get", "update"}, pw, dgst); err != nil {
 		return err
 	}
 
-	if needsDriver && dv != "" {
-		// this pretty much never works, is it even worth having?
-		// better approach could be to try to create another chroot/container that is built with same kernel packages as the host
-		// could nvidia-headless-no-dkms- be reusable
-		if err := run(ctx, []string{"apt-get", "install", "-y", "nvidia-driver-" + dv}, pw, dgst); err != nil {
-			return err
-		}
-		_, err := os.Stat("/proc/driver/nvidia")
-		if err != nil {
-			return errors.Wrapf(err, "failed to install NVIDIA kernel module. Please install NVIDIA drivers manually")
-		}
-	}
+	return run(ctx, []string{"apt-get", "install", "-y", "--no-install-recommends",
+		"libnvidia-compute-" + dv,
+		"libnvidia-extra-" + dv,
+		"libnvidia-gl-" + dv,
+		"nvidia-utils-" + dv,
+	}, pw, dgst)
+}
 
-	pkgs := []string{
-		"nvidia-container-toolkit-base",
-	}
-	if dv != "" {
-		pkgs = append(pkgs, []string{
-			"libnvidia-compute-" + dv,
-			"libnvidia-extra-" + dv,
-			"libnvidia-gl-" + dv,
-			"nvidia-utils-" + dv,
-		}...)
-	}
+func installContainerToolkit(ctx context.Context, pw progress.Writer, dgst digest.Digest) error {
+	aptURL := "https://nvidia.github.io/libnvidia-container/"
+	keyTarget := "/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg"
 
-	if err := run(ctx, append([]string{"apt-get", "install", "-y", "--no-install-recommends"}, pkgs...), pw, dgst); err != nil {
-		return err
-	}
+	if _, err := os.Stat(keyTarget); err != nil {
+		fmt.Fprintf(newStream(pw, 2, dgst), "Downloading NVIDIA container toolkit GPG key\n")
 
-	if err := os.MkdirAll("/etc/cdi", 0700); err != nil {
-		return errors.Wrapf(err, "failed to create /etc/cdi")
-	}
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, aptURL+"gpgkey", nil)
+		if err != nil {
+			return errors.Wrapf(err, "failed to create request for NVIDIA container toolkit GPG key")
+		}
 
-	buf := &bytes.Buffer{}
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			return errors.Wrapf(err, "failed to download NVIDIA container toolkit GPG key")
+		}
 
-	cmd := exec.CommandContext(ctx, "nvidia-ctk", "cdi", "generate")
-	cmd.Stdout = buf
-	cmd.Stderr = newStream(pw, 2, dgst)
-	if err := cmd.Run(); err != nil {
-		return errors.Wrapf(err, "failed to generate CDI spec")
+		cmd := exec.CommandContext(ctx, "gpg", "--dearmor", "-o", keyTarget)
+		cmd.Stdin = resp.Body
+		cmd.Stderr = newStream(pw, 2, dgst)
+		if err := cmd.Run(); err != nil {
+			return errors.Wrapf(err, "failed to install NVIDIA container toolkit GPG key")
+		}
+		resp.Body.Close()
 	}
 
-	if len(buf.Bytes()) == 0 {
-		return errors.Errorf("nvidia-ctk output is empty")
+	if err := os.WriteFile("/etc/apt/sources.list.d/nvidia-container-toolkit.list", []byte("deb [signed-by="+keyTarget+"] "+aptURL+"stable/deb/$(ARCH) /"), 0644); err != nil {
+		return errors.Wrapf(err, "failed to add NVIDIA container toolkit apt repo")
 	}
 
-	if err := os.WriteFile("/etc/cdi/nvidia.yaml", buf.Bytes(), 0644); err != nil {
-		return errors.Wrapf(err, "failed to write /etc/cdi/nvidia.yaml")
+	if err := run(ctx, []string{"apt-get", "update"}, pw, dgst); err != nil {
+		return err
 	}
 
-	return nil
-}
-
-func run(ctx context.Context, args []string, pw progress.Writer, dgst digest.Digest) error {
-	fmt.Fprintf(newStream(pw, 2, dgst), "> %s\n", strings.Join(args, " "))
-	cmd := exec.CommandContext(ctx, args[0], args[1:]...) //nolint:gosec
-	cmd.Stderr = newStream(pw, 2, dgst)
-	cmd.Stdout = newStream(pw, 1, dgst)
-	return cmd.Run()
+	return run(ctx, []string{"apt-get", "install", "-y", "nvidia-container-toolkit-base"}, pw, dgst)
 }
 
 func readVersion() (string, error) {
@@ -326,3 +355,13 @@ func hasWSLGPU() bool {
 	_, err := os.Stat("/dev/dxg")
 	return err == nil
 }
+
+func hasDriversInstalled() bool {
+	// Check for libcuda in the standard locations to confirm NVIDIA GPU drivers
+	for _, p := range libcudaGlobs {
+		if matches, err := filepath.Glob(p); err == nil && len(matches) > 0 {
+			return true
+		}
+	}
+	return false
+}