Skip to content

Commit 6b037a0

Browse files
elezarelezar
authored
Merge pull request #917 from elezar/remove-nvidia-container-runtime-hook-in-cdi-mode
Remove nvidia container runtime hook in cdi mode
2 parents b170a35 + 9eccc16 commit 6b037a0

File tree

5 files changed

+191
-77
lines changed

5 files changed

+191
-77
lines changed

‎internal/modifier/csv.go

+1-11
Original file line numberDiff line numberDiff line change
@@ -68,20 +68,10 @@ func NewCSVModifier(logger logger.Interface, cfg *config.Config, container image
6868
return nil, fmt.Errorf("failed to get CDI spec: %v", err)
6969
}
7070

71-
cdiModifier, err := cdi.New(
71+
return cdi.New(
7272
cdi.WithLogger(logger),
7373
cdi.WithSpec(spec.Raw()),
7474
)
75-
if err != nil {
76-
return nil, fmt.Errorf("failed to construct CDI modifier: %v", err)
77-
}
78-
79-
modifiers := Merge(
80-
nvidiaContainerRuntimeHookRemover{logger},
81-
cdiModifier,
82-
)
83-
84-
return modifiers, nil
8575
}
8676

8777
func checkRequirements(logger logger.Interface, image image.CUDA) error {

‎internal/modifier/csv_test.go

-64
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ package modifier
1919
import (
2020
"testing"
2121

22-
"github.com/opencontainers/runtime-spec/specs-go"
2322
testlog "github.com/sirupsen/logrus/hooks/test"
2423
"github.com/stretchr/testify/require"
2524

@@ -74,66 +73,3 @@ func TestNewCSVModifier(t *testing.T) {
7473
})
7574
}
7675
}
77-
78-
func TestCSVModifierRemovesHook(t *testing.T) {
79-
logger, _ := testlog.NewNullLogger()
80-
81-
testCases := []struct {
82-
description string
83-
spec *specs.Spec
84-
expectedError error
85-
expectedSpec *specs.Spec
86-
}{
87-
{
88-
description: "modification removes existing nvidia-container-runtime-hook",
89-
spec: &specs.Spec{
90-
Hooks: &specs.Hooks{
91-
Prestart: []specs.Hook{
92-
{
93-
Path: "/path/to/nvidia-container-runtime-hook",
94-
Args: []string{"/path/to/nvidia-container-runtime-hook", "prestart"},
95-
},
96-
},
97-
},
98-
},
99-
expectedSpec: &specs.Spec{
100-
Hooks: &specs.Hooks{
101-
Prestart: []specs.Hook{},
102-
},
103-
},
104-
},
105-
{
106-
description: "modification removes existing nvidia-container-toolkit",
107-
spec: &specs.Spec{
108-
Hooks: &specs.Hooks{
109-
Prestart: []specs.Hook{
110-
{
111-
Path: "/path/to/nvidia-container-toolkit",
112-
Args: []string{"/path/to/nvidia-container-toolkit", "prestart"},
113-
},
114-
},
115-
},
116-
},
117-
expectedSpec: &specs.Spec{
118-
Hooks: &specs.Hooks{
119-
Prestart: []specs.Hook{},
120-
},
121-
},
122-
},
123-
}
124-
125-
for _, tc := range testCases {
126-
t.Run(tc.description, func(t *testing.T) {
127-
m := nvidiaContainerRuntimeHookRemover{logger: logger}
128-
129-
err := m.Modify(tc.spec)
130-
if tc.expectedError != nil {
131-
require.Error(t, err)
132-
} else {
133-
require.NoError(t, err)
134-
}
135-
136-
require.Empty(t, tc.spec.Hooks.Prestart)
137-
})
138-
}
139-
}

‎internal/modifier/hook_remover.go

+7
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,13 @@ type nvidiaContainerRuntimeHookRemover struct {
3333

3434
var _ oci.SpecModifier = (*nvidiaContainerRuntimeHookRemover)(nil)
3535

36+
// NewNvidiaContainerRuntimeHookRemover creates a modifier that removes any NVIDIA Container Runtime hooks from the provided spec.
37+
func NewNvidiaContainerRuntimeHookRemover(logger logger.Interface) oci.SpecModifier {
38+
return nvidiaContainerRuntimeHookRemover{
39+
logger: logger,
40+
}
41+
}
42+
3643
// Modify removes any NVIDIA Container Runtime hooks from the provided spec
3744
func (m nvidiaContainerRuntimeHookRemover) Modify(spec *specs.Spec) error {
3845
if spec == nil {

‎internal/runtime/runtime_factory.go

+4-2
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Sp
8585
switch modifierType {
8686
case "mode":
8787
modifiers = append(modifiers, modeModifier)
88+
case "nvidia-hook-remover":
89+
modifiers = append(modifiers, modifier.NewNvidiaContainerRuntimeHookRemover(logger))
8890
case "graphics":
8991
graphicsModifier, err := modifier.NewGraphicsModifier(logger, cfg, image, driver)
9092
if err != nil {
@@ -121,10 +123,10 @@ func supportedModifierTypes(mode string) []string {
121123
switch mode {
122124
case "cdi":
123125
// For CDI mode we make no additional modifications.
124-
return []string{"mode"}
126+
return []string{"nvidia-hook-remover", "mode"}
125127
case "csv":
126128
// For CSV mode we support mode and feature-gated modification.
127-
return []string{"mode", "feature-gated"}
129+
return []string{"nvidia-hook-remover", "mode", "feature-gated"}
128130
default:
129131
return []string{"mode", "graphics", "feature-gated"}
130132
}

‎internal/runtime/runtime_factory_test.go

+179
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import (
3030

3131
"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
3232
"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/root"
33+
"github.com/NVIDIA/nvidia-container-toolkit/internal/oci"
3334
"github.com/NVIDIA/nvidia-container-toolkit/internal/test"
3435
)
3536

@@ -165,3 +166,181 @@ func TestFactoryMethod(t *testing.T) {
165166
})
166167
}
167168
}
169+
170+
func TestNewSpecModifier(t *testing.T) {
171+
logger, _ := testlog.NewNullLogger()
172+
driver := root.New(
173+
root.WithDriverRoot("/nvidia/driver/root"),
174+
)
175+
testCases := []struct {
176+
description string
177+
config *config.Config
178+
spec *specs.Spec
179+
expectedSpec *specs.Spec
180+
}{
181+
{
182+
description: "csv mode removes nvidia-container-runtime-hook",
183+
config: &config.Config{
184+
NVIDIAContainerRuntimeConfig: config.RuntimeConfig{
185+
Mode: "csv",
186+
},
187+
},
188+
spec: &specs.Spec{
189+
Hooks: &specs.Hooks{
190+
Prestart: []specs.Hook{
191+
{
192+
Path: "/path/to/nvidia-container-runtime-hook",
193+
Args: []string{"/path/to/nvidia-container-runtime-hook", "prestart"},
194+
},
195+
},
196+
},
197+
},
198+
expectedSpec: &specs.Spec{
199+
Hooks: &specs.Hooks{
200+
Prestart: nil,
201+
},
202+
},
203+
},
204+
{
205+
description: "csv mode removes nvidia-container-toolkit",
206+
config: &config.Config{
207+
NVIDIAContainerRuntimeConfig: config.RuntimeConfig{
208+
Mode: "csv",
209+
},
210+
},
211+
spec: &specs.Spec{
212+
Hooks: &specs.Hooks{
213+
Prestart: []specs.Hook{
214+
{
215+
Path: "/path/to/nvidia-container-toolkit",
216+
Args: []string{"/path/to/nvidia-container-toolkit", "prestart"},
217+
},
218+
},
219+
},
220+
},
221+
expectedSpec: &specs.Spec{
222+
Hooks: &specs.Hooks{
223+
Prestart: nil,
224+
},
225+
},
226+
},
227+
{
228+
description: "cdi mode removes nvidia-container-runtime-hook",
229+
config: &config.Config{
230+
NVIDIAContainerRuntimeConfig: config.RuntimeConfig{
231+
Mode: "cdi",
232+
},
233+
},
234+
spec: &specs.Spec{
235+
Hooks: &specs.Hooks{
236+
Prestart: []specs.Hook{
237+
{
238+
Path: "/path/to/nvidia-container-runtime-hook",
239+
Args: []string{"/path/to/nvidia-container-runtime-hook", "prestart"},
240+
},
241+
},
242+
},
243+
},
244+
expectedSpec: &specs.Spec{
245+
Hooks: &specs.Hooks{
246+
Prestart: nil,
247+
},
248+
},
249+
},
250+
{
251+
description: "cdi mode removes nvidia-container-toolkit",
252+
config: &config.Config{
253+
NVIDIAContainerRuntimeConfig: config.RuntimeConfig{
254+
Mode: "cdi",
255+
},
256+
},
257+
spec: &specs.Spec{
258+
Hooks: &specs.Hooks{
259+
Prestart: []specs.Hook{
260+
{
261+
Path: "/path/to/nvidia-container-toolkit",
262+
Args: []string{"/path/to/nvidia-container-toolkit", "prestart"},
263+
},
264+
},
265+
},
266+
},
267+
expectedSpec: &specs.Spec{
268+
Hooks: &specs.Hooks{
269+
Prestart: nil,
270+
},
271+
},
272+
},
273+
{
274+
description: "legacy mode keeps nvidia-container-runtime-hook",
275+
config: &config.Config{
276+
NVIDIAContainerRuntimeConfig: config.RuntimeConfig{
277+
Mode: "legacy",
278+
},
279+
},
280+
spec: &specs.Spec{
281+
Hooks: &specs.Hooks{
282+
Prestart: []specs.Hook{
283+
{
284+
Path: "/path/to/nvidia-container-runtime-hook",
285+
Args: []string{"/path/to/nvidia-container-runtime-hook", "prestart"},
286+
},
287+
},
288+
},
289+
},
290+
expectedSpec: &specs.Spec{
291+
Hooks: &specs.Hooks{
292+
Prestart: []specs.Hook{
293+
{
294+
Path: "/path/to/nvidia-container-runtime-hook",
295+
Args: []string{"/path/to/nvidia-container-runtime-hook", "prestart"},
296+
},
297+
},
298+
},
299+
},
300+
},
301+
{
302+
description: "legacy mode keeps nvidia-container-toolkit",
303+
config: &config.Config{
304+
NVIDIAContainerRuntimeConfig: config.RuntimeConfig{
305+
Mode: "legacy",
306+
},
307+
},
308+
spec: &specs.Spec{
309+
Hooks: &specs.Hooks{
310+
Prestart: []specs.Hook{
311+
{
312+
Path: "/path/to/nvidia-container-toolkit",
313+
Args: []string{"/path/to/nvidia-container-toolkit", "prestart"},
314+
},
315+
},
316+
},
317+
},
318+
expectedSpec: &specs.Spec{
319+
Hooks: &specs.Hooks{
320+
Prestart: []specs.Hook{
321+
{
322+
Path: "/path/to/nvidia-container-toolkit",
323+
Args: []string{"/path/to/nvidia-container-toolkit", "prestart"},
324+
},
325+
},
326+
},
327+
},
328+
},
329+
}
330+
331+
for _, tc := range testCases {
332+
t.Run(tc.description, func(t *testing.T) {
333+
spec := &oci.SpecMock{
334+
LoadFunc: func() (*specs.Spec, error) {
335+
return tc.spec, nil
336+
},
337+
}
338+
m, err := newSpecModifier(logger, tc.config, spec, driver)
339+
require.NoError(t, err)
340+
341+
err = m.Modify(tc.spec)
342+
require.NoError(t, err)
343+
require.EqualValues(t, tc.expectedSpec, tc.spec)
344+
})
345+
}
346+
}

0 commit comments

Comments
 (0)