From 5bedbc2b50bff202436c8e6c5a68d6e9ec0e1dfa Mon Sep 17 00:00:00 2001 From: Evan Lezar Date: Fri, 27 Sep 2024 10:36:39 +0200 Subject: [PATCH] Convert docker to runtime package Signed-off-by: Evan Lezar --- tools/container/docker/docker.go | 212 ------------------ tools/container/nvidia-toolkit/run.go | 46 +--- tools/container/runtime/docker/docker.go | 92 ++++++++ .../{ => runtime}/docker/docker_linux.go | 2 +- .../{ => runtime}/docker/docker_other.go | 2 +- .../{ => runtime}/docker/docker_test.go | 24 +- tools/container/runtime/runtime.go | 32 ++- 7 files changed, 136 insertions(+), 274 deletions(-) delete mode 100644 tools/container/docker/docker.go create mode 100644 tools/container/runtime/docker/docker.go rename tools/container/{ => runtime}/docker/docker_linux.go (99%) rename tools/container/{ => runtime}/docker/docker_other.go (98%) rename tools/container/{ => runtime}/docker/docker_test.go (96%) diff --git a/tools/container/docker/docker.go b/tools/container/docker/docker.go deleted file mode 100644 index 9ae92f86..00000000 --- a/tools/container/docker/docker.go +++ /dev/null @@ -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 = "" - 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 = "" - 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) -} diff --git a/tools/container/nvidia-toolkit/run.go b/tools/container/nvidia-toolkit/run.go index d17421b3..65a807ec 100644 --- a/tools/container/nvidia-toolkit/run.go +++ b/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") diff --git a/tools/container/runtime/docker/docker.go b/tools/container/runtime/docker/docker.go new file mode 100644 index 00000000..c3a07cca --- /dev/null +++ b/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) +} diff --git a/tools/container/docker/docker_linux.go b/tools/container/runtime/docker/docker_linux.go similarity index 99% rename from tools/container/docker/docker_linux.go rename to tools/container/runtime/docker/docker_linux.go index a4db4aae..2a0f9729 100644 --- a/tools/container/docker/docker_linux.go +++ b/tools/container/runtime/docker/docker_linux.go @@ -14,7 +14,7 @@ # limitations under the License. **/ -package main +package docker import ( "fmt" diff --git a/tools/container/docker/docker_other.go b/tools/container/runtime/docker/docker_other.go similarity index 98% rename from tools/container/docker/docker_other.go rename to tools/container/runtime/docker/docker_other.go index 5078d649..7b3ab618 100644 --- a/tools/container/docker/docker_other.go +++ b/tools/container/runtime/docker/docker_other.go @@ -17,7 +17,7 @@ # limitations under the License. **/ -package main +package docker import ( "errors" diff --git a/tools/container/docker/docker_test.go b/tools/container/runtime/docker/docker_test.go similarity index 96% rename from tools/container/docker/docker_test.go rename to tools/container/runtime/docker/docker_test.go index 20c66d0d..03699b40 100644 --- a/tools/container/docker/docker_test.go +++ b/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{ - Options: container.Options{ - RuntimeName: tc.runtimeName, - RuntimeDir: runtimeDir, - SetAsDefault: tc.setAsDefault, - }, + o := &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{ - Options: container.Options{ - RuntimeName: tc.runtimeName, - RuntimeDir: runtimeDir, - SetAsDefault: tc.setAsDefault, - }, + o := &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) diff --git a/tools/container/runtime/runtime.go b/tools/container/runtime/runtime.go index 9f35e885..d02aceb3 100644 --- a/tools/container/runtime/runtime.go +++ b/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 { - // We set this option here to ensure that it is avalable in future calls. +func ValidateOptions(opts *Options, runtime string, toolkitRoot string) error { + // 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) } - 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) } - return fmt.Errorf("undefined runtime %v", runtime) }