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

Merge branch 'refactor-envvar-devices' into 'main'

Browse Source 
Add DevicesFromEnvvars function to CUDA image abstraction

See merge request nvidia/container-toolkit/container-toolkit!178
This commit is contained in:
Evan Lezar
2022-07-04 08:47:28 +00:00
parent 877083f091 4e08ec2405
commit 89824849d3
6 changed files with 79 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
internal/config/image/cuda_image.go
View File

@@ -112,6 +112,36 @@ func (i CUDA) HasDisableRequire() bool {
return false
}
// DevicesFromEnvvars returns the devices requested by the image through environment variables
func (i CUDA) DevicesFromEnvvars(envVars ...string) []string {
// Grab a reference to devices from the first envvar
// in the list that actually exists in the environment.
var devices *string
for _, envVar := range envVars {
if devs, ok := i[envVar]; ok {
devices = &devs
break
}
}
// Environment variable unset with legacy image: default to "all".
if devices == nil && i.IsLegacy() {
return []string{"all"}
}
// Environment variable unset or empty or "void": return nil
if devices == nil || len(*devices) == 0 || *devices == "void" {
return nil
}
// Environment variable set to "none": reset to "".
if *devices == "none" {
return []string{""}
}
return strings.Split(*devices, ",")
}
func (i CUDA) legacyVersion() (string, error) {
majorMinor, err := parseMajorMinorVersion(i[envCUDAVersion])
if err != nil {

25
internal/modifier/csv.go
View File

@@ -50,10 +50,13 @@ func NewCSVModifier(logger *logrus.Logger, cfg *config.Config, ociSpec oci.Spec)
return nil, fmt.Errorf("failed to load OCI spec: %v", err)
}
// We check whether a modification is required and return a nil modifier if this is not the case.
visibleDevices, exists := ociSpec.LookupEnv(visibleDevicesEnvvar)
if !exists || visibleDevices == "" || visibleDevices == visibleDevicesVoid {
logger.Infof("No modification required: %v=%v (exists=%v)", visibleDevicesEnvvar, visibleDevices, exists)
image, err := image.NewCUDAImageFromSpec(rawSpec)
if err != nil {
return nil, err
}
if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices) == 0 {
logger.Infof("No modification required; no devices requested")
return nil, nil
}
logger.Infof("Constructing modifier from config: %+v", *cfg)
@@ -63,14 +66,7 @@ func NewCSVModifier(logger *logrus.Logger, cfg *config.Config, ociSpec oci.Spec)
NVIDIAContainerToolkitCLIExecutablePath: cfg.NVIDIACTKConfig.Path,
}
// TODO: Once the devices have been encapsulated in the CUDA image, this can be moved to before the
// visible devices are checked.
image, err := image.NewCUDAImageFromSpec(rawSpec)
if err != nil {
return nil, err
}
if err := checkRequirements(logger, &image); err != nil {
if err := checkRequirements(logger, image); err != nil {
return nil, fmt.Errorf("requirements not met: %v", err)
}
@@ -79,8 +75,7 @@ func NewCSVModifier(logger *logrus.Logger, cfg *config.Config, ociSpec oci.Spec)
return nil, fmt.Errorf("failed to get list of CSV files: %v", err)
}
nvidiaRequireJetpack, _ := ociSpec.LookupEnv(nvidiaRequireJetpackEnvvar)
if nvidiaRequireJetpack != "csv-mounts=all" {
if nvidiaRequireJetpack, _ := image[nvidiaRequireJetpackEnvvar]; nvidiaRequireJetpack != "csv-mounts=all" {
csvFiles = csv.BaseFilesOnly(csvFiles)
}
@@ -114,7 +109,7 @@ func NewCSVModifier(logger *logrus.Logger, cfg *config.Config, ociSpec oci.Spec)
return modifiers, nil
}
func checkRequirements(logger *logrus.Logger, image *image.CUDA) error {
func checkRequirements(logger *logrus.Logger, image image.CUDA) error {
if image.HasDisableRequire() {
// TODO: We could print the real value here instead
logger.Debugf("NVIDIA_DISABLE_REQUIRE=%v; skipping requirement checks", true)

16
internal/modifier/gds.go
View File

@@ -20,6 +20,7 @@ import (
"fmt"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
"github.com/NVIDIA/nvidia-container-toolkit/internal/oci"
"github.com/sirupsen/logrus"
@@ -32,19 +33,22 @@ const (
// NewGDSModifier creates the modifiers for GDS devices.
// If the spec does not contain the NVIDIA_GDS=enabled environment variable no changes are made.
func NewGDSModifier(logger *logrus.Logger, cfg *config.Config, ociSpec oci.Spec) (oci.SpecModifier, error) {
_, err := ociSpec.Load()
rawSpec, err := ociSpec.Load()
if err != nil {
return nil, fmt.Errorf("failed to load OCI spec: %v", err)
}
// We check whether a modification is required and return a nil modifier if this is not the case.
visibleDevices, exists := ociSpec.LookupEnv(visibleDevicesEnvvar)
if !exists || visibleDevices == "" || visibleDevices == visibleDevicesVoid {
logger.Infof("No modification required: %v=%v (exists=%v)", visibleDevicesEnvvar, visibleDevices, exists)
image, err := image.NewCUDAImageFromSpec(rawSpec)
if err != nil {
return nil, err
}
if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices) == 0 {
logger.Infof("No modification required; no devices requested")
return nil, nil
}
if gds, _ := ociSpec.LookupEnv(nvidiaGDSEnvvar); gds != "enabled" {
if gds, _ := image[nvidiaGDSEnvvar]; gds != "enabled" {
return nil, nil
}

16
internal/modifier/mofed.go
View File

@@ -20,6 +20,7 @@ import (
"fmt"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
"github.com/NVIDIA/nvidia-container-toolkit/internal/oci"
"github.com/sirupsen/logrus"
@@ -32,19 +33,22 @@ const (
// NewMOFEDModifier creates the modifiers for MOFED devices.
// If the spec does not contain the NVIDIA_MOFED=enabled environment variable no changes are made.
func NewMOFEDModifier(logger *logrus.Logger, cfg *config.Config, ociSpec oci.Spec) (oci.SpecModifier, error) {
_, err := ociSpec.Load()
rawSpec, err := ociSpec.Load()
if err != nil {
return nil, fmt.Errorf("failed to load OCI spec: %v", err)
}
// We check whether a modification is required and return a nil modifier if this is not the case.
visibleDevices, exists := ociSpec.LookupEnv(visibleDevicesEnvvar)
if !exists || visibleDevices == "" || visibleDevices == visibleDevicesVoid {
logger.Infof("No modification required: %v=%v (exists=%v)", visibleDevicesEnvvar, visibleDevices, exists)
image, err := image.NewCUDAImageFromSpec(rawSpec)
if err != nil {
return nil, err
}
if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices) == 0 {
logger.Infof("No modification required; no devices requested")
return nil, nil
}
if mofed, _ := ociSpec.LookupEnv(nvidiaMOFEDEnvvar); mofed != "enabled" {
if mofed, _ := image[nvidiaMOFEDEnvvar]; mofed != "enabled" {
return nil, nil
}