Use logger in toolkit install

tools/container/nvidia-toolkit/run.go

@@ -8,10 +8,10 @@ import (
 	"strings"
 	"syscall"
 
-	log "github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v2"
 	"golang.org/x/sys/unix"
 
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/runtime"
 	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/toolkit"
 )
@@ -51,12 +51,42 @@ func (o options) toolkitRoot() string {
 var Version = "development"
 
 func main() {
-	remainingArgs, root, err := ParseArgs(os.Args)
+	logger := logger.New()
+
+	remainingArgs, root, err := ParseArgs(logger, os.Args)
 	if err != nil {
-		log.Errorf("Error: unable to parse arguments: %v", err)
+		logger.Errorf("Error: unable to parse arguments: %v", err)
 		os.Exit(1)
 	}
 
+	c := new(logger, root)
+
+	// Run the CLI
+	logger.Infof("Starting %v", c.Name)
+	if err := c.Run(remainingArgs); err != nil {
+		logger.Errorf("error running nvidia-toolkit: %v", err)
+		os.Exit(1)
+	}
+
+	logger.Infof("Completed %v", c.Name)
+}
+
+type app struct {
+	logger      logger.Interface
+	defaultRoot string
+
+	toolkit *toolkit.Installer
+}
+
+func new(logger logger.Interface, defaultRoot string) *cli.App {
+	a := app{
+		logger:      logger,
+		defaultRoot: defaultRoot,
+	}
+	return a.build()
+}
+
+func (a app) build() *cli.App {
 	options := options{
 		toolkitOptions: toolkit.Options{},
 	}
@@ -68,10 +98,11 @@ func main() {
 	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 {
-		return validateFlags(ctx, &options)
+		a.init(&options)
+		return a.validateFlags(ctx, &options)
 	}
 	c.Action = func(ctx *cli.Context) error {
-		return Run(ctx, &options)
+		return a.Run(ctx, &options)
 	}
 
 	// Setup flags for the CLI
@@ -102,7 +133,7 @@ func main() {
 		},
 		&cli.StringFlag{
 			Name:        "root",
-			Value:       root,
+			Value:       a.defaultRoot,
 			Usage:       "the folder where the NVIDIA Container Toolkit is to be installed. It will be installed to `ROOT`/toolkit",
 			Destination: &options.root,
 			EnvVars:     []string{"ROOT"},
@@ -119,21 +150,28 @@ func main() {
 	c.Flags = append(c.Flags, toolkit.Flags(&options.toolkitOptions)...)
 	c.Flags = append(c.Flags, runtime.Flags(&options.runtimeOptions)...)
 
-	// Run the CLI
-	log.Infof("Starting %v", c.Name)
-	if err := c.Run(remainingArgs); err != nil {
-		log.Errorf("error running nvidia-toolkit: %v", err)
-		os.Exit(1)
-	}
+	return c
+}
 
-	log.Infof("Completed %v", c.Name)
+func (a *app) init(o *options) {
+	a.toolkit = toolkit.NewInstaller(
+		toolkit.WithLogger(a.logger),
+		toolkit.WithToolkitRoot(o.toolkitRoot()),
+	)
 }
 
-func validateFlags(_ *cli.Context, o *options) error {
+func (a *app) validateFlags(_ *cli.Context, o *options) error {
+	if o.root == "" {
+		return fmt.Errorf("the install root must be specified")
+	}
+	if _, exists := availableRuntimes[o.runtime]; !exists {
+		return fmt.Errorf("unknown runtime: %v", o.runtime)
+	}
 	if filepath.Base(o.pidFile) != toolkitPidFilename {
 		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 {
@@ -143,17 +181,12 @@ func validateFlags(_ *cli.Context, o *options) error {
 }
 
 // Run runs the core logic of the CLI
-func Run(c *cli.Context, o *options) error {
-	err := verifyFlags(o)
-	if err != nil {
-		return fmt.Errorf("unable to verify flags: %v", err)
-	}
-
-	err = initialize(o.pidFile)
+func (a *app) Run(c *cli.Context, o *options) error {
+	err := a.initialize(o.pidFile)
 	if err != nil {
 		return fmt.Errorf("unable to initialize: %v", err)
 	}
-	defer shutdown(o.pidFile)
+	defer a.shutdown(o.pidFile)
 
 	if len(o.toolkitOptions.ContainerRuntimeRuntimes.Value()) == 0 {
 		lowlevelRuntimePaths, err := runtime.GetLowlevelRuntimePaths(&o.runtimeOptions, o.runtime)
@@ -164,7 +197,12 @@ func 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)
 	}
@@ -175,7 +213,7 @@ func Run(c *cli.Context, o *options) error {
 	}
 
 	if !o.noDaemon {
-		err = waitForSignal()
+		err = a.waitForSignal()
 		if err != nil {
 			return fmt.Errorf("unable to wait for signal: %v", err)
 		}
@@ -191,8 +229,8 @@ func Run(c *cli.Context, o *options) error {
 
 // ParseArgs checks if a single positional argument was defined and extracts this the root.
 // If no positional arguments are defined, it is assumed that the root is specified as a flag.
-func ParseArgs(args []string) ([]string, string, error) {
-	log.Infof("Parsing arguments")
+func ParseArgs(logger logger.Interface, args []string) ([]string, string, error) {
+	logger.Infof("Parsing arguments")
 
 	if len(args) < 2 {
 		return args, "", nil
@@ -217,20 +255,8 @@ func ParseArgs(args []string) ([]string, string, error) {
 	return nil, "", fmt.Errorf("unexpected positional argument(s) %v", args[2:lastPositionalArg+1])
 }
 
-func verifyFlags(o *options) error {
-	log.Infof("Verifying Flags")
-	if o.root == "" {
-		return fmt.Errorf("the install root must be specified")
-	}
-
-	if _, exists := availableRuntimes[o.runtime]; !exists {
-		return fmt.Errorf("unknown runtime: %v", o.runtime)
-	}
-	return nil
-}
-
-func initialize(pidFile string) error {
-	log.Infof("Initializing")
+func (a *app) initialize(pidFile string) error {
+	a.logger.Infof("Initializing")
 
 	if dir := filepath.Dir(pidFile); dir != "" {
 		err := os.MkdirAll(dir, 0755)
@@ -246,8 +272,8 @@ func initialize(pidFile string) error {
 
 	err = unix.Flock(int(f.Fd()), unix.LOCK_EX|unix.LOCK_NB)
 	if err != nil {
-		log.Warningf("Unable to get exclusive lock on '%v'", pidFile)
-		log.Warningf("This normally means an instance of the NVIDIA toolkit Container is already running, aborting")
+		a.logger.Warningf("Unable to get exclusive lock on '%v'", pidFile)
+		a.logger.Warningf("This normally means an instance of the NVIDIA toolkit Container is already running, aborting")
 		return fmt.Errorf("unable to get flock on pidfile: %v", err)
 	}
 
@@ -264,27 +290,27 @@ func initialize(pidFile string) error {
 		case <-waitingForSignal:
 			signalReceived <- true
 		default:
-			log.Infof("Signal received, exiting early")
-			shutdown(pidFile)
+			a.logger.Infof("Signal received, exiting early")
+			a.shutdown(pidFile)
 			os.Exit(0)
 		}
 	}()
 
 	return nil
 }
 
-func waitForSignal() error {
-	log.Infof("Waiting for signal")
+func (a *app) waitForSignal() error {
+	a.logger.Infof("Waiting for signal")
 	waitingForSignal <- true
 	<-signalReceived
 	return nil
 }
 
-func shutdown(pidFile string) {
-	log.Infof("Shutting Down")
+func (a *app) shutdown(pidFile string) {
+	a.logger.Infof("Shutting Down")
 
 	err := os.Remove(pidFile)
 	if err != nil {
-		log.Warningf("Unable to remove pidfile: %v", err)
+		a.logger.Warningf("Unable to remove pidfile: %v", err)
 	}
 }

tools/container/nvidia-toolkit/run_test.go

@@ -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
@@ -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 {

tools/container/toolkit/executable.go

@@ -23,8 +23,6 @@ import (
 	"path/filepath"
 	"sort"
 	"strings"
-
-	log "github.com/sirupsen/logrus"
 )
 
 type executableTarget struct {
@@ -33,6 +31,7 @@ type executableTarget struct {
 }
 
 type executable struct {
+	fileInstaller
 	source   string
 	target   executableTarget
 	env      map[string]string
@@ -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
 }

tools/container/toolkit/executable_test.go

@@ -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"
@@ -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)
 
@@ -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)
@@ -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",

tools/container/toolkit/options.go

@@ -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
+	}
+}