Skip to content

Commit 95b458f

Browse files
elezarelezar
committed
Move tests
Signed-off-by: Evan Lezar <[email protected]>
1 parent 5496a0f commit 95b458f

File tree

5 files changed

+344
-351
lines changed

5 files changed

+344
-351
lines changed

‎cmd/nvidia-container-runtime-hook/container_config.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -265,7 +265,7 @@ func (hookConfig *hookConfig) getContainerConfig() (config containerConfig) {
265265
image.WithDisableRequire(hookConfig.DisableRequire),
266266
image.WithAcceptDeviceListAsVolumeMounts(hookConfig.AcceptDeviceListAsVolumeMounts),
267267
image.WithAcceptEnvvarUnprivileged(hookConfig.AcceptEnvvarUnprivileged),
268-
image.WithAdditionalVisibleDevicesEnvVars(hookConfig.getSwarmResourceEnvvars()...),
268+
image.WithPreferredVisibleDevicesEnvVars(hookConfig.getSwarmResourceEnvvars()...),
269269
)
270270
if err != nil {
271271
log.Panicln(err)

‎cmd/nvidia-container-runtime-hook/container_config_test.go‎

Lines changed: 1 addition & 341 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,8 @@
11
package main
22

33
import (
4-
"path/filepath"
54
"testing"
65

7-
"github.com/opencontainers/runtime-spec/specs-go"
86
"github.com/stretchr/testify/require"
97

108
"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
@@ -480,7 +478,7 @@ func TestGetNvidiaConfig(t *testing.T) {
480478
image, _ := image.New(
481479
image.WithEnvMap(tc.env),
482480
image.WithPrivileged(tc.privileged),
483-
image.WithAdditionalVisibleDevicesEnvVars(tc.hookConfig.getSwarmResourceEnvvars()...),
481+
image.WithPreferredVisibleDevicesEnvVars(tc.hookConfig.getSwarmResourceEnvvars()...),
484482
)
485483

486484
// Wrap the call to getNvidiaConfig() in a closure.
@@ -521,344 +519,6 @@ func TestGetNvidiaConfig(t *testing.T) {
521519
}
522520
}
523521

524-
func TestDeviceListSourcePriority(t *testing.T) {
525-
var tests = []struct {
526-
description string
527-
mountDevices []specs.Mount
528-
envvarDevices string
529-
privileged bool
530-
acceptUnprivileged bool
531-
acceptMounts bool
532-
expectedDevices []string
533-
}{
534-
{
535-
description: "Mount devices, unprivileged, no accept unprivileged",
536-
mountDevices: []specs.Mount{
537-
{
538-
Source: "/dev/null",
539-
Destination: filepath.Join(image.DeviceListAsVolumeMountsRoot, "GPU0"),
540-
},
541-
{
542-
Source: "/dev/null",
543-
Destination: filepath.Join(image.DeviceListAsVolumeMountsRoot, "GPU1"),
544-
},
545-
},
546-
envvarDevices: "GPU2,GPU3",
547-
privileged: false,
548-
acceptUnprivileged: false,
549-
acceptMounts: true,
550-
expectedDevices: []string{"GPU0", "GPU1"},
551-
},
552-
{
553-
description: "No mount devices, unprivileged, no accept unprivileged",
554-
mountDevices: nil,
555-
envvarDevices: "GPU0,GPU1",
556-
privileged: false,
557-
acceptUnprivileged: false,
558-
acceptMounts: true,
559-
expectedDevices: nil,
560-
},
561-
{
562-
description: "No mount devices, privileged, no accept unprivileged",
563-
mountDevices: nil,
564-
envvarDevices: "GPU0,GPU1",
565-
privileged: true,
566-
acceptUnprivileged: false,
567-
acceptMounts: true,
568-
expectedDevices: []string{"GPU0", "GPU1"},
569-
},
570-
{
571-
description: "No mount devices, unprivileged, accept unprivileged",
572-
mountDevices: nil,
573-
envvarDevices: "GPU0,GPU1",
574-
privileged: false,
575-
acceptUnprivileged: true,
576-
acceptMounts: true,
577-
expectedDevices: []string{"GPU0", "GPU1"},
578-
},
579-
{
580-
description: "Mount devices, unprivileged, accept unprivileged, no accept mounts",
581-
mountDevices: []specs.Mount{
582-
{
583-
Source: "/dev/null",
584-
Destination: filepath.Join(image.DeviceListAsVolumeMountsRoot, "GPU0"),
585-
},
586-
{
587-
Source: "/dev/null",
588-
Destination: filepath.Join(image.DeviceListAsVolumeMountsRoot, "GPU1"),
589-
},
590-
},
591-
envvarDevices: "GPU2,GPU3",
592-
privileged: false,
593-
acceptUnprivileged: true,
594-
acceptMounts: false,
595-
expectedDevices: []string{"GPU2", "GPU3"},
596-
},
597-
{
598-
description: "Mount devices, unprivileged, no accept unprivileged, no accept mounts",
599-
mountDevices: []specs.Mount{
600-
{
601-
Source: "/dev/null",
602-
Destination: filepath.Join(image.DeviceListAsVolumeMountsRoot, "GPU0"),
603-
},
604-
{
605-
Source: "/dev/null",
606-
Destination: filepath.Join(image.DeviceListAsVolumeMountsRoot, "GPU1"),
607-
},
608-
},
609-
envvarDevices: "GPU2,GPU3",
610-
privileged: false,
611-
acceptUnprivileged: false,
612-
acceptMounts: false,
613-
expectedDevices: nil,
614-
},
615-
}
616-
for _, tc := range tests {
617-
t.Run(tc.description, func(t *testing.T) {
618-
// Wrap the call to getDevices() in a closure.
619-
var devices []string
620-
getDevices := func() {
621-
image, _ := image.New(
622-
image.WithEnvMap(
623-
map[string]string{
624-
image.EnvVarNvidiaVisibleDevices: tc.envvarDevices,
625-
},
626-
),
627-
image.WithMounts(tc.mountDevices),
628-
image.WithPrivileged(tc.privileged),
629-
image.WithAcceptDeviceListAsVolumeMounts(tc.acceptMounts),
630-
image.WithAcceptEnvvarUnprivileged(tc.acceptUnprivileged),
631-
)
632-
devices = image.VisibleDevices()
633-
}
634-
635-
// For all other tests, just grab the devices and check the results
636-
getDevices()
637-
638-
require.Equal(t, tc.expectedDevices, devices)
639-
})
640-
}
641-
}
642-
643-
func TestGetDevicesFromEnvvar(t *testing.T) {
644-
envDockerResourceGPUs := "DOCKER_RESOURCE_GPUS"
645-
gpuID := "GPU-12345"
646-
anotherGPUID := "GPU-67890"
647-
thirdGPUID := "MIG-12345"
648-
649-
var tests = []struct {
650-
description string
651-
swarmResourceEnvvars []string
652-
env map[string]string
653-
expectedDevices []string
654-
}{
655-
{
656-
description: "empty env returns nil for non-legacy image",
657-
},
658-
{
659-
description: "blank NVIDIA_VISIBLE_DEVICES returns nil for non-legacy image",
660-
env: map[string]string{
661-
image.EnvVarNvidiaVisibleDevices: "",
662-
},
663-
},
664-
{
665-
description: "'void' NVIDIA_VISIBLE_DEVICES returns nil for non-legacy image",
666-
env: map[string]string{
667-
image.EnvVarNvidiaVisibleDevices: "void",
668-
},
669-
},
670-
{
671-
description: "'none' NVIDIA_VISIBLE_DEVICES returns empty for non-legacy image",
672-
env: map[string]string{
673-
image.EnvVarNvidiaVisibleDevices: "none",
674-
},
675-
expectedDevices: []string{""},
676-
},
677-
{
678-
description: "NVIDIA_VISIBLE_DEVICES set returns value for non-legacy image",
679-
env: map[string]string{
680-
image.EnvVarNvidiaVisibleDevices: gpuID,
681-
},
682-
expectedDevices: []string{gpuID},
683-
},
684-
{
685-
description: "NVIDIA_VISIBLE_DEVICES set returns value for legacy image",
686-
env: map[string]string{
687-
image.EnvVarNvidiaVisibleDevices: gpuID,
688-
image.EnvVarCudaVersion: "legacy",
689-
},
690-
expectedDevices: []string{gpuID},
691-
},
692-
{
693-
description: "empty env returns all for legacy image",
694-
env: map[string]string{
695-
image.EnvVarCudaVersion: "legacy",
696-
},
697-
expectedDevices: []string{"all"},
698-
},
699-
// Add the `DOCKER_RESOURCE_GPUS` envvar and ensure that this is ignored when
700-
// not enabled
701-
{
702-
description: "missing NVIDIA_VISIBLE_DEVICES returns nil for non-legacy image",
703-
env: map[string]string{
704-
envDockerResourceGPUs: anotherGPUID,
705-
},
706-
},
707-
{
708-
description: "blank NVIDIA_VISIBLE_DEVICES returns nil for non-legacy image",
709-
env: map[string]string{
710-
image.EnvVarNvidiaVisibleDevices: "",
711-
envDockerResourceGPUs: anotherGPUID,
712-
},
713-
},
714-
{
715-
description: "'void' NVIDIA_VISIBLE_DEVICES returns nil for non-legacy image",
716-
env: map[string]string{
717-
image.EnvVarNvidiaVisibleDevices: "void",
718-
envDockerResourceGPUs: anotherGPUID,
719-
},
720-
},
721-
{
722-
description: "'none' NVIDIA_VISIBLE_DEVICES returns empty for non-legacy image",
723-
env: map[string]string{
724-
image.EnvVarNvidiaVisibleDevices: "none",
725-
envDockerResourceGPUs: anotherGPUID,
726-
},
727-
expectedDevices: []string{""},
728-
},
729-
{
730-
description: "NVIDIA_VISIBLE_DEVICES set returns value for non-legacy image",
731-
env: map[string]string{
732-
image.EnvVarNvidiaVisibleDevices: gpuID,
733-
envDockerResourceGPUs: anotherGPUID,
734-
},
735-
expectedDevices: []string{gpuID},
736-
},
737-
{
738-
description: "NVIDIA_VISIBLE_DEVICES set returns value for legacy image",
739-
env: map[string]string{
740-
image.EnvVarNvidiaVisibleDevices: gpuID,
741-
envDockerResourceGPUs: anotherGPUID,
742-
image.EnvVarCudaVersion: "legacy",
743-
},
744-
expectedDevices: []string{gpuID},
745-
},
746-
{
747-
description: "empty env returns all for legacy image",
748-
env: map[string]string{
749-
envDockerResourceGPUs: anotherGPUID,
750-
image.EnvVarCudaVersion: "legacy",
751-
},
752-
expectedDevices: []string{"all"},
753-
},
754-
// Add the `DOCKER_RESOURCE_GPUS` envvar and ensure that this is selected when
755-
// enabled
756-
{
757-
description: "empty env returns nil for non-legacy image",
758-
swarmResourceEnvvars: []string{envDockerResourceGPUs},
759-
},
760-
{
761-
description: "blank DOCKER_RESOURCE_GPUS returns nil for non-legacy image",
762-
swarmResourceEnvvars: []string{envDockerResourceGPUs},
763-
env: map[string]string{
764-
envDockerResourceGPUs: "",
765-
},
766-
},
767-
{
768-
description: "'void' DOCKER_RESOURCE_GPUS returns nil for non-legacy image",
769-
swarmResourceEnvvars: []string{envDockerResourceGPUs},
770-
env: map[string]string{
771-
envDockerResourceGPUs: "void",
772-
},
773-
},
774-
{
775-
description: "'none' DOCKER_RESOURCE_GPUS returns empty for non-legacy image",
776-
swarmResourceEnvvars: []string{envDockerResourceGPUs},
777-
env: map[string]string{
778-
envDockerResourceGPUs: "none",
779-
},
780-
expectedDevices: []string{""},
781-
},
782-
{
783-
description: "DOCKER_RESOURCE_GPUS set returns value for non-legacy image",
784-
swarmResourceEnvvars: []string{envDockerResourceGPUs},
785-
env: map[string]string{
786-
envDockerResourceGPUs: gpuID,
787-
},
788-
expectedDevices: []string{gpuID},
789-
},
790-
{
791-
description: "DOCKER_RESOURCE_GPUS set returns value for legacy image",
792-
swarmResourceEnvvars: []string{envDockerResourceGPUs},
793-
env: map[string]string{
794-
envDockerResourceGPUs: gpuID,
795-
image.EnvVarCudaVersion: "legacy",
796-
},
797-
expectedDevices: []string{gpuID},
798-
},
799-
{
800-
description: "DOCKER_RESOURCE_GPUS is selected if present",
801-
swarmResourceEnvvars: []string{envDockerResourceGPUs},
802-
env: map[string]string{
803-
envDockerResourceGPUs: anotherGPUID,
804-
},
805-
expectedDevices: []string{anotherGPUID},
806-
},
807-
{
808-
description: "DOCKER_RESOURCE_GPUS overrides NVIDIA_VISIBLE_DEVICES if present",
809-
swarmResourceEnvvars: []string{envDockerResourceGPUs},
810-
env: map[string]string{
811-
image.EnvVarNvidiaVisibleDevices: gpuID,
812-
envDockerResourceGPUs: anotherGPUID,
813-
},
814-
expectedDevices: []string{anotherGPUID},
815-
},
816-
{
817-
description: "DOCKER_RESOURCE_GPUS_ADDITIONAL overrides NVIDIA_VISIBLE_DEVICES if present",
818-
swarmResourceEnvvars: []string{"DOCKER_RESOURCE_GPUS_ADDITIONAL"},
819-
env: map[string]string{
820-
image.EnvVarNvidiaVisibleDevices: gpuID,
821-
"DOCKER_RESOURCE_GPUS_ADDITIONAL": anotherGPUID,
822-
},
823-
expectedDevices: []string{anotherGPUID},
824-
},
825-
{
826-
description: "All available swarm resource envvars are selected and override NVIDIA_VISIBLE_DEVICES if present",
827-
swarmResourceEnvvars: []string{"DOCKER_RESOURCE_GPUS", "DOCKER_RESOURCE_GPUS_ADDITIONAL"},
828-
env: map[string]string{
829-
image.EnvVarNvidiaVisibleDevices: gpuID,
830-
"DOCKER_RESOURCE_GPUS": thirdGPUID,
831-
"DOCKER_RESOURCE_GPUS_ADDITIONAL": anotherGPUID,
832-
},
833-
expectedDevices: []string{thirdGPUID, anotherGPUID},
834-
},
835-
{
836-
description: "DOCKER_RESOURCE_GPUS_ADDITIONAL or DOCKER_RESOURCE_GPUS override NVIDIA_VISIBLE_DEVICES if present",
837-
swarmResourceEnvvars: []string{"DOCKER_RESOURCE_GPUS", "DOCKER_RESOURCE_GPUS_ADDITIONAL"},
838-
env: map[string]string{
839-
image.EnvVarNvidiaVisibleDevices: gpuID,
840-
"DOCKER_RESOURCE_GPUS_ADDITIONAL": anotherGPUID,
841-
},
842-
expectedDevices: []string{anotherGPUID},
843-
},
844-
}
845-
846-
for _, tc := range tests {
847-
t.Run(tc.description, func(t *testing.T) {
848-
image, _ := image.New(
849-
image.WithEnvMap(tc.env),
850-
image.WithPrivileged(true),
851-
image.WithAcceptDeviceListAsVolumeMounts(false),
852-
image.WithAcceptEnvvarUnprivileged(false),
853-
image.WithAdditionalVisibleDevicesEnvVars(tc.swarmResourceEnvvars...),
854-
)
855-
856-
devices := image.VisibleDevices()
857-
require.EqualValues(t, tc.expectedDevices, devices)
858-
})
859-
}
860-
}
861-
862522
func TestGetDriverCapabilities(t *testing.T) {
863523

864524
supportedCapabilities := "compute,display,utility,video"

‎internal/config/image/builder.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -138,14 +138,14 @@ func WithPrivileged(isPrivileged bool) Option {
138138
}
139139
}
140140

141-
// WithAdditionalVisibleDevicesEnvVars sets the visible devices environment variables.
141+
// WithPreferredVisibleDevicesEnvVars sets the visible devices environment variables.
142142
// For each envvar passed:
143143
// - Split on commas
144144
// - Trim spaces and ignore empty strings
145145
// - Concatenate all of this
146-
func WithAdditionalVisibleDevicesEnvVars(visibleDevicesEnvVars ...string) Option {
146+
func WithPreferredVisibleDevicesEnvVars(preferredVisibleDeviceEnvVars ...string) Option {
147147
return func(b *builder) error {
148-
b.additionalVisibleDevicesEnvVars = visibleDevicesEnvVars
148+
b.preferredVisibleDeviceEnvVars = preferredVisibleDeviceEnvVars
149149
return nil
150150
}
151151
}

0 commit comments

Comments
 (0)