Convert docker to runtime package

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

tools/container/docker/docker.go

@@ -1,212 +0,0 @@
-/**
-# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-*/
-
-package main
-
-import (
-	"fmt"
-	"os"
-
-	log "github.com/sirupsen/logrus"
-	cli "github.com/urfave/cli/v2"
-
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/info"
-	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/docker"
-	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
-)
-
-const (
-	defaultConfig       = "/etc/docker/daemon.json"
-	defaultSocket       = "/var/run/docker.sock"
-	defaultSetAsDefault = true
-	// defaultRuntimeName specifies the NVIDIA runtime to be use as the default runtime if setting the default runtime is enabled
-	defaultRuntimeName   = "nvidia"
-	defaultRestartMode   = "signal"
-	defaultHostRootMount = "/host"
-)
-
-// options stores the configuration from the command line or environment variables
-type options struct {
-	container.Options
-}
-
-func main() {
-	options := options{}
-
-	// Create the top-level CLI
-	c := cli.NewApp()
-	c.Name = "docker"
-	c.Usage = "Update docker config with the nvidia runtime"
-	c.Version = info.GetVersionString()
-
-	// Create the 'setup' subcommand
-	setup := cli.Command{}
-	setup.Name = "setup"
-	setup.Usage = "Trigger docker config to be updated"
-	setup.ArgsUsage = "<runtime_dirname>"
-	setup.Action = func(c *cli.Context) error {
-		return Setup(c, &options)
-	}
-	setup.Before = func(c *cli.Context) error {
-		return container.ParseArgs(c, &options.Options)
-	}
-
-	// Create the 'cleanup' subcommand
-	cleanup := cli.Command{}
-	cleanup.Name = "cleanup"
-	cleanup.Usage = "Trigger any updates made to docker config to be undone"
-	cleanup.ArgsUsage = "<runtime_dirname>"
-	cleanup.Action = func(c *cli.Context) error {
-		return Cleanup(c, &options)
-	}
-	cleanup.Before = func(c *cli.Context) error {
-		return container.ParseArgs(c, &options.Options)
-	}
-
-	// Register the subcommands with the top-level CLI
-	c.Commands = []*cli.Command{
-		&setup,
-		&cleanup,
-	}
-
-	// Setup common flags across both subcommands. All subcommands get the same
-	// set of flags even if they don't use some of them. This is so that we
-	// only require the user to specify one set of flags for both 'startup'
-	// and 'cleanup' to simplify things.
-	commonFlags := []cli.Flag{
-		&cli.StringFlag{
-			Name:        "config",
-			Usage:       "Path to docker config file",
-			Value:       defaultConfig,
-			Destination: &options.Config,
-			EnvVars:     []string{"RUNTIME_CONFIG", "DOCKER_CONFIG"},
-		},
-		&cli.StringFlag{
-			Name:        "socket",
-			Usage:       "Path to the docker socket file",
-			Value:       defaultSocket,
-			Destination: &options.Socket,
-			EnvVars:     []string{"RUNTIME_SOCKET", "DOCKER_SOCKET"},
-		},
-		&cli.StringFlag{
-			Name:        "restart-mode",
-			Usage:       "Specify how docker should be restarted; If 'none' is selected it will not be restarted [signal | systemd | none ]",
-			Value:       defaultRestartMode,
-			Destination: &options.RestartMode,
-			EnvVars:     []string{"RUNTIME_RESTART_MODE", "DOCKER_RESTART_MODE"},
-		},
-		&cli.StringFlag{
-			Name:        "host-root",
-			Usage:       "Specify the path to the host root to be used when restarting docker using systemd",
-			Value:       defaultHostRootMount,
-			Destination: &options.HostRootMount,
-			EnvVars:     []string{"HOST_ROOT_MOUNT"},
-			// Restart using systemd is currently not supported.
-			// We hide this option for the time being.
-			Hidden: true,
-		},
-		&cli.StringFlag{
-			Name:        "runtime-name",
-			Aliases:     []string{"nvidia-runtime-name", "runtime-class"},
-			Usage:       "Specify the name of the `nvidia` runtime. If set-as-default is selected, the runtime is used as the default runtime.",
-			Value:       defaultRuntimeName,
-			Destination: &options.RuntimeName,
-			EnvVars:     []string{"NVIDIA_RUNTIME_NAME", "DOCKER_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.",
-			Value:       defaultSetAsDefault,
-			Destination: &options.SetAsDefault,
-			EnvVars:     []string{"NVIDIA_RUNTIME_SET_AS_DEFAULT", "DOCKER_SET_AS_DEFAULT"},
-			Hidden:      true,
-		},
-	}
-
-	// Update the subcommand flags with the common subcommand flags
-	setup.Flags = append([]cli.Flag{}, commonFlags...)
-	cleanup.Flags = append([]cli.Flag{}, commonFlags...)
-
-	// Run the top-level CLI
-	if err := c.Run(os.Args); err != nil {
-		log.Errorf("Error running docker configuration: %v", err)
-		os.Exit(1)
-	}
-}
-
-// Setup updates docker configuration to include the nvidia runtime and reloads it
-func Setup(c *cli.Context, o *options) error {
-	log.Infof("Starting 'setup' for %v", c.App.Name)
-
-	cfg, err := docker.New(
-		docker.WithPath(o.Config),
-	)
-	if err != nil {
-		return fmt.Errorf("unable to load config: %v", err)
-	}
-
-	err = o.Configure(cfg)
-	if err != nil {
-		return fmt.Errorf("unable to configure docker: %v", err)
-	}
-
-	err = RestartDocker(o)
-	if err != nil {
-		return fmt.Errorf("unable to restart docker: %v", err)
-	}
-
-	log.Infof("Completed 'setup' for %v", c.App.Name)
-
-	return nil
-}
-
-// Cleanup reverts docker configuration to remove the nvidia runtime and reloads it
-func Cleanup(c *cli.Context, o *options) error {
-	log.Infof("Starting 'cleanup' for %v", c.App.Name)
-
-	cfg, err := docker.New(
-		docker.WithPath(o.Config),
-	)
-	if err != nil {
-		return fmt.Errorf("unable to load config: %v", err)
-	}
-
-	err = o.Unconfigure(cfg)
-	if err != nil {
-		return fmt.Errorf("unable to unconfigure docker: %v", err)
-	}
-
-	err = RestartDocker(o)
-	if err != nil {
-		return fmt.Errorf("unable to signal docker: %v", err)
-	}
-
-	log.Infof("Completed 'cleanup' for %v", c.App.Name)
-
-	return nil
-}
-
-// RestartDocker restarts docker depending on the value of restartModeFlag
-func RestartDocker(o *options) error {
-	return o.Restart("docker", SignalDocker)
-}

tools/container/nvidia-toolkit/run.go

@@ -3,7 +3,6 @@ package main
 import (
 	"fmt"
 	"os"
-	"os/exec"
 	"os/signal"
 	"path/filepath"
 	"strings"
@@ -92,6 +91,7 @@ func main() {
 			Destination: &options.runtime,
 			EnvVars:     []string{"RUNTIME"},
 		},
+		// TODO: Remove runtime-args
 		&cli.StringFlag{
 			Name:        "runtime-args",
 			Aliases:     []string{"u"},
@@ -136,7 +136,7 @@ func validateFlags(_ *cli.Context, o *options) error {
 	if err := toolkit.ValidateOptions(&o.toolkitOptions, o.toolkitRoot()); err != nil {
 		return err
 	}
-	if err := runtime.ValidateOptions(&o.runtimeOptions, o.toolkitRoot()); err != nil {
+	if err := runtime.ValidateOptions(&o.runtimeOptions, o.runtime, o.toolkitRoot()); err != nil {
 		return err
 	}
 	return nil
@@ -160,7 +160,7 @@ func Run(c *cli.Context, o *options) error {
 		return fmt.Errorf("unable to install toolkit: %v", err)
 	}
 
-	err = setupRuntime(o)
+	err = runtime.Setup(c, &o.runtimeOptions, o.runtime)
 	if err != nil {
 		return fmt.Errorf("unable to setup runtime: %v", err)
 	}
@@ -171,7 +171,7 @@ func Run(c *cli.Context, o *options) error {
 			return fmt.Errorf("unable to wait for signal: %v", err)
 		}
 
-		err = cleanupRuntime(o)
+		err = runtime.Cleanup(c, &o.runtimeOptions, o.runtime)
 		if err != nil {
 			return fmt.Errorf("unable to cleanup runtime: %v", err)
 		}
@@ -264,25 +264,6 @@ func initialize(pidFile string) error {
 	return nil
 }
 
-func setupRuntime(o *options) error {
-	toolkitDir := filepath.Join(o.root, toolkitSubDir)
-
-	log.Infof("Setting up runtime")
-
-	cmdline := fmt.Sprintf("%v setup %v %v\n", o.runtime, o.runtimeArgs, toolkitDir)
-
-	//nolint:gosec // TODO: Can we harden this so that there is less risk of command injection
-	cmd := exec.Command("sh", "-c", cmdline)
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-	err := cmd.Run()
-	if err != nil {
-		return fmt.Errorf("error running %v command: %v", o.runtime, err)
-	}
-
-	return nil
-}
-
 func waitForSignal() error {
 	log.Infof("Waiting for signal")
 	waitingForSignal <- true
@@ -290,25 +271,6 @@ func waitForSignal() error {
 	return nil
 }
 
-func cleanupRuntime(o *options) error {
-	toolkitDir := filepath.Join(o.root, toolkitSubDir)
-
-	log.Infof("Cleaning up Runtime")
-
-	cmdline := fmt.Sprintf("%v cleanup %v %v\n", o.runtime, o.runtimeArgs, toolkitDir)
-
-	//nolint:gosec // TODO: Can we harden this so that there is less risk of command injection
-	cmd := exec.Command("sh", "-c", cmdline)
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-	err := cmd.Run()
-	if err != nil {
-		return fmt.Errorf("error running %v command: %v", o.runtime, err)
-	}
-
-	return nil
-}
-
 func shutdown(pidFile string) {
 	log.Infof("Shutting Down")
 

tools/container/runtime/docker/docker.go

@@ -0,0 +1,92 @@
+/**
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+*/
+
+package docker
+
+import (
+	"fmt"
+
+	log "github.com/sirupsen/logrus"
+	cli "github.com/urfave/cli/v2"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/docker"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
+)
+
+const (
+	Name = "docker"
+
+	DefaultConfig      = "/etc/docker/daemon.json"
+	DefaultSocket      = "/var/run/docker.sock"
+	DefaultRestartMode = "signal"
+)
+
+// Setup updates docker configuration to include the nvidia runtime and reloads it
+func Setup(c *cli.Context, o *container.Options) error {
+	log.Infof("Starting 'setup' for %v", c.App.Name)
+
+	cfg, err := docker.New(
+		docker.WithPath(o.Config),
+	)
+	if err != nil {
+		return fmt.Errorf("unable to load config: %v", err)
+	}
+
+	err = o.Configure(cfg)
+	if err != nil {
+		return fmt.Errorf("unable to configure docker: %v", err)
+	}
+
+	err = RestartDocker(o)
+	if err != nil {
+		return fmt.Errorf("unable to restart docker: %v", err)
+	}
+
+	log.Infof("Completed 'setup' for %v", c.App.Name)
+
+	return nil
+}
+
+// Cleanup reverts docker configuration to remove the nvidia runtime and reloads it
+func Cleanup(c *cli.Context, o *container.Options) error {
+	log.Infof("Starting 'cleanup' for %v", c.App.Name)
+
+	cfg, err := docker.New(
+		docker.WithPath(o.Config),
+	)
+	if err != nil {
+		return fmt.Errorf("unable to load config: %v", err)
+	}
+
+	err = o.Unconfigure(cfg)
+	if err != nil {
+		return fmt.Errorf("unable to unconfigure docker: %v", err)
+	}
+
+	err = RestartDocker(o)
+	if err != nil {
+		return fmt.Errorf("unable to signal docker: %v", err)
+	}
+
+	log.Infof("Completed 'cleanup' for %v", c.App.Name)
+
+	return nil
+}
+
+// RestartDocker restarts docker depending on the value of restartModeFlag
+func RestartDocker(o *container.Options) error {
+	return o.Restart("docker", SignalDocker)
+}

tools/container/runtime/docker/docker_linux.go

@@ -14,7 +14,7 @@
 # limitations under the License.
 **/
 
-package main
+package docker
 
 import (
 	"fmt"

tools/container/runtime/docker/docker_other.go

@@ -17,7 +17,7 @@
 # limitations under the License.
 **/
 
-package main
+package docker
 
 import (
 	"errors"

tools/container/runtime/docker/docker_test.go

@@ -14,7 +14,7 @@
 # limitations under the License.
 */
 
-package main
+package docker
 
 import (
 	"encoding/json"
@@ -52,12 +52,10 @@ func TestUpdateConfigDefaultRuntime(t *testing.T) {
 	}
 
 	for i, tc := range testCases {
-		o := &options{
+		o := &container.Options{
-			Options: container.Options{
 			RuntimeName:  tc.runtimeName,
 			RuntimeDir:   runtimeDir,
 			SetAsDefault: tc.setAsDefault,
-			},
 		}
 
 		config := docker.Config(map[string]interface{}{})
@@ -238,12 +236,10 @@ func TestUpdateConfig(t *testing.T) {
 	for i, tc := range testCases {
 		tc := tc
 
-		o := &options{
+		o := &container.Options{
-			Options: container.Options{
 			RuntimeName:  tc.runtimeName,
 			RuntimeDir:   runtimeDir,
 			SetAsDefault: tc.setAsDefault,
-			},
 		}
 
 		err := o.UpdateConfig(&tc.config)
@@ -365,7 +361,7 @@ func TestRevertConfig(t *testing.T) {
 
 	for i, tc := range testCases {
 		tc := tc
-		o := &options{}
+		o := &container.Options{}
 		err := o.RevertConfig(&tc.config)
 
 		require.NoError(t, err, "%d: %v", i, tc)

tools/container/runtime/runtime.go

@@ -22,6 +22,7 @@ import (
 	"github.com/urfave/cli/v2"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
+	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/runtime/docker"
 )
 
 const (
@@ -89,20 +90,43 @@ func Flags(opts *Options) []cli.Flag {
 }
 
 // ValidateOptions checks whether the specified options are valid
-func ValidateOptions(opts *Options, toolkitRoot string) error {
+func ValidateOptions(opts *Options, runtime string, toolkitRoot string) error {
-	// We set this option here to ensure that it is avalable in future calls.
+	// We set this option here to ensure that it is available in future calls.
 	opts.RuntimeDir = toolkitRoot
+
+	// Apply the runtime-specific config changes.
+	switch runtime {
+	case docker.Name:
+		if opts.Config == runtimeSpecificDefault {
+			opts.Config = docker.DefaultConfig
+		}
+		if opts.Socket == runtimeSpecificDefault {
+			opts.Socket = docker.DefaultSocket
+		}
+		if opts.RestartMode == runtimeSpecificDefault {
+			opts.RestartMode = docker.DefaultRestartMode
+		}
+	default:
+		return fmt.Errorf("undefined runtime %v", runtime)
+	}
+
 	return nil
 }
 
 func Setup(c *cli.Context, opts *Options, runtime string) error {
 	switch runtime {
-	}
+	case docker.Name:
+		return docker.Setup(c, &opts.Options)
+	default:
 		return fmt.Errorf("undefined runtime %v", runtime)
+	}
 }
 
 func Cleanup(c *cli.Context, opts *Options, runtime string) error {
 	switch runtime {
-	}
+	case docker.Name:
+		return docker.Cleanup(c, &opts.Options)
+	default:
 		return fmt.Errorf("undefined runtime %v", runtime)
+	}
 }