From e89d94be48369d0cdac042814b56bf12d24223e4 Mon Sep 17 00:00:00 2001 From: Ayush Ranjan Date: Wed, 17 Jan 2024 11:07:13 -0800 Subject: [PATCH] Refactor nvproxy to expose useful API for supported driver versions. - Replaced GetSupportedDriversAndChecksums() with ForEachSupportDriver(), LatestDriver() and GetDriverChecksum(). This API is more efficient. GetSupportedDriversAndChecksums() was allocating a map with all versions and checksums while most callers did not want checksum. - Add validate_checksum command to tools/gpu:main which validates that abis map has valid and correct checksums. The checksum command was fixed to print the checksum for the provided version (as the description suggested). PiperOrigin-RevId: 599230148 --- BUILD | 2 +- pkg/sentry/devices/nvproxy/nvproxy_test.go | 6 +-- pkg/sentry/devices/nvproxy/version.go | 53 +++++++++++++------ runsc/cmd/nvproxy/list_supported_drivers.go | 4 +- tools/gpu/drivers/install_driver.go | 56 +++++++++------------ tools/gpu/drivers/install_driver_test.go | 39 +++++++------- tools/gpu/main.go | 56 ++++++++++++++------- 7 files changed, 126 insertions(+), 90 deletions(-) diff --git a/BUILD b/BUILD index ce145ba36..21d5e51c7 100644 --- a/BUILD +++ b/BUILD @@ -1,8 +1,8 @@ -load("@rules_license//rules:license.bzl", "license") load("//tools:defs.bzl", "build_test", "gazelle", "go_path") load("//tools/nogo:defs.bzl", "nogo_config") load("//tools/yamltest:defs.bzl", "yaml_test") load("//website:defs.bzl", "doc") +load("@rules_license//rules:license.bzl", "license") package( default_applicable_licenses = ["//:license"], diff --git a/pkg/sentry/devices/nvproxy/nvproxy_test.go b/pkg/sentry/devices/nvproxy/nvproxy_test.go index 8c3079c4e..2d3e7270b 100644 --- a/pkg/sentry/devices/nvproxy/nvproxy_test.go +++ b/pkg/sentry/devices/nvproxy/nvproxy_test.go @@ -39,11 +39,11 @@ func TestNVOS21ParamsSize(t *testing.T) { // 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. +// `make sudo TARGETS=//tools/gpu:main ARGS="validate_checksum"`and fix mismatches. func TestAllSupportedHashesPresent(t *testing.T) { Init() - for version, checksum := range GetSupportedDriversAndChecksums() { - if checksum == "" { + for version, abi := range abis { + if abi.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 210ea774a..3b3fd0ed4 100644 --- a/pkg/sentry/devices/nvproxy/version.go +++ b/pkg/sentry/devices/nvproxy/version.go @@ -71,25 +71,25 @@ 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 +// isGreaterThan returns true if v is greater than other. +// 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 { +func (v DriverVersion) isGreaterThan(other DriverVersion) bool { switch { case v.major > other.major: - return v + return true case other.major > v.major: - return other + return false case v.minor > other.minor: - return v + return true case other.minor > v.minor: - return other + return false case v.patch > other.patch: - return v + return true case other.patch > v.patch: - return other + return false default: - return v + return true } } @@ -138,7 +138,7 @@ var abisOnce sync.Once // 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. +// Run `make sudo TARGETS=//tools/gpu:main ARGS="checksum --version={}"` to get checksum. func addDriverABI(major, minor, patch int, runfileChecksum string, cons driverABIFunc) driverABIFunc { if abis == nil { abis = make(map[DriverVersion]abiConAndChecksum) @@ -348,11 +348,32 @@ func Init() { }) } -// GetSupportedDriversAndChecksums returns supported driver ABIs. -func GetSupportedDriversAndChecksums() map[DriverVersion]string { - versions := make(map[DriverVersion]string, len(abis)) +// ForEachSupportDriver calls f on all supported drivers. +// Precondition: Init() must have been called. +func ForEachSupportDriver(f func(version DriverVersion, checksum string)) { for version, abi := range abis { - versions[version] = abi.checksum + f(version, abi.checksum) } - return versions +} + +// LatestDriver returns the latest supported driver. +// Precondition: Init() must have been called. +func LatestDriver() DriverVersion { + var ret DriverVersion + for version := range abis { + if version.isGreaterThan(ret) { + ret = version + } + } + return ret +} + +// ExpectedDriverChecksum returns the expected checksum for a given version. +// Precondition: Init() must have been called. +func ExpectedDriverChecksum(version DriverVersion) (string, bool) { + abi, ok := abis[version] + if !ok { + return "", false + } + return abi.checksum, true } diff --git a/runsc/cmd/nvproxy/list_supported_drivers.go b/runsc/cmd/nvproxy/list_supported_drivers.go index ed211ee14..3fc43ae41 100644 --- a/runsc/cmd/nvproxy/list_supported_drivers.go +++ b/runsc/cmd/nvproxy/list_supported_drivers.go @@ -53,9 +53,9 @@ func (*listSupportedDrivers) Execute(ctx context.Context, f *flag.FlagSet, args return subcommands.ExitUsageError } - for version, _ := range nvproxy.GetSupportedDriversAndChecksums() { + nvproxy.ForEachSupportDriver(func(version nvproxy.DriverVersion, _ string) { fmt.Println(version) - } + }) return subcommands.ExitSuccess } diff --git a/tools/gpu/drivers/install_driver.go b/tools/gpu/drivers/install_driver.go index eb1f3dd14..47cd84247 100644 --- a/tools/gpu/drivers/install_driver.go +++ b/tools/gpu/drivers/install_driver.go @@ -43,25 +43,23 @@ func init() { 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 + expectedChecksumFunc func(nvproxy.DriverVersion) (string, bool) + getCurrentDriverFunc func() (nvproxy.DriverVersion, error) + downloadFunc func(context.Context, string) (io.ReadCloser, error) + installFunc 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, + expectedChecksumFunc: nvproxy.ExpectedDriverChecksum, + getCurrentDriverFunc: getCurrentDriver, + downloadFunc: DownloadDriver, + installFunc: installDriver, } switch { case latest: - for v := range ret.getSupportedDriverFunc() { - ret.requestedVersion = v.IsGreaterThan(ret.requestedVersion) - } + ret.requestedVersion = nvproxy.LatestDriver() default: d, err := nvproxy.DriverVersionFrom(requestedVersion) if err != nil { @@ -77,8 +75,7 @@ func NewInstaller(requestedVersion string, latest bool) (*Installer, error) { // 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 { + if _, ok := i.expectedChecksumFunc(i.requestedVersion); !ok { return fmt.Errorf("requested driver %q is not supported", i.requestedVersion) } @@ -86,7 +83,7 @@ func (i *Installer) MaybeInstall(ctx context.Context) error { if err != nil { log.Warningf("failed to get current driver: %v", err) } - if existingDriver.Equals(driver) { + if existingDriver.Equals(i.requestedVersion) { log.Infof("Driver already installed: %s", i.requestedVersion) return nil } @@ -100,7 +97,7 @@ func (i *Installer) MaybeInstall(ctx context.Context) error { } log.Infof("Downloading driver: %s", i.requestedVersion) - reader, err := i.downloadFunction(ctx, i.requestedVersion.String()) + reader, err := i.downloadFunc(ctx, i.requestedVersion.String()) if err != nil { return fmt.Errorf("failed to download driver: %w", err) } @@ -110,7 +107,7 @@ func (i *Installer) MaybeInstall(ctx context.Context) error { return fmt.Errorf("failed to open driver file: %w", err) } defer os.Remove(f.Name()) - if err := i.writeAndCheck(f, reader, driver); err != nil { + if err := i.writeAndCheck(f, reader); err != nil { f.Close() return fmt.Errorf("writeAndCheck: %w", err) } @@ -122,7 +119,7 @@ func (i *Installer) MaybeInstall(ctx context.Context) error { } log.Infof("Driver downloaded: %s", i.requestedVersion) log.Infof("Installing driver: %s", i.requestedVersion) - if err := i.installFunction(f.Name()); err != nil { + if err := i.installFunc(f.Name()); err != nil { return fmt.Errorf("failed to install driver: %w", err) } log.Infof("Installation Complete!") @@ -140,7 +137,7 @@ func (i *Installer) uninstallDriver(ctx context.Context, driverVersion string) e return nil } -func (i *Installer) writeAndCheck(f *os.File, reader io.ReadCloser, driverVersion nvproxy.DriverVersion) error { +func (i *Installer) writeAndCheck(f *os.File, reader io.ReadCloser) error { checksum := sha256.New() buf := make([]byte, 1024*1024) for { @@ -159,9 +156,12 @@ func (i *Installer) writeAndCheck(f *os.File, reader io.ReadCloser, driverVersio } } gotChecksum := fmt.Sprintf("%x", checksum.Sum(nil)) - wantChecksum := i.getSupportedDriverFunc()[driverVersion] + wantChecksum, ok := i.expectedChecksumFunc(i.requestedVersion) + if !ok { + return fmt.Errorf("requested driver %q is not supported", i.requestedVersion) + } if gotChecksum != wantChecksum { - return fmt.Errorf("driver %q checksum mismatch: got %q, want %q", driverVersion, gotChecksum, wantChecksum) + return fmt.Errorf("driver %q checksum mismatch: got %q, want %q", i.requestedVersion, gotChecksum, wantChecksum) } return nil } @@ -183,15 +183,6 @@ func getCurrentDriver() (nvproxy.DriverVersion, error) { 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(outfile string) error { @@ -205,11 +196,10 @@ func ListSupportedDrivers(outfile string) error { out = f } - supportedDrivers := nvproxy.GetSupportedDriversAndChecksums() - list := make([]string, 0, len(supportedDrivers)) - for version := range supportedDrivers { + var list []string + nvproxy.ForEachSupportDriver(func(version nvproxy.DriverVersion, checksum string) { list = append(list, version.String()) - } + }) if _, err := out.WriteString(strings.Join(list, " ") + "\n"); err != nil { return fmt.Errorf("failed to write to outfile: %w", err) } diff --git a/tools/gpu/drivers/install_driver_test.go b/tools/gpu/drivers/install_driver_test.go index a89091a7b..6f2b6f065 100644 --- a/tools/gpu/drivers/install_driver_test.go +++ b/tools/gpu/drivers/install_driver_test.go @@ -34,10 +34,15 @@ func TestVersionInstalled(t *testing.T) { 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, + requestedVersion: version, + expectedChecksumFunc: func(v nvproxy.DriverVersion) (string, bool) { + if v == version { + return checksum, true + } + return "", false + }, + getCurrentDriverFunc: getFunction, + downloadFunc: downloadFunction, } if err := installer.MaybeInstall(ctx); err != nil { t.Fatalf("Installation failed: %v", err) @@ -50,10 +55,8 @@ func TestVersionNotSupported(t *testing.T) { 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", - } + expectedChecksumFunc: func(v nvproxy.DriverVersion) (string, bool) { + return "", false }, } err := installer.MaybeInstall(ctx) @@ -74,12 +77,13 @@ func TestShaMismatch(t *testing.T) { getCurrentDriverFunc: func() (nvproxy.DriverVersion, error) { return nvproxy.DriverVersion{}, nil }, - getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string { - return map[nvproxy.DriverVersion]string{ - version: "mismatched checksum", + expectedChecksumFunc: func(v nvproxy.DriverVersion) (string, bool) { + if v == version { + return "mismatch", true } + return "", false }, - downloadFunction: func(context.Context, string) (io.ReadCloser, error) { + downloadFunc: func(context.Context, string) (io.ReadCloser, error) { reader := bytes.NewReader([]byte("some content")) return io.NopCloser(reader), nil }, @@ -104,16 +108,17 @@ func TestDriverInstalls(t *testing.T) { getCurrentDriverFunc: func() (nvproxy.DriverVersion, error) { return nvproxy.DriverVersion{}, nil }, - getSupportedDriverFunc: func() map[nvproxy.DriverVersion]string { - return map[nvproxy.DriverVersion]string{ - version: checksum, + expectedChecksumFunc: func(v nvproxy.DriverVersion) (string, bool) { + if v == version { + return checksum, true } + return "", false }, - downloadFunction: func(context.Context, string) (io.ReadCloser, error) { + downloadFunc: func(context.Context, string) (io.ReadCloser, error) { reader := bytes.NewReader(content) return io.NopCloser(reader), nil }, - installFunction: func(_ string) error { + installFunc: func(_ string) error { return nil }, } diff --git a/tools/gpu/main.go b/tools/gpu/main.go index dd24e66b3..41a8a6bbf 100644 --- a/tools/gpu/main.go +++ b/tools/gpu/main.go @@ -27,12 +27,14 @@ import ( ) 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" + installCmdStr = "install" + installDescription = "installs a driver on the host machine" + checksumCmdStr = "checksum" + checksumDescription = "computes the sha256 checksum for a given driver version" + validateChecksumCmdStr = "validate_checksum" + validateChecksumDescription = "validates the checksum of all supported drivers" + listCmdStr = "list" + listDescription = "lists the supported drivers" ) var ( @@ -41,17 +43,22 @@ var ( latest = installCmd.Bool("latest", false, "install the latest supported driver") version = installCmd.String("version", "", "version of the driver") + // Computes the sha256 checksum for a given driver's .run file from the nvidia site. + checksumCmd = flag.NewFlagSet(checksumCmdStr, flag.ContinueOnError) + checksumVersion = checksumCmd.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) + validateChecksumCmd = flag.NewFlagSet(validateChecksumCmdStr, flag.ContinueOnError) // The list command returns the list of supported drivers from this tool. listCmd = flag.NewFlagSet(listCmdStr, flag.ContinueOnError) outfile = listCmd.String("outfile", "", "if set, write the list output to this file") commandSet = map[*flag.FlagSet]string{ - installCmd: installDescription, - checksumCmd: checksumDescription, - listCmd: listDescription, + installCmd: installDescription, + checksumCmd: checksumDescription, + validateChecksumCmd: validateChecksumDescription, + listCmd: listDescription, } ) @@ -61,7 +68,7 @@ func printUsage() { Available commands:` fmt.Println(usage) - for _, f := range []*flag.FlagSet{installCmd, checksumCmd, listCmd} { + for _, f := range []*flag.FlagSet{installCmd, checksumCmd, validateChecksumCmd, listCmd} { fmt.Printf("%s %s\n", f.Name(), commandSet[f]) f.PrintDefaults() } @@ -73,6 +80,7 @@ func main() { printUsage() os.Exit(1) } + nvproxy.Init() switch os.Args[1] { case installCmdStr: if err := installCmd.Parse(os.Args[2:]); err != nil { @@ -94,18 +102,30 @@ func main() { os.Exit(1) } - for version, storedChecksum := range nvproxy.GetSupportedDriversAndChecksums() { - checksum, err := drivers.ChecksumDriver(ctx, version.String()) + checksum, err := drivers.ChecksumDriver(ctx, *checksumVersion) + if err != nil { + log.Warningf("Failed to compute checksum: %v", err) + os.Exit(1) + } + fmt.Printf("Checksum: %q\n", checksum) + case validateChecksumCmdStr: + if err := validateChecksumCmd.Parse(os.Args[2:]); err != nil { + log.Warningf("%s failed with: %v", validateChecksumCmdStr, err) + os.Exit(1) + } + + nvproxy.ForEachSupportDriver(func(version nvproxy.DriverVersion, checksum string) { + wantChecksum, err := drivers.ChecksumDriver(ctx, version.String()) if err != nil { log.Warningf("error on version %q: %v", version.String(), err) - continue + return } - if checksum != storedChecksum { - log.Warningf("Checksum Mismatch on driver %q got: %q want: %q", version.String(), storedChecksum, checksum) - continue + if checksum != wantChecksum { + log.Warningf("Checksum mismatch on driver %q got: %q want: %q", version.String(), checksum, wantChecksum) + return } 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)