Skip to content

Commit

Permalink
Add imex mode to CDI spec generation
Browse files Browse the repository at this point in the history
This change adds a imex mode to CDI spec generation. This mode detected
generates CDI specifications for existing IMEX channels. By default these
devices have the fully qualified CDI device names:

nvidia.com/imex-channel=<ID>

Signed-off-by: Evan Lezar <[email protected]>
  • Loading branch information
elezar committed Nov 25, 2024
1 parent 4e5f02d commit 3384ce2
Show file tree
Hide file tree
Showing 8 changed files with 211 additions and 0 deletions.
118 changes: 118 additions & 0 deletions pkg/nvcdi/lib-imex.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
/**
# Copyright 2024 NVIDIA CORPORATION
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
**/

package nvcdi

import (
"fmt"
"path/filepath"
"strconv"
"strings"

"tags.cncf.io/container-device-interface/pkg/cdi"
"tags.cncf.io/container-device-interface/specs-go"

"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"

"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
)

type imexlib nvcdilib

var _ Interface = (*imexlib)(nil)

const (
classImexChannel = "imex-channel"
)

// GetSpec returns a CDI spec for all available IMEX channels.
func (l *imexlib) GetSpec() (spec.Interface, error) {
return nil, nil
}

// GetAllDeviceSpecs returns the device specs for all available devices.
func (l *imexlib) GetAllDeviceSpecs() ([]specs.Device, error) {
channelsDiscoverer := discover.NewCharDeviceDiscoverer(
l.logger,
l.devRoot,
[]string{"/dev/nvidia-caps-imex-channels/channel*"},
)

channels, err := channelsDiscoverer.Devices()
if err != nil {
return nil, err
}

var channelIDs []string
for _, channel := range channels {
channelIDs = append(channelIDs, filepath.Base(channel.Path))
}

return l.GetDeviceSpecsByID(channelIDs...)
}

// GetCommonEdits returns an empty set of edits for IMEX devices.
func (l *imexlib) GetCommonEdits() (*cdi.ContainerEdits, error) {
return edits.FromDiscoverer(discover.None{})
}

// GetDeviceSpecsByID returns the CDI device specs for the IMEX channels specified.
func (l *imexlib) GetDeviceSpecsByID(ids ...string) ([]specs.Device, error) {
var deviceSpecs []specs.Device
for _, id := range ids {
trimmed := strings.TrimPrefix(id, "channel")
_, err := strconv.ParseUint(trimmed, 10, 64)
if err != nil {
return nil, fmt.Errorf("invalid channel ID %v: %w", id, err)
}
path := "/dev/nvidia-caps-imex-channels/channel" + trimmed
deviceSpec := specs.Device{
Name: trimmed,
ContainerEdits: specs.ContainerEdits{
DeviceNodes: []*specs.DeviceNode{
{
Path: path,
HostPath: filepath.Join(l.devRoot, path),
},
},
},
}
deviceSpecs = append(deviceSpecs, deviceSpec)
}
return deviceSpecs, nil
}

// GetGPUDeviceEdits is unsupported for the imexlib specs
func (l *imexlib) GetGPUDeviceEdits(device.Device) (*cdi.ContainerEdits, error) {
return nil, fmt.Errorf("GetGPUDeviceEdits is not supported")
}

// GetGPUDeviceSpecs is unsupported for the imexlib specs
func (l *imexlib) GetGPUDeviceSpecs(int, device.Device) ([]specs.Device, error) {
return nil, fmt.Errorf("GetGPUDeviceSpecs is not supported")
}

// GetMIGDeviceEdits is unsupported for the imexlib specs
func (l *imexlib) GetMIGDeviceEdits(device.Device, device.MigDevice) (*cdi.ContainerEdits, error) {
return nil, fmt.Errorf("GetMIGDeviceEdits is not supported")
}

// GetMIGDeviceSpecs is unsupported for the imexlib specs
func (l *imexlib) GetMIGDeviceSpecs(int, device.Device, int, device.MigDevice) ([]specs.Device, error) {
return nil, fmt.Errorf("GetMIGDeviceSpecs is not supported")
}
80 changes: 80 additions & 0 deletions pkg/nvcdi/lib-imex_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
/**
# Copyright 2024 NVIDIA CORPORATION
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
**/

package nvcdi

import (
"bytes"
"path/filepath"
"strings"
"testing"

testlog "github.com/sirupsen/logrus/hooks/test"
"github.com/stretchr/testify/require"

"github.com/NVIDIA/nvidia-container-toolkit/internal/test"
)

func TestImexMode(t *testing.T) {
t.Setenv("__NVCT_TESTING_DEVICES_ARE_FILES", "true")

logger, _ := testlog.NewNullLogger()

moduleRoot, err := test.GetModuleRoot()
require.NoError(t, err)
hostRoot := filepath.Join(moduleRoot, "testdata", "lookup", "rootfs-1")

expectedSpec := `---
cdiVersion: 0.5.0
containerEdits:
env:
- NVIDIA_VISIBLE_DEVICES=void
devices:
- containerEdits:
deviceNodes:
- hostPath: {{ .hostRoot }}/dev/nvidia-caps-imex-channels/channel0
path: /dev/nvidia-caps-imex-channels/channel0
name: "0"
- containerEdits:
deviceNodes:
- hostPath: {{ .hostRoot }}/dev/nvidia-caps-imex-channels/channel1
path: /dev/nvidia-caps-imex-channels/channel1
name: "1"
- containerEdits:
deviceNodes:
- hostPath: {{ .hostRoot }}/dev/nvidia-caps-imex-channels/channel2047
path: /dev/nvidia-caps-imex-channels/channel2047
name: "2047"
kind: nvidia.com/imex-channel
`
expectedSpec = strings.ReplaceAll(expectedSpec, "{{ .hostRoot }}", hostRoot)

lib, err := New(
WithLogger(logger),
WithMode(ModeImex),
WithDriverRoot(hostRoot),
)
require.NoError(t, err)

spec, err := lib.GetSpec()
require.NoError(t, err)

var b bytes.Buffer

_, err = spec.WriteTo(&b)
require.NoError(t, err)
require.Equal(t, expectedSpec, b.String())
}
5 changes: 5 additions & 0 deletions pkg/nvcdi/lib.go
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,11 @@ func New(opts ...Option) (Interface, error) {
l.class = "mofed"
}
lib = (*mofedlib)(l)
case ModeImex:
if l.class == "" {
l.class = classImexChannel
}
lib = (*imexlib)(l)
default:
return nil, fmt.Errorf("unknown mode %q", l.mode)
}
Expand Down
2 changes: 2 additions & 0 deletions pkg/nvcdi/mode.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ const (
// ModeCSV configures the CDI spec generator to generate a spec based on the contents of CSV
// mountspec files.
ModeCSV = Mode("csv")
// ModeImex configures the CDI spec generated to generate a spec for the available IMEX channels.
ModeImex = Mode("imex")
)

type modeConstraint interface {
Expand Down
Empty file.
Empty file.
Empty file.
6 changes: 6 additions & 0 deletions tools/container/toolkit/toolkit_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,12 @@ devices:
path: /dev/nvidia0
- hostPath: /host/driver/root/dev/nvidiactl
path: /dev/nvidiactl
- hostPath: /host/driver/root/dev/nvidia-caps-imex-channels/channel0
path: /dev/nvidia-caps-imex-channels/channel0
- hostPath: /host/driver/root/dev/nvidia-caps-imex-channels/channel1
path: /dev/nvidia-caps-imex-channels/channel1
- hostPath: /host/driver/root/dev/nvidia-caps-imex-channels/channel2047
path: /dev/nvidia-caps-imex-channels/channel2047
name: all
kind: example.com/class
`,
Expand Down

0 comments on commit 3384ce2

Please sign in to comment.