Add nvidia installer tool for installing NVIDIA drivers in buildkite tests.

PiperOrigin-RevId: 574371646

.buildkite/pipeline.yaml

@@ -168,6 +168,16 @@ steps:
       - git checkout go && git clean -xf .
       - GOOS=linux GOARCH=mips go build $$PACKAGES
 
+  # GPU workflow.
+  - <<: *common
+    <<: *source_test
+    label: "GPU tests"
+    commands:
+      - uname -r
+      - gcc --version
+      - make sudo TARGETS=//tools/gpu:main ARGS="install --latest" || cat /var/log/nvidia-installer.log
+    agents:
+      queue: gpu
   # Release workflow.
   - <<: *common
     <<: *source_test

pkg/sentry/devices/nvproxy/BUILD

@@ -27,7 +27,10 @@ go_library(
         "uvm_unsafe.go",
         "version.go",
     ],
-    visibility = ["//pkg/sentry:internal"],
+    visibility = [
+        "//pkg/sentry:internal",
+        "//tools/gpu:__subpackages__",
+    ],
     deps = [
         "//pkg/abi/linux",
         "//pkg/abi/nvgpu",

pkg/sentry/devices/nvproxy/nvproxy.go

@@ -39,7 +39,8 @@ func Register(vfsObj *vfs.VirtualFilesystem, uvmDevMajor uint32) error {
 	if err != nil {
 		return fmt.Errorf("failed to get Nvidia driver version: %w", err)
 	}
-	version, err := driverVersionFrom(versionStr)
+	log.Debugf("NVIDIA driver version: %s", versionStr)
+	version, err := DriverVersionFrom(versionStr)
 	if err != nil {
 		return fmt.Errorf("failed to parse Nvidia driver version %s: %w", versionStr, err)
 	}
@@ -47,10 +48,9 @@ func Register(vfsObj *vfs.VirtualFilesystem, uvmDevMajor uint32) error {
 	if !ok {
 		return fmt.Errorf("unsupported Nvidia driver version: %s", versionStr)
 	}
-	log.Infof("Nvidia driver version: %s", versionStr)
 	nvp := &nvproxy{
 		objsLive: make(map[nvgpu.Handle]*object),
-		abi:      abiCons(),
+		abi:      abiCons.cons(),
 	}
 	for minor := uint32(0); minor <= nvgpu.NV_CONTROL_DEVICE_MINOR; minor++ {
 		if err := vfsObj.RegisterDevice(vfs.CharDevice, nvgpu.NV_MAJOR_DEVICE_NUMBER, minor, &frontendDevice{

pkg/sentry/devices/nvproxy/nvproxy_test.go

@@ -23,8 +23,8 @@ import (
 func TestInit(t *testing.T) {
 	// Test that initializing all driverABI works (does not panic or anything).
 	Init()
-	for _, cons := range abis {
-		cons()
+	for _, abi := range abis {
+		abi.cons()
 	}
 }
 
@@ -36,3 +36,15 @@ func TestNVOS21ParamsSize(t *testing.T) {
 		t.Errorf("SizeofNVOS21ParametersV535(%#08x) != SizeofNVOS21Parameters(%#08x)", nvgpu.SizeofNVOS21ParametersV535, nvgpu.SizeofNVOS21Parameters)
 	}
 }
+
+// TestAllSupportedHashesPresent tests that all the supported versions in nvproxy have hash entries
+// in this tool's map. If you're here because of failures run:
+// `make sudo TARGETS=//tools/gpu:main ARGS="checksum"`and fix mismatches in supported drivers.
+func TestAllSupportedHashesPresent(t *testing.T) {
+	Init()
+	for version, checksum := range GetSupportedDriversAndChecksums() {
+		if checksum == "" {
+			t.Errorf("unexpected empty value for driver %q", version.String())
+		}
+	}
+}

pkg/sentry/devices/nvproxy/version.go

@@ -23,40 +23,74 @@ import (
 	"gvisor.dev/gvisor/pkg/sync"
 )
 
-type driverVersion struct {
+// DriverVersion represents a NVIDIA driver version patch release.
+type DriverVersion struct {
 	major int
 	minor int
 	patch int
 }
 
-func driverVersionFrom(version string) (driverVersion, error) {
+// NewDriverVersion returns a new driver version.
+func NewDriverVersion(major, minor, patch int) DriverVersion {
+	return DriverVersion{major, minor, patch}
+}
+
+// DriverVersionFrom returns a DriverVersion from a string.
+func DriverVersionFrom(version string) (DriverVersion, error) {
 	parts := strings.Split(version, ".")
 	if len(parts) != 3 {
-		return driverVersion{}, fmt.Errorf("invalid format of version string %q", version)
+		return DriverVersion{}, fmt.Errorf("invalid format of version string %q", version)
 	}
 	var (
-		res driverVersion
+		res DriverVersion
 		err error
 	)
 	res.major, err = strconv.Atoi(parts[0])
 	if err != nil {
-		return driverVersion{}, fmt.Errorf("invalid format for major version %q: %v", version, err)
+		return DriverVersion{}, fmt.Errorf("invalid format for major version %q: %v", version, err)
 	}
 	res.minor, err = strconv.Atoi(parts[1])
 	if err != nil {
-		return driverVersion{}, fmt.Errorf("invalid format for minor version %q: %v", version, err)
+		return DriverVersion{}, fmt.Errorf("invalid format for minor version %q: %v", version, err)
 	}
 	res.patch, err = strconv.Atoi(parts[2])
 	if err != nil {
-		return driverVersion{}, fmt.Errorf("invalid format for patch version %q: %v", version, err)
+		return DriverVersion{}, fmt.Errorf("invalid format for patch version %q: %v", version, err)
 	}
 	return res, nil
 }
 
-func (v driverVersion) String() string {
+func (v DriverVersion) String() string {
 	return fmt.Sprintf("%02d.%02d.%02d", v.major, v.minor, v.patch)
 }
 
+// Equals returns true if the two driver versions are equal.
+func (v DriverVersion) Equals(other DriverVersion) bool {
+	return v.major == other.major && v.minor == other.minor && v.patch == other.patch
+}
+
+// IsGreaterThan returns the "greater" driver version.
+// IsGreaterThan returns true if v is more recent than other, assuming v and other are on the same
+// dev branch.
+func (v DriverVersion) IsGreaterThan(other DriverVersion) DriverVersion {
+	switch {
+	case v.major > other.major:
+		return v
+	case other.major > v.major:
+		return other
+	case v.minor > other.minor:
+		return v
+	case other.minor > v.minor:
+		return other
+	case v.patch > other.patch:
+		return v
+	case other.patch > v.patch:
+		return other
+	default:
+		return v
+	}
+}
+
 type frontendIoctlHandler func(fi *frontendIoctlState) (uintptr, error)
 type controlCmdHandler func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS54Parameters) (uintptr, error)
 type allocationClassHandler func(fi *frontendIoctlState, ioctlParams *nvgpu.NVOS64ParametersV535, isNVOS64 bool) (uintptr, error)
@@ -66,6 +100,13 @@ type uvmIoctlHandler func(ui *uvmIoctlState) (uintptr, error)
 // This indirection exists to avoid memory usage from unused driver ABIs.
 type driverABIFunc func() *driverABI
 
+// abiConAndChecksum couples the driver's abiConstructor to the SHA256 checksum of its linux .run
+// driver installer file from NVIDIA.
+type abiConAndChecksum struct {
+	cons     driverABIFunc
+	checksum string
+}
+
 // driverABI defines the Nvidia kernel driver ABI proxied at a given version.
 //
 // The Nvidia driver's ioctl interface branches widely at various places in the
@@ -89,21 +130,27 @@ type driverABI struct {
 
 // abis is a global map containing all supported Nvidia driver ABIs. This is
 // initialized on Init() and is immutable henceforth.
-var abis map[driverVersion]driverABIFunc
+var abis map[DriverVersion]abiConAndChecksum
 var abisOnce sync.Once
 
-func addDriverABI(major, minor, patch int, cons driverABIFunc) driverABIFunc {
+// Note: runfileChecksum is the checksum of the .run file of the driver installer for linux from
+// nvidia.
+// To add a new version, add in support as normal and add the "addDriverABI" call for your version.
+// Run `make sudo TARGETS=//tools/gpu:main ARGS="checksum"` and fill in mismatches.
+func addDriverABI(major, minor, patch int, runfileChecksum string, cons driverABIFunc) driverABIFunc {
 	if abis == nil {
-		abis = make(map[driverVersion]driverABIFunc)
+		abis = make(map[DriverVersion]abiConAndChecksum)
 	}
-	abis[driverVersion{major, minor, patch}] = cons
+	version := NewDriverVersion(major, minor, patch)
+	abis[version] = abiConAndChecksum{cons: cons, checksum: runfileChecksum}
 	return cons
 }
 
 // Init initializes abis global map.
 func Init() {
 	abisOnce.Do(func() {
-		v525_60_13 := addDriverABI(525, 60, 13, func() *driverABI {
+		v525_60_13Checksum := "dce1c184f9f038be72237ccd29c66bb151077f6037f1c158c83d582bd2dba8ca"
+		v525_60_13 := addDriverABI(525, 60, 13, v525_60_13Checksum, func() *driverABI {
 			// 525.60.13 is the earliest driver version supported by nvproxy. Since
 			// there is no parent to inherit from, the driverABI needs to be constructed
 			// with the entirety of the nvproxy functionality at this version.
@@ -257,8 +304,11 @@ func Init() {
 
 		// The following versions do not exist on the main branch. They branched off
 		// the main branch at 525.89.02.
-		v525_105_17 := addDriverABI(525, 105, 17, v525_89_02)
-		_ = addDriverABI(525, 125, 06, v525_105_17)
+		v525_105_17Checksum := "c635a21a282c9b53485f19ebb64a0f4b536a968b94d4d97629e0bc547a58142a"
+		v525_105_17 := addDriverABI(525, 105, 17, v525_105_17Checksum, v525_89_02)
+
+		v525_125_06Checksum := "b5275689f4a833c37a507717ac8f0ee2f1f5cd2b7e236ffa70aad8dfb7455b9d"
+		_ = addDriverABI(525, 125, 06, v525_125_06Checksum, v525_105_17)
 
 		// v535.43.02 is an intermediate unqualified version from the main branch.
 		v535_43_02 := func() *driverABI {
@@ -270,7 +320,18 @@ func Init() {
 			return abi
 		}
 
-		v535_54_03 := addDriverABI(535, 54, 03, v535_43_02)
-		_ = addDriverABI(535, 104, 05, v535_54_03)
+		v535_54_03Checksum := "454764f57ea1b9e19166a370f78be10e71f0626438fb197f726dc3caf05b4082"
+		v535_54_03 := addDriverABI(535, 54, 03, v535_54_03Checksum, v535_43_02)
+		v535_104_05Checksum := "2f9d609d1da770beee757636635c46e7ed8253ade887b87c7a5482e33fcbedc9"
+		_ = addDriverABI(535, 104, 05, v535_104_05Checksum, v535_54_03)
 	})
 }
+
+// GetSupportedDriversAndChecksums returns supported driver ABIs.
+func GetSupportedDriversAndChecksums() map[DriverVersion]string {
+	versions := make(map[DriverVersion]string, len(abis))
+	for version, abi := range abis {
+		versions[version] = abi.checksum
+	}
+	return versions
+}

tools/gpu/BUILD

@@ -0,0 +1,17 @@
+load("//tools:defs.bzl", "go_binary")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+go_binary(
+    name = "main",
+    srcs = ["main.go"],
+    deps = [
+        "//pkg/log",
+        "//pkg/sentry/devices/nvproxy",
+        "//runsc/flag",
+        "//tools/gpu/drivers",
+    ],
+)

tools/gpu/drivers/BUILD

@@ -0,0 +1,23 @@
+load("//tools:defs.bzl", "go_library", "go_test")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    licenses = ["notice"],
+)
+
+go_library(
+    name = "drivers",
+    srcs = ["install_driver.go"],
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/log",
+        "//pkg/sentry/devices/nvproxy",
+    ],
+)
+
+go_test(
+    name = "drivers_test",
+    srcs = ["install_driver_test.go"],
+    library = ":drivers",
+    deps = ["//pkg/sentry/devices/nvproxy"],
+)

tools/gpu/drivers/install_driver.go

@@ -0,0 +1,272 @@
+// Copyright 2023 The gVisor Authors.
+//
+// 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 drivers contains methods to download and install drivers.
+package drivers
+
+import (
+	"context"
+	"crypto/sha256"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"os/exec"
+	"strings"
+
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy"
+)
+
+const (
+	nvidiaSMIPath       = "/usr/bin/nvidia-smi"
+	nvidiaUninstallPath = "/usr/bin/nvidia-uninstall"
+	nvidiaBaseURL       = "https://us.download.nvidia.com/tesla/"
+)
+
+func init() {
+	nvproxy.Init()
+}
+
+// Installer handles the logic to install drivers.
+type Installer struct {
+	requestedVersion nvproxy.DriverVersion
+	// include functions so they can be mocked in tests.
+	getSupportedDriverFunc func() map[nvproxy.DriverVersion]string
+	getCurrentDriverFunc   func() (nvproxy.DriverVersion, error)
+	downloadFunction       func(context.Context, string) (io.ReadCloser, error)
+	installFunction        func(string) error
+}
+
+// NewInstaller returns a driver installer instance.
+func NewInstaller(requestedVersion string, latest bool) (*Installer, error) {
+	ret := &Installer{
+		getSupportedDriverFunc: nvproxy.GetSupportedDriversAndChecksums,
+		getCurrentDriverFunc:   getCurrentDriver,
+		downloadFunction:       DownloadDriver,
+		installFunction:        installDriver,
+	}
+	switch {
+	case latest:
+		for v := range ret.getSupportedDriverFunc() {
+			ret.requestedVersion = v.IsGreaterThan(ret.requestedVersion)
+		}
+	default:
+		d, err := nvproxy.DriverVersionFrom(requestedVersion)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse requested driver version: %w", err)
+		}
+		ret.requestedVersion = d
+	}
+
+	return ret, nil
+}
+
+// MaybeInstall installs a driver if 1) no driver is present on the system already or 2) the
+// driver currently installed does not match the requested version.
+func (i *Installer) MaybeInstall(ctx context.Context) error {
+	// If we don't support the driver, don't attempt to install it.
+	driver, supported := i.getRequestedDriver()
+	if !supported {
+		return fmt.Errorf("requested driver %q is not supported", i.requestedVersion)
+	}
+
+	existingDriver, err := i.getCurrentDriverFunc()
+	if err != nil {
+		log.Warningf("failed to get current driver: %v", err)
+	}
+	if existingDriver.Equals(driver) {
+		log.Infof("Driver already installed: %s", i.requestedVersion)
+		return nil
+	}
+
+	if !existingDriver.Equals(nvproxy.DriverVersion{}) {
+		log.Infof("Uninstalling driver: %s", existingDriver)
+		if err := i.uninstallDriver(ctx, existingDriver.String()); err != nil {
+			return fmt.Errorf("failed to uninstall driver: %w", err)
+		}
+		log.Infof("Driver uninstalled: %s", i.requestedVersion)
+	}
+
+	log.Infof("Downloading driver: %s", i.requestedVersion)
+	reader, err := i.downloadFunction(ctx, i.requestedVersion.String())
+	if err != nil {
+		return fmt.Errorf("failed to download driver: %w", err)
+	}
+
+	f, err := os.CreateTemp("", "")
+	if err != nil {
+		return fmt.Errorf("failed to open driver file: %w", err)
+	}
+	defer os.Remove(f.Name())
+	if err := i.writeAndCheck(f, reader, driver); err != nil {
+		f.Close()
+		return fmt.Errorf("writeAndCheck: %w", err)
+	}
+	if err := f.Chmod(0755); err != nil {
+		return fmt.Errorf("failed to chmod: %w", err)
+	}
+	if err := f.Close(); err != nil {
+		return fmt.Errorf("failed to close driver file: %w", err)
+	}
+	log.Infof("Driver downloaded: %s", i.requestedVersion)
+	log.Infof("Installing driver: %s", i.requestedVersion)
+	if err := i.installFunction(f.Name()); err != nil {
+		return fmt.Errorf("failed to install driver: %w", err)
+	}
+	log.Infof("Installation Complete!")
+	return nil
+}
+
+func (i *Installer) uninstallDriver(ctx context.Context, driverVersion string) error {
+	exec.Command(nvidiaUninstallPath, "-s", driverVersion)
+	cmd := exec.Command(nvidiaUninstallPath, "-s")
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("failed to run nvidia-uninstall: %w", err)
+	}
+	return nil
+}
+
+func (i *Installer) writeAndCheck(f *os.File, reader io.ReadCloser, driverVersion nvproxy.DriverVersion) error {
+	checksum := sha256.New()
+	buf := make([]byte, 1024*1024)
+	for {
+		n, err := reader.Read(buf[0:])
+		if err != nil && err != io.EOF {
+			return fmt.Errorf("failed to read: %w", err)
+		}
+		if n == 0 || err == io.EOF {
+			break
+		}
+		if _, err := checksum.Write(buf[:n]); err != nil {
+			return fmt.Errorf("failed to write: %w", err)
+		}
+		if _, err := f.Write(buf[:n]); err != nil {
+			return fmt.Errorf("failed to write: %w", err)
+		}
+	}
+	gotChecksum := fmt.Sprintf("%x", checksum.Sum(nil))
+	wantChecksum := i.getSupportedDriverFunc()[driverVersion]
+	if gotChecksum != wantChecksum {
+		return fmt.Errorf("driver %q checksum mismatch: got %q, want %q", driverVersion, gotChecksum, wantChecksum)
+	}
+	return nil
+}
+
+func getCurrentDriver() (nvproxy.DriverVersion, error) {
+	_, err := os.Stat(nvidiaSMIPath)
+	// If the nvidia-smi executable does not exist, then we don't have a driver installed.
+	if os.IsNotExist(err) {
+		return nvproxy.DriverVersion{}, fmt.Errorf("nvidia-smi does not exist at path: %q", nvidiaSMIPath)
+	}
+	if err != nil {
+		return nvproxy.DriverVersion{}, fmt.Errorf("failed to stat nvidia-smi: %w", err)
+	}
+	out, err := exec.Command(nvidiaSMIPath, []string{"--query-gpu", "driver_version", "--format=csv,noheader"}...).CombinedOutput()
+	if err != nil {
+		log.Warningf("failed to run nvidia-smi: %v", err)
+		return nvproxy.DriverVersion{}, fmt.Errorf("failed to run nvidia-smi: %w", err)
+	}
+	return nvproxy.DriverVersionFrom(strings.TrimSpace(string(out)))
+}
+
+func (i *Installer) getRequestedDriver() (nvproxy.DriverVersion, bool) {
+	for version := range i.getSupportedDriverFunc() {
+		if version == i.requestedVersion {
+			return version, true
+		}
+	}
+	return nvproxy.DriverVersion{}, false
+}
+
+// ListSupportedDrivers prints the driver to stderr in a format that can be
+// consumed by the Makefile to iterate tests across drivers.
+func ListSupportedDrivers() {
+	supportedDrivers := nvproxy.GetSupportedDriversAndChecksums()
+	list := make([]string, 0, len(supportedDrivers))
+	for version := range nvproxy.GetSupportedDriversAndChecksums() {
+		list = append(list, version.String())
+	}
+	fmt.Println(strings.Join(list, " "))
+}
+
+// ChecksumDriver downloads and returns the SHA265 checksum of the driver.
+func ChecksumDriver(ctx context.Context, driverVersion string) (string, error) {
+	f, err := DownloadDriver(ctx, driverVersion)
+	if err != nil {
+		return "", fmt.Errorf("failed to download driver: %w", err)
+	}
+	checksum := sha256.New()
+	for {
+		n, err := io.Copy(checksum, f)
+		if err == io.EOF || n == 0 {
+			break
+		}
+		if err != nil {
+			return "", fmt.Errorf("failed to copy driver: %w", err)
+		}
+	}
+	return fmt.Sprintf("%x", checksum.Sum(nil)), nil
+}
+
+// DownloadDriver downloads the requested driver and returns the binary as a []byte so it can be
+// checked before written to disk.
+func DownloadDriver(ctx context.Context, driverVersion string) (io.ReadCloser, error) {
+	url := fmt.Sprintf("%s%s/NVIDIA-Linux-x86_64-%s.run", nvidiaBaseURL, driverVersion, driverVersion)
+	resp, err := http.Get(url)
+	if err != nil {
+		return nil, fmt.Errorf("failed to download driver: %w", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("failed to download driver with status: %w", err)
+	}
+	return resp.Body, nil
+}
+
+func installDriver(driverPath string) error {
+	// Certain VMs can be broken if we attempt to install drivers on them. Do a simple check of the
+	// PCI device to ensure we have a GPU attached.
+	out, err := exec.Command("lspci").CombinedOutput()
+	if err != nil {
+		return fmt.Errorf("failed to run lspci: %w out: %s", err, string(out))
+	}
+	if !strings.Contains(string(out), "NVIDIA") {
+		return fmt.Errorf("no NVIDIA PCI device on host:\n%s", string(out))
+	}
+
+	driverArgs := strings.Split("--dkms -a -s --no-drm --install-libglvnd", " ")
+	cmd := exec.Command(driverPath, driverArgs...)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	/*
+		cmd.Env = append(os.Environ(),
+			"IGNORE_CC_MISMATCH=1",
+			"LLVM=1",
+			"LLVM_IS=1",
+		)
+	*/
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("failed to run nvidia-install: %w out: %s", err, string(out))
+	}
+
+	cmd = exec.Command(nvidiaSMIPath)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("failed to run nvidia-install: %w out: %s", err, string(out))
+	}
+	return nil
+}

tools/gpu/drivers/install_driver_test.go

@@ -0,0 +1,123 @@
+// Copyright 2023 The gVisor Authors.
+//
+// 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 drivers
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"fmt"
+	"io"
+	"strings"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy"
+)
+
+// TestVersionInstalled tests when the version is already installed.
+func TestVersionInstalled(t *testing.T) {
+	ctx := context.Background()
+	versionContent := []byte("some cool content")
+	checksum := fmt.Sprintf("%x", sha256.Sum256(versionContent))
+	version := nvproxy.NewDriverVersion(1, 2, 3)
+	getFunction := func() (nvproxy.DriverVersion, error) { return version, nil }
+	downloadFunction := func(context.Context, string) (io.ReadCloser, error) { return nil, fmt.Errorf("should not get here") }
+	installer := &Installer{
+		requestedVersion:       version,
+		getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string { return map[nvproxy.DriverVersion]string{version: checksum} },
+		getCurrentDriverFunc:   getFunction,
+		downloadFunction:       downloadFunction,
+	}
+	if err := installer.MaybeInstall(ctx); err != nil {
+		t.Fatalf("Installation failed: %v", err)
+	}
+}
+
+// TestVersionNotSupported tests when the version is not supported.
+func TestVersionNotSupported(t *testing.T) {
+	ctx := context.Background()
+	unsupportedVersion := nvproxy.NewDriverVersion(1, 2, 3)
+	installer := &Installer{
+		requestedVersion: unsupportedVersion,
+		getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string {
+			return map[nvproxy.DriverVersion]string{
+				nvproxy.NewDriverVersion(2, 3, 4): "other version",
+			}
+		},
+	}
+	err := installer.MaybeInstall(ctx)
+	if err == nil {
+		t.Fatalf("Installation succeeded, want error")
+	}
+	if !strings.Contains(err.Error(), "not supported") {
+		t.Errorf("Installation failed, want error containing 'not supported' got: %s", err.Error())
+	}
+}
+
+// TestShaMismatch tests when a checksum of a driver doesn't match what's in the map.
+func TestShaMismatch(t *testing.T) {
+	ctx := context.Background()
+	version := nvproxy.NewDriverVersion(1, 2, 3)
+	installer := &Installer{
+		requestedVersion: version,
+		getCurrentDriverFunc: func() (nvproxy.DriverVersion, error) {
+			return nvproxy.DriverVersion{}, nil
+		},
+		getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string {
+			return map[nvproxy.DriverVersion]string{
+				version: "mismatched checksum",
+			}
+		},
+		downloadFunction: func(context.Context, string) (io.ReadCloser, error) {
+			reader := bytes.NewReader([]byte("some content"))
+			return io.NopCloser(reader), nil
+		},
+	}
+	err := installer.MaybeInstall(ctx)
+	if err == nil {
+		t.Fatalf("Installation succeeded, want error")
+	}
+	if !strings.Contains(err.Error(), "checksum mismatch") {
+		t.Errorf("Installation failed, want error containing 'mismatch checksum' got: %s", err.Error())
+	}
+}
+
+// TestDriverInstalls tests the successful installation of a driver.
+func TestDriverInstalls(t *testing.T) {
+	ctx := context.Background()
+	content := []byte("some content")
+	checksum := fmt.Sprintf("%x", sha256.Sum256(content))
+	version := nvproxy.NewDriverVersion(1, 2, 3)
+	installer := &Installer{
+		requestedVersion: version,
+		getCurrentDriverFunc: func() (nvproxy.DriverVersion, error) {
+			return nvproxy.DriverVersion{}, nil
+		},
+		getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string {
+			return map[nvproxy.DriverVersion]string{
+				version: checksum,
+			}
+		},
+		downloadFunction: func(context.Context, string) (io.ReadCloser, error) {
+			reader := bytes.NewReader(content)
+			return io.NopCloser(reader), nil
+		},
+		installFunction: func(_ string) error {
+			return nil
+		},
+	}
+	if err := installer.MaybeInstall(ctx); err != nil {
+		t.Fatalf("Installation failed: %v", err)
+	}
+}

tools/gpu/main.go

@@ -0,0 +1,118 @@
+// Copyright 2023 The gVisor Authors.
+//
+// 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 downloads and installs drivers.
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy"
+	"gvisor.dev/gvisor/runsc/flag"
+	"gvisor.dev/gvisor/tools/gpu/drivers"
+)
+
+const (
+	installCmdStr       = "install"
+	installDescription  = "installs a driver on the host machine"
+	checksumCmdStr      = "checksum"
+	checksumDescription = "computes the sha256 checksum for a given driver version"
+	listCmdStr          = "list"
+	listDescription     = "lists the supported drivers"
+)
+
+var (
+	// Install installs a give driver on the host machine.
+	installCmd = flag.NewFlagSet(installCmdStr, flag.ContinueOnError)
+	latest     = installCmd.Bool("latest", false, "install the latest supported driver")
+	version    = installCmd.String("version", "", "version of the driver")
+
+	// Validates all supported driver's checksums of each driver's .run file from the nvidia site.
+	checksumCmd = flag.NewFlagSet(checksumCmdStr, flag.ContinueOnError)
+
+	// The list command returns the list of supported drivers from this tool.
+	listCmd = flag.NewFlagSet(listCmdStr, flag.ContinueOnError)
+
+	commandSet = map[*flag.FlagSet]string{
+		installCmd:  installDescription,
+		checksumCmd: checksumDescription,
+		listCmd:     listDescription,
+	}
+)
+
+// printUsage prints the top level usage string.
+func printUsage() {
+	usage := `Usage: main <command> <flags> ...
+
+Available commands:`
+	fmt.Println(usage)
+	for _, f := range []*flag.FlagSet{installCmd, checksumCmd, listCmd} {
+		fmt.Printf("%s	%s\n", f.Name(), commandSet[f])
+		f.PrintDefaults()
+	}
+}
+
+func main() {
+	ctx := context.Background()
+	if len(os.Args) < 2 {
+		printUsage()
+		os.Exit(1)
+	}
+	switch os.Args[1] {
+	case installCmdStr:
+		if err := installCmd.Parse(os.Args[2:]); err != nil {
+			log.Warningf("%s failed with: %v", installCmdStr, err)
+			os.Exit(1)
+		}
+		installer, err := drivers.NewInstaller(*version, *latest)
+		if err != nil {
+			log.Warningf("Failed to create installer: %v", err.Error())
+			os.Exit(1)
+		}
+		if err := installer.MaybeInstall(ctx); err != nil {
+			log.Warningf("Failed to install driver: %v", err.Error())
+			os.Exit(1)
+		}
+	case checksumCmdStr:
+		if err := checksumCmd.Parse(os.Args[2:]); err != nil {
+			log.Warningf("%s failed with: %v", checksumCmdStr, err)
+			os.Exit(1)
+		}
+
+		for version, storedChecksum := range nvproxy.GetSupportedDriversAndChecksums() {
+			checksum, err := drivers.ChecksumDriver(ctx, version.String())
+			if err != nil {
+				log.Warningf("error on version %q: %v", version.String(), err)
+				continue
+			}
+			if checksum != storedChecksum {
+				log.Warningf("Checksum Mismatch on driver %q got: %q want: %q", version.String(), checksum, storedChecksum)
+				continue
+			}
+			log.Infof("Checksum matched on driver %q.", version.String())
+		}
+	case listCmdStr:
+		if err := listCmd.Parse(os.Args[2:]); err != nil {
+			log.Warningf("%s failed with: %v", listCmdStr, err)
+			os.Exit(1)
+		}
+		drivers.ListSupportedDrivers()
+	default:
+		printUsage()
+		os.Exit(1)
+	}
+}