mirror of
https://github.com/NVIDIA/nvidia-container-toolkit
synced 2024-11-22 08:18:32 +00:00
Add --discovery-mode to nvidia-ctk cdi generate command
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:
parent
5103adab89
commit
20d6e9af04
@ -36,6 +36,8 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
discoveryModeNVML = "nvml"
|
||||
|
||||
formatJSON = "json"
|
||||
formatYAML = "yaml"
|
||||
)
|
||||
@ -50,6 +52,7 @@ type config struct {
|
||||
deviceNameStrategy string
|
||||
driverRoot string
|
||||
nvidiaCTKPath string
|
||||
discoveryMode string
|
||||
}
|
||||
|
||||
// NewCommand constructs a generate-cdi command with the specified logger
|
||||
@ -88,6 +91,12 @@ func (m command) build() *cli.Command {
|
||||
Value: formatYAML,
|
||||
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{
|
||||
Name: "device-name-strategy",
|
||||
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)
|
||||
}
|
||||
|
||||
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)
|
||||
if err != nil {
|
||||
return err
|
||||
@ -229,6 +245,7 @@ func (m command) generateSpec(cfg *config) (*specs.Spec, error) {
|
||||
nvcdi.WithDeviceNamer(deviceNamer),
|
||||
nvcdi.WithDeviceLib(devicelib),
|
||||
nvcdi.WithNvmlLib(nvmllib),
|
||||
nvcdi.WithMode(string(cfg.discoveryMode)),
|
||||
)
|
||||
|
||||
deviceSpecs, err := cdilib.GetAllDeviceSpecs()
|
||||
@ -298,3 +315,5 @@ func createParentDirsIfRequired(filename string) error {
|
||||
}
|
||||
return os.MkdirAll(dir, 0755)
|
||||
}
|
||||
|
||||
type discoveryMode string
|
||||
|
Loading…
Reference in New Issue
Block a user