@@ -9,9 +9,10 @@ import (
99 "path/filepath"
1010 "strings"
1111
12- "github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
1312 "github.com/opencontainers/runtime-spec/specs-go"
1413 "golang.org/x/mod/semver"
14+
15+ "github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
1516)
1617
1718const (
@@ -22,6 +23,7 @@ const (
2223 envNVVisibleDevices = "NVIDIA_VISIBLE_DEVICES"
2324 envNVMigConfigDevices = "NVIDIA_MIG_CONFIG_DEVICES"
2425 envNVMigMonitorDevices = "NVIDIA_MIG_MONITOR_DEVICES"
26+ envNVImexChannels = "NVIDIA_IMEX_CHANNELS"
2527 envNVDriverCapabilities = "NVIDIA_DRIVER_CAPABILITIES"
2628)
2729
@@ -37,6 +39,7 @@ type nvidiaConfig struct {
3739 Devices string
3840 MigConfigDevices string
3941 MigMonitorDevices string
42+ ImexChannels string
4043 DriverCapabilities string
4144 // Requirements defines the requirements DSL for the container to run.
4245 // This is empty if no specific requirements are needed, or if requirements are
@@ -271,6 +274,13 @@ func getMigMonitorDevices(env map[string]string) *string {
271274 return nil
272275}
273276
277+ func getImexChannels (env map [string ]string ) * string {
278+ if chans , ok := env [envNVImexChannels ]; ok {
279+ return & chans
280+ }
281+ return nil
282+ }
283+
274284func (c * HookConfig ) getDriverCapabilities (env map [string ]string , legacyImage bool ) image.DriverCapabilities {
275285 // We use the default driver capabilities by default. This is filtered to only include the
276286 // supported capabilities
@@ -324,6 +334,11 @@ func getNvidiaConfig(hookConfig *HookConfig, image image.CUDA, mounts []Mount, p
324334 log .Panicln ("cannot set MIG_MONITOR_DEVICES in non privileged container" )
325335 }
326336
337+ var imexChannels string
338+ if c := getImexChannels (image ); c != nil {
339+ imexChannels = * c
340+ }
341+
327342 driverCapabilities := hookConfig .getDriverCapabilities (image , legacyImage ).String ()
328343
329344 requirements , err := image .GetRequirements ()
@@ -335,6 +350,7 @@ func getNvidiaConfig(hookConfig *HookConfig, image image.CUDA, mounts []Mount, p
335350 Devices : devices ,
336351 MigConfigDevices : migConfigDevices ,
337352 MigMonitorDevices : migMonitorDevices ,
353+ ImexChannels : imexChannels ,
338354 DriverCapabilities : driverCapabilities ,
339355 Requirements : requirements ,
340356 }
0 commit comments