Refactor toolking to setup and cleanup configs

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2023-06-09 17:17:54 +02:00
parent 23041be511
commit 9f5c82420a
7 changed files with 276 additions and 367 deletions

View File

@ -0,0 +1,123 @@
/**
# 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 container
import (
"fmt"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator"
"github.com/sirupsen/logrus"
"github.com/urfave/cli/v2"
)
// Options defines the shared options for the CLIs to configure containers runtimes.
type Options struct {
Config string
Socket string
RuntimeName string
RuntimeDir string
SetAsDefault bool
RestartMode string
HostRootMount string
}
// ParseArgs parses the command line arguments to the CLI
func ParseArgs(c *cli.Context, o *Options) error {
if o.RuntimeDir != "" {
logrus.Debug("Runtime directory already set; ignoring arguments")
return nil
}
args := c.Args()
logrus.Infof("Parsing arguments: %v", args.Slice())
if c.NArg() != 1 {
return fmt.Errorf("incorrect number of arguments")
}
o.RuntimeDir = args.Get(0)
logrus.Infof("Successfully parsed arguments")
return nil
}
// Configure applies the options to the specified config
func (o Options) Configure(cfg engine.Interface) error {
err := o.UpdateConfig(cfg)
if err != nil {
return fmt.Errorf("unable to update config: %v", err)
}
return o.flush(cfg)
}
// Unconfigure removes the options from the specified config
func (o Options) Unconfigure(cfg engine.Interface) error {
err := o.RevertConfig(cfg)
if err != nil {
return fmt.Errorf("unable to update config: %v", err)
}
return o.flush(cfg)
}
// flush flushes the specified config to disk
func (o Options) flush(cfg engine.Interface) error {
logrus.Infof("Flushing 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 {
logrus.Infof("Config file is empty, removed")
}
return nil
}
// UpdateConfig updates the specified config to include the nvidia runtimes
func (o Options) UpdateConfig(cfg engine.Interface) 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", name, err)
}
}
return nil
}
// RevertConfig reverts the specified config to remove the nvidia runtimes
func (o Options) RevertConfig(cfg engine.Interface) 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)
}
}
return nil
}

View File

@ -21,6 +21,7 @@ import (
"testing" "testing"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd" "github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd"
"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
"github.com/pelletier/go-toml" "github.com/pelletier/go-toml"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
@ -31,7 +32,7 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
testCases := []struct { testCases := []struct {
legacyConfig bool legacyConfig bool
setAsDefault bool setAsDefault bool
runtimeClass string runtimeName string
expectedDefaultRuntimeName interface{} expectedDefaultRuntimeName interface{}
expectedDefaultRuntimeBinary interface{} expectedDefaultRuntimeBinary interface{}
}{ }{
@ -51,14 +52,14 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
{ {
legacyConfig: true, legacyConfig: true,
setAsDefault: true, setAsDefault: true,
runtimeClass: "NAME", runtimeName: "NAME",
expectedDefaultRuntimeName: nil, expectedDefaultRuntimeName: nil,
expectedDefaultRuntimeBinary: "/test/runtime/dir/nvidia-container-runtime", expectedDefaultRuntimeBinary: "/test/runtime/dir/nvidia-container-runtime",
}, },
{ {
legacyConfig: true, legacyConfig: true,
setAsDefault: true, setAsDefault: true,
runtimeClass: "nvidia-experimental", runtimeName: "nvidia-experimental",
expectedDefaultRuntimeName: nil, expectedDefaultRuntimeName: nil,
expectedDefaultRuntimeBinary: "/test/runtime/dir/nvidia-container-runtime.experimental", expectedDefaultRuntimeBinary: "/test/runtime/dir/nvidia-container-runtime.experimental",
}, },
@ -77,14 +78,14 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
{ {
legacyConfig: false, legacyConfig: false,
setAsDefault: true, setAsDefault: true,
runtimeClass: "NAME", runtimeName: "NAME",
expectedDefaultRuntimeName: "NAME", expectedDefaultRuntimeName: "NAME",
expectedDefaultRuntimeBinary: nil, expectedDefaultRuntimeBinary: nil,
}, },
{ {
legacyConfig: false, legacyConfig: false,
setAsDefault: true, setAsDefault: true,
runtimeClass: "nvidia-experimental", runtimeName: "nvidia-experimental",
expectedDefaultRuntimeName: "nvidia-experimental", expectedDefaultRuntimeName: "nvidia-experimental",
expectedDefaultRuntimeBinary: nil, expectedDefaultRuntimeBinary: nil,
}, },
@ -93,11 +94,13 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
for i, tc := range testCases { for i, tc := range testCases {
t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
o := &options{ o := &options{
useLegacyConfig: tc.legacyConfig, Options: container.Options{
setAsDefault: tc.setAsDefault, RuntimeName: tc.runtimeName,
runtimeClass: tc.runtimeClass, RuntimeDir: runtimeDir,
SetAsDefault: tc.setAsDefault,
},
runtimeType: runtimeType, runtimeType: runtimeType,
runtimeDir: runtimeDir, useLegacyConfig: tc.legacyConfig,
} }
config, err := toml.TreeFromMap(map[string]interface{}{}) config, err := toml.TreeFromMap(map[string]interface{}{})
@ -109,7 +112,7 @@ func TestUpdateV1ConfigDefaultRuntime(t *testing.T) {
RuntimeType: runtimeType, RuntimeType: runtimeType,
} }
err = UpdateConfig(v1, o) err = o.UpdateConfig(v1)
require.NoError(t, err, "%d: %v", i, tc) require.NoError(t, err, "%d: %v", i, tc)
defaultRuntimeName := v1.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"}) defaultRuntimeName := v1.GetPath([]string{"plugins", "cri", "containerd", "default_runtime_name"})
@ -138,11 +141,11 @@ func TestUpdateV1Config(t *testing.T) {
const runtimeDir = "/test/runtime/dir" const runtimeDir = "/test/runtime/dir"
testCases := []struct { testCases := []struct {
runtimeClass string runtimeName string
expectedConfig map[string]interface{} expectedConfig map[string]interface{}
}{ }{
{ {
runtimeClass: "nvidia", runtimeName: "nvidia",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(1), "version": int64(1),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -200,7 +203,7 @@ func TestUpdateV1Config(t *testing.T) {
}, },
}, },
{ {
runtimeClass: "NAME", runtimeName: "NAME",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(1), "version": int64(1),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -258,7 +261,7 @@ func TestUpdateV1Config(t *testing.T) {
}, },
}, },
{ {
runtimeClass: "nvidia-experimental", runtimeName: "nvidia-experimental",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(1), "version": int64(1),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -320,9 +323,10 @@ func TestUpdateV1Config(t *testing.T) {
for i, tc := range testCases { for i, tc := range testCases {
t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
o := &options{ o := &options{
runtimeClass: tc.runtimeClass, Options: container.Options{
runtimeType: runtimeType, RuntimeName: tc.runtimeName,
runtimeDir: runtimeDir, RuntimeDir: runtimeDir,
},
} }
config, err := toml.TreeFromMap(map[string]interface{}{}) config, err := toml.TreeFromMap(map[string]interface{}{})
@ -335,7 +339,7 @@ func TestUpdateV1Config(t *testing.T) {
ContainerAnnotations: []string{"cdi.k8s.io/*"}, ContainerAnnotations: []string{"cdi.k8s.io/*"},
} }
err = UpdateConfig(v1, o) err = o.UpdateConfig(v1)
require.NoError(t, err) require.NoError(t, err)
expected, err := toml.TreeFromMap(tc.expectedConfig) expected, err := toml.TreeFromMap(tc.expectedConfig)
@ -350,11 +354,11 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
const runtimeDir = "/test/runtime/dir" const runtimeDir = "/test/runtime/dir"
testCases := []struct { testCases := []struct {
runtimeClass string runtimeName string
expectedConfig map[string]interface{} expectedConfig map[string]interface{}
}{ }{
{ {
runtimeClass: "nvidia", runtimeName: "nvidia",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(1), "version": int64(1),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -426,7 +430,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
}, },
}, },
{ {
runtimeClass: "NAME", runtimeName: "NAME",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(1), "version": int64(1),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -498,7 +502,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
}, },
}, },
{ {
runtimeClass: "nvidia-experimental", runtimeName: "nvidia-experimental",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(1), "version": int64(1),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -574,9 +578,10 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
for i, tc := range testCases { for i, tc := range testCases {
t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
o := &options{ o := &options{
runtimeClass: tc.runtimeClass, Options: container.Options{
runtimeType: runtimeType, RuntimeName: tc.runtimeName,
runtimeDir: runtimeDir, RuntimeDir: runtimeDir,
},
} }
config, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary")) config, err := toml.TreeFromMap(runcConfigMapV1("/runc-binary"))
@ -589,7 +594,7 @@ func TestUpdateV1ConfigWithRuncPresent(t *testing.T) {
ContainerAnnotations: []string{"cdi.k8s.io/*"}, ContainerAnnotations: []string{"cdi.k8s.io/*"},
} }
err = UpdateConfig(v1, o) err = o.UpdateConfig(v1)
require.NoError(t, err) require.NoError(t, err)
expected, err := toml.TreeFromMap(tc.expectedConfig) expected, err := toml.TreeFromMap(tc.expectedConfig)
@ -653,7 +658,9 @@ func TestRevertV1Config(t *testing.T) {
for i, tc := range testCases { for i, tc := range testCases {
t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
o := &options{ o := &options{
runtimeClass: "nvidia", Options: container.Options{
RuntimeName: "nvidia",
},
} }
config, err := toml.TreeFromMap(tc.config) config, err := toml.TreeFromMap(tc.config)
@ -668,7 +675,7 @@ func TestRevertV1Config(t *testing.T) {
RuntimeType: runtimeType, RuntimeType: runtimeType,
} }
err = RevertConfig(v1, o) err = o.RevertConfig(v1)
require.NoError(t, err, "%d: %v", i, tc) require.NoError(t, err, "%d: %v", i, tc)
configContents, _ := toml.Marshal(config) configContents, _ := toml.Marshal(config)

View File

@ -21,6 +21,7 @@ import (
"testing" "testing"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd" "github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd"
"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
"github.com/pelletier/go-toml" "github.com/pelletier/go-toml"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
@ -34,38 +35,38 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
testCases := []struct { testCases := []struct {
setAsDefault bool setAsDefault bool
runtimeClass string runtimeName string
expectedDefaultRuntimeName interface{} expectedDefaultRuntimeName interface{}
}{ }{
{}, {},
{ {
setAsDefault: false, setAsDefault: false,
runtimeClass: "nvidia", runtimeName: "nvidia",
expectedDefaultRuntimeName: nil, expectedDefaultRuntimeName: nil,
}, },
{ {
setAsDefault: false, setAsDefault: false,
runtimeClass: "NAME", runtimeName: "NAME",
expectedDefaultRuntimeName: nil, expectedDefaultRuntimeName: nil,
}, },
{ {
setAsDefault: false, setAsDefault: false,
runtimeClass: "nvidia-experimental", runtimeName: "nvidia-experimental",
expectedDefaultRuntimeName: nil, expectedDefaultRuntimeName: nil,
}, },
{ {
setAsDefault: true, setAsDefault: true,
runtimeClass: "nvidia", runtimeName: "nvidia",
expectedDefaultRuntimeName: "nvidia", expectedDefaultRuntimeName: "nvidia",
}, },
{ {
setAsDefault: true, setAsDefault: true,
runtimeClass: "NAME", runtimeName: "NAME",
expectedDefaultRuntimeName: "NAME", expectedDefaultRuntimeName: "NAME",
}, },
{ {
setAsDefault: true, setAsDefault: true,
runtimeClass: "nvidia-experimental", runtimeName: "nvidia-experimental",
expectedDefaultRuntimeName: "nvidia-experimental", expectedDefaultRuntimeName: "nvidia-experimental",
}, },
} }
@ -73,9 +74,11 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
for i, tc := range testCases { for i, tc := range testCases {
t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
o := &options{ o := &options{
setAsDefault: tc.setAsDefault, Options: container.Options{
runtimeClass: tc.runtimeClass, RuntimeName: tc.runtimeName,
runtimeDir: runtimeDir, RuntimeDir: runtimeDir,
SetAsDefault: tc.setAsDefault,
},
} }
config, err := toml.TreeFromMap(map[string]interface{}{}) config, err := toml.TreeFromMap(map[string]interface{}{})
@ -86,7 +89,7 @@ func TestUpdateV2ConfigDefaultRuntime(t *testing.T) {
RuntimeType: runtimeType, RuntimeType: runtimeType,
} }
err = UpdateConfig(v2, o) err = o.UpdateConfig(v2)
require.NoError(t, err) require.NoError(t, err)
defaultRuntimeName := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"}) defaultRuntimeName := config.GetPath([]string{"plugins", "io.containerd.grpc.v1.cri", "containerd", "default_runtime_name"})
@ -100,11 +103,11 @@ func TestUpdateV2Config(t *testing.T) {
const expectedVersion = int64(2) const expectedVersion = int64(2)
testCases := []struct { testCases := []struct {
runtimeClass string runtimeName string
expectedConfig map[string]interface{} expectedConfig map[string]interface{}
}{ }{
{ {
runtimeClass: "nvidia", runtimeName: "nvidia",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(2), "version": int64(2),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -158,7 +161,7 @@ func TestUpdateV2Config(t *testing.T) {
}, },
}, },
{ {
runtimeClass: "NAME", runtimeName: "NAME",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(2), "version": int64(2),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -212,7 +215,7 @@ func TestUpdateV2Config(t *testing.T) {
}, },
}, },
{ {
runtimeClass: "nvidia-experimental", runtimeName: "nvidia-experimental",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(2), "version": int64(2),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -270,9 +273,11 @@ func TestUpdateV2Config(t *testing.T) {
for i, tc := range testCases { for i, tc := range testCases {
t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
o := &options{ o := &options{
runtimeClass: tc.runtimeClass, Options: container.Options{
RuntimeName: tc.runtimeName,
RuntimeDir: runtimeDir,
},
runtimeType: runtimeType, runtimeType: runtimeType,
runtimeDir: runtimeDir,
} }
config, err := toml.TreeFromMap(map[string]interface{}{}) config, err := toml.TreeFromMap(map[string]interface{}{})
@ -284,7 +289,7 @@ func TestUpdateV2Config(t *testing.T) {
ContainerAnnotations: []string{"cdi.k8s.io/*"}, ContainerAnnotations: []string{"cdi.k8s.io/*"},
} }
err = UpdateConfig(v2, o) err = o.UpdateConfig(v2)
require.NoError(t, err) require.NoError(t, err)
expected, err := toml.TreeFromMap(tc.expectedConfig) expected, err := toml.TreeFromMap(tc.expectedConfig)
@ -300,11 +305,11 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
const runtimeDir = "/test/runtime/dir" const runtimeDir = "/test/runtime/dir"
testCases := []struct { testCases := []struct {
runtimeClass string runtimeName string
expectedConfig map[string]interface{} expectedConfig map[string]interface{}
}{ }{
{ {
runtimeClass: "nvidia", runtimeName: "nvidia",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(2), "version": int64(2),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -372,7 +377,7 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
}, },
}, },
{ {
runtimeClass: "NAME", runtimeName: "NAME",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(2), "version": int64(2),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -440,7 +445,7 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
}, },
}, },
{ {
runtimeClass: "nvidia-experimental", runtimeName: "nvidia-experimental",
expectedConfig: map[string]interface{}{ expectedConfig: map[string]interface{}{
"version": int64(2), "version": int64(2),
"plugins": map[string]interface{}{ "plugins": map[string]interface{}{
@ -512,9 +517,10 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
for i, tc := range testCases { for i, tc := range testCases {
t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
o := &options{ o := &options{
runtimeClass: tc.runtimeClass, Options: container.Options{
runtimeType: runtimeType, RuntimeName: tc.runtimeName,
runtimeDir: runtimeDir, RuntimeDir: runtimeDir,
},
} }
config, err := toml.TreeFromMap(runcConfigMapV2("/runc-binary")) config, err := toml.TreeFromMap(runcConfigMapV2("/runc-binary"))
@ -526,7 +532,7 @@ func TestUpdateV2ConfigWithRuncPresent(t *testing.T) {
ContainerAnnotations: []string{"cdi.k8s.io/*"}, ContainerAnnotations: []string{"cdi.k8s.io/*"},
} }
err = UpdateConfig(v2, o) err = o.UpdateConfig(v2)
require.NoError(t, err) require.NoError(t, err)
expected, err := toml.TreeFromMap(tc.expectedConfig) expected, err := toml.TreeFromMap(tc.expectedConfig)
@ -585,7 +591,9 @@ func TestRevertV2Config(t *testing.T) {
for i, tc := range testCases { for i, tc := range testCases {
t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
o := &options{ o := &options{
runtimeClass: "nvidia", Options: container.Options{
RuntimeName: "nvidia",
},
} }
config, err := toml.TreeFromMap(tc.config) config, err := toml.TreeFromMap(tc.config)
@ -599,7 +607,7 @@ func TestRevertV2Config(t *testing.T) {
RuntimeType: runtimeType, RuntimeType: runtimeType,
} }
err = RevertConfig(v2, o) err = o.RevertConfig(v2)
require.NoError(t, err) require.NoError(t, err)
configContents, _ := toml.Marshal(config) configContents, _ := toml.Marshal(config)

View File

@ -24,9 +24,8 @@ import (
"syscall" "syscall"
"time" "time"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd" "github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd"
"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator" "github.com/NVIDIA/nvidia-container-toolkit/tools/container"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
cli "github.com/urfave/cli/v2" cli "github.com/urfave/cli/v2"
) )
@ -63,13 +62,7 @@ var nvidiaRuntimeBinaries = map[string]string{
// options stores the configuration from the command line or environment variables // options stores the configuration from the command line or environment variables
type options struct { type options struct {
config string container.Options
socket string
runtimeClass string
runtimeDir string
setAsDefault bool
restartMode string
hostRootMount string
// containerd-specific options // containerd-specific options
useLegacyConfig bool useLegacyConfig bool
@ -96,7 +89,7 @@ func main() {
return Setup(c, &options) return Setup(c, &options)
} }
setup.Before = func(c *cli.Context) error { setup.Before = func(c *cli.Context) error {
return ParseArgs(c, &options) return container.ParseArgs(c, &options.Options)
} }
// Create the 'cleanup' subcommand // Create the 'cleanup' subcommand
@ -108,7 +101,7 @@ func main() {
return Cleanup(c, &options) return Cleanup(c, &options)
} }
cleanup.Before = func(c *cli.Context) error { cleanup.Before = func(c *cli.Context) error {
return ParseArgs(c, &options) return container.ParseArgs(c, &options.Options)
} }
// Register the subcommands with the top-level CLI // Register the subcommands with the top-level CLI
@ -126,21 +119,21 @@ func main() {
Name: "config", Name: "config",
Usage: "Path to the containerd config file", Usage: "Path to the containerd config file",
Value: defaultConfig, Value: defaultConfig,
Destination: &options.config, Destination: &options.Config,
EnvVars: []string{"CONTAINERD_CONFIG", "RUNTIME_CONFIG"}, EnvVars: []string{"CONTAINERD_CONFIG", "RUNTIME_CONFIG"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "socket", Name: "socket",
Usage: "Path to the containerd socket file", Usage: "Path to the containerd socket file",
Value: defaultSocket, Value: defaultSocket,
Destination: &options.socket, Destination: &options.Socket,
EnvVars: []string{"CONTAINERD_SOCKET", "RUNTIME_SOCKET"}, EnvVars: []string{"CONTAINERD_SOCKET", "RUNTIME_SOCKET"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "restart-mode", Name: "restart-mode",
Usage: "Specify how containerd should be restarted; If 'none' is selected, it will not be restarted [signal | systemd | none]", Usage: "Specify how containerd should be restarted; If 'none' is selected, it will not be restarted [signal | systemd | none]",
Value: defaultRestartMode, Value: defaultRestartMode,
Destination: &options.restartMode, Destination: &options.RestartMode,
EnvVars: []string{"CONTAINERD_RESTART_MODE", "RUNTIME_RESTART_MODE"}, EnvVars: []string{"CONTAINERD_RESTART_MODE", "RUNTIME_RESTART_MODE"},
}, },
&cli.StringFlag{ &cli.StringFlag{
@ -148,21 +141,21 @@ func main() {
Aliases: []string{"runtime-name", "nvidia-runtime-name"}, Aliases: []string{"runtime-name", "nvidia-runtime-name"},
Usage: "The name of the runtime class to set for the nvidia-container-runtime", Usage: "The name of the runtime class to set for the nvidia-container-runtime",
Value: defaultRuntimeClass, Value: defaultRuntimeClass,
Destination: &options.runtimeClass, Destination: &options.RuntimeName,
EnvVars: []string{"CONTAINERD_RUNTIME_CLASS", "NVIDIA_RUNTIME_NAME"}, EnvVars: []string{"CONTAINERD_RUNTIME_CLASS", "NVIDIA_RUNTIME_NAME"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "nvidia-runtime-dir", Name: "nvidia-runtime-dir",
Aliases: []string{"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", 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, Destination: &options.RuntimeDir,
EnvVars: []string{"NVIDIA_RUNTIME_DIR"}, EnvVars: []string{"NVIDIA_RUNTIME_DIR"},
}, },
&cli.BoolFlag{ &cli.BoolFlag{
Name: "set-as-default", Name: "set-as-default",
Usage: "Set nvidia-container-runtime as the default runtime", Usage: "Set nvidia-container-runtime as the default runtime",
Value: defaultSetAsDefault, Value: defaultSetAsDefault,
Destination: &options.setAsDefault, Destination: &options.SetAsDefault,
EnvVars: []string{"CONTAINERD_SET_AS_DEFAULT", "NVIDIA_RUNTIME_SET_AS_DEFAULT"}, EnvVars: []string{"CONTAINERD_SET_AS_DEFAULT", "NVIDIA_RUNTIME_SET_AS_DEFAULT"},
Hidden: true, Hidden: true,
}, },
@ -170,7 +163,7 @@ func main() {
Name: "host-root", Name: "host-root",
Usage: "Specify the path to the host root to be used when restarting containerd using systemd", Usage: "Specify the path to the host root to be used when restarting containerd using systemd",
Value: defaultHostRootMount, Value: defaultHostRootMount,
Destination: &options.hostRootMount, Destination: &options.HostRootMount,
EnvVars: []string{"HOST_ROOT_MOUNT"}, EnvVars: []string{"HOST_ROOT_MOUNT"},
}, },
&cli.BoolFlag{ &cli.BoolFlag{
@ -208,7 +201,7 @@ func Setup(c *cli.Context, o *options) error {
log.Infof("Starting 'setup' for %v", c.App.Name) log.Infof("Starting 'setup' for %v", c.App.Name)
cfg, err := containerd.New( cfg, err := containerd.New(
containerd.WithPath(o.config), containerd.WithPath(o.Config),
containerd.WithRuntimeType(o.runtimeType), containerd.WithRuntimeType(o.runtimeType),
containerd.WithUseLegacyConfig(o.useLegacyConfig), containerd.WithUseLegacyConfig(o.useLegacyConfig),
containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...), containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...),
@ -217,18 +210,9 @@ func Setup(c *cli.Context, o *options) error {
return fmt.Errorf("unable to load config: %v", err) return fmt.Errorf("unable to load config: %v", err)
} }
err = UpdateConfig(cfg, o) err = o.Configure(cfg)
if err != nil { if err != nil {
return fmt.Errorf("unable to update config: %v", err) return fmt.Errorf("unable to configure containerd: %v", err)
}
log.Infof("Flushing containerd 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 = RestartContainerd(o) err = RestartContainerd(o)
@ -246,7 +230,7 @@ func Cleanup(c *cli.Context, o *options) error {
log.Infof("Starting 'cleanup' for %v", c.App.Name) log.Infof("Starting 'cleanup' for %v", c.App.Name)
cfg, err := containerd.New( cfg, err := containerd.New(
containerd.WithPath(o.config), containerd.WithPath(o.Config),
containerd.WithRuntimeType(o.runtimeType), containerd.WithRuntimeType(o.runtimeType),
containerd.WithUseLegacyConfig(o.useLegacyConfig), containerd.WithUseLegacyConfig(o.useLegacyConfig),
containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...), containerd.WithContainerAnnotations(o.containerAnnotationsFromCDIPrefixes()...),
@ -255,18 +239,9 @@ func Cleanup(c *cli.Context, o *options) error {
return fmt.Errorf("unable to load config: %v", err) return fmt.Errorf("unable to load config: %v", err)
} }
err = RevertConfig(cfg, o) err = o.Unconfigure(cfg)
if err != nil { if err != nil {
return fmt.Errorf("unable to update config: %v", err) return fmt.Errorf("unable to unconfigure containerd: %v", err)
}
log.Infof("Flushing containerd 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 = RestartContainerd(o) err = RestartContainerd(o)
@ -279,65 +254,11 @@ func Cleanup(c *cli.Context, o *options) error {
return nil return nil
} }
// 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
}
// UpdateConfig updates the containerd config to include the nvidia-container-runtime
func UpdateConfig(cfg engine.Interface, o *options) error {
runtimes := operator.GetRuntimes(
operator.WithNvidiaRuntimeName(o.runtimeClass),
operator.WithSetAsDefault(o.setAsDefault),
operator.WithRoot(o.runtimeDir),
)
for class, runtime := range runtimes {
err := cfg.AddRuntime(class, runtime.Path, runtime.SetAsDefault)
if err != nil {
return fmt.Errorf("unable to update config for runtime class '%v': %v", class, err)
}
}
return nil
}
// RevertConfig reverts the containerd config to remove the nvidia-container-runtime
func RevertConfig(cfg engine.Interface, o *options) error {
runtimes := operator.GetRuntimes(
operator.WithNvidiaRuntimeName(o.runtimeClass),
operator.WithSetAsDefault(o.setAsDefault),
operator.WithRoot(o.runtimeDir),
)
for class := range runtimes {
err := cfg.RemoveRuntime(class)
if err != nil {
return fmt.Errorf("unable to revert config for runtime class '%v': %v", class, err)
}
}
return nil
}
// RestartContainerd restarts containerd depending on the value of restartModeFlag // RestartContainerd restarts containerd depending on the value of restartModeFlag
func RestartContainerd(o *options) error { func RestartContainerd(o *options) error {
switch o.restartMode { switch o.RestartMode {
case restartModeNone: case restartModeNone:
log.Warnf("Skipping sending signal to containerd due to --restart-mode=%v", o.restartMode) log.Warnf("Skipping sending signal to containerd due to --restart-mode=%v", o.RestartMode)
return nil return nil
case restartModeSignal: case restartModeSignal:
err := SignalContainerd(o) err := SignalContainerd(o)
@ -345,9 +266,9 @@ func RestartContainerd(o *options) error {
return fmt.Errorf("unable to signal containerd: %v", err) return fmt.Errorf("unable to signal containerd: %v", err)
} }
case restartModeSystemd: case restartModeSystemd:
return RestartContainerdSystemd(o.hostRootMount) return RestartContainerdSystemd(o.HostRootMount)
default: default:
return fmt.Errorf("Invalid restart mode specified: %v", o.restartMode) return fmt.Errorf("Invalid restart mode specified: %v", o.RestartMode)
} }
return nil return nil
@ -359,7 +280,7 @@ func SignalContainerd(o *options) error {
// Wrap the logic to perform the SIGHUP in a function so we can retry it on failure // Wrap the logic to perform the SIGHUP in a function so we can retry it on failure
retriable := func() error { retriable := func() error {
conn, err := net.Dial("unix", o.socket) conn, err := net.Dial("unix", o.Socket)
if err != nil { if err != nil {
return fmt.Errorf("unable to dial: %v", err) return fmt.Errorf("unable to dial: %v", err)
} }

View File

@ -24,9 +24,8 @@ import (
"path/filepath" "path/filepath"
"github.com/NVIDIA/nvidia-container-toolkit/internal/config" "github.com/NVIDIA/nvidia-container-toolkit/internal/config"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/crio" "github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/crio"
"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator" "github.com/NVIDIA/nvidia-container-toolkit/tools/container"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
cli "github.com/urfave/cli/v2" cli "github.com/urfave/cli/v2"
) )
@ -52,13 +51,7 @@ const (
// options stores the configuration from the command linek or environment variables // options stores the configuration from the command linek or environment variables
type options struct { type options struct {
config string container.Options
socket string
runtimeClass string
runtimeDir string
setAsDefault bool
restartMode string
hostRootMount string
configMode string configMode string
@ -86,7 +79,7 @@ func main() {
return Setup(c, &options) return Setup(c, &options)
} }
setup.Before = func(c *cli.Context) error { setup.Before = func(c *cli.Context) error {
return ParseArgs(c, &options) return container.ParseArgs(c, &options.Options)
} }
// Create the 'cleanup' subcommand // Create the 'cleanup' subcommand
@ -97,7 +90,7 @@ func main() {
return Cleanup(c, &options) return Cleanup(c, &options)
} }
cleanup.Before = func(c *cli.Context) error { cleanup.Before = func(c *cli.Context) error {
return ParseArgs(c, &options) return container.ParseArgs(c, &options.Options)
} }
// Register the subcommands with the top-level CLI // Register the subcommands with the top-level CLI
@ -115,14 +108,14 @@ func main() {
Name: "config", Name: "config",
Usage: "Path to the cri-o config file", Usage: "Path to the cri-o config file",
Value: defaultConfig, Value: defaultConfig,
Destination: &options.config, Destination: &options.Config,
EnvVars: []string{"CRIO_CONFIG", "RUNTIME_CONFIG"}, EnvVars: []string{"CRIO_CONFIG", "RUNTIME_CONFIG"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "socket", Name: "socket",
Usage: "Path to the crio socket file", Usage: "Path to the crio socket file",
Value: "", Value: "",
Destination: &options.socket, Destination: &options.Socket,
EnvVars: []string{"CRIO_SOCKET", "RUNTIME_SOCKET"}, EnvVars: []string{"CRIO_SOCKET", "RUNTIME_SOCKET"},
Hidden: true, Hidden: true,
}, },
@ -130,7 +123,7 @@ func main() {
Name: "restart-mode", Name: "restart-mode",
Usage: "Specify how cri-o should be restarted; If 'none' is selected, it will not be restarted [systemd | none]", Usage: "Specify how cri-o should be restarted; If 'none' is selected, it will not be restarted [systemd | none]",
Value: defaultRestartMode, Value: defaultRestartMode,
Destination: &options.restartMode, Destination: &options.RestartMode,
EnvVars: []string{"CRIO_RESTART_MODE", "RUNTIME_RESTART_MODE"}, EnvVars: []string{"CRIO_RESTART_MODE", "RUNTIME_RESTART_MODE"},
}, },
&cli.StringFlag{ &cli.StringFlag{
@ -138,21 +131,21 @@ func main() {
Aliases: []string{"runtime-name", "nvidia-runtime-name"}, Aliases: []string{"runtime-name", "nvidia-runtime-name"},
Usage: "The name of the runtime class to set for the nvidia-container-runtime", Usage: "The name of the runtime class to set for the nvidia-container-runtime",
Value: defaultRuntimeClass, Value: defaultRuntimeClass,
Destination: &options.runtimeClass, Destination: &options.RuntimeName,
EnvVars: []string{"CRIO_RUNTIME_CLASS", "NVIDIA_RUNTIME_NAME"}, EnvVars: []string{"CRIO_RUNTIME_CLASS", "NVIDIA_RUNTIME_NAME"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "nvidia-runtime-dir", Name: "nvidia-runtime-dir",
Aliases: []string{"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", 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, Destination: &options.RuntimeDir,
EnvVars: []string{"NVIDIA_RUNTIME_DIR"}, EnvVars: []string{"NVIDIA_RUNTIME_DIR"},
}, },
&cli.BoolFlag{ &cli.BoolFlag{
Name: "set-as-default", Name: "set-as-default",
Usage: "Set nvidia-container-runtime as the default runtime", Usage: "Set nvidia-container-runtime as the default runtime",
Value: defaultSetAsDefault, Value: defaultSetAsDefault,
Destination: &options.setAsDefault, Destination: &options.SetAsDefault,
EnvVars: []string{"CRIO_SET_AS_DEFAULT", "NVIDIA_RUNTIME_SET_AS_DEFAULT"}, EnvVars: []string{"CRIO_SET_AS_DEFAULT", "NVIDIA_RUNTIME_SET_AS_DEFAULT"},
Hidden: true, Hidden: true,
}, },
@ -160,7 +153,7 @@ func main() {
Name: "host-root", Name: "host-root",
Usage: "Specify the path to the host root to be used when restarting crio using systemd", Usage: "Specify the path to the host root to be used when restarting crio using systemd",
Value: defaultHostRootMount, Value: defaultHostRootMount,
Destination: &options.hostRootMount, Destination: &options.HostRootMount,
EnvVars: []string{"HOST_ROOT_MOUNT"}, EnvVars: []string{"HOST_ROOT_MOUNT"},
}, },
&cli.StringFlag{ &cli.StringFlag{
@ -223,7 +216,7 @@ func setupHook(o *options) error {
} }
hookPath := getHookPath(o.hooksDir, o.hookFilename) hookPath := getHookPath(o.hooksDir, o.hookFilename)
err = createHook(o.runtimeDir, hookPath) err = createHook(o.RuntimeDir, hookPath)
if err != nil { if err != nil {
return fmt.Errorf("error creating hook: %v", err) return fmt.Errorf("error creating hook: %v", err)
} }
@ -236,24 +229,15 @@ func setupConfig(o *options) error {
log.Infof("Updating config file") log.Infof("Updating config file")
cfg, err := crio.New( cfg, err := crio.New(
crio.WithPath(o.config), crio.WithPath(o.Config),
) )
if err != nil { if err != nil {
return fmt.Errorf("unable to load config: %v", err) return fmt.Errorf("unable to load config: %v", err)
} }
err = UpdateConfig(cfg, o) err = o.Configure(cfg)
if err != nil { if err != nil {
return fmt.Errorf("unable to update config: %v", err) return fmt.Errorf("unable to configure cri-o: %v", err)
}
log.Infof("Flushing cri-o 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 = RestartCrio(o) err = RestartCrio(o)
@ -296,24 +280,15 @@ func cleanupConfig(o *options) error {
log.Infof("Reverting config file modifications") log.Infof("Reverting config file modifications")
cfg, err := crio.New( cfg, err := crio.New(
crio.WithPath(o.config), crio.WithPath(o.Config),
) )
if err != nil { if err != nil {
return fmt.Errorf("unable to load config: %v", err) return fmt.Errorf("unable to load config: %v", err)
} }
err = RevertConfig(cfg, o) err = o.Unconfigure(cfg)
if err != nil { if err != nil {
return fmt.Errorf("unable to update config: %v", err) return fmt.Errorf("unable to unconfigure cri-o: %v", err)
}
log.Infof("Flushing cri-o 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 = RestartCrio(o) err = RestartCrio(o)
@ -324,27 +299,6 @@ func cleanupConfig(o *options) error {
return nil return nil
} }
// 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
}
func createHook(toolkitDir string, hookPath string) error { func createHook(toolkitDir string, hookPath string) error {
hook, err := os.Create(hookPath) hook, err := os.Create(hookPath)
if err != nil { if err != nil {
@ -385,49 +339,16 @@ func generateOciHook(toolkitDir string) podmanHook {
return hook return hook
} }
// UpdateConfig updates the cri-o config to include the NVIDIA Container Runtime
func UpdateConfig(cfg engine.Interface, o *options) error {
runtimes := operator.GetRuntimes(
operator.WithNvidiaRuntimeName(o.runtimeClass),
operator.WithSetAsDefault(o.setAsDefault),
operator.WithRoot(o.runtimeDir),
)
for class, runtime := range runtimes {
err := cfg.AddRuntime(class, runtime.Path, runtime.SetAsDefault)
if err != nil {
return fmt.Errorf("unable to update config for runtime class '%v': %v", class, err)
}
}
return nil
}
// RevertConfig reverts the cri-o config to remove the NVIDIA Container Runtime
func RevertConfig(cfg engine.Interface, o *options) error {
runtimes := operator.GetRuntimes(
operator.WithNvidiaRuntimeName(o.runtimeClass),
operator.WithSetAsDefault(o.setAsDefault),
operator.WithRoot(o.runtimeDir),
)
for class := range runtimes {
err := cfg.RemoveRuntime(class)
if err != nil {
return fmt.Errorf("unable to revert config for runtime class '%v': %v", class, err)
}
}
return nil
}
// RestartCrio restarts crio depending on the value of restartModeFlag // RestartCrio restarts crio depending on the value of restartModeFlag
func RestartCrio(o *options) error { func RestartCrio(o *options) error {
switch o.restartMode { switch o.RestartMode {
case restartModeNone: case restartModeNone:
log.Warnf("Skipping restart of crio due to --restart-mode=%v", o.restartMode) log.Warnf("Skipping restart of crio due to --restart-mode=%v", o.RestartMode)
return nil return nil
case restartModeSystemd: case restartModeSystemd:
return RestartCrioSystemd(o.hostRootMount) return RestartCrioSystemd(o.HostRootMount)
default: default:
return fmt.Errorf("invalid restart mode specified: %v", o.restartMode) return fmt.Errorf("invalid restart mode specified: %v", o.RestartMode)
} }
} }

View File

@ -23,9 +23,8 @@ import (
"syscall" "syscall"
"time" "time"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/docker" "github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/docker"
"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator" "github.com/NVIDIA/nvidia-container-toolkit/tools/container"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
cli "github.com/urfave/cli/v2" cli "github.com/urfave/cli/v2"
) )
@ -56,13 +55,7 @@ const (
// options stores the configuration from the command line or environment variables // options stores the configuration from the command line or environment variables
type options struct { type options struct {
config string container.Options
socket string
runtimeName string
runtimeDir string
setAsDefault bool
restartMode string
hostRootMount string
} }
func main() { func main() {
@ -83,7 +76,7 @@ func main() {
return Setup(c, &options) return Setup(c, &options)
} }
setup.Before = func(c *cli.Context) error { setup.Before = func(c *cli.Context) error {
return ParseArgs(c, &options) return container.ParseArgs(c, &options.Options)
} }
// Create the 'cleanup' subcommand // Create the 'cleanup' subcommand
@ -95,7 +88,7 @@ func main() {
return Cleanup(c, &options) return Cleanup(c, &options)
} }
cleanup.Before = func(c *cli.Context) error { cleanup.Before = func(c *cli.Context) error {
return ParseArgs(c, &options) return container.ParseArgs(c, &options.Options)
} }
// Register the subcommands with the top-level CLI // Register the subcommands with the top-level CLI
@ -113,28 +106,28 @@ func main() {
Name: "config", Name: "config",
Usage: "Path to docker config file", Usage: "Path to docker config file",
Value: defaultConfig, Value: defaultConfig,
Destination: &options.config, Destination: &options.Config,
EnvVars: []string{"DOCKER_CONFIG", "RUNTIME_CONFIG"}, EnvVars: []string{"DOCKER_CONFIG", "RUNTIME_CONFIG"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "socket", Name: "socket",
Usage: "Path to the docker socket file", Usage: "Path to the docker socket file",
Value: defaultSocket, Value: defaultSocket,
Destination: &options.socket, Destination: &options.Socket,
EnvVars: []string{"DOCKER_SOCKET", "RUNTIME_SOCKET"}, EnvVars: []string{"DOCKER_SOCKET", "RUNTIME_SOCKET"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "restart-mode", Name: "restart-mode",
Usage: "Specify how docker should be restarted; If 'none' is selected it will not be restarted [signal | none]", Usage: "Specify how docker should be restarted; If 'none' is selected it will not be restarted [signal | none]",
Value: defaultRestartMode, Value: defaultRestartMode,
Destination: &options.restartMode, Destination: &options.RestartMode,
EnvVars: []string{"DOCKER_RESTART_MODE", "RUNTIME_RESTART_MODE"}, EnvVars: []string{"DOCKER_RESTART_MODE", "RUNTIME_RESTART_MODE"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "host-root", Name: "host-root",
Usage: "Specify the path to the host root to be used when restarting docker using systemd", Usage: "Specify the path to the host root to be used when restarting docker using systemd",
Value: defaultHostRootMount, Value: defaultHostRootMount,
Destination: &options.hostRootMount, Destination: &options.HostRootMount,
EnvVars: []string{"HOST_ROOT_MOUNT"}, EnvVars: []string{"HOST_ROOT_MOUNT"},
// Restart using systemd is currently not supported. // Restart using systemd is currently not supported.
// We hide this option for the time being. // We hide this option for the time being.
@ -145,21 +138,21 @@ func main() {
Aliases: []string{"runtime-class", "nvidia-runtime-name"}, Aliases: []string{"runtime-class", "nvidia-runtime-name"},
Usage: "Specify the name of the `nvidia` runtime. If set-as-default is selected, the runtime is used as the default runtime.", Usage: "Specify the name of the `nvidia` runtime. If set-as-default is selected, the runtime is used as the default runtime.",
Value: defaultRuntimeName, Value: defaultRuntimeName,
Destination: &options.runtimeName, Destination: &options.RuntimeName,
EnvVars: []string{"DOCKER_RUNTIME_NAME", "NVIDIA_RUNTIME_NAME"}, EnvVars: []string{"DOCKER_RUNTIME_NAME", "NVIDIA_RUNTIME_NAME"},
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "nvidia-runtime-dir", Name: "nvidia-runtime-dir",
Aliases: []string{"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", 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, Destination: &options.RuntimeDir,
EnvVars: []string{"NVIDIA_RUNTIME_DIR"}, EnvVars: []string{"NVIDIA_RUNTIME_DIR"},
}, },
&cli.BoolFlag{ &cli.BoolFlag{
Name: "set-as-default", 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", 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",
Value: defaultSetAsDefault, Value: defaultSetAsDefault,
Destination: &options.setAsDefault, Destination: &options.SetAsDefault,
EnvVars: []string{"DOCKER_SET_AS_DEFAULT", "NVIDIA_RUNTIME_SET_AS_DEFAULT"}, EnvVars: []string{"DOCKER_SET_AS_DEFAULT", "NVIDIA_RUNTIME_SET_AS_DEFAULT"},
Hidden: true, Hidden: true,
}, },
@ -181,21 +174,15 @@ func Setup(c *cli.Context, o *options) error {
log.Infof("Starting 'setup' for %v", c.App.Name) log.Infof("Starting 'setup' for %v", c.App.Name)
cfg, err := docker.New( cfg, err := docker.New(
docker.WithPath(o.config), docker.WithPath(o.Config),
) )
if err != nil { if err != nil {
return fmt.Errorf("unable to load config: %v", err) return fmt.Errorf("unable to load config: %v", err)
} }
err = UpdateConfig(cfg, o) err = o.Configure(cfg)
if err != nil { if err != nil {
return fmt.Errorf("unable to update config: %v", err) return fmt.Errorf("unable to configure docker: %v", err)
}
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)
} }
err = RestartDocker(o) err = RestartDocker(o)
@ -213,24 +200,15 @@ func Cleanup(c *cli.Context, o *options) error {
log.Infof("Starting 'cleanup' for %v", c.App.Name) log.Infof("Starting 'cleanup' for %v", c.App.Name)
cfg, err := docker.New( cfg, err := docker.New(
docker.WithPath(o.config), docker.WithPath(o.Config),
) )
if err != nil { if err != nil {
return fmt.Errorf("unable to load config: %v", err) return fmt.Errorf("unable to load config: %v", err)
} }
err = RevertConfig(cfg, o) err = o.Unconfigure(cfg)
if err != nil { if err != nil {
return fmt.Errorf("unable to update config: %v", err) return fmt.Errorf("unable to unconfigure docker: %v", err)
}
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) err = RestartDocker(o)
@ -243,73 +221,18 @@ func Cleanup(c *cli.Context, o *options) error {
return nil return nil
} }
// 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
}
// UpdateConfig updates the docker config to include the nvidia runtimes
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", name, err)
}
}
return nil
}
// RevertConfig reverts the docker config to remove the nvidia runtime
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)
}
}
return nil
}
// RestartDocker restarts docker depending on the value of restartModeFlag // RestartDocker restarts docker depending on the value of restartModeFlag
func RestartDocker(o *options) error { func RestartDocker(o *options) error {
switch o.restartMode { switch o.RestartMode {
case restartModeNone: case restartModeNone:
log.Warnf("Skipping sending signal to docker due to --restart-mode=%v", o.restartMode) log.Warnf("Skipping sending signal to docker due to --restart-mode=%v", o.RestartMode)
case restartModeSignal: case restartModeSignal:
err := SignalDocker(o.socket) err := SignalDocker(o.Socket)
if err != nil { if err != nil {
return fmt.Errorf("unable to signal docker: %v", err) return fmt.Errorf("unable to signal docker: %v", err)
} }
default: default:
return fmt.Errorf("invalid restart mode specified: %v", o.restartMode) return fmt.Errorf("invalid restart mode specified: %v", o.RestartMode)
} }
return nil return nil

View File

@ -21,6 +21,7 @@ import (
"testing" "testing"
"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/docker" "github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/docker"
"github.com/NVIDIA/nvidia-container-toolkit/tools/container"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
@ -56,14 +57,16 @@ func TestUpdateConfigDefaultRuntime(t *testing.T) {
for i, tc := range testCases { for i, tc := range testCases {
o := &options{ o := &options{
setAsDefault: tc.setAsDefault, Options: container.Options{
runtimeName: tc.runtimeName, RuntimeName: tc.runtimeName,
runtimeDir: runtimeDir, RuntimeDir: runtimeDir,
SetAsDefault: tc.setAsDefault,
},
} }
config := docker.Config(map[string]interface{}{}) config := docker.Config(map[string]interface{}{})
err := UpdateConfig(&config, o) err := o.UpdateConfig(&config)
require.NoError(t, err, "%d: %v", i, tc) require.NoError(t, err, "%d: %v", i, tc)
defaultRuntimeName := config["default-runtime"] defaultRuntimeName := config["default-runtime"]
@ -314,13 +317,15 @@ func TestUpdateConfig(t *testing.T) {
} }
for i, tc := range testCases { for i, tc := range testCases {
options := &options{ o := &options{
setAsDefault: tc.setAsDefault, Options: container.Options{
runtimeName: tc.runtimeName, RuntimeName: tc.runtimeName,
runtimeDir: runtimeDir, RuntimeDir: runtimeDir,
SetAsDefault: tc.setAsDefault,
},
} }
err := UpdateConfig(&tc.config, options) err := o.UpdateConfig(&tc.config)
require.NoError(t, err, "%d: %v", i, tc) require.NoError(t, err, "%d: %v", i, tc)
configContent, err := json.MarshalIndent(tc.config, "", " ") configContent, err := json.MarshalIndent(tc.config, "", " ")
@ -457,7 +462,8 @@ func TestRevertConfig(t *testing.T) {
} }
for i, tc := range testCases { for i, tc := range testCases {
err := RevertConfig(&tc.config, &options{}) o := &options{}
err := o.RevertConfig(&tc.config)
require.NoError(t, err, "%d: %v", i, tc) require.NoError(t, err, "%d: %v", i, tc)