diff --git a/tools/container/containerd/containerd_test.go b/tools/container/containerd/containerd_test.go deleted file mode 100644 index 748219b3..00000000 --- a/tools/container/containerd/containerd_test.go +++ /dev/null @@ -1,106 +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 ( - "testing" - - "github.com/stretchr/testify/require" -) - -func TestOptions(t *testing.T) { - testCases := []struct { - options options - expectedDefaultRuntime string - expectedRuntimeBinaries map[string]string - }{ - { - expectedRuntimeBinaries: map[string]string{ - "nvidia": "nvidia-container-runtime", - "nvidia-experimental": "nvidia-container-runtime.experimental", - }, - }, - { - options: options{ - setAsDefault: true, - }, - expectedDefaultRuntime: "nvidia", - expectedRuntimeBinaries: map[string]string{ - "nvidia": "nvidia-container-runtime", - "nvidia-experimental": "nvidia-container-runtime.experimental", - }, - }, - { - options: options{ - setAsDefault: true, - runtimeClass: "nvidia", - }, - expectedDefaultRuntime: "nvidia", - expectedRuntimeBinaries: map[string]string{ - "nvidia": "nvidia-container-runtime", - "nvidia-experimental": "nvidia-container-runtime.experimental", - }, - }, - { - options: options{ - setAsDefault: true, - runtimeClass: "NAME", - }, - expectedDefaultRuntime: "NAME", - expectedRuntimeBinaries: map[string]string{ - "NAME": "nvidia-container-runtime", - "nvidia-experimental": "nvidia-container-runtime.experimental", - }, - }, - { - options: options{ - setAsDefault: false, - runtimeClass: "NAME", - }, - expectedRuntimeBinaries: map[string]string{ - "NAME": "nvidia-container-runtime", - "nvidia-experimental": "nvidia-container-runtime.experimental", - }, - }, - { - options: options{ - setAsDefault: true, - runtimeClass: "nvidia-experimental", - }, - expectedDefaultRuntime: "nvidia-experimental", - expectedRuntimeBinaries: map[string]string{ - "nvidia": "nvidia-container-runtime", - "nvidia-experimental": "nvidia-container-runtime.experimental", - }, - }, - { - options: options{ - setAsDefault: false, - runtimeClass: "nvidia-experimental", - }, - expectedRuntimeBinaries: map[string]string{ - "nvidia": "nvidia-container-runtime", - "nvidia-experimental": "nvidia-container-runtime.experimental", - }, - }, - } - - for i, tc := range testCases { - require.Equal(t, tc.expectedDefaultRuntime, tc.options.getDefaultRuntime(), "%d: %v", i, tc) - require.EqualValues(t, tc.expectedRuntimeBinaries, tc.options.getRuntimeBinaries(), "%d: %v", i, tc) - } -} diff --git a/tools/container/operator/operator.go b/tools/container/operator/operator.go new file mode 100644 index 00000000..b52e97ad --- /dev/null +++ b/tools/container/operator/operator.go @@ -0,0 +1,130 @@ +/** +# 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 operator + +import "path/filepath" + +const ( + defaultRuntimeName = "nvidia" + experimentalRuntimeName = "nvidia-experimental" + + defaultRoot = "/usr/bin" +) + +// Runtime defines a runtime to be configured. +// The path and whether the runtime is the default runtime can be specfied +type Runtime struct { + name string + Path string + SetAsDefault bool +} + +// Runtimes defines a set of runtimes to be configure for use in the GPU Operator +type Runtimes map[string]Runtime + +type config struct { + root string + nvidiaRuntimeName string + setAsDefault bool +} + +// GetRuntimes returns the set of runtimes to be configured for use with the GPU Operator. +func GetRuntimes(opts ...Option) Runtimes { + c := &config{} + for _, opt := range opts { + opt(c) + } + + if c.root == "" { + c.root = defaultRoot + } + if c.nvidiaRuntimeName == "" { + c.nvidiaRuntimeName = defaultRuntimeName + } + + runtimes := make(Runtimes) + runtimes.add(c.nvidiaRuntime()) + + modes := []string{"experimental"} + for _, mode := range modes { + runtimes.add(c.modeRuntime(mode)) + } + return runtimes +} + +// DefaultRuntimeName returns the name of the default runtime. +func (r Runtimes) DefaultRuntimeName() string { + for _, runtime := range r { + if runtime.SetAsDefault { + return runtime.name + } + } + return "" +} + +// Add a runtime to the set of runtimes. +func (r *Runtimes) add(runtime Runtime) { + (*r)[runtime.name] = runtime +} + +// nvidiaRuntime creates a runtime that corresponds to the nvidia runtime. +// If the nvidiaRuntimeName is specified, this name us used unless this is exactly equal to nvidia-experimental. +func (c config) nvidiaRuntime() Runtime { + name := c.nvidiaRuntimeName + if name == experimentalRuntimeName { + name = defaultRuntimeName + } + return c.newRuntime(name, "nvidia-container-runtime") +} + +// modeRuntime creates a runtime for the specified mode. +func (c config) modeRuntime(mode string) Runtime { + return c.newRuntime("nvidia-"+mode, "nvidia-container-runtime."+mode) +} + +// newRuntime creates a runtime based on the configuration +func (c config) newRuntime(name string, binary string) Runtime { + return Runtime{ + name: name, + Path: filepath.Join(c.root, binary), + SetAsDefault: c.setAsDefault && name == c.nvidiaRuntimeName, + } +} + +// Option is a functional option for configuring set of runtimes. +type Option func(*config) + +// WithRoot sets the root directory for the runtime binaries. +func WithRoot(root string) Option { + return func(c *config) { + c.root = root + } +} + +// WithNvidiaRuntimeName sets the name of the nvidia runtime. +func WithNvidiaRuntimeName(name string) Option { + return func(c *config) { + c.nvidiaRuntimeName = name + } +} + +// WithSetAsDefault sets the default runtime to the nvidia runtime. +func WithSetAsDefault(set bool) Option { + return func(c *config) { + c.setAsDefault = set + } +} diff --git a/tools/container/operator/operator_test.go b/tools/container/operator/operator_test.go new file mode 100644 index 00000000..07417fb9 --- /dev/null +++ b/tools/container/operator/operator_test.go @@ -0,0 +1,151 @@ +/** +# 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 operator + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestOptions(t *testing.T) { + testCases := []struct { + setAsDefault bool + nvidiaRuntimeName string + root string + expectedDefaultRuntime string + expectedRuntimes Runtimes + }{ + { + expectedRuntimes: Runtimes{ + "nvidia": Runtime{ + name: "nvidia", + Path: "/usr/bin/nvidia-container-runtime", + }, + "nvidia-experimental": Runtime{ + name: "nvidia-experimental", + Path: "/usr/bin/nvidia-container-runtime.experimental", + }, + }, + }, + { + setAsDefault: true, + expectedDefaultRuntime: "nvidia", + expectedRuntimes: Runtimes{ + "nvidia": Runtime{ + name: "nvidia", + Path: "/usr/bin/nvidia-container-runtime", + SetAsDefault: true, + }, + "nvidia-experimental": Runtime{ + name: "nvidia-experimental", + Path: "/usr/bin/nvidia-container-runtime.experimental", + }, + }, + }, + { + setAsDefault: true, + nvidiaRuntimeName: "nvidia", + expectedDefaultRuntime: "nvidia", + expectedRuntimes: Runtimes{ + "nvidia": Runtime{ + name: "nvidia", + Path: "/usr/bin/nvidia-container-runtime", + SetAsDefault: true, + }, + "nvidia-experimental": Runtime{ + name: "nvidia-experimental", + Path: "/usr/bin/nvidia-container-runtime.experimental", + }, + }, + }, + { + setAsDefault: true, + nvidiaRuntimeName: "NAME", + expectedDefaultRuntime: "NAME", + expectedRuntimes: Runtimes{ + "NAME": Runtime{ + name: "NAME", + Path: "/usr/bin/nvidia-container-runtime", + SetAsDefault: true, + }, + "nvidia-experimental": Runtime{ + name: "nvidia-experimental", + Path: "/usr/bin/nvidia-container-runtime.experimental", + }, + }, + }, + { + setAsDefault: false, + nvidiaRuntimeName: "NAME", + expectedRuntimes: Runtimes{ + "NAME": Runtime{ + name: "NAME", + Path: "/usr/bin/nvidia-container-runtime", + }, + "nvidia-experimental": Runtime{ + name: "nvidia-experimental", + Path: "/usr/bin/nvidia-container-runtime.experimental", + }, + }, + }, + { + setAsDefault: true, + nvidiaRuntimeName: "nvidia-experimental", + expectedDefaultRuntime: "nvidia-experimental", + expectedRuntimes: Runtimes{ + "nvidia": Runtime{ + name: "nvidia", + Path: "/usr/bin/nvidia-container-runtime", + }, + "nvidia-experimental": Runtime{ + name: "nvidia-experimental", + Path: "/usr/bin/nvidia-container-runtime.experimental", + SetAsDefault: true, + }, + }, + }, + { + setAsDefault: false, + nvidiaRuntimeName: "nvidia-experimental", + expectedRuntimes: Runtimes{ + "nvidia": Runtime{ + name: "nvidia", + Path: "/usr/bin/nvidia-container-runtime", + }, + "nvidia-experimental": Runtime{ + name: "nvidia-experimental", + Path: "/usr/bin/nvidia-container-runtime.experimental", + }, + }, + }, + } + + for i, tc := range testCases { + t.Run(fmt.Sprintf("%d", i), func(t *testing.T) { + runtimes := GetRuntimes( + WithNvidiaRuntimeName(tc.nvidiaRuntimeName), + WithSetAsDefault(tc.setAsDefault), + WithRoot(tc.root), + ) + + require.EqualValues(t, tc.expectedRuntimes, runtimes) + require.Equal(t, tc.expectedDefaultRuntime, runtimes.DefaultRuntimeName()) + }) + } +}