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

Add mode constants to nvcdi

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2023-02-20 16:27:34 +02:00
parent b76808dbd5
commit 2680c45811
3 changed files with 20 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -36,10 +36,6 @@ import (
)
const (
discoveryModeAuto = "auto"
discoveryModeNVML = "nvml"
discoveryModeWSL = "wsl"
formatJSON = "json"
formatYAML = "yaml"
@ -97,8 +93,8 @@ func (m command) build() *cli.Command {
},
&cli.StringFlag{
Name: "discovery-mode",
Usage: "The mode to use when discovering the available entities. One of [auto | nvml | wsl]. I mode is set to 'auto' the mode will be determined based on the system configuration.",
Value: discoveryModeAuto,
Usage: "The mode to use when discovering the available entities. One of [auto | nvml | wsl]. If mode is set to 'auto' the mode will be determined based on the system configuration.",
Value: nvcdi.ModeAuto,
Destination: &cfg.discoveryMode,
},
&cli.StringFlag{
@ -133,9 +129,9 @@ func (m command) validateFlags(r *cli.Context, cfg *config) error {
cfg.discoveryMode = strings.ToLower(cfg.discoveryMode)
switch cfg.discoveryMode {
case discoveryModeAuto:
case discoveryModeNVML:
case discoveryModeWSL:
case nvcdi.ModeAuto:
case nvcdi.ModeNvml:
case nvcdi.ModeWsl:
default:
return fmt.Errorf("invalid discovery mode: %v", cfg.discoveryMode)
}

9
pkg/nvcdi/api.go
View File

@ -22,6 +22,15 @@ import (
"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device"
)
const (
// ModeAuto configures the CDI spec generator to automatically detect the system configuration
ModeAuto = "auto"
// ModeNvml configures the CDI spec generator to use the NVML library.
ModeNvml = "nvml"
// ModeWsl configures the CDI spec generator to generate a WSL spec.
ModeWsl = "wsl"
)
// Interface defines the API for the nvcdi package
type Interface interface {
GetCommonEdits() (*cdi.ContainerEdits, error)

12
pkg/nvcdi/lib.go
View File

@ -42,7 +42,7 @@ func New(opts ...Option) Interface {
opt(l)
}
if l.mode == "" {
l.mode = "auto"
l.mode = ModeAuto
}
if l.logger == nil {
l.logger = logrus.StandardLogger()
@ -61,7 +61,7 @@ func New(opts ...Option) Interface {
}
switch l.resolveMode() {
case "nvml":
case ModeNvml:
if l.nvmllib == nil {
l.nvmllib = nvml.New()
}
@ -70,7 +70,7 @@ func New(opts ...Option) Interface {
}
return (*nvmllib)(l)
case "wsl":
case ModeWsl:
return (*wsllib)(l)
}
@ -80,7 +80,7 @@ func New(opts ...Option) Interface {
// resolveMode resolves the mode for CDI spec generation based on the current system.
func (l *nvcdilib) resolveMode() (rmode string) {
if l.mode != "auto" {
if l.mode != ModeAuto {
return l.mode
}
defer func() {
@ -91,8 +91,8 @@ func (l *nvcdilib) resolveMode() (rmode string) {
l.logger.Debugf("Is WSL-based system? %v: %v", isWSL, reason)
if isWSL {
return "wsl"
return ModeWsl
}
return "nvml"
return ModeNvml
}