Add runtimeDir as argument

Thsi change adds the --nvidia-runtime-dir as a command line
argument when configuring container runtimes in the toolkit container.
This removes the need to set it via the command line.

Signed-off-by: Evan Lezar <elezar@nvidia.com>

tools/container/containerd/containerd.go

@@ -66,12 +66,14 @@ type options struct {
 	config        string
 	socket        string
 	runtimeClass  string
-	runtimeType     string
+	runtimeDir    string
 	setAsDefault  bool
 	restartMode   string
 	hostRootMount string
-	runtimeDir      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

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

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