diff --git a/.buildkite/pipeline.yaml b/.buildkite/pipeline.yaml index 7de1373ee..966174445 100644 --- a/.buildkite/pipeline.yaml +++ b/.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 diff --git a/pkg/sentry/devices/nvproxy/BUILD b/pkg/sentry/devices/nvproxy/BUILD index 4951df448..4fc677f55 100644 --- a/pkg/sentry/devices/nvproxy/BUILD +++ b/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", diff --git a/pkg/sentry/devices/nvproxy/nvproxy.go b/pkg/sentry/devices/nvproxy/nvproxy.go index bd58cc4d3..a92c7edd2 100644 --- a/pkg/sentry/devices/nvproxy/nvproxy.go +++ b/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{ diff --git a/pkg/sentry/devices/nvproxy/nvproxy_test.go b/pkg/sentry/devices/nvproxy/nvproxy_test.go index b0e483b65..8c3079c4e 100644 --- a/pkg/sentry/devices/nvproxy/nvproxy_test.go +++ b/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()) + } + } +} diff --git a/pkg/sentry/devices/nvproxy/version.go b/pkg/sentry/devices/nvproxy/version.go index 15cd8ea5f..b05e9518d 100644 --- a/pkg/sentry/devices/nvproxy/version.go +++ b/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 +} diff --git a/tools/gpu/BUILD b/tools/gpu/BUILD new file mode 100644 index 000000000..8a9ed4e09 --- /dev/null +++ b/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", + ], +) diff --git a/tools/gpu/drivers/BUILD b/tools/gpu/drivers/BUILD new file mode 100644 index 000000000..11f75eb7e --- /dev/null +++ b/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"], +) diff --git a/tools/gpu/drivers/install_driver.go b/tools/gpu/drivers/install_driver.go new file mode 100644 index 000000000..48e98336b --- /dev/null +++ b/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 +} diff --git a/tools/gpu/drivers/install_driver_test.go b/tools/gpu/drivers/install_driver_test.go new file mode 100644 index 000000000..a89091a7b --- /dev/null +++ b/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) + } +} diff --git a/tools/gpu/main.go b/tools/gpu/main.go new file mode 100644 index 000000000..d2e5fec84 --- /dev/null +++ b/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 ... + +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) + } +}