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

Move toolkit options to struct

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2022-07-25 10:01:33 +02:00
parent 94e650c518
commit 35982e51bf
1 changed files with 26 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
tools/container/toolkit/toolkit.go
View File

@ -40,12 +40,18 @@ const (
) )
var toolkitDirArg string var toolkitDirArg string
var nvidiaDriverRootFlag string
var nvidiaContainerRuntimeDebugFlag string type options struct {
var nvidiaContainerRuntimeLogLevelFlag string DriverRoot string
var nvidiaContainerCLIDebugFlag string ContainerRuntimeDebug string
ContainerRuntimeLogLevel string
ContainerCLIDebug string
}
func main() { func main() {
opts := options{}
// Create the top-level CLI // Create the top-level CLI
c := cli.NewApp() c := cli.NewApp()
c.Name = "toolkit" c.Name = "toolkit"
@ -58,7 +64,9 @@ func main() {
install.Usage = "Install the components of the NVIDIA container toolkit" install.Usage = "Install the components of the NVIDIA container toolkit"
install.ArgsUsage = "<toolkit_directory>" install.ArgsUsage = "<toolkit_directory>"
install.Before = parseArgs install.Before = parseArgs
install.Action = Install install.Action = func(c *cli.Context) error {
return Install(c, &opts)
}
// Create the 'delete' command // Create the 'delete' command
delete := cli.Command{} delete := cli.Command{}
@ -78,24 +86,24 @@ func main() {
&cli.StringFlag{ &cli.StringFlag{
Name: "nvidia-driver-root", Name: "nvidia-driver-root",
Value: DefaultNvidiaDriverRoot, Value: DefaultNvidiaDriverRoot,
Destination: &nvidiaDriverRootFlag, Destination: &opts.DriverRoot,
EnvVars: []string{"NVIDIA_DRIVER_ROOT"}, EnvVars: []string{"NVIDIA_DRIVER_ROOT"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "nvidia-container-runtime-debug", Name: "nvidia-container-runtime-debug",
Usage: "Specify the location of the debug log file for the NVIDIA Container Runtime", Usage: "Specify the location of the debug log file for the NVIDIA Container Runtime",
Destination: &nvidiaContainerRuntimeDebugFlag, Destination: &opts.ContainerRuntimeDebug,
EnvVars: []string{"NVIDIA_CONTAINER_RUNTIME_DEBUG"}, EnvVars: []string{"NVIDIA_CONTAINER_RUNTIME_DEBUG"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "nvidia-container-runtime-debug-log-level", Name: "nvidia-container-runtime-debug-log-level",
Destination: &nvidiaContainerRuntimeLogLevelFlag, Destination: &opts.ContainerRuntimeLogLevel,
EnvVars: []string{"NVIDIA_CONTAINER_RUNTIME_LOG_LEVEL"}, EnvVars: []string{"NVIDIA_CONTAINER_RUNTIME_LOG_LEVEL"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "nvidia-container-cli-debug", Name: "nvidia-container-cli-debug",
Usage: "Specify the location of the debug log file for the NVIDIA Container CLI", Usage: "Specify the location of the debug log file for the NVIDIA Container CLI",
Destination: &nvidiaContainerCLIDebugFlag, Destination: &opts.ContainerCLIDebug,
EnvVars: []string{"NVIDIA_CONTAINER_CLI_DEBUG"}, EnvVars: []string{"NVIDIA_CONTAINER_CLI_DEBUG"},
}, },
} }
@ -135,7 +143,7 @@ func Delete(cli *cli.Context) error {
// Install installs the components of the NVIDIA container toolkit. // Install installs the components of the NVIDIA container toolkit.
// Any existing installation is removed. // Any existing installation is removed.
func Install(cli *cli.Context) error { func Install(cli *cli.Context, opts *options) error {
log.Infof("Installing NVIDIA container toolkit to '%v'", toolkitDirArg) log.Infof("Installing NVIDIA container toolkit to '%v'", toolkitDirArg)
log.Infof("Removing existing NVIDIA container toolkit installation") log.Infof("Removing existing NVIDIA container toolkit installation")
@ -157,7 +165,7 @@ func Install(cli *cli.Context) error {
return fmt.Errorf("error installing NVIDIA container library: %v", err) return fmt.Errorf("error installing NVIDIA container library: %v", err)
} }
err = installContainerRuntimes(toolkitDirArg, nvidiaDriverRootFlag) err = installContainerRuntimes(toolkitDirArg, opts.DriverRoot)
if err != nil { if err != nil {
return fmt.Errorf("error installing NVIDIA container runtime: %v", err) return fmt.Errorf("error installing NVIDIA container runtime: %v", err)
} }
@ -172,7 +180,7 @@ func Install(cli *cli.Context) error {
return fmt.Errorf("error installing NVIDIA container runtime hook: %v", err) return fmt.Errorf("error installing NVIDIA container runtime hook: %v", err)
} }
err = installToolkitConfig(toolkitConfigPath, nvidiaDriverRootFlag, nvidiaContainerCliExecutable) err = installToolkitConfig(toolkitConfigPath, nvidiaContainerCliExecutable, opts)
if err != nil { if err != nil {
return fmt.Errorf("error installing NVIDIA container toolkit config: %v", err) return fmt.Errorf("error installing NVIDIA container toolkit config: %v", err)
} }
@ -230,7 +238,7 @@ func installLibrary(libName string, toolkitDir string) error {
// installToolkitConfig installs the config file for the NVIDIA container toolkit ensuring // installToolkitConfig installs the config file for the NVIDIA container toolkit ensuring
// that the settings are updated to match the desired install and nvidia driver directories. // that the settings are updated to match the desired install and nvidia driver directories.
func installToolkitConfig(toolkitConfigPath string, nvidiaDriverDir string, nvidiaContainerCliExecutablePath string) error { func installToolkitConfig(toolkitConfigPath string, nvidiaContainerCliExecutablePath string, opts *options) error {
log.Infof("Installing NVIDIA container toolkit config '%v'", toolkitConfigPath) log.Infof("Installing NVIDIA container toolkit config '%v'", toolkitConfigPath)
config, err := toml.LoadFile(nvidiaContainerToolkitConfigSource) config, err := toml.LoadFile(nvidiaContainerToolkitConfigSource)
@ -253,17 +261,17 @@ func installToolkitConfig(toolkitConfigPath string, nvidiaDriverDir string, nvid
ldconfigPath := fmt.Sprintf("%s", config.GetPath(nvidiaContainerCliKey("ldconfig"))) ldconfigPath := fmt.Sprintf("%s", config.GetPath(nvidiaContainerCliKey("ldconfig")))
// Use the driver run root as the root: // Use the driver run root as the root:
driverLdconfigPath := "@" + filepath.Join(nvidiaDriverDir, strings.TrimPrefix(ldconfigPath, "@/")) driverLdconfigPath := "@" + filepath.Join(opts.DriverRoot, strings.TrimPrefix(ldconfigPath, "@/"))
config.SetPath(nvidiaContainerCliKey("root"), nvidiaDriverDir) config.SetPath(nvidiaContainerCliKey("root"), opts.DriverRoot)
config.SetPath(nvidiaContainerCliKey("path"), nvidiaContainerCliExecutablePath) config.SetPath(nvidiaContainerCliKey("path"), nvidiaContainerCliExecutablePath)
config.SetPath(nvidiaContainerCliKey("ldconfig"), driverLdconfigPath) config.SetPath(nvidiaContainerCliKey("ldconfig"), driverLdconfigPath)
// Set the debug options if selected // Set the debug options if selected
debugOptions := map[string]string{ debugOptions := map[string]string{
"nvidia-container-runtime.debug": nvidiaContainerRuntimeDebugFlag, "nvidia-container-runtime.debug": opts.ContainerRuntimeDebug,
"nvidia-container-runtime.log-level": nvidiaContainerRuntimeLogLevelFlag, "nvidia-container-runtime.log-level": opts.ContainerRuntimeLogLevel,
"nvidia-container-cli.debug": nvidiaContainerCLIDebugFlag, "nvidia-container-cli.debug": opts.ContainerCLIDebug,
} }
for key, value := range debugOptions { for key, value := range debugOptions {
if value == "" { if value == "" {