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

Use CUDA.DevicesFromEnvvar to check if modifications are required

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar
2022-06-29 17:33:57 +02:00
parent 925c348565
commit 4e08ec2405
3 changed files with 30 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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
}