visor/runsc/specutils/nvidia.go

// 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 specutils

import (
	"fmt"
	"path/filepath"
	"regexp"
	"slices"
	"strconv"
	"strings"

	specs "github.com/opencontainers/runtime-spec/specs-go"
	"gvisor.dev/gvisor/pkg/abi/nvgpu"
	"gvisor.dev/gvisor/pkg/log"
	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy/nvconf"
	"gvisor.dev/gvisor/runsc/config"
)

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 returns true if nvproxy should be enabled.
func NVProxyEnabled(spec *specs.Spec, conf *config.Config) bool {
	return conf.NVProxy || AnnotationToBool(spec, AnnotationNVProxy) || GPUFunctionalityRequested(spec, conf)
}

// GPUFunctionalityRequested returns true if the container should have access
// to GPU functionality based on the spec.
func GPUFunctionalityRequested(spec *specs.Spec, conf *config.Config) bool {
	// In GKE, the nvidia_gpu device plugin injects NVIDIA devices into
	// spec.Linux.Devices when GPUs are allocated to a container.
	if spec.Linux != nil {
		for _, dev := range spec.Linux.Devices {
			if dev.Path == "/dev/nvidiactl" {
				return true
			}
		}
	}
	return GPUFunctionalityRequestedViaHook(spec, conf)
}

// GPUFunctionalityRequestedViaHook returns true if the container should have
// access to GPU functionality configured via nvidia-container-runtime-hook.
// This hook is used by:
// - Docker when using `--gpus` flag from the CLI.
// - nvidia-container-runtime when using its legacy mode.
func GPUFunctionalityRequestedViaHook(spec *specs.Spec, conf *config.Config) bool {
	kind := findNvidiaHook(spec, conf)
	if kind == nvidiaHookNone {
		return false
	}
	// In Docker mode, GPU access is only requested if NVIDIA_VISIBLE_DEVICES is
	// non-empty and set to a value that doesn't mean "no GPU".
	if spec.Process == nil {
		return false
	}
	nvd := nvidiaVisibleDevsEnvVar(spec.Process.Env, kind)
	// 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"
}

type nvidiaHookKind uint8

const (
	nvidiaHookNone nvidiaHookKind = iota
	nvidiaHookLegacy
	nvidiaHookCDI
)

func findNvidiaHook(spec *specs.Spec, conf *config.Config) nvidiaHookKind {
	if spec.Hooks != nil {
		for _, h := range spec.Hooks.CreateContainer {
			if strings.HasSuffix(h.Path, "/nvidia-cdi-hook") {
				return nvidiaHookCDI
			}
		}
		for _, h := range spec.Hooks.Prestart {
			if strings.HasSuffix(h.Path, "/nvidia-container-runtime-hook") {
				return nvidiaHookLegacy
			}
		}
	}
	if conf.NVProxyDocker {
		// This has the effect of injecting the nvidia-container-runtime-hook.
		return nvidiaHookLegacy
	}
	return nvidiaHookNone
}

func nvidiaVisibleDevsEnvVar(envs []string, kind nvidiaHookKind) string {
	if kind != nvidiaHookCDI {
		nvd, _ := EnvVar(envs, nvidiaVisibleDevsEnv)
		return nvd
	}
	// After
	// https://github.com/NVIDIA/nvidia-container-toolkit/commit/88ad42ccd16a3ba6fd7867511d4070aa8111b8fa,
	// when Docker communicates with nvidia-container-toolkit via CDI, the
	// latter always appends a "NVIDIA_VISIBLE_DEVICES=void" environment
	// variable to the container spec, which we want to ignore.
	const prefix = nvidiaVisibleDevsEnv + "="
	nvds := make([]string, 0, 2)
	for _, env := range envs {
		if strings.HasPrefix(env, prefix) {
			nvds = append(nvds, env[len(prefix):])
		}
	}
	if len(nvds) == 0 {
		return ""
	}
	if len(nvds) > 1 && nvds[len(nvds)-1] == "void" {
		return nvds[len(nvds)-2]
	}
	return nvds[len(nvds)-1]
}

// ParseNvidiaVisibleDevices parses NVIDIA_VISIBLE_DEVICES env var and returns
// the devices specified in it. This can be passed to nvidia-container-cli.
//
// Precondition: GPUFunctionalityRequestedViaHook(spec, conf).
func ParseNvidiaVisibleDevices(spec *specs.Spec, conf *config.Config) (string, error) {
	nvd := nvidiaVisibleDevsEnvVar(spec.Process.Env, findNvidiaHook(spec, conf))
	if nvd == "none" {
		return "", nil
	}
	if nvd == "all" {
		return "all", nil
	}

	for _, gpuDev := range strings.Split(nvd, ",") {
		// Validate gpuDev. We only support the following formats for now:
		// * GPU indices (e.g. 0,1,2)
		// * GPU UUIDs (e.g. GPU-fef8089b)
		//
		// We do not support MIG devices yet.
		if strings.HasPrefix(gpuDev, "GPU-") {
			continue
		}
		_, err := strconv.ParseUint(gpuDev, 10, 32)
		if err != nil {
			return "", fmt.Errorf("invalid %q in NVIDIA_VISIBLE_DEVICES %q: %w", gpuDev, nvd, err)
		}
	}
	return nvd, nil
}

// NVProxyDriverCapsAllowed returns the driver capabilities allowed by the
// configuration, irrespective of what a container requests.
// This should be used to determine the bounding set of driver capabilities
// that a container can request.
func NVProxyDriverCapsAllowed(conf *config.Config) (nvconf.DriverCaps, error) {
	// Construct the set of allowed driver capabilities.
	allowedDriverCaps, hasAll, err := nvconf.DriverCapsFromString(conf.NVProxyAllowedDriverCapabilities)
	if err != nil {
		return 0, fmt.Errorf("invalid set of allowed NVIDIA driver capabilities %q: %w", conf.NVProxyAllowedDriverCapabilities, err)
	}
	// allowedDriverCaps is already a subset of `nvconf.SupportedDriverCaps`
	// as this was checked by `config.Config.validate`.
	if hasAll {
		allowedDriverCaps |= nvconf.AllContainerDriverCaps
	}
	// GPUDirect RDMA requires the privileged CapRDMA capability to export GPU
	// memory to a dma-buf fd. There is no container-facing driver-capability
	// flag for it, so enable it implicitly whenever RDMA passthrough is on.
	if conf.RDMAProxy {
		allowedDriverCaps |= nvconf.CapRDMA
	}
	return allowedDriverCaps, 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) {
	allowedDriverCaps, err := NVProxyDriverCapsAllowed(conf)
	if err != nil {
		return 0, err
	}

	// 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 & allowedDriverCaps, nil
	}
	if len(driverCapsEnvStr) == 0 {
		// Empty. Fallback to nvconf.DefaultDriverCaps.
		return nvconf.DefaultDriverCaps & allowedDriverCaps, nil
	}
	envDriverCaps, enableAll, err := nvconf.DriverCapsFromString(driverCapsEnvStr)
	if err != nil {
		return 0, fmt.Errorf("invalid set of requested NVIDIA driver capabilities %q: %w", driverCapsEnvStr, err)
	}
	if enableAll {
		// The "all" keyword here is confusing but we need to match the behavior of
		// nvidia-container-toolkit:cmd/nvidia-container-runtime-hook/container_config.go:getDriverCapabilities.
		// If the environment variable contains "all", the intuitive thing to do
		// would be to expand it to mean "all allowed capabilities". Rather, in
		// nvidia-container-toolkit, it means "expand to the entire set of
		// capabilities, then silently drop all disallowed or unsupported
		// capabilities from this set".
		// We aim to be drop-in compatible with this behavior so we need to
		// implement it too, but log a warning when these behaviors result in a
		// different outcome.
		if intuitiveCaps := envDriverCaps | allowedDriverCaps; intuitiveCaps != allowedDriverCaps {
			log.Warningf("Container requested NVIDIA driver capabilities %q; this expands to %v which is a larger set than allowed capabilities (%v). The extra capabilities (%v) will be dropped.", driverCapsEnvStr, intuitiveCaps, allowedDriverCaps, intuitiveCaps&^allowedDriverCaps)
		}
		return allowedDriverCaps, nil
	}
	// Intersect what's requested with what's allowed.
	if driverCaps := allowedDriverCaps & envDriverCaps; driverCaps != envDriverCaps {
		return 0, fmt.Errorf(`disallowed driver capabilities requested: "%v" (allowed "%v"), update --nvproxy-allowed-driver-capabilities to allow them`, envDriverCaps, driverCaps)
	}
	return envDriverCaps, 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
}

// removeNvproxyFrontendDevices removes regular Nvidia frontend devices from
// both o and n, since we handle these by remapping.
func removeNvproxyFrontendDevices(field, cName string, o, n []specs.LinuxDevice) ([]specs.LinuxDevice, []specs.LinuxDevice, error) {
	nvidiaDeviceRegex := regexp.MustCompile(`^nvidia(\d+)$`)
	shouldDelete := func(dev specs.LinuxDevice) bool {
		basename := filepath.Base(dev.Path)
		ms := nvidiaDeviceRegex.FindStringSubmatch(basename)
		if ms == nil {
			return false
		}
		minor, err := strconv.ParseUint(ms[1], 10, 32)
		if err != nil {
			return false
		}
		if minor > nvgpu.NV_MINOR_DEVICE_NUMBER_REGULAR_MAX {
			return false
		}
		return dev.Major == nvgpu.NV_MAJOR_DEVICE_NUMBER && dev.Minor == int64(minor)
	}
	o = slices.DeleteFunc(slices.Clone(o), shouldDelete)
	n = slices.DeleteFunc(slices.Clone(n), shouldDelete)
	return o, n, nil
}