UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2024-11-25 21:39:10 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add --discovery-mode to nvidia-ctk cdi generate command

Browse Source 
This change adds --discovery-mode flag to the nvidia-ctk cdi generate
command and plumbs this through to the CDI API.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2023-02-08 15:29:35 +01:00
parent 5103adab89
commit 20d6e9af04
1 changed files with 19 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -36,6 +36,8 @@ import (
) )
const ( const (
discoveryModeNVML = "nvml"
formatJSON = "json" formatJSON = "json"
formatYAML = "yaml" formatYAML = "yaml"
) )
@ -50,6 +52,7 @@ type config struct {
deviceNameStrategy string deviceNameStrategy string
driverRoot string driverRoot string
nvidiaCTKPath string nvidiaCTKPath string
discoveryMode string
} }
// NewCommand constructs a generate-cdi command with the specified logger // NewCommand constructs a generate-cdi command with the specified logger
@ -88,6 +91,12 @@ func (m command) build() *cli.Command {
Value: formatYAML, Value: formatYAML,
Destination: &cfg.format, Destination: &cfg.format,
}, },
&cli.StringFlag{
Name: "discovery-mode",
Usage: "The mode to use when discovering the available entities. One of [nvml]",
Value: discoveryModeNVML,
Destination: &cfg.discoveryMode,
},
&cli.StringFlag{ &cli.StringFlag{
Name: "device-name-strategy", Name: "device-name-strategy",
Usage: "Specify the strategy for generating device names. One of [index | uuid | type-index]", Usage: "Specify the strategy for generating device names. One of [index | uuid | type-index]",
@ -118,6 +127,13 @@ func (m command) validateFlags(r *cli.Context, cfg *config) error {
return fmt.Errorf("invalid output format: %v", cfg.format) return fmt.Errorf("invalid output format: %v", cfg.format)
} }
cfg.discoveryMode = strings.ToLower(cfg.discoveryMode)
switch cfg.discoveryMode {
case discoveryModeNVML:
default:
return fmt.Errorf("invalid discovery mode: %v", cfg.discoveryMode)
}
_, err := nvcdi.NewDeviceNamer(cfg.deviceNameStrategy) _, err := nvcdi.NewDeviceNamer(cfg.deviceNameStrategy)
if err != nil { if err != nil {
return err return err
@ -229,6 +245,7 @@ func (m command) generateSpec(cfg *config) (*specs.Spec, error) {
nvcdi.WithDeviceNamer(deviceNamer), nvcdi.WithDeviceNamer(deviceNamer),
nvcdi.WithDeviceLib(devicelib), nvcdi.WithDeviceLib(devicelib),
nvcdi.WithNvmlLib(nvmllib), nvcdi.WithNvmlLib(nvmllib),
nvcdi.WithMode(string(cfg.discoveryMode)),
) )
deviceSpecs, err := cdilib.GetAllDeviceSpecs() deviceSpecs, err := cdilib.GetAllDeviceSpecs()
@ -298,3 +315,5 @@ func createParentDirsIfRequired(filename string) error {
} }
return os.MkdirAll(dir, 0755) return os.MkdirAll(dir, 0755)
} }
type discoveryMode string