Pass the NVIDIA_DRIVER_CAPABILITIES env var to nvidia-container-cli.

runsc attempts to emulate nvidia-container-runtime-hook. But it was always
passing "--compute --utility" as driver capability flags to
`nvidia-container-cli configure` command.

Fix runsc to emulate nvidia-container-runtime-hook correctly by parsing
NVIDIA_DRIVER_CAPABILITIES and converting that comma-separated list to flags.

This is in preparation for adding support for non-compute GPU workloads in
nvproxy :)

Updates #9452
Updates #10856

PiperOrigin-RevId: 671644915

pkg/sentry/devices/nvproxy/nvconf/BUILD

@@ -0,0 +1,16 @@
+load("//tools:defs.bzl", "go_library")
+
+package(default_applicable_licenses = ["//:license"])
+
+licenses(["notice"])
+
+go_library(
+    name = "nvconf",
+    srcs = [
+        "caps.go",
+        "nvconf.go",
+    ],
+    visibility = [
+        "//pkg/sentry:internal",
+    ],
+)

pkg/sentry/devices/nvproxy/nvconf/caps.go

@@ -0,0 +1,107 @@
+// Copyright 2024 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 nvconf
+
+import "strings"
+
+// DriverCap is a GPU driver capability (like compute, graphics, etc.).
+type DriverCap string
+
+// Driver capabilities understood by nvproxy.
+const (
+	// AllCap is a special value that means all supported driver capabilities.
+	AllCap DriverCap = "all"
+
+	Compat32Cap DriverCap = "compat32"
+	ComputeCap  DriverCap = "compute"
+	DisplayCap  DriverCap = "display"
+	GraphicsCap DriverCap = "graphics"
+	NGXCap      DriverCap = "ngx"
+	UtilityCap  DriverCap = "utility"
+	VideoCap    DriverCap = "video"
+)
+
+// ToFlag converts the driver capability to a flag for nvidia-container-cli.
+// See nvidia-container-toolkit/cmd/nvidia-container-runtime-hook/capabilities.go:capabilityToCLI().
+func (c DriverCap) ToFlag() string {
+	return "--" + string(c)
+}
+
+// DriverCaps is a set of GPU driver capabilities.
+type DriverCaps map[DriverCap]struct{}
+
+// DefaultDriverCaps is the set of driver capabilities that are enabled by
+// default in the absence of any other configuration. See
+// nvidia-container-toolkit/internal/config/image/capabilities.go:DefaultDriverCapabilities.
+var DefaultDriverCaps = DriverCaps{
+	ComputeCap: struct{}{},
+	UtilityCap: struct{}{},
+}
+
+// SupportedDriverCaps is the set of driver capabilities that are supported by
+// nvproxy. Similar to
+// nvidia-container-toolkit/internal/config/image/capabilities.go:SupportedDriverCapabilities.
+var SupportedDriverCaps = DriverCaps{
+	ComputeCap: struct{}{},
+	UtilityCap: struct{}{},
+}
+
+// KnownDriverCapValues is the set of understood driver capability values.
+var KnownDriverCapValues = DriverCaps{
+	Compat32Cap: struct{}{},
+	ComputeCap:  struct{}{},
+	DisplayCap:  struct{}{},
+	GraphicsCap: struct{}{},
+	NGXCap:      struct{}{},
+	UtilityCap:  struct{}{},
+	VideoCap:    struct{}{},
+}
+
+// DriverCapsFromString constructs NvidiaDriverCaps from a comma-separated list
+// of driver capabilities.
+func DriverCapsFromString(caps string) DriverCaps {
+	res := make(DriverCaps)
+	for _, cap := range strings.Split(caps, ",") {
+		trimmed := strings.TrimSpace(cap)
+		if len(trimmed) == 0 {
+			continue
+		}
+		res[DriverCap(trimmed)] = struct{}{}
+	}
+	return res
+}
+
+// HasAll returns true if the set of driver capabilities contains "all".
+func (c DriverCaps) HasAll() bool {
+	_, ok := c[AllCap]
+	return ok
+}
+
+// Intersect returns the intersection of two sets of driver capabilities.
+func (c DriverCaps) Intersect(c2 DriverCaps) DriverCaps {
+	if c2.HasAll() {
+		return c
+	}
+	if c.HasAll() {
+		return c2
+	}
+	res := make(DriverCaps)
+	for cap := range c2 {
+		if _, ok := c[cap]; ok {
+			res[cap] = struct{}{}
+		}
+	}
+	return res
+}

pkg/sentry/devices/nvproxy/nvconf/nvconf.go

@@ -0,0 +1,21 @@
+// Copyright 2024 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 nvconf provides configuration structures and utilities for nvproxy.
+//
+// This package is separate from the main nvproxy package to allow reading and
+// working with nvproxy configuration without importing the full nvproxy
+// package and its dependencies. This is useful for tools and applications that
+// only need to interact with the configuration.
+package nvconf

runsc/config/BUILD

@@ -18,6 +18,7 @@ go_library(
     deps = [
         "//pkg/log",
         "//pkg/refs",
+        "//pkg/sentry/devices/nvproxy/nvconf",
         "//pkg/sentry/watchdog",
         "//runsc/flag",
         "//runsc/version",

runsc/config/config.go

@@ -27,6 +27,7 @@ import (
 
 	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/pkg/refs"
+	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy/nvconf"
 	"gvisor.dev/gvisor/pkg/sentry/watchdog"
 	"gvisor.dev/gvisor/runsc/flag"
 	"gvisor.dev/gvisor/runsc/version"
@@ -325,6 +326,10 @@ type Config struct {
 	// the latest supported NVIDIA driver ABI.
 	NVProxyDriverVersion string `flag:"nvproxy-driver-version"`
 
+	// NVProxyAllowUnsupportedCapabilities is a comma-separated list of driver
+	// capabilities that are allowed to be requested by the container.
+	NVProxyAllowedDriverCapabilities string `flag:"nvproxy-allowed-driver-capabilities"`
+
 	// TPUProxy enables support for TPUs.
 	TPUProxy bool `flag:"tpuproxy"`
 
@@ -408,6 +413,18 @@ func (c *Config) validate() error {
 	if len(c.ProfilingMetrics) > 0 && len(c.ProfilingMetricsLog) == 0 {
 		return fmt.Errorf("profiling-metrics flag requires defining a profiling-metrics-log for output")
 	}
+	for _, cap := range strings.Split(c.NVProxyAllowedDriverCapabilities, ",") {
+		nvcap := nvconf.DriverCap(cap)
+		if nvcap == nvconf.AllCap {
+			continue
+		}
+		if _, ok := nvconf.KnownDriverCapValues[nvcap]; !ok {
+			return fmt.Errorf("nvproxy-allowed-driver-capabilities contains invalid capability %q", cap)
+		}
+		if _, ok := nvconf.SupportedDriverCaps[nvcap]; !ok {
+			log.Warningf("nvproxy-allowed-driver-capabilities contains unsupported capability %q", cap)
+		}
+	}
 	return nil
 }
 

runsc/config/flags.go

@@ -134,6 +134,7 @@ func RegisterFlags(flagSet *flag.FlagSet) {
 	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.String("nvproxy-allowed-driver-capabilities", "utility,compute", "Comma separated list of NVIDIA driver capabilities that are allowed to be requested by the container. If 'all' is specified here, it is resolved to all driver capabilities supported in nvproxy. If 'all' is requested by the container, it is resolved to this list.")
 	flagSet.Bool("tpuproxy", false, "EXPERIMENTAL: enable support for TPU device passthrough.")
 
 	// Test flags, not to be used outside tests, ever.

runsc/container/container.go

@@ -2011,12 +2011,20 @@ func nvproxySetupAfterGoferUserns(spec *specs.Spec, conf *config.Config, goferCm
 			"configure",
 			fmt.Sprintf("--ldconfig=@%s", ldconfigPath),
 			"--no-cgroups", // runsc doesn't configure device cgroups yet
-			"--utility",
-			"--compute",
 			fmt.Sprintf("--pid=%d", goferCmd.Process.Pid),
 			fmt.Sprintf("--device=%s", devices),
-			spec.Root.Path,
 		}
+		// Pass driver capabilities specified via NVIDIA_DRIVER_CAPABILITIES as flags. See
+		// nvidia-container-toolkit/cmd/nvidia-container-runtime-hook/main.go:doPrestart().
+		driverCaps, err := specutils.NVProxyDriverCapsFromEnv(spec, conf)
+		if err != nil {
+			return fmt.Errorf("failed to get driver capabilities: %w", err)
+		}
+		for cap := range driverCaps {
+			argv = append(argv, cap.ToFlag())
+		}
+		// Add rootfs path as the final argument.
+		argv = append(argv, spec.Root.Path)
 		log.Debugf("Executing %q", argv)
 		var stdout, stderr strings.Builder
 		cmd := exec.Cmd{

runsc/specutils/BUILD

@@ -19,7 +19,7 @@ go_library(
         "//pkg/abi/linux",
         "//pkg/bits",
         "//pkg/log",
-        "//pkg/sentry/devices/tpuproxy/vfio",
+        "//pkg/sentry/devices/nvproxy/nvconf",
         "//pkg/sentry/kernel/auth",
         "//runsc/config",
         "//runsc/flag",
@@ -36,5 +36,9 @@ go_test(
     size = "small",
     srcs = ["specutils_test.go"],
     library = ":specutils",
-    deps = ["@com_github_opencontainers_runtime_spec//specs-go:go_default_library"],
+    deps = [
+        "//pkg/sentry/devices/nvproxy/nvconf",
+        "//runsc/config",
+        "@com_github_opencontainers_runtime_spec//specs-go:go_default_library",
+    ],
 )

runsc/specutils/nvidia.go

@@ -20,13 +20,26 @@ import (
 	"strings"
 
 	specs "github.com/opencontainers/runtime-spec/specs-go"
+	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy/nvconf"
 	"gvisor.dev/gvisor/runsc/config"
 )
 
-const nvdEnvVar = "NVIDIA_VISIBLE_DEVICES"
-
-// AnnotationNVProxy enables nvproxy.
-const AnnotationNVProxy = "dev.gvisor.internal.nvproxy"
+const (
+	// NVIDIA_VISIBLE_DEVICES environment variable controls which GPUs are
+	// visible and accessible to the container.
+	nvidiaVisibleDevsEnv = "NVIDIA_VISIBLE_DEVICES"
+	// NVIDIA_DRIVER_CAPABILITIES environment variable allows to fine-tune which
+	// NVIDIA driver components are mounted and accessible within a container.
+	nvidiaDriverCapsEnv = "NVIDIA_DRIVER_CAPABILITIES"
+	// CUDA_VERSION environment variable indicates the version of the CUDA
+	// toolkit installed on in the container image.
+	cudaVersionEnv = "CUDA_VERSION"
+	// NVIDIA_REQUIRE_CUDA environment variable indicates the CUDA toolkit
+	// version that a container needs.
+	requireCudaEnv = "NVIDIA_REQUIRE_CUDA"
+	// AnnotationNVProxy enables nvproxy.
+	AnnotationNVProxy = "dev.gvisor.internal.nvproxy"
+)
 
 // NVProxyEnabled checks both the nvproxy annotation and conf.NVProxy to see if nvproxy is enabled.
 func NVProxyEnabled(spec *specs.Spec, conf *config.Config) bool {
@@ -78,7 +91,7 @@ func gpuFunctionalityRequestedViaHook(spec *specs.Spec, conf *config.Config) boo
 	if spec.Process == nil {
 		return false
 	}
-	nvd, _ := EnvVar(spec.Process.Env, nvdEnvVar)
+	nvd, _ := EnvVar(spec.Process.Env, nvidiaVisibleDevsEnv)
 	// A value of "none" means "no GPU device, but still access to driver
 	// functionality", so it is not a value we check for here.
 	return nvd != "" && nvd != "void"
@@ -105,7 +118,7 @@ func isNvidiaHookPresent(spec *specs.Spec, conf *config.Config) bool {
 //
 // Precondition: conf.NVProxyDocker && GPUFunctionalityRequested(spec, conf).
 func ParseNvidiaVisibleDevices(spec *specs.Spec) (string, error) {
-	nvd, _ := EnvVar(spec.Process.Env, nvdEnvVar)
+	nvd, _ := EnvVar(spec.Process.Env, nvidiaVisibleDevsEnv)
 	if nvd == "none" {
 		return "", nil
 	}
@@ -130,3 +143,49 @@ func ParseNvidiaVisibleDevices(spec *specs.Spec) (string, error) {
 	}
 	return nvd, nil
 }
+
+// NVProxyDriverCapsFromEnv returns the driver capabilities requested by the
+// application via the NVIDIA_DRIVER_CAPABILITIES env var. See
+// nvidia-container-toolkit/cmd/nvidia-container-runtime-hook/container_config.go:getDriverCapabilities().
+func NVProxyDriverCapsFromEnv(spec *specs.Spec, conf *config.Config) (nvconf.DriverCaps, error) {
+	// Construct the set of allowed driver capabilities.
+	allowedDriverCaps := nvconf.DriverCapsFromString(conf.NVProxyAllowedDriverCapabilities)
+	// Resolve "all" to nvconf.SupportedDriverCaps.
+	if allowedDriverCaps.HasAll() {
+		delete(allowedDriverCaps, nvconf.AllCap)
+		for cap := range nvconf.SupportedDriverCaps {
+			allowedDriverCaps[cap] = struct{}{}
+		}
+	}
+
+	// Extract the set of driver capabilities requested by the application.
+	driverCapsEnvStr, ok := EnvVar(spec.Process.Env, nvidiaDriverCapsEnv)
+	if !ok {
+		// Nothing requested. Fallback to default configurations.
+		if IsLegacyCudaImage(spec) {
+			return allowedDriverCaps, nil
+		}
+		return nvconf.DefaultDriverCaps, nil
+	}
+	if len(driverCapsEnvStr) == 0 {
+		// Empty. Fallback to nvconf.DefaultDriverCaps.
+		return nvconf.DefaultDriverCaps, nil
+	}
+	envDriverCaps := nvconf.DriverCapsFromString(driverCapsEnvStr)
+
+	// Intersect what's requested with what's allowed. Since allowedDriverCaps
+	// doesn't contain "all", the result will also not contain "all".
+	driverCaps := allowedDriverCaps.Intersect(envDriverCaps)
+	if !envDriverCaps.HasAll() && len(driverCaps) != len(envDriverCaps) {
+		return nil, fmt.Errorf("disallowed driver capabilities requested: '%v' (allowed '%v'), update --nvproxy-allowed-driver-capabilities to allow them", envDriverCaps, driverCaps)
+	}
+	return driverCaps, nil
+}
+
+// IsLegacyCudaImage returns true if spec represents a legacy CUDA image.
+// See nvidia-container-toolkit/internal/config/image/cuda_image.go:IsLegacy().
+func IsLegacyCudaImage(spec *specs.Spec) bool {
+	cudaVersion, _ := EnvVar(spec.Process.Env, cudaVersionEnv)
+	requireCuda, _ := EnvVar(spec.Process.Env, requireCudaEnv)
+	return len(cudaVersion) > 0 && len(requireCuda) == 0
+}

runsc/specutils/specutils.go

@@ -35,7 +35,6 @@ import (
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/bits"
 	"gvisor.dev/gvisor/pkg/log"
-	"gvisor.dev/gvisor/pkg/sentry/devices/tpuproxy/vfio"
 	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
 	"gvisor.dev/gvisor/runsc/config"
 	"gvisor.dev/gvisor/runsc/flag"
@@ -577,7 +576,7 @@ func TPUProxyIsEnabled(spec *specs.Spec, conf *config.Config) bool {
 // VFIOFunctionalityRequested returns true if the container should have access
 // to VFIO functionality.
 func VFIOFunctionalityRequested(dev *specs.LinuxDevice) bool {
-	return strings.HasPrefix(dev.Path, filepath.Dir(vfio.VFIOPath))
+	return strings.HasPrefix(dev.Path, "/dev/vfio")
 }
 
 // AcceleratorFunctionalityRequested returns true if the container should have