diff --git a/tools/container/containerd/containerd.go b/tools/container/containerd/containerd.go index 2a41ec13..045459e6 100644 --- a/tools/container/containerd/containerd.go +++ b/tools/container/containerd/containerd.go @@ -63,15 +63,17 @@ var nvidiaRuntimeBinaries = map[string]string{ // options stores the configuration from the command line or environment variables type options struct { - config string - socket string - runtimeClass string - runtimeType string - setAsDefault bool - restartMode string - hostRootMount string - runtimeDir string + config string + socket string + runtimeClass string + runtimeDir string + setAsDefault bool + restartMode string + hostRootMount string + + // containerd-specific options useLegacyConfig bool + runtimeType string ContainerRuntimeModesCDIAnnotationPrefixes cli.StringSlice } @@ -93,6 +95,9 @@ func main() { setup.Action = func(c *cli.Context) error { return Setup(c, &options) } + setup.Before = func(c *cli.Context) error { + return ParseArgs(c, &options) + } // Create the 'cleanup' subcommand cleanup := cli.Command{} @@ -102,6 +107,9 @@ func main() { cleanup.Action = func(c *cli.Context) error { return Cleanup(c, &options) } + cleanup.Before = func(c *cli.Context) error { + return ParseArgs(c, &options) + } // Register the subcommands with the top-level CLI c.Commands = []*cli.Command{ @@ -143,6 +151,13 @@ func main() { Destination: &options.runtimeClass, EnvVars: []string{"CONTAINERD_RUNTIME_CLASS", "NVIDIA_RUNTIME_NAME"}, }, + &cli.StringFlag{ + Name: "nvidia-runtime-dir", + Aliases: []string{"runtime-dir"}, + Usage: "The path where the nvidia-container-runtime binaries are located. If this is not specified, the first argument will be used instead", + Destination: &options.runtimeDir, + EnvVars: []string{"NVIDIA_RUNTIME_DIR"}, + }, &cli.BoolFlag{ Name: "set-as-default", Usage: "Set nvidia-container-runtime as the default runtime", @@ -192,12 +207,6 @@ func main() { func Setup(c *cli.Context, o *options) error { log.Infof("Starting 'setup' for %v", c.App.Name) - runtimeDir, err := ParseArgs(c) - if err != nil { - return fmt.Errorf("unable to parse args: %v", err) - } - o.runtimeDir = runtimeDir - cfg, err := containerd.New( containerd.WithPath(o.config), containerd.WithRuntimeType(o.runtimeType), @@ -236,11 +245,6 @@ func Setup(c *cli.Context, o *options) error { func Cleanup(c *cli.Context, o *options) error { log.Infof("Starting 'cleanup' for %v", c.App.Name) - _, err := ParseArgs(c) - if err != nil { - return fmt.Errorf("unable to parse args: %v", err) - } - cfg, err := containerd.New( containerd.WithPath(o.config), containerd.WithRuntimeType(o.runtimeType), @@ -276,17 +280,24 @@ func Cleanup(c *cli.Context, o *options) error { } // ParseArgs parses the command line arguments to the CLI -func ParseArgs(c *cli.Context) (string, error) { +func ParseArgs(c *cli.Context, o *options) error { + if o.runtimeDir != "" { + log.Debug("Runtime directory already set") + return nil + } + args := c.Args() log.Infof("Parsing arguments: %v", args.Slice()) - if args.Len() != 1 { - return "", fmt.Errorf("incorrect number of arguments") + if c.NArg() != 1 { + return fmt.Errorf("incorrect number of arguments") } - runtimeDir := args.Get(0) + + o.runtimeDir = args.Get(0) + log.Infof("Successfully parsed arguments") - return runtimeDir, nil + return nil } // UpdateConfig updates the containerd config to include the nvidia-container-runtime diff --git a/tools/container/crio/crio.go b/tools/container/crio/crio.go index a55f005d..1e14d87d 100644 --- a/tools/container/crio/crio.go +++ b/tools/container/crio/crio.go @@ -52,18 +52,19 @@ const ( // options stores the configuration from the command linek or environment variables type options struct { - configMode string - - hooksDir string - hookFilename string - runtimeDir string - config string socket string runtimeClass string + runtimeDir string setAsDefault bool restartMode string hostRootMount string + + configMode string + + // hook-specific options + hooksDir string + hookFilename string } func main() { @@ -95,6 +96,10 @@ func main() { cleanup.Action = func(c *cli.Context) error { return Cleanup(c, &options) } + cleanup.Before = func(c *cli.Context) error { + return ParseArgs(c, &options) + } + // Register the subcommands with the top-level CLI c.Commands = []*cli.Command{ &setup, @@ -136,6 +141,13 @@ func main() { Destination: &options.runtimeClass, EnvVars: []string{"CRIO_RUNTIME_CLASS", "NVIDIA_RUNTIME_NAME"}, }, + &cli.StringFlag{ + Name: "nvidia-runtime-dir", + Aliases: []string{"runtime-dir"}, + Usage: "The path where the nvidia-container-runtime binaries are located. If this is not specified, the first argument will be used instead", + Destination: &options.runtimeDir, + EnvVars: []string{"NVIDIA_RUNTIME_DIR"}, + }, &cli.BoolFlag{ Name: "set-as-default", Usage: "Set nvidia-container-runtime as the default runtime", @@ -314,13 +326,20 @@ func cleanupConfig(o *options) error { // ParseArgs parses the command line arguments to the CLI func ParseArgs(c *cli.Context, o *options) error { + if o.runtimeDir != "" { + log.Debug("Runtime directory already set") + return nil + } + args := c.Args() log.Infof("Parsing arguments: %v", args.Slice()) if c.NArg() != 1 { return fmt.Errorf("incorrect number of arguments") } + o.runtimeDir = args.Get(0) + log.Infof("Successfully parsed arguments") return nil diff --git a/tools/container/docker/docker.go b/tools/container/docker/docker.go index 942c3892..c7f69cf7 100644 --- a/tools/container/docker/docker.go +++ b/tools/container/docker/docker.go @@ -54,19 +54,13 @@ const ( socketMessageToGetPID = "GET /info HTTP/1.0\r\n\r\n" ) -// nvidiaRuntimeBinaries defines a map of runtime names to binary names -var nvidiaRuntimeBinaries = map[string]string{ - nvidiaRuntimeName: nvidiaRuntimeBinary, - nvidiaExperimentalRuntimeName: nvidiaExperimentalRuntimeBinary, -} - // options stores the configuration from the command line or environment variables type options struct { config string socket string runtimeName string - setAsDefault bool runtimeDir string + setAsDefault bool restartMode string hostRootMount string } @@ -88,6 +82,9 @@ func main() { setup.Action = func(c *cli.Context) error { return Setup(c, &options) } + setup.Before = func(c *cli.Context) error { + return ParseArgs(c, &options) + } // Create the 'cleanup' subcommand cleanup := cli.Command{} @@ -97,6 +94,9 @@ func main() { cleanup.Action = func(c *cli.Context) error { return Cleanup(c, &options) } + cleanup.Before = func(c *cli.Context) error { + return ParseArgs(c, &options) + } // Register the subcommands with the top-level CLI c.Commands = []*cli.Command{ @@ -148,6 +148,13 @@ func main() { Destination: &options.runtimeName, EnvVars: []string{"DOCKER_RUNTIME_NAME", "NVIDIA_RUNTIME_NAME"}, }, + &cli.StringFlag{ + Name: "nvidia-runtime-dir", + Aliases: []string{"runtime-dir"}, + Usage: "The path where the nvidia-container-runtime binaries are located. If this is not specified, the first argument will be used instead", + Destination: &options.runtimeDir, + EnvVars: []string{"NVIDIA_RUNTIME_DIR"}, + }, &cli.BoolFlag{ Name: "set-as-default", Usage: "Set the `nvidia` runtime as the default runtime. If --runtime-name is specified as `nvidia-experimental` the experimental runtime is set as the default runtime instead", @@ -173,12 +180,6 @@ func main() { func Setup(c *cli.Context, o *options) error { log.Infof("Starting 'setup' for %v", c.App.Name) - runtimeDir, err := ParseArgs(c) - if err != nil { - return fmt.Errorf("unable to parse args: %v", err) - } - o.runtimeDir = runtimeDir - cfg, err := docker.New( docker.WithPath(o.config), ) @@ -211,11 +212,6 @@ func Setup(c *cli.Context, o *options) error { func Cleanup(c *cli.Context, o *options) error { log.Infof("Starting 'cleanup' for %v", c.App.Name) - _, err := ParseArgs(c) - if err != nil { - return fmt.Errorf("unable to parse args: %v", err) - } - cfg, err := docker.New( docker.WithPath(o.config), ) @@ -248,17 +244,24 @@ func Cleanup(c *cli.Context, o *options) error { } // ParseArgs parses the command line arguments to the CLI -func ParseArgs(c *cli.Context) (string, error) { +func ParseArgs(c *cli.Context, o *options) error { + if o.runtimeDir != "" { + log.Debug("Runtime directory already set") + return nil + } + args := c.Args() log.Infof("Parsing arguments: %v", args.Slice()) - if args.Len() != 1 { - return "", fmt.Errorf("incorrect number of arguments") + if c.NArg() != 1 { + return fmt.Errorf("incorrect number of arguments") } - runtimeDir := args.Get(0) + + o.runtimeDir = args.Get(0) + log.Infof("Successfully parsed arguments") - return runtimeDir, nil + return nil } // UpdateConfig updates the docker config to include the nvidia runtimes