From eaf3f3d7f21c44dfa0f24bfd65bc0ffa48c70f2d Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Mon, 15 Apr 2024 19:32:42 -0700
Subject: [PATCH] Add a flag to specify the NVIDIA driver ABI version to use.

This is useful for testing new NVIDIA driver versions before they are
officially released, and in general to just try out if gVisor works with
the currently-installed NVIDIA driver version regardless of whether it is
an exact match with what gVisor supports.

PiperOrigin-RevId: 625166530
---
 runsc/boot/loader.go     |  6 ++++--
 runsc/config/config.go   |  6 ++++++
 runsc/config/flags.go    |  1 +
 runsc/sandbox/sandbox.go | 19 +++++++++++++++++--
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/runsc/boot/loader.go b/runsc/boot/loader.go
index a9c214e41..28bd17bb2 100644
--- a/runsc/boot/loader.go
+++ b/runsc/boot/loader.go
@@ -131,7 +131,8 @@ type containerInfo struct {
 	// nvidiaUVMDevMajor is the device major number used for nvidia-uvm.
 	nvidiaUVMDevMajor uint32
 
-	// nvidiaDriverVersion is the Nvidia driver version on the host.
+	// nvidiaDriverVersion is the NVIDIA driver ABI version to use for
+	// communicating with NVIDIA devices on the host.
 	nvidiaDriverVersion string
 }
 
@@ -292,7 +293,8 @@ type Args struct {
 	// ProfileOpts contains the set of profiles to enable and the
 	// corresponding FDs where profile data will be written.
 	ProfileOpts profile.Opts
-	// NvidiaDriverVersion is the Nvidia driver version on the host.
+	// NvidiaDriverVersion is the NVIDIA driver ABI version to use for
+	// communicating with NVIDIA devices on the host.
 	NvidiaDriverVersion string
 }
 
diff --git a/runsc/config/config.go b/runsc/config/config.go
index acd6c70b8..cc1c88890 100644
--- a/runsc/config/config.go
+++ b/runsc/config/config.go
@@ -310,6 +310,12 @@ type Config struct {
 	// effect of injecting nvidia-container-runtime-hook as a prestart hook.
 	NVProxyDocker bool `flag:"nvproxy-docker"`
 
+	// NVProxyDriverVersion is the version of the NVIDIA driver ABI to use.
+	// If empty, it is autodetected from the installed NVIDIA driver.
+	// It can also be set to the special value "latest" to force the use of
+	// the latest supported NVIDIA driver ABI.
+	NVProxyDriverVersion string `flag:"nvproxy-driver-version"`
+
 	// TPUProxy enables support for TPUs.
 	TPUProxy bool `flag:"tpuproxy"`
 
diff --git a/runsc/config/flags.go b/runsc/config/flags.go
index db49aaf88..39052472b 100644
--- a/runsc/config/flags.go
+++ b/runsc/config/flags.go
@@ -128,6 +128,7 @@ func RegisterFlags(flagSet *flag.FlagSet) {
 	// Flags that control sandbox runtime behavior: accelerator related.
 	flagSet.Bool("nvproxy", false, "EXPERIMENTAL: enable support for Nvidia GPUs")
 	flagSet.Bool("nvproxy-docker", false, "DEPRECATED: use nvidia-container-runtime or `docker run --gpus` directly. Or manually add nvidia-container-runtime-hook as a prestart hook and set up NVIDIA_VISIBLE_DEVICES container environment variable.")
+	flagSet.String("nvproxy-driver-version", "", "NVIDIA driver ABI version to use. If empty, autodetect installed driver version. The special value 'latest' may also be used to use the latest ABI.")
 	flagSet.Bool("tpuproxy", false, "EXPERIMENTAL: enable support for TPU device passthrough.")
 
 	// Test flags, not to be used outside tests, ever.
diff --git a/runsc/sandbox/sandbox.go b/runsc/sandbox/sandbox.go
index dcf0e2ac7..02b766cd4 100644
--- a/runsc/sandbox/sandbox.go
+++ b/runsc/sandbox/sandbox.go
@@ -841,11 +841,11 @@ func (s *Sandbox) createSandboxProcess(conf *config.Config, args *Args, startSyn
 	}
 
 	if specutils.NVProxyEnabled(args.Spec, conf) {
-		nvidiaDriverVersion, err := nvproxy.HostDriverVersion()
+		version, err := getNvproxyDriverVersion(conf)
 		if err != nil {
 			return fmt.Errorf("failed to get Nvidia driver version: %w", err)
 		}
-		cmd.Args = append(cmd.Args, "--nvidia-driver-version="+nvidiaDriverVersion)
+		cmd.Args = append(cmd.Args, "--nvidia-driver-version="+version)
 	}
 
 	// Joins the network namespace if network is enabled. the sandbox talks
@@ -1556,6 +1556,21 @@ func deviceFileForPlatform(name, devicePath string) (*os.File, error) {
 	return f, nil
 }
 
+// getNvproxyDriverVersion returns the NVIDIA driver ABI version to use by
+// nvproxy.
+func getNvproxyDriverVersion(conf *config.Config) (string, error) {
+	switch conf.NVProxyDriverVersion {
+	case "":
+		return nvproxy.HostDriverVersion()
+	case "latest":
+		nvproxy.Init()
+		return nvproxy.LatestDriver().String(), nil
+	default:
+		version, err := nvproxy.DriverVersionFrom(conf.NVProxyDriverVersion)
+		return version.String(), err
+	}
+}
+
 // checkBinaryPermissions verifies that the required binary bits are set on
 // the runsc executable.
 func checkBinaryPermissions(conf *config.Config) error {