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

Add WSL2 discovery and spec generation

Browse Source 
These changes add a wsl discovery mode to the nvidia-ctk cdi generate command.

If wsl mode is enabled, the driver store for the available devices is used as
the source for discovered entities.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar
2023-02-06 18:53:24 +01:00
parent d226925fe7
commit b21dc929ef
6 changed files with 375 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
cmd/nvidia-ctk/cdi/generate/generate.go
View File

@@ -37,9 +37,12 @@ import (
const (
discoveryModeNVML = "nvml"
discoveryModeWSL = "wsl"
formatJSON = "json"
formatYAML = "yaml"
allDeviceName = "all"
)
type command struct {
@@ -93,7 +96,7 @@ func (m command) build() *cli.Command {
},
&cli.StringFlag{
Name: "discovery-mode",
Usage: "The mode to use when discovering the available entities. One of [nvml]",
Usage: "The mode to use when discovering the available entities. One of [nvml | wsl]",
Value: discoveryModeNVML,
Destination: &cfg.discoveryMode,
},
@@ -130,6 +133,7 @@ func (m command) validateFlags(r *cli.Context, cfg *config) error {
cfg.discoveryMode = strings.ToLower(cfg.discoveryMode)
switch cfg.discoveryMode {
case discoveryModeNVML:
case discoveryModeWSL:
default:
return fmt.Errorf("invalid discovery mode: %v", cfg.discoveryMode)
}
@@ -252,10 +256,20 @@ func (m command) generateSpec(cfg *config) (*specs.Spec, error) {
if err != nil {
return nil, fmt.Errorf("failed to create device CDI specs: %v", err)
}
allDevice := createAllDevice(deviceSpecs)
deviceSpecs = append(deviceSpecs, allDevice)
var hasAll bool
for _, deviceSpec := range deviceSpecs {
if deviceSpec.Name == allDeviceName {
hasAll = true
break
}
}
if !hasAll {
allDevice, err := MergeDeviceSpecs(deviceSpecs, allDeviceName)
if err != nil {
return nil, fmt.Errorf("failed to create CDI specification for %q device: %v", allDeviceName, err)
}
deviceSpecs = append(deviceSpecs, allDevice)
}
commonEdits, err := cdilib.GetCommonEdits()
if err != nil {
@@ -287,22 +301,32 @@ func (m command) generateSpec(cfg *config) (*specs.Spec, error) {
return &spec, nil
}
// createAllDevice creates an 'all' device which combines the edits from the previous devices
func createAllDevice(deviceSpecs []specs.Device) specs.Device {
edits := edits.NewContainerEdits()
// MergeDeviceSpecs creates a device with the specified name which combines the edits from the previous devices.
// If a device of the specified name already exists, an error is returned.
func MergeDeviceSpecs(deviceSpecs []specs.Device, mergedDeviceName string) (specs.Device, error) {
if err := cdi.ValidateDeviceName(mergedDeviceName); err != nil {
return specs.Device{}, fmt.Errorf("invalid device name %q: %v", mergedDeviceName, err)
}
for _, d := range deviceSpecs {
if d.Name == mergedDeviceName {
return specs.Device{}, fmt.Errorf("device %q already exists", mergedDeviceName)
}
}
mergedEdits := edits.NewContainerEdits()
for _, d := range deviceSpecs {
edit := cdi.ContainerEdits{
ContainerEdits: &d.ContainerEdits,
}
edits.Append(&edit)
mergedEdits.Append(&edit)
}
all := specs.Device{
Name: "all",
ContainerEdits: *edits.ContainerEdits,
merged := specs.Device{
Name: mergedDeviceName,
ContainerEdits: *mergedEdits.ContainerEdits,
}
return all
return merged, nil
}
// createParentDirsIfRequired creates the parent folders of the specified path if requried.
@@ -315,5 +339,3 @@ func createParentDirsIfRequired(filename string) error {
}
return os.MkdirAll(dir, 0755)
}
type discoveryMode string