UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-06-26 18:18:24 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

[no-relnote] Use image.CUDA to extract visible devices

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar
2024-10-14 15:06:06 +02:00
parent 1991b3ef2a
commit 92df542f2f
11 changed files with 313 additions and 337 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
internal/modifier/cdi.go
View File

@@ -90,11 +90,9 @@ func getDevicesFromSpec(logger logger.Interface, ociSpec oci.Spec, cfg *config.C
}
}
envDevices := container.DevicesFromEnvvars(visibleDevicesEnvvar)
var devices []string
seen := make(map[string]bool)
for _, name := range envDevices.List() {
for _, name := range container.VisibleDevicesFromEnvVar() {
if !parser.IsQualifiedName(name) {
name = fmt.Sprintf("%s=%s", cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.DefaultKind, name)
}

15
internal/modifier/csv.go
View File

@@ -30,23 +30,16 @@ import (
"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi"
)
const (
visibleDevicesEnvvar = "NVIDIA_VISIBLE_DEVICES"
visibleDevicesVoid = "void"
nvidiaRequireJetpackEnvvar = "NVIDIA_REQUIRE_JETPACK"
)
// NewCSVModifier creates a modifier that applies modications to an OCI spec if required by the runtime wrapper.
// The modifications are defined by CSV MountSpecs.
func NewCSVModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) {
if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices.List()) == 0 {
func NewCSVModifier(logger logger.Interface, cfg *config.Config, container image.CUDA) (oci.SpecModifier, error) {
if devices := container.VisibleDevicesFromEnvVar(); len(devices) == 0 {
logger.Infof("No modification required; no devices requested")
return nil, nil
}
logger.Infof("Constructing modifier from config: %+v", *cfg)
if err := checkRequirements(logger, image); err != nil {
if err := checkRequirements(logger, container); err != nil {
return nil, fmt.Errorf("requirements not met: %v", err)
}
@@ -55,7 +48,7 @@ func NewCSVModifier(logger logger.Interface, cfg *config.Config, image image.CUD
return nil, fmt.Errorf("failed to get list of CSV files: %v", err)
}
if image.Getenv(nvidiaRequireJetpackEnvvar) != "csv-mounts=all" {
if container.Getenv(image.EnvVarNvidiaRequireJetpack) != "csv-mounts=all" {
csvFiles = csv.BaseFilesOnly(csvFiles)
}

2
internal/modifier/gated.go
View File

@@ -36,7 +36,7 @@ import (
//
// If not devices are selected, no changes are made.
func NewFeatureGatedModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) {
if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices.List()) == 0 {
if devices := image.VisibleDevicesFromEnvVar(); len(devices) == 0 {
logger.Infof("No modification required; no devices requested")
return nil, nil
}

8
internal/modifier/graphics.go
View File

@@ -29,8 +29,8 @@ import (
// NewGraphicsModifier constructs a modifier that injects graphics-related modifications into an OCI runtime specification.
// The value of the NVIDIA_DRIVER_CAPABILITIES environment variable is checked to determine if this modification should be made.
func NewGraphicsModifier(logger logger.Interface, cfg *config.Config, image image.CUDA, driver *root.Driver) (oci.SpecModifier, error) {
if required, reason := requiresGraphicsModifier(image); !required {
func NewGraphicsModifier(logger logger.Interface, cfg *config.Config, containerImage image.CUDA, driver *root.Driver) (oci.SpecModifier, error) {
if required, reason := requiresGraphicsModifier(containerImage); !required {
logger.Infof("No graphics modifier required: %v", reason)
return nil, nil
}
@@ -50,7 +50,7 @@ func NewGraphicsModifier(logger logger.Interface, cfg *config.Config, image imag
devRoot := driver.Root
drmNodes, err := discover.NewDRMNodesDiscoverer(
logger,
image.DevicesFromEnvvars(visibleDevicesEnvvar),
containerImage.DevicesFromEnvvars(image.EnvVarNvidiaVisibleDevices),
devRoot,
nvidiaCDIHookPath,
)
@@ -67,7 +67,7 @@ func NewGraphicsModifier(logger logger.Interface, cfg *config.Config, image imag
// requiresGraphicsModifier determines whether a graphics modifier is required.
func requiresGraphicsModifier(cudaImage image.CUDA) (bool, string) {
if devices := cudaImage.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices.List()) == 0 {
if devices := cudaImage.VisibleDevicesFromEnvVar(); len(devices) == 0 {
return false, "no devices requested"
}