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

Return empty requirements if NVIDIA_DISABLE_REQUIRE is true

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar
2023-07-07 12:34:43 +02:00
parent 7451e6eb75
commit 1081cecea9
7 changed files with 104 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
cmd/nvidia-container-runtime-hook/container_config.go
View File

@@ -38,8 +38,10 @@ type nvidiaConfig struct {
MigConfigDevices string
MigMonitorDevices string
DriverCapabilities string
Requirements []string
DisableRequire bool
// Requirements defines the requirements DSL for the container to run.
// This is empty if no specific requirements are needed, or if requirements are
// explicitly disabled.
Requirements []string
}
type containerConfig struct {
@@ -327,15 +329,12 @@ func getNvidiaConfig(hookConfig *HookConfig, image image.CUDA, mounts []Mount, p
log.Panicln("failed to get requirements", err)
}
disableRequire := image.HasDisableRequire()
return &nvidiaConfig{
Devices: devices,
MigConfigDevices: migConfigDevices,
MigMonitorDevices: migMonitorDevices,
DriverCapabilities: driverCapabilities,
Requirements: requirements,
DisableRequire: disableRequire,
}
}
@@ -353,7 +352,10 @@ func getContainerConfig(hook HookConfig) (config containerConfig) {
s := loadSpec(path.Join(b, "config.json"))
image, err := image.NewCUDAImageFromEnv(s.Process.Env)
image, err := image.New(
image.WithEnv(s.Process.Env),
image.WithDisableRequire(hook.DisableRequire),
)
if err != nil {
log.Panicln(err)
}

25
cmd/nvidia-container-runtime-hook/container_config_test.go
View File

@@ -40,7 +40,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "all",
DriverCapabilities: allDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -54,7 +53,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "all",
DriverCapabilities: allDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -86,7 +84,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "",
DriverCapabilities: allDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -100,7 +97,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: allDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -115,7 +111,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: defaultDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -130,7 +125,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: allDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -145,7 +139,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: "video,display",
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -162,7 +155,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: "video,display",
Requirements: []string{"cuda>=9.0", "req0=true", "req1=false"},
DisableRequire: false,
},
},
{
@@ -179,8 +171,7 @@ func TestGetNvidiaConfig(t *testing.T) {
expectedConfig: &nvidiaConfig{
Devices: "gpu0,gpu1",
DriverCapabilities: "video,display",
Requirements: []string{"cuda>=9.0", "req0=true", "req1=false"},
DisableRequire: true,
Requirements: []string{},
},
},
{
@@ -211,7 +202,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "all",
DriverCapabilities: defaultDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -243,7 +233,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "",
DriverCapabilities: defaultDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -257,7 +246,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: defaultDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -272,7 +260,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: defaultDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -287,7 +274,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: allDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -302,7 +288,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: "video,display",
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -319,7 +304,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "gpu0,gpu1",
DriverCapabilities: "video,display",
Requirements: []string{"cuda>=9.0", "req0=true", "req1=false"},
DisableRequire: false,
},
},
{
@@ -336,8 +320,7 @@ func TestGetNvidiaConfig(t *testing.T) {
expectedConfig: &nvidiaConfig{
Devices: "gpu0,gpu1",
DriverCapabilities: "video,display",
Requirements: []string{"cuda>=9.0", "req0=true", "req1=false"},
DisableRequire: true,
Requirements: []string{},
},
},
{
@@ -351,7 +334,6 @@ func TestGetNvidiaConfig(t *testing.T) {
Devices: "all",
DriverCapabilities: defaultDriverCapabilities.String(),
Requirements: []string{},
DisableRequire: false,
},
},
{
@@ -367,7 +349,6 @@ func TestGetNvidiaConfig(t *testing.T) {
MigConfigDevices: "mig0,mig1",
DriverCapabilities: defaultDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -393,7 +374,6 @@ func TestGetNvidiaConfig(t *testing.T) {
MigMonitorDevices: "mig0,mig1",
DriverCapabilities: defaultDriverCapabilities.String(),
Requirements: []string{"cuda>=9.0"},
DisableRequire: false,
},
},
{
@@ -525,7 +505,6 @@ func TestGetNvidiaConfig(t *testing.T) {
require.Equal(t, tc.expectedConfig.DriverCapabilities, config.DriverCapabilities)
require.ElementsMatch(t, tc.expectedConfig.Requirements, config.Requirements)
require.Equal(t, tc.expectedConfig.DisableRequire, config.DisableRequire)
})
}
}

6
cmd/nvidia-container-runtime-hook/main.go
View File

@@ -137,10 +137,8 @@ func doPrestart() {
args = append(args, capabilityToCLI(cap))
}
if !hook.DisableRequire && !nvidia.DisableRequire {
for _, req := range nvidia.Requirements {
args = append(args, fmt.Sprintf("--require=%s", req))
}
for _, req := range nvidia.Requirements {
args = append(args, fmt.Sprintf("--require=%s", req))
}
args = append(args, fmt.Sprintf("--pid=%s", strconv.FormatUint(uint64(container.Pid), 10)))