@@ -25,64 +25,34 @@ import (
2525 "tags.cncf.io/container-device-interface/pkg/cdi"
2626 "tags.cncf.io/container-device-interface/specs-go"
2727
28- "github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
29-
3028 "github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
3129 "github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
32- "github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
3330)
3431
3532type imexlib nvcdilib
3633
37- var _ Interface = (* imexlib )(nil )
34+ var _ wrapped = (* imexlib )(nil )
3835
3936const (
4037 classImexChannel = "imex-channel"
4138)
4239
43- // GetSpec should not be called for imexlib.
44- func (l * imexlib ) GetSpec (... string ) (spec.Interface , error ) {
45- return nil , fmt .Errorf ("unexpected call to imexlib.GetSpec()" )
46- }
47-
48- // GetAllDeviceSpecs returns the device specs for all available devices.
49- func (l * imexlib ) GetAllDeviceSpecs () ([]specs.Device , error ) {
50- channelsDiscoverer := discover .NewCharDeviceDiscoverer (
51- l .logger ,
52- l .devRoot ,
53- []string {"/dev/nvidia-caps-imex-channels/channel*" },
54- )
55-
56- channels , err := channelsDiscoverer .Devices ()
57- if err != nil {
58- return nil , err
59- }
60-
61- var channelIDs []string
62- for _ , channel := range channels {
63- channelIDs = append (channelIDs , filepath .Base (channel .Path ))
64- }
65-
66- return l .GetDeviceSpecsByID (channelIDs ... )
67- }
68-
6940// GetCommonEdits returns an empty set of edits for IMEX devices.
7041func (l * imexlib ) GetCommonEdits () (* cdi.ContainerEdits , error ) {
7142 return edits .FromDiscoverer (discover.None {})
7243}
7344
7445// GetDeviceSpecsByID returns the CDI device specs for the IMEX channels specified.
7546func (l * imexlib ) GetDeviceSpecsByID (ids ... string ) ([]specs.Device , error ) {
47+ channelsIDs , err := l .getChannelIDs (ids ... )
48+ if err != nil {
49+ return nil , err
50+ }
7651 var deviceSpecs []specs.Device
77- for _ , id := range ids {
78- trimmed := strings .TrimPrefix (id , "channel" )
79- _ , err := strconv .ParseUint (trimmed , 10 , 64 )
80- if err != nil {
81- return nil , fmt .Errorf ("invalid channel ID %v: %w" , id , err )
82- }
83- path := "/dev/nvidia-caps-imex-channels/channel" + trimmed
52+ for _ , id := range channelsIDs {
53+ path := "/dev/nvidia-caps-imex-channels/channel" + id
8454 deviceSpec := specs.Device {
85- Name : trimmed ,
55+ Name : id ,
8656 ContainerEdits : specs.ContainerEdits {
8757 DeviceNodes : []* specs.DeviceNode {
8858 {
@@ -97,22 +67,40 @@ func (l *imexlib) GetDeviceSpecsByID(ids ...string) ([]specs.Device, error) {
9767 return deviceSpecs , nil
9868}
9969
100- // GetGPUDeviceEdits is unsupported for the imexlib specs
101- func (l * imexlib ) GetGPUDeviceEdits (device.Device ) (* cdi.ContainerEdits , error ) {
102- return nil , fmt .Errorf ("GetGPUDeviceEdits is not supported" )
70+ func (l * imexlib ) getChannelIDs (ids ... string ) ([]string , error ) {
71+ var channelIDs []string
72+ for _ , id := range ids {
73+ trimmed := strings .TrimPrefix (id , "channel" )
74+ if trimmed == "all" {
75+ return l .getAllChannelIDs ()
76+ }
77+ _ , err := strconv .ParseUint (trimmed , 10 , 64 )
78+ if err != nil {
79+ return nil , fmt .Errorf ("invalid channel ID %v: %w" , id , err )
80+ }
81+ channelIDs = append (channelIDs , trimmed )
82+ }
83+ return channelIDs , nil
10384}
10485
105- // GetGPUDeviceSpecs is unsupported for the imexlib specs
106- func (l * imexlib ) GetGPUDeviceSpecs (int , device.Device ) ([]specs.Device , error ) {
107- return nil , fmt .Errorf ("GetGPUDeviceSpecs is not supported" )
108- }
86+ // getAllChannelIDs returns the device IDs for all available IMEX channels.
87+ func (l * imexlib ) getAllChannelIDs () ([]string , error ) {
88+ channelsDiscoverer := discover .NewCharDeviceDiscoverer (
89+ l .logger ,
90+ l .devRoot ,
91+ []string {"/dev/nvidia-caps-imex-channels/channel*" },
92+ )
10993
110- // GetMIGDeviceEdits is unsupported for the imexlib specs
111- func (l * imexlib ) GetMIGDeviceEdits (device.Device , device.MigDevice ) (* cdi.ContainerEdits , error ) {
112- return nil , fmt .Errorf ("GetMIGDeviceEdits is not supported" )
113- }
94+ channels , err := channelsDiscoverer .Devices ()
95+ if err != nil {
96+ return nil , err
97+ }
98+
99+ var channelIDs []string
100+ for _ , channel := range channels {
101+ channelID := filepath .Base (channel .Path )
102+ channelIDs = append (channelIDs , strings .TrimPrefix (channelID , "channel" ))
103+ }
114104
115- // GetMIGDeviceSpecs is unsupported for the imexlib specs
116- func (l * imexlib ) GetMIGDeviceSpecs (int , device.Device , int , device.MigDevice ) ([]specs.Device , error ) {
117- return nil , fmt .Errorf ("GetMIGDeviceSpecs is not supported" )
105+ return channelIDs , nil
118106}
0 commit comments