Move runtime flags to struct

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2022-08-26 14:59:23 +02:00
parent 503ed96275
commit 8d623967ed

View File

@ -31,14 +31,21 @@ var waitingForSignal = make(chan bool, 1)
var signalReceived = make(chan bool, 1) var signalReceived = make(chan bool, 1)
var destinationArg string var destinationArg string
var noDaemonFlag bool
var runtimeFlag string // options stores the command line arguments
var runtimeArgsFlag string type options struct {
noDaemon bool
runtime string
runtimeArgs string
}
// Version defines the CLI version. This is set at build time using LD FLAGS // Version defines the CLI version. This is set at build time using LD FLAGS
var Version = "development" var Version = "development"
func main() { func main() {
options := options{}
// Create the top-level CLI // Create the top-level CLI
c := cli.NewApp() c := cli.NewApp()
c.Name = "nvidia-toolkit" c.Name = "nvidia-toolkit"
@ -46,7 +53,9 @@ func main() {
c.UsageText = "DESTINATION [-n | --no-daemon] [-r | --runtime] [-u | --runtime-args]" c.UsageText = "DESTINATION [-n | --no-daemon] [-r | --runtime] [-u | --runtime-args]"
c.Description = "DESTINATION points to the host path underneath which the nvidia-container-toolkit should be installed.\nIt will be installed at ${DESTINATION}/toolkit" c.Description = "DESTINATION points to the host path underneath which the nvidia-container-toolkit should be installed.\nIt will be installed at ${DESTINATION}/toolkit"
c.Version = Version c.Version = Version
c.Action = Run c.Action = func(ctx *cli.Context) error {
return Run(ctx, &options)
}
// Setup flags for the CLI // Setup flags for the CLI
c.Flags = []cli.Flag{ c.Flags = []cli.Flag{
@ -54,7 +63,7 @@ func main() {
Name: "no-daemon", Name: "no-daemon",
Aliases: []string{"n"}, Aliases: []string{"n"},
Usage: "terminate immediatly after setting up the runtime. Note that no cleanup will be performed", Usage: "terminate immediatly after setting up the runtime. Note that no cleanup will be performed",
Destination: &noDaemonFlag, Destination: &options.noDaemon,
EnvVars: []string{"NO_DAEMON"}, EnvVars: []string{"NO_DAEMON"},
}, },
&cli.StringFlag{ &cli.StringFlag{
@ -62,7 +71,7 @@ func main() {
Aliases: []string{"r"}, Aliases: []string{"r"},
Usage: "the runtime to setup on this node. One of {'docker', 'crio', 'containerd'}", Usage: "the runtime to setup on this node. One of {'docker', 'crio', 'containerd'}",
Value: defaultRuntime, Value: defaultRuntime,
Destination: &runtimeFlag, Destination: &options.runtime,
EnvVars: []string{"RUNTIME"}, EnvVars: []string{"RUNTIME"},
}, },
&cli.StringFlag{ &cli.StringFlag{
@ -70,7 +79,7 @@ func main() {
Aliases: []string{"u"}, Aliases: []string{"u"},
Usage: "arguments to pass to 'docker', 'crio', or 'containerd' setup command", Usage: "arguments to pass to 'docker', 'crio', or 'containerd' setup command",
Value: defaultRuntimeArgs, Value: defaultRuntimeArgs,
Destination: &runtimeArgsFlag, Destination: &options.runtimeArgs,
EnvVars: []string{"RUNTIME_ARGS"}, EnvVars: []string{"RUNTIME_ARGS"},
}, },
} }
@ -93,8 +102,8 @@ func main() {
} }
// Run runs the core logic of the CLI // Run runs the core logic of the CLI
func Run(c *cli.Context) error { func Run(c *cli.Context, o *options) error {
err := verifyFlags() err := verifyFlags(o)
if err != nil { if err != nil {
return fmt.Errorf("unable to verify flags: %v", err) return fmt.Errorf("unable to verify flags: %v", err)
} }
@ -110,18 +119,18 @@ func Run(c *cli.Context) error {
return fmt.Errorf("unable to install toolkit: %v", err) return fmt.Errorf("unable to install toolkit: %v", err)
} }
err = setupRuntime() err = setupRuntime(o)
if err != nil { if err != nil {
return fmt.Errorf("unable to setup runtime: %v", err) return fmt.Errorf("unable to setup runtime: %v", err)
} }
if !noDaemonFlag { if !o.noDaemon {
err = waitForSignal() err = waitForSignal()
if err != nil { if err != nil {
return fmt.Errorf("unable to wait for signal: %v", err) return fmt.Errorf("unable to wait for signal: %v", err)
} }
err = cleanupRuntime() err = cleanupRuntime(o)
if err != nil { if err != nil {
return fmt.Errorf("unable to cleanup runtime: %v", err) return fmt.Errorf("unable to cleanup runtime: %v", err)
} }
@ -166,10 +175,10 @@ func ParseArgs(args []string) ([]string, error) {
return append([]string{args[0]}, args[numPositionalArgs:]...), nil return append([]string{args[0]}, args[numPositionalArgs:]...), nil
} }
func verifyFlags() error { func verifyFlags(o *options) error {
log.Infof("Verifying Flags") log.Infof("Verifying Flags")
if _, exists := availableRuntimes[runtimeFlag]; !exists { if _, exists := availableRuntimes[o.runtime]; !exists {
return fmt.Errorf("unknown runtime: %v", runtimeFlag) return fmt.Errorf("unknown runtime: %v", o.runtime)
} }
return nil return nil
} }
@ -232,19 +241,19 @@ func installToolkit() error {
return nil return nil
} }
func setupRuntime() error { func setupRuntime(o *options) error {
toolkitDir := filepath.Join(destinationArg, toolkitSubDir) toolkitDir := filepath.Join(destinationArg, toolkitSubDir)
log.Infof("Setting up runtime") log.Infof("Setting up runtime")
cmdline := fmt.Sprintf("%v setup %v %v\n", runtimeFlag, runtimeArgsFlag, toolkitDir) cmdline := fmt.Sprintf("%v setup %v %v\n", o.runtime, o.runtimeArgs, toolkitDir)
cmd := exec.Command("sh", "-c", cmdline) cmd := exec.Command("sh", "-c", cmdline)
cmd.Stdout = os.Stdout cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr cmd.Stderr = os.Stderr
err := cmd.Run() err := cmd.Run()
if err != nil { if err != nil {
return fmt.Errorf("error running %v command: %v", runtimeFlag, err) return fmt.Errorf("error running %v command: %v", o.runtime, err)
} }
return nil return nil
@ -257,19 +266,19 @@ func waitForSignal() error {
return nil return nil
} }
func cleanupRuntime() error { func cleanupRuntime(o *options) error {
toolkitDir := filepath.Join(destinationArg, toolkitSubDir) toolkitDir := filepath.Join(destinationArg, toolkitSubDir)
log.Infof("Cleaning up Runtime") log.Infof("Cleaning up Runtime")
cmdline := fmt.Sprintf("%v cleanup %v %v\n", runtimeFlag, runtimeArgsFlag, toolkitDir) cmdline := fmt.Sprintf("%v cleanup %v %v\n", o.runtime, o.runtimeArgs, toolkitDir)
cmd := exec.Command("sh", "-c", cmdline) cmd := exec.Command("sh", "-c", cmdline)
cmd.Stdout = os.Stdout cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr cmd.Stderr = os.Stderr
err := cmd.Run() err := cmd.Run()
if err != nil { if err != nil {
return fmt.Errorf("error running %v command: %v", runtimeFlag, err) return fmt.Errorf("error running %v command: %v", o.runtime, err)
} }
return nil return nil