From 0e6dc3f7eafd5f424769b3128c8ad560f89f31fb Mon Sep 17 00:00:00 2001 From: Evan Lezar Date: Fri, 14 Jan 2022 16:29:20 +0100 Subject: [PATCH] Move docker config handling to internal package In preparation for adding a command to the nvidia-ctk CLI to modify the docker config, this change refactors load, update, and flush logic from the toolkit container docker CLI to an internal package. Signed-off-by: Evan Lezar --- internal/config/docker/docker.go | 115 +++++++++++++ internal/config/docker/docker_test.go | 228 ++++++++++++++++++++++++++ tools/container/docker/docker.go | 78 +-------- 3 files changed, 348 insertions(+), 73 deletions(-) create mode 100644 internal/config/docker/docker.go create mode 100644 internal/config/docker/docker_test.go diff --git a/internal/config/docker/docker.go b/internal/config/docker/docker.go new file mode 100644 index 00000000..efe65aef --- /dev/null +++ b/internal/config/docker/docker.go @@ -0,0 +1,115 @@ +/** +# Copyright (c) 2021-2022, 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 ( + "bytes" + "encoding/json" + "fmt" + "io/ioutil" + "os" + + log "github.com/sirupsen/logrus" +) + +// LoadConfig loads the docker config from disk +func LoadConfig(configFilePath string) (map[string]interface{}, error) { + log.Infof("Loading docker config from %v", configFilePath) + + info, err := os.Stat(configFilePath) + if os.IsExist(err) && info.IsDir() { + return nil, fmt.Errorf("config file is a directory") + } + + cfg := make(map[string]interface{}) + + if os.IsNotExist(err) { + log.Infof("Config file does not exist, creating new one") + return cfg, nil + } + + readBytes, err := ioutil.ReadFile(configFilePath) + if err != nil { + return nil, fmt.Errorf("unable to read config: %v", err) + } + + reader := bytes.NewReader(readBytes) + if err := json.NewDecoder(reader).Decode(&cfg); err != nil { + return nil, err + } + + log.Infof("Successfully loaded config") + return cfg, nil +} + +// UpdateConfig updates the docker config to include the nvidia runtimes +func UpdateConfig(config map[string]interface{}, defaultRuntime string, newRuntimes map[string]interface{}) error { + if defaultRuntime != "" { + config["default-runtime"] = defaultRuntime + } + + // Read the existing runtimes + runtimes := make(map[string]interface{}) + if _, exists := config["runtimes"]; exists { + runtimes = config["runtimes"].(map[string]interface{}) + } + + // Add / update the runtime definitions + for name, rt := range newRuntimes { + runtimes[name] = rt + } + + // Update the runtimes definition + if len(runtimes) > 0 { + config["runtimes"] = runtimes + } + return nil +} + +// FlushConfig flushes the updated/reverted config out to disk +func FlushConfig(cfg map[string]interface{}, configFilePath string) error { + log.Infof("Flushing docker config to %v", configFilePath) + + output, err := json.MarshalIndent(cfg, "", " ") + if err != nil { + return fmt.Errorf("unable to convert to JSON: %v", err) + } + + switch len(output) { + case 0: + err := os.Remove(configFilePath) + if err != nil { + return fmt.Errorf("unable to remove empty file: %v", err) + } + log.Infof("Config empty, removing file") + default: + f, err := os.Create(configFilePath) + if err != nil { + return fmt.Errorf("unable to open %v for writing: %v", configFilePath, err) + } + defer f.Close() + + _, err = f.WriteString(string(output)) + if err != nil { + return fmt.Errorf("unable to write output: %v", err) + } + } + + log.Infof("Successfully flushed config") + + return nil +} diff --git a/internal/config/docker/docker_test.go b/internal/config/docker/docker_test.go new file mode 100644 index 00000000..9c4e961e --- /dev/null +++ b/internal/config/docker/docker_test.go @@ -0,0 +1,228 @@ +/** +# Copyright (c) 2021-2022, 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 ( + "encoding/json" + "fmt" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestUpdateConfigDefaultRuntime(t *testing.T) { + testCases := []struct { + config map[string]interface{} + defaultRuntime string + runtimeName string + expectedDefaultRuntimeName interface{} + }{ + { + defaultRuntime: "", + expectedDefaultRuntimeName: nil, + }, + { + defaultRuntime: "NAME", + expectedDefaultRuntimeName: "NAME", + }, + { + config: map[string]interface{}{ + "default-runtime": "ALREADY_SET", + }, + defaultRuntime: "", + expectedDefaultRuntimeName: "ALREADY_SET", + }, + { + config: map[string]interface{}{ + "default-runtime": "ALREADY_SET", + }, + defaultRuntime: "NAME", + expectedDefaultRuntimeName: "NAME", + }, + } + + for i, tc := range testCases { + t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) { + if tc.config == nil { + tc.config = make(map[string]interface{}) + } + err := UpdateConfig(tc.config, tc.defaultRuntime, nil) + require.NoError(t, err) + + defaultRuntimeName := tc.config["default-runtime"] + require.EqualValues(t, tc.expectedDefaultRuntimeName, defaultRuntimeName) + }) + } +} + +func TestUpdateConfigRuntimes(t *testing.T) { + testCases := []struct { + config map[string]interface{} + runtimes map[string]interface{} + expectedConfig map[string]interface{} + }{ + { + config: map[string]interface{}{}, + runtimes: map[string]interface{}{ + "runtime1": map[string]interface{}{ + "path": "/test/runtime/dir/runtime1", + "args": []string{}, + }, + "runtime2": map[string]interface{}{ + "path": "/test/runtime/dir/runtime2", + "args": []string{}, + }, + }, + expectedConfig: map[string]interface{}{ + "runtimes": map[string]interface{}{ + "runtime1": map[string]interface{}{ + "path": "/test/runtime/dir/runtime1", + "args": []string{}, + }, + "runtime2": map[string]interface{}{ + "path": "/test/runtime/dir/runtime2", + "args": []string{}, + }, + }, + }, + }, + { + config: map[string]interface{}{ + "runtimes": map[string]interface{}{ + "runtime1": map[string]interface{}{ + "path": "runtime1", + "args": []string{}, + }, + }, + }, + runtimes: map[string]interface{}{ + "runtime1": map[string]interface{}{ + "path": "/test/runtime/dir/runtime1", + "args": []string{}, + }, + "runtime2": map[string]interface{}{ + "path": "/test/runtime/dir/runtime2", + "args": []string{}, + }, + }, + expectedConfig: map[string]interface{}{ + "runtimes": map[string]interface{}{ + "runtime1": map[string]interface{}{ + "path": "/test/runtime/dir/runtime1", + "args": []string{}, + }, + "runtime2": map[string]interface{}{ + "path": "/test/runtime/dir/runtime2", + "args": []string{}, + }, + }, + }, + }, + { + config: map[string]interface{}{ + "runtimes": map[string]interface{}{ + "not-nvidia": map[string]interface{}{ + "path": "some-other-path", + "args": []string{}, + }, + }, + }, + runtimes: map[string]interface{}{ + "runtime1": map[string]interface{}{ + "path": "/test/runtime/dir/runtime1", + "args": []string{}, + }, + }, + expectedConfig: map[string]interface{}{ + "runtimes": map[string]interface{}{ + "not-nvidia": map[string]interface{}{ + "path": "some-other-path", + "args": []string{}, + }, + "runtime1": map[string]interface{}{ + "path": "/test/runtime/dir/runtime1", + "args": []string{}, + }, + }, + }, + }, + { + config: map[string]interface{}{ + "exec-opts": []string{"native.cgroupdriver=systemd"}, + "log-driver": "json-file", + "log-opts": map[string]string{ + "max-size": "100m", + }, + "storage-driver": "overlay2", + }, + runtimes: map[string]interface{}{ + "runtime1": map[string]interface{}{ + "path": "/test/runtime/dir/runtime1", + "args": []string{}, + }, + }, + expectedConfig: map[string]interface{}{ + "exec-opts": []string{"native.cgroupdriver=systemd"}, + "log-driver": "json-file", + "log-opts": map[string]string{ + "max-size": "100m", + }, + "storage-driver": "overlay2", + "runtimes": map[string]interface{}{ + "runtime1": map[string]interface{}{ + "path": "/test/runtime/dir/runtime1", + "args": []string{}, + }, + }, + }, + }, + { + config: map[string]interface{}{ + "exec-opts": []string{"native.cgroupdriver=systemd"}, + "log-driver": "json-file", + "log-opts": map[string]string{ + "max-size": "100m", + }, + "storage-driver": "overlay2", + }, + expectedConfig: map[string]interface{}{ + "exec-opts": []string{"native.cgroupdriver=systemd"}, + "log-driver": "json-file", + "log-opts": map[string]string{ + "max-size": "100m", + }, + "storage-driver": "overlay2", + }, + }, + } + + for i, tc := range testCases { + t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) { + err := UpdateConfig(tc.config, "", tc.runtimes) + require.NoError(t, err) + + configContent, err := json.MarshalIndent(tc.config, "", " ") + require.NoError(t, err) + + expectedContent, err := json.MarshalIndent(tc.expectedConfig, "", " ") + require.NoError(t, err) + + require.EqualValues(t, string(expectedContent), string(configContent)) + }) + + } +} diff --git a/tools/container/docker/docker.go b/tools/container/docker/docker.go index 21866a1b..ce6e7dca 100644 --- a/tools/container/docker/docker.go +++ b/tools/container/docker/docker.go @@ -17,16 +17,14 @@ package main import ( - "bytes" - "encoding/json" "fmt" - "io/ioutil" "net" "os" "path/filepath" "syscall" "time" + "github.com/NVIDIA/nvidia-container-toolkit/internal/config/docker" log "github.com/sirupsen/logrus" cli "github.com/urfave/cli/v2" ) @@ -247,52 +245,15 @@ func ParseArgs(c *cli.Context) (string, error) { // LoadConfig loads the docker config from disk func LoadConfig(config string) (map[string]interface{}, error) { - log.Infof("Loading config: %v", config) - - info, err := os.Stat(config) - if os.IsExist(err) && info.IsDir() { - return nil, fmt.Errorf("config file is a directory") - } - - cfg := make(map[string]interface{}) - - if os.IsNotExist(err) { - log.Infof("Config file does not exist, creating new one") - return cfg, nil - } - - readBytes, err := ioutil.ReadFile(config) - if err != nil { - return nil, fmt.Errorf("unable to read config: %v", err) - } - - reader := bytes.NewReader(readBytes) - if err := json.NewDecoder(reader).Decode(&cfg); err != nil { - return nil, err - } - - log.Infof("Successfully loaded config") - return cfg, nil + return docker.LoadConfig(config) } // UpdateConfig updates the docker config to include the nvidia runtimes func UpdateConfig(config map[string]interface{}, o *options) error { defaultRuntime := o.getDefaultRuntime() - if defaultRuntime != "" { - config["default-runtime"] = defaultRuntime - } + runtimes := o.runtimes() - runtimes := make(map[string]interface{}) - if _, exists := config["runtimes"]; exists { - runtimes = config["runtimes"].(map[string]interface{}) - } - - for name, rt := range o.runtimes() { - runtimes[name] = rt - } - - config["runtimes"] = runtimes - return nil + return docker.UpdateConfig(config, defaultRuntime, runtimes) } //RevertConfig reverts the docker config to remove the nvidia runtime @@ -320,36 +281,7 @@ func RevertConfig(config map[string]interface{}) error { // FlushConfig flushes the updated/reverted config out to disk func FlushConfig(cfg map[string]interface{}, config string) error { - log.Infof("Flushing config") - - output, err := json.MarshalIndent(cfg, "", " ") - if err != nil { - return fmt.Errorf("unable to convert to JSON: %v", err) - } - - switch len(output) { - case 0: - err := os.Remove(config) - if err != nil { - return fmt.Errorf("unable to remove empty file: %v", err) - } - log.Infof("Config empty, removing file") - default: - f, err := os.Create(config) - if err != nil { - return fmt.Errorf("unable to open %v for writing: %v", config, err) - } - defer f.Close() - - _, err = f.WriteString(string(output)) - if err != nil { - return fmt.Errorf("unable to write output: %v", err) - } - } - - log.Infof("Successfully flushed config") - - return nil + return docker.FlushConfig(cfg, config) } // RestartDocker restarts docker depending on the value of restartModeFlag