UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-06-25 17:47:35 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Construct CUDA image object once

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2023-06-05 16:01:20 +02:00
parent 2da32970b9
commit d92300506c
4 changed files with 17 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
internal/modifier/gds.go
View File

@ -32,17 +32,7 @@ const (
// NewGDSModifier creates the modifiers for GDS devices. // NewGDSModifier creates the modifiers for GDS devices.
// If the spec does not contain the NVIDIA_GDS=enabled environment variable no changes are made. // If the spec does not contain the NVIDIA_GDS=enabled environment variable no changes are made.
func NewGDSModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Spec) (oci.SpecModifier, error) { func NewGDSModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) {
rawSpec, err := ociSpec.Load()
if err != nil {
return nil, fmt.Errorf("failed to load OCI spec: %v", err)
}
image, err := image.NewCUDAImageFromSpec(rawSpec)
if err != nil {
return nil, err
}
if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices.List()) == 0 { if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices.List()) == 0 {
logger.Infof("No modification required; no devices requested") logger.Infof("No modification required; no devices requested")
return nil, nil return nil, nil

12
internal/modifier/graphics.go
View File

@ -28,17 +28,7 @@ import (
// NewGraphicsModifier constructs a modifier that injects graphics-related modifications into an OCI runtime specification. // 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. // 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, ociSpec oci.Spec) (oci.SpecModifier, error) { func NewGraphicsModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) {
rawSpec, err := ociSpec.Load()
if err != nil {
return nil, fmt.Errorf("failed to load OCI spec: %v", err)
}
image, err := image.NewCUDAImageFromSpec(rawSpec)
if err != nil {
return nil, err
}
if required, reason := requiresGraphicsModifier(image); !required { if required, reason := requiresGraphicsModifier(image); !required {
logger.Infof("No graphics modifier required: %v", reason) logger.Infof("No graphics modifier required: %v", reason)
return nil, nil return nil, nil

12
internal/modifier/mofed.go
View File

@ -32,17 +32,7 @@ const (
// NewMOFEDModifier creates the modifiers for MOFED devices. // NewMOFEDModifier creates the modifiers for MOFED devices.
// If the spec does not contain the NVIDIA_MOFED=enabled environment variable no changes are made. // If the spec does not contain the NVIDIA_MOFED=enabled environment variable no changes are made.
func NewMOFEDModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Spec) (oci.SpecModifier, error) { func NewMOFEDModifier(logger logger.Interface, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) {
rawSpec, err := ociSpec.Load()
if err != nil {
return nil, fmt.Errorf("failed to load OCI spec: %v", err)
}
image, err := image.NewCUDAImageFromSpec(rawSpec)
if err != nil {
return nil, err
}
if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices.List()) == 0 { if devices := image.DevicesFromEnvvars(visibleDevicesEnvvar); len(devices.List()) == 0 {
logger.Infof("No modification required; no devices requested") logger.Infof("No modification required; no devices requested")
return nil, nil return nil, nil

17
internal/runtime/runtime_factory.go
View File

@ -20,6 +20,7 @@ import (
"fmt" "fmt"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config" "github.com/NVIDIA/nvidia-container-toolkit/internal/config"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
"github.com/NVIDIA/nvidia-container-toolkit/internal/info" "github.com/NVIDIA/nvidia-container-toolkit/internal/info"
"github.com/NVIDIA/nvidia-container-toolkit/internal/logger" "github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
"github.com/NVIDIA/nvidia-container-toolkit/internal/modifier" "github.com/NVIDIA/nvidia-container-toolkit/internal/modifier"
@ -61,6 +62,16 @@ func newNVIDIAContainerRuntime(logger logger.Interface, cfg *config.Config, argv
// newSpecModifier is a factory method that creates constructs an OCI spec modifer based on the provided config. // newSpecModifier is a factory method that creates constructs an OCI spec modifer based on the provided config.
func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Spec, argv []string) (oci.SpecModifier, error) { func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Spec, argv []string) (oci.SpecModifier, error) {
rawSpec, err := ociSpec.Load()
if err != nil {
return nil, fmt.Errorf("failed to load OCI spec: %v", err)
}
image, err := image.NewCUDAImageFromSpec(rawSpec)
if err != nil {
return nil, err
}
mode := info.ResolveAutoMode(logger, cfg.NVIDIAContainerRuntimeConfig.Mode) mode := info.ResolveAutoMode(logger, cfg.NVIDIAContainerRuntimeConfig.Mode)
modeModifier, err := newModeModifier(logger, mode, cfg, ociSpec, argv) modeModifier, err := newModeModifier(logger, mode, cfg, ociSpec, argv)
if err != nil { if err != nil {
@ -71,17 +82,17 @@ func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Sp
return modeModifier, nil return modeModifier, nil
} }
graphicsModifier, err := modifier.NewGraphicsModifier(logger, cfg, ociSpec) graphicsModifier, err := modifier.NewGraphicsModifier(logger, cfg, image)
if err != nil { if err != nil {
return nil, err return nil, err
} }
gdsModifier, err := modifier.NewGDSModifier(logger, cfg, ociSpec) gdsModifier, err := modifier.NewGDSModifier(logger, cfg, image)
if err != nil { if err != nil {
return nil, err return nil, err
} }
mofedModifier, err := modifier.NewMOFEDModifier(logger, cfg, ociSpec) mofedModifier, err := modifier.NewMOFEDModifier(logger, cfg, image)
if err != nil { if err != nil {
return nil, err return nil, err
} }