Skip to content

Commit 48c9397

Browse files
committed
[no-relnote] Move nvcdi wrapper to separate file
Signed-off-by: Evan Lezar <[email protected]>
1 parent bd8ecb1 commit 48c9397

File tree

2 files changed

+72
-44
lines changed

2 files changed

+72
-44
lines changed

pkg/nvcdi/lib.go

Lines changed: 0 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -24,27 +24,15 @@ import (
2424
"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
2525
"github.com/NVIDIA/go-nvlib/pkg/nvlib/info"
2626
"github.com/NVIDIA/go-nvml/pkg/nvml"
27-
"tags.cncf.io/container-device-interface/pkg/cdi"
2827

29-
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
3028
"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
3129
"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/cuda"
3230
"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/root"
3331
"github.com/NVIDIA/nvidia-container-toolkit/internal/nvsandboxutils"
3432
"github.com/NVIDIA/nvidia-container-toolkit/internal/platform-support/tegra/csv"
35-
"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
3633
"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/transform"
3734
)
3835

39-
type wrapper struct {
40-
Interface
41-
42-
vendor string
43-
class string
44-
45-
mergedDeviceOptions []transform.MergedDeviceOption
46-
}
47-
4836
type nvcdilib struct {
4937
logger logger.Interface
5038
nvmllib nvml.Interface
@@ -177,38 +165,6 @@ func New(opts ...Option) (Interface, error) {
177165
return &w, nil
178166
}
179167

180-
// GetSpec combines the device specs and common edits from the wrapped Interface to a single spec.Interface.
181-
func (l *wrapper) GetSpec() (spec.Interface, error) {
182-
deviceSpecs, err := l.GetAllDeviceSpecs()
183-
if err != nil {
184-
return nil, err
185-
}
186-
187-
edits, err := l.GetCommonEdits()
188-
if err != nil {
189-
return nil, err
190-
}
191-
192-
return spec.New(
193-
spec.WithDeviceSpecs(deviceSpecs),
194-
spec.WithEdits(*edits.ContainerEdits),
195-
spec.WithVendor(l.vendor),
196-
spec.WithClass(l.class),
197-
spec.WithMergedDeviceOptions(l.mergedDeviceOptions...),
198-
)
199-
}
200-
201-
// GetCommonEdits returns the wrapped edits and adds additional edits on top.
202-
func (m *wrapper) GetCommonEdits() (*cdi.ContainerEdits, error) {
203-
edits, err := m.Interface.GetCommonEdits()
204-
if err != nil {
205-
return nil, err
206-
}
207-
edits.Env = append(edits.Env, image.EnvVarNvidiaVisibleDevices+"=void")
208-
209-
return edits, nil
210-
}
211-
212168
// resolveMode resolves the mode for CDI spec generation based on the current system.
213169
func (l *nvcdilib) resolveMode() (rmode string) {
214170
if l.mode != ModeAuto {

pkg/nvcdi/wrapper.go

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
/**
2+
# Copyright 2024 NVIDIA CORPORATION
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
**/
16+
17+
package nvcdi
18+
19+
import (
20+
"tags.cncf.io/container-device-interface/pkg/cdi"
21+
"tags.cncf.io/container-device-interface/specs-go"
22+
23+
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
24+
"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
25+
"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/transform"
26+
)
27+
28+
type wrapper struct {
29+
Interface
30+
31+
vendor string
32+
class string
33+
34+
mergedDeviceOptions []transform.MergedDeviceOption
35+
}
36+
37+
// GetSpec combines the device specs and common edits from the wrapped Interface to a single spec.Interface.
38+
func (l *wrapper) GetSpec() (spec.Interface, error) {
39+
deviceSpecs, err := l.GetAllDeviceSpecs()
40+
if err != nil {
41+
return nil, err
42+
}
43+
44+
edits, err := l.GetCommonEdits()
45+
if err != nil {
46+
return nil, err
47+
}
48+
49+
return spec.New(
50+
spec.WithDeviceSpecs(deviceSpecs),
51+
spec.WithEdits(*edits.ContainerEdits),
52+
spec.WithVendor(l.vendor),
53+
spec.WithClass(l.class),
54+
spec.WithMergedDeviceOptions(l.mergedDeviceOptions...),
55+
)
56+
}
57+
58+
// GetAllDeviceSpecs returns the device specs for all available devices.
59+
func (l *wrapper) GetAllDeviceSpecs() ([]specs.Device, error) {
60+
return l.Interface.GetAllDeviceSpecs()
61+
}
62+
63+
// GetCommonEdits returns the wrapped edits and adds additional edits on top.
64+
func (m *wrapper) GetCommonEdits() (*cdi.ContainerEdits, error) {
65+
edits, err := m.Interface.GetCommonEdits()
66+
if err != nil {
67+
return nil, err
68+
}
69+
edits.Env = append(edits.Env, image.EnvVarNvidiaVisibleDevices+"=void")
70+
71+
return edits, nil
72+
}

0 commit comments

Comments
 (0)