UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2024-11-22 00:08:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Migrate docker config to engine.Interface

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2023-02-23 15:43:15 +02:00
parent e5bb4d2718
commit 9fff19da23
6 changed files with 221 additions and 181 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
cmd/nvidia-ctk/runtime/configure/configure.go
View File

@ -127,13 +127,14 @@ func (m command) configureDocker(c *cli.Context, config *config) error {
configFilePath = defaultDockerConfigFilePath
}
cfg, err := docker.LoadConfig(configFilePath)
cfg, err := docker.New(
docker.WithPath(configFilePath),
)
if err != nil {
return fmt.Errorf("unable to load config: %v", err)
}
err = docker.UpdateConfig(
cfg,
err = cfg.AddRuntime(
config.nvidiaOptions.RuntimeName,
config.nvidiaOptions.RuntimePath,
config.nvidiaOptions.SetAsDefault,
@ -150,12 +151,16 @@ func (m command) configureDocker(c *cli.Context, config *config) error {
os.Stdout.WriteString(fmt.Sprintf("%s\n", output))
return nil
}
err = docker.FlushConfig(cfg, configFilePath)
n, err := cfg.Save(configFilePath)
if err != nil {
return fmt.Errorf("unable to flush config: %v", err)
}
m.logger.Infof("Wrote updated config to %v", configFilePath)
if n == 0 {
m.logger.Infof("Removed empty config from %v", configFilePath)
} else {
m.logger.Infof("Wrote updated config to %v", configFilePath)
}
m.logger.Infof("It is recommended that the docker daemon be restarted.")
return nil

145
internal/config/engine/docker/docker.go
View File

@ -17,47 +17,39 @@
package docker
import (
"bytes"
"encoding/json"
"fmt"
"io/ioutil"
"os"
log "github.com/sirupsen/logrus"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
)
// LoadConfig loads the docker config from disk
func LoadConfig(configFilePath string) (map[string]interface{}, error) {
log.Infof("Loading docker config from %v", configFilePath)
const (
defaultDockerRuntime = "runc"
)
info, err := os.Stat(configFilePath)
if os.IsExist(err) && info.IsDir() {
return nil, fmt.Errorf("config file is a directory")
// Config defines a docker config file.
// TODO: This should not be public, but we need to access it from the tests in tools/container/docker
type Config map[string]interface{}
// New creates a docker config with the specified options
func New(opts ...Option) (engine.Interface, error) {
b := &builder{}
for _, opt := range opts {
opt(b)
}
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
return b.build()
}
// UpdateConfig updates the docker config to include the nvidia runtimes
func UpdateConfig(config map[string]interface{}, runtimeName string, runtimePath string, setAsDefault bool) error {
// AddRuntime adds a new runtime to the docker config
func (c *Config) AddRuntime(name string, path string, setAsDefault bool) error {
if c == nil {
return fmt.Errorf("config is nil")
}
config := *c
// Read the existing runtimes
runtimes := make(map[string]interface{})
if _, exists := config["runtimes"]; exists {
@ -65,53 +57,84 @@ func UpdateConfig(config map[string]interface{}, runtimeName string, runtimePath
}
// Add / update the runtime definitions
runtimes[runtimeName] = map[string]interface{}{
"path": runtimePath,
runtimes[name] = map[string]interface{}{
"path": path,
"args": []string{},
}
// Update the runtimes definition
if len(runtimes) > 0 {
config["runtimes"] = runtimes
}
config["runtimes"] = runtimes
if setAsDefault {
config["default-runtime"] = runtimeName
config["default-runtime"] = name
}
*c = config
return nil
}
// DefaultRuntime returns the default runtime for the docker config
func (c Config) DefaultRuntime() string {
r, ok := c["default-runtime"].(string)
if !ok {
return ""
}
return r
}
// RemoveRuntime removes a runtime from the docker config
func (c *Config) RemoveRuntime(name string) error {
if c == nil {
return nil
}
config := *c
if _, exists := config["default-runtime"]; exists {
defaultRuntime := config["default-runtime"].(string)
if defaultRuntime == name {
config["default-runtime"] = defaultDockerRuntime
}
}
if _, exists := config["runtimes"]; exists {
runtimes := config["runtimes"].(map[string]interface{})
delete(runtimes, name)
if len(runtimes) == 0 {
delete(config, "runtimes")
}
}
*c = config
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, "", " ")
// Save writes the config to the specified path
func (c Config) Save(path string) (int64, error) {
output, err := json.MarshalIndent(c, "", " ")
if err != nil {
return fmt.Errorf("unable to convert to JSON: %v", err)
return 0, fmt.Errorf("unable to convert to JSON: %v", err)
}
switch len(output) {
case 0:
err := os.Remove(configFilePath)
if len(output) == 0 {
err := os.Remove(path)
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)
return 0, fmt.Errorf("unable to remove empty file: %v", err)
}
return 0, nil
}
log.Infof("Successfully flushed config")
f, err := os.Create(path)
if err != nil {
return 0, fmt.Errorf("unable to open %v for writing: %v", path, err)
}
defer f.Close()
return nil
n, err := f.WriteString(string(output))
if err != nil {
return 0, fmt.Errorf("unable to write output: %v", err)
}
return int64(n), nil
}

8
internal/config/engine/docker/docker_test.go
View File

@ -26,7 +26,7 @@ import (
func TestUpdateConfigDefaultRuntime(t *testing.T) {
testCases := []struct {
config map[string]interface{}
config Config
runtimeName string
setAsDefault bool
expectedDefaultRuntimeName interface{}
@ -63,7 +63,7 @@ func TestUpdateConfigDefaultRuntime(t *testing.T) {
if tc.config == nil {
tc.config = make(map[string]interface{})
}
err := UpdateConfig(tc.config, tc.runtimeName, "", tc.setAsDefault)
err := tc.config.AddRuntime(tc.runtimeName, "", tc.setAsDefault)
require.NoError(t, err)
defaultRuntimeName := tc.config["default-runtime"]
@ -74,7 +74,7 @@ func TestUpdateConfigDefaultRuntime(t *testing.T) {
func TestUpdateConfigRuntimes(t *testing.T) {
testCases := []struct {
config map[string]interface{}
config Config
runtimes map[string]string
expectedConfig map[string]interface{}
}{
@ -198,7 +198,7 @@ func TestUpdateConfigRuntimes(t *testing.T) {
for i, tc := range testCases {
t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
for runtimeName, runtimePath := range tc.runtimes {
err := UpdateConfig(tc.config, runtimeName, runtimePath, false)
err := tc.config.AddRuntime(runtimeName, runtimePath, false)
require.NoError(t, err)
}

80
internal/config/engine/docker/option.go Normal file
View File

@ -0,0 +1,80 @@
/**
# Copyright (c) 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"
)
type builder struct {
path string
}
// Option defines a function that can be used to configure the config builder
type Option func(*builder)
// WithPath sets the path for the config builder
func WithPath(path string) Option {
return func(b *builder) {
b.path = path
}
}
func (b *builder) build() (*Config, error) {
if b.path == "" {
empty := make(Config)
return &empty, nil
}
return loadConfig(b.path)
}
// loadConfig loads the docker config from disk
func loadConfig(configFilePath string) (*Config, 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(Config)
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
}

100
tools/container/docker/docker.go
View File

@ -20,11 +20,12 @@ import (
"fmt"
"net"
"os"
"path/filepath"
"syscall"
"time"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/docker"
"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator"
log "github.com/sirupsen/logrus"
cli "github.com/urfave/cli/v2"
)
@ -170,7 +171,9 @@ func Setup(c *cli.Context, o *options) error {
}
o.runtimeDir = runtimeDir
cfg, err := LoadConfig(o.config)
cfg, err := docker.New(
docker.WithPath(o.config),
)
if err != nil {
return fmt.Errorf("unable to load config: %v", err)
}
@ -180,7 +183,8 @@ func Setup(c *cli.Context, o *options) error {
return fmt.Errorf("unable to update config: %v", err)
}
err = FlushConfig(cfg, o.config)
log.Infof("Flushing docker config to %v", o.config)
_, err = cfg.Save(o.config)
if err != nil {
return fmt.Errorf("unable to flush config: %v", err)
}
@ -204,20 +208,26 @@ func Cleanup(c *cli.Context, o *options) error {
return fmt.Errorf("unable to parse args: %v", err)
}
cfg, err := LoadConfig(o.config)
cfg, err := docker.New(
docker.WithPath(o.config),
)
if err != nil {
return fmt.Errorf("unable to load config: %v", err)
}
err = RevertConfig(cfg)
err = RevertConfig(cfg, o)
if err != nil {
return fmt.Errorf("unable to update config: %v", err)
}
err = FlushConfig(cfg, o.config)
log.Infof("Flushing docker config to %v", o.config)
n, err := cfg.Save(o.config)
if err != nil {
return fmt.Errorf("unable to flush config: %v", err)
}
if n == 0 {
log.Infof("Config file is empty, removed")
}
err = RestartDocker(o)
if err != nil {
@ -243,18 +253,17 @@ func ParseArgs(c *cli.Context) (string, error) {
return runtimeDir, nil
}
// LoadConfig loads the docker config from disk
func LoadConfig(config string) (map[string]interface{}, error) {
return docker.LoadConfig(config)
}
// UpdateConfig updates the docker config to include the nvidia runtimes
func UpdateConfig(config map[string]interface{}, o *options) error {
for runtimeName, runtimePath := range o.getRuntimeBinaries() {
setAsDefault := runtimeName == o.getDefaultRuntime()
err := docker.UpdateConfig(config, runtimeName, runtimePath, setAsDefault)
func UpdateConfig(cfg engine.Interface, o *options) error {
runtimes := operator.GetRuntimes(
operator.WithNvidiaRuntimeName(o.runtimeName),
operator.WithSetAsDefault(o.setAsDefault),
operator.WithRoot(o.runtimeDir),
)
for name, runtime := range runtimes {
err := cfg.AddRuntime(name, runtime.Path, runtime.SetAsDefault)
if err != nil {
return fmt.Errorf("failed to update runtime %q: %v", runtimeName, err)
return fmt.Errorf("failed to update runtime %q: %v", name, err)
}
}
@ -262,33 +271,22 @@ func UpdateConfig(config map[string]interface{}, o *options) error {
}
// RevertConfig reverts the docker config to remove the nvidia runtime
func RevertConfig(config map[string]interface{}) error {
if _, exists := config["default-runtime"]; exists {
defaultRuntime := config["default-runtime"].(string)
if _, exists := nvidiaRuntimeBinaries[defaultRuntime]; exists {
config["default-runtime"] = defaultDockerRuntime
func RevertConfig(cfg engine.Interface, o *options) error {
runtimes := operator.GetRuntimes(
operator.WithNvidiaRuntimeName(o.runtimeName),
operator.WithSetAsDefault(o.setAsDefault),
operator.WithRoot(o.runtimeDir),
)
for name := range runtimes {
err := cfg.RemoveRuntime(name)
if err != nil {
return fmt.Errorf("failed to remove runtime %q: %v", name, err)
}
}
if _, exists := config["runtimes"]; exists {
runtimes := config["runtimes"].(map[string]interface{})
for name := range nvidiaRuntimeBinaries {
delete(runtimes, name)
}
if len(runtimes) == 0 {
delete(config, "runtimes")
}
}
return nil
}
// FlushConfig flushes the updated/reverted config out to disk
func FlushConfig(cfg map[string]interface{}, config string) error {
return docker.FlushConfig(cfg, config)
}
// RestartDocker restarts docker depending on the value of restartModeFlag
func RestartDocker(o *options) error {
switch o.restartMode {
@ -385,31 +383,3 @@ func SignalDocker(socket string) error {
return nil
}
// getDefaultRuntime returns the default runtime for the configured options.
// If the configuration is invalid or the default runtimes should not be set
// the empty string is returned.
func (o options) getDefaultRuntime() string {
if o.setAsDefault == false {
return ""
}
return o.runtimeName
}
// getRuntimeBinaries returns a map of runtime names to binary paths. This includes the
// renaming of the `nvidia` runtime as per the --runtime-class command line flag.
func (o options) getRuntimeBinaries() map[string]string {
runtimeBinaries := make(map[string]string)
for rt, bin := range nvidiaRuntimeBinaries {
runtime := rt
if o.runtimeName != "" && o.runtimeName != nvidiaExperimentalRuntimeName && runtime == defaultRuntimeName {
runtime = o.runtimeName
}
runtimeBinaries[runtime] = filepath.Join(o.runtimeDir, bin)
}
return runtimeBinaries
}

54
tools/container/docker/docker_test.go
View File

@ -20,6 +20,7 @@ import (
"encoding/json"
"testing"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config/engine/docker"
"github.com/stretchr/testify/require"
)
@ -60,9 +61,9 @@ func TestUpdateConfigDefaultRuntime(t *testing.T) {
runtimeDir: runtimeDir,
}
config := map[string]interface{}{}
config := docker.Config(map[string]interface{}{})
err := UpdateConfig(config, o)
err := UpdateConfig(&config, o)
require.NoError(t, err, "%d: %v", i, tc)
defaultRuntimeName := config["default-runtime"]
@ -74,7 +75,7 @@ func TestUpdateConfig(t *testing.T) {
const runtimeDir = "/test/runtime/dir"
testCases := []struct {
config map[string]interface{}
config docker.Config
setAsDefault bool
runtimeName string
expectedConfig map[string]interface{}
@ -254,7 +255,8 @@ func TestUpdateConfig(t *testing.T) {
runtimeName: tc.runtimeName,
runtimeDir: runtimeDir,
}
err := UpdateConfig(tc.config, options)
err := UpdateConfig(&tc.config, options)
require.NoError(t, err, "%d: %v", i, tc)
configContent, err := json.MarshalIndent(tc.config, "", " ")
@ -269,7 +271,7 @@ func TestUpdateConfig(t *testing.T) {
func TestRevertConfig(t *testing.T) {
testCases := []struct {
config map[string]interface{}
config docker.Config
expectedConfig map[string]interface{}
}{
{
@ -368,7 +370,7 @@ func TestRevertConfig(t *testing.T) {
}
for i, tc := range testCases {
err := RevertConfig(tc.config)
err := RevertConfig(&tc.config, &options{})
require.NoError(t, err, "%d: %v", i, tc)
@ -381,43 +383,3 @@ func TestRevertConfig(t *testing.T) {
require.EqualValues(t, string(expectedContent), string(configContent), "%d: %v", i, tc)
}
}
func TestFlagsDefaultRuntime(t *testing.T) {
testCases := []struct {
setAsDefault bool
runtimeName string
expected string
}{
{
expected: "",
},
{
runtimeName: "not-bool",
expected: "",
},
{
setAsDefault: false,
runtimeName: "nvidia",
expected: "",
},
{
setAsDefault: true,
runtimeName: "nvidia",
expected: "nvidia",
},
{
setAsDefault: true,
runtimeName: "nvidia-experimental",
expected: "nvidia-experimental",
},
}
for i, tc := range testCases {
f := options{
setAsDefault: tc.setAsDefault,
runtimeName: tc.runtimeName,
}
require.Equal(t, tc.expected, f.getDefaultRuntime(), "%d: %v", i, tc)
}
}