Skip to content

Commit

Permalink
[no-relnote] Use logger in toolkit installation
Browse files Browse the repository at this point in the history
Signed-off-by: Evan Lezar <[email protected]>
  • Loading branch information
elezar committed Oct 30, 2024
1 parent 830904f commit dbed2e9
Show file tree
Hide file tree
Showing 9 changed files with 217 additions and 103 deletions.
19 changes: 17 additions & 2 deletions tools/container/nvidia-toolkit/run.go
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,8 @@ func main() {
type app struct {
logger logger.Interface
defaultRoot string

toolkit *toolkit.Installer
}

func new(logger logger.Interface, defaultRoot string) *cli.App {
Expand All @@ -96,6 +98,7 @@ func (a app) build() *cli.App {
c.Description = "DESTINATION points to the host path underneath which the nvidia-container-toolkit should be installed.\nIt will be installed at ${DESTINATION}/toolkit"
c.Version = Version
c.Before = func(ctx *cli.Context) error {
a.init(&options)
return a.validateFlags(ctx, &options)
}
c.Action = func(ctx *cli.Context) error {
Expand Down Expand Up @@ -150,6 +153,13 @@ func (a app) build() *cli.App {
return c
}

func (a *app) init(o *options) {
a.toolkit = toolkit.NewInstaller(
toolkit.WithLogger(a.logger),
toolkit.WithToolkitRoot(o.toolkitRoot()),
)
}

func (a *app) validateFlags(_ *cli.Context, o *options) error {
if o.root == "" {
return fmt.Errorf("the install root must be specified")
Expand All @@ -161,7 +171,7 @@ func (a *app) validateFlags(_ *cli.Context, o *options) error {
return fmt.Errorf("invalid toolkit.pid path %v", o.pidFile)
}

if err := toolkit.ValidateOptions(&o.toolkitOptions, o.toolkitRoot()); err != nil {
if err := a.toolkit.ValidateOptions(&o.toolkitOptions); err != nil {
return err
}
if err := runtime.ValidateOptions(&o.runtimeOptions, o.runtime, o.toolkitRoot()); err != nil {
Expand All @@ -187,7 +197,12 @@ func (a *app) Run(c *cli.Context, o *options) error {

o.toolkitOptions.ContainerRuntimeRuntimes = *cli.NewStringSlice(lowlevelRuntimePaths...)
}
err = toolkit.Install(c, &o.toolkitOptions, "", o.toolkitRoot())

installer := toolkit.NewInstaller(
toolkit.WithLogger(a.logger),
toolkit.WithToolkitRoot(o.toolkitRoot()),
)
err = installer.Install(c, &o.toolkitOptions)
if err != nil {
return fmt.Errorf("unable to install toolkit: %v", err)
}
Expand Down
4 changes: 3 additions & 1 deletion tools/container/nvidia-toolkit/run_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,12 @@ import (
"fmt"
"testing"

testlog "github.com/sirupsen/logrus/hooks/test"
"github.com/stretchr/testify/require"
)

func TestParseArgs(t *testing.T) {
logger, _ := testlog.NewNullLogger()
testCases := []struct {
args []string
expectedRemaining []string
Expand Down Expand Up @@ -70,7 +72,7 @@ func TestParseArgs(t *testing.T) {

for i, tc := range testCases {
t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
remaining, root, err := ParseArgs(tc.args)
remaining, root, err := ParseArgs(logger, tc.args)
if tc.expectedError != nil {
require.EqualError(t, err, tc.expectedError.Error())
} else {
Expand Down
11 changes: 5 additions & 6 deletions tools/container/toolkit/executable.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,6 @@ import (
"path/filepath"
"sort"
"strings"

log "github.com/sirupsen/logrus"
)

type executableTarget struct {
Expand All @@ -33,6 +31,7 @@ type executableTarget struct {
}

type executable struct {
fileInstaller
source string
target executableTarget
env map[string]string
Expand All @@ -43,21 +42,21 @@ type executable struct {
// install installs an executable component of the NVIDIA container toolkit. The source executable
// is copied to a `.real` file and a wapper is created to set up the environment as required.
func (e executable) install(destFolder string) (string, error) {
log.Infof("Installing executable '%v' to %v", e.source, destFolder)
e.logger.Infof("Installing executable '%v' to %v", e.source, destFolder)

dotfileName := e.dotfileName()

installedDotfileName, err := installFileToFolderWithName(destFolder, dotfileName, e.source)
installedDotfileName, err := e.installFileToFolderWithName(destFolder, dotfileName, e.source)
if err != nil {
return "", fmt.Errorf("error installing file '%v' as '%v': %v", e.source, dotfileName, err)
}
log.Infof("Installed '%v'", installedDotfileName)
e.logger.Infof("Installed '%v'", installedDotfileName)

wrapperFilename, err := e.installWrapper(destFolder, installedDotfileName)
if err != nil {
return "", fmt.Errorf("error wrapping '%v': %v", installedDotfileName, err)
}
log.Infof("Installed wrapper '%v'", wrapperFilename)
e.logger.Infof("Installed wrapper '%v'", wrapperFilename)

return wrapperFilename, nil
}
Expand Down
10 changes: 10 additions & 0 deletions tools/container/toolkit/executable_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,13 @@ import (
"strings"
"testing"

testlog "github.com/sirupsen/logrus/hooks/test"
"github.com/stretchr/testify/require"
)

func TestWrapper(t *testing.T) {
logger, _ := testlog.NewNullLogger()

const shebang = "#! /bin/sh"
const destFolder = "/dest/folder"
const dotfileName = "source.real"
Expand Down Expand Up @@ -98,6 +101,8 @@ func TestWrapper(t *testing.T) {
for i, tc := range testCases {
buf := &bytes.Buffer{}

tc.e.logger = logger

err := tc.e.writeWrapperTo(buf, destFolder, dotfileName)
require.NoError(t, err)

Expand All @@ -107,6 +112,8 @@ func TestWrapper(t *testing.T) {
}

func TestInstallExecutable(t *testing.T) {
logger, _ := testlog.NewNullLogger()

inputFolder, err := os.MkdirTemp("", "")
require.NoError(t, err)
defer os.RemoveAll(inputFolder)
Expand All @@ -121,6 +128,9 @@ func TestInstallExecutable(t *testing.T) {
require.NoError(t, sourceFile.Close())

e := executable{
fileInstaller: fileInstaller{
logger: logger,
},
source: source,
target: executableTarget{
dotfileName: "input.real",
Expand Down
40 changes: 40 additions & 0 deletions tools/container/toolkit/options.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
/**
# Copyright 2024 NVIDIA CORPORATION
#
# 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 toolkit

import "github.com/NVIDIA/nvidia-container-toolkit/internal/logger"

// An Option provides a mechanism to configure an Installer.
type Option func(*Installer)

func WithLogger(logger logger.Interface) Option {
return func(i *Installer) {
i.logger = logger
}
}

func WithToolkitRoot(toolkitRoot string) Option {
return func(i *Installer) {
i.toolkitRoot = toolkitRoot
}
}

func WithSourceRoot(sourceRoot string) Option {
return func(i *Installer) {
i.sourceRoot = sourceRoot
}
}
19 changes: 10 additions & 9 deletions tools/container/toolkit/runtime.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,10 @@ const (

// installContainerRuntimes sets up the NVIDIA container runtimes, copying the executables
// and implementing the required wrapper
func installContainerRuntimes(sourceRoot string, toolkitDir string) error {
func (t *Installer) installContainerRuntimes(toolkitDir string) error {
runtimes := operator.GetRuntimes()
for _, runtime := range runtimes {
r := newNvidiaContainerRuntimeInstaller(filepath.Join(sourceRoot, runtime.Path))
r := t.newNvidiaContainerRuntimeInstaller(runtime.Path)

_, err := r.install(toolkitDir)
if err != nil {
Expand All @@ -46,17 +46,17 @@ func installContainerRuntimes(sourceRoot string, toolkitDir string) error {
// This installer will copy the specified source executable to the toolkit directory.
// The executable is copied to a file with the same name as the source, but with a ".real" suffix and a wrapper is
// created to allow for the configuration of the runtime environment.
func newNvidiaContainerRuntimeInstaller(source string) *executable {
func (t *Installer) newNvidiaContainerRuntimeInstaller(source string) *executable {
wrapperName := filepath.Base(source)
dotfileName := wrapperName + ".real"
target := executableTarget{
dotfileName: dotfileName,
wrapperName: wrapperName,
}
return newRuntimeInstaller(source, target, nil)
return t.newRuntimeInstaller(source, target, nil)
}

func newRuntimeInstaller(source string, target executableTarget, env map[string]string) *executable {
func (t *Installer) newRuntimeInstaller(source string, target executableTarget, env map[string]string) *executable {
preLines := []string{
"",
"cat /proc/modules | grep -e \"^nvidia \" >/dev/null 2>&1",
Expand All @@ -74,10 +74,11 @@ func newRuntimeInstaller(source string, target executableTarget, env map[string]
}

r := executable{
source: source,
target: target,
env: runtimeEnv,
preLines: preLines,
fileInstaller: t.fileInstaller,
source: source,
target: target,
env: runtimeEnv,
preLines: preLines,
}

return &r
Expand Down
9 changes: 8 additions & 1 deletion tools/container/toolkit/runtime_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,18 @@ import (
"strings"
"testing"

testlog "github.com/sirupsen/logrus/hooks/test"
"github.com/stretchr/testify/require"
)

func TestNvidiaContainerRuntimeInstallerWrapper(t *testing.T) {
r := newNvidiaContainerRuntimeInstaller(nvidiaContainerRuntimeSource)
logger, _ := testlog.NewNullLogger()
i := Installer{
fileInstaller: fileInstaller{
logger: logger,
},
}
r := i.newNvidiaContainerRuntimeInstaller(nvidiaContainerRuntimeSource)

const shebang = "#! /bin/sh"
const destFolder = "/dest/folder"
Expand Down
Loading

0 comments on commit dbed2e9

Please sign in to comment.