Consolidate code to find nvidia-ctk

Signed-off-by: Evan Lezar <elezar@nvidia.com>

cmd/nvidia-ctk/cdi/generate/generate.go

@@ -25,7 +25,6 @@ import (
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
 	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi"
 	specs "github.com/container-orchestrated-devices/container-device-interface/specs-go"
 	"github.com/sirupsen/logrus"
@@ -36,9 +35,6 @@ import (
 )
 
 const (
-	nvidiaCTKExecutable      = "nvidia-ctk"
-	nvidiaCTKDefaultFilePath = "/usr/bin/" + nvidiaCTKExecutable
-
 	formatJSON = "json"
 	formatYAML = "yaml"
 )
@@ -168,30 +164,6 @@ func (m command) run(c *cli.Context, cfg *config) error {
 	return nil
 }
 
-// findNvidiaCTK locates the nvidia-ctk executable to be used in hooks.
-// If an override is specified, this is used instead.
-func (m command) findNvidiaCTK(override string) string {
-	if override != "" {
-		m.logger.Debugf("Using specified NVIDIA Container Toolkit CLI path %v", override)
-		return override
-	}
-
-	lookup := lookup.NewExecutableLocator(m.logger, "")
-	hookPath := nvidiaCTKDefaultFilePath
-	targets, err := lookup.Locate(nvidiaCTKExecutable)
-	if err != nil {
-		m.logger.Warnf("Failed to locate %v: %v", nvidiaCTKExecutable, err)
-	} else if len(targets) == 0 {
-		m.logger.Warnf("%v not found", nvidiaCTKExecutable)
-	} else {
-		m.logger.Debugf("Found %v candidates: %v", nvidiaCTKExecutable, targets)
-		hookPath = targets[0]
-	}
-	m.logger.Debugf("Using NVIDIA Container Toolkit CLI path %v", hookPath)
-
-	return hookPath
-}
-
 func formatFromFilename(filename string) string {
 	ext := filepath.Ext(filename)
 	switch strings.ToLower(ext) {
@@ -230,7 +202,7 @@ func (m command) generateSpec(root string, nvidiaCTKPath string) (*specs.Spec, e
 
 	devicelib := device.New(device.WithNvml(nvmllib))
 
-	useNvidiaCTKPath := m.findNvidiaCTK(nvidiaCTKPath)
+	useNvidiaCTKPath := discover.FindNvidiaCTK(m.logger, nvidiaCTKPath)
 
 	deviceSpecs, err := m.generateDeviceSpecs(devicelib, root, useNvidiaCTKPath)
 	if err != nil {

internal/discover/graphics.go

@@ -96,7 +96,6 @@ func NewGraphicsMountsDiscoverer(logger *logrus.Logger, root string) (Discover,
 type drmDevicesByPath struct {
 	None
 	logger        *logrus.Logger
-	lookup        lookup.Locator
 	nvidiaCTKPath string
 	root          string
 	devicesFrom   Discover
@@ -106,8 +105,7 @@ type drmDevicesByPath struct {
 func newCreateDRMByPathSymlinks(logger *logrus.Logger, devices Discover, cfg *Config) Discover {
 	d := drmDevicesByPath{
 		logger:        logger,
-		lookup:        lookup.NewExecutableLocator(logger, cfg.Root),
+		nvidiaCTKPath: FindNvidiaCTK(logger, cfg.NvidiaCTKPath),
-		nvidiaCTKPath: cfg.NvidiaCTKPath,
 		root:          cfg.Root,
 		devicesFrom:   devices,
 	}
@@ -132,26 +130,14 @@ func (d drmDevicesByPath) Hooks() ([]Hook, error) {
 		return nil, nil
 	}
 
-	hookPath := nvidiaCTKDefaultFilePath
+	args := []string{d.nvidiaCTKPath, "hook", "create-symlinks"}
-	targets, err := d.lookup.Locate(d.nvidiaCTKPath)
-	if err != nil {
-		d.logger.Warnf("Failed to locate %v: %v", d.nvidiaCTKPath, err)
-	} else if len(targets) == 0 {
-		d.logger.Warnf("%v not found", d.nvidiaCTKPath)
-	} else {
-		d.logger.Debugf("Found %v candidates: %v", d.nvidiaCTKPath, targets)
-		hookPath = targets[0]
-	}
-	d.logger.Debugf("Using NVIDIA Container Toolkit CLI path %v", hookPath)
-
-	args := []string{hookPath, "hook", "create-symlinks"}
 	for _, l := range links {
 		args = append(args, "--link", l)
 	}
 
 	h := Hook{
 		Lifecycle: cdi.CreateContainerHook,
-		Path:      hookPath,
+		Path:      d.nvidiaCTKPath,
 		Args:      args,
 	}
 

internal/discover/hooks.go

@@ -0,0 +1,63 @@
+/**
+# 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 discover
+
+import (
+	"path/filepath"
+
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
+	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi"
+	"github.com/sirupsen/logrus"
+)
+
+const (
+	nvidiaCTKExecutable      = "nvidia-ctk"
+	nvidiaCTKDefaultFilePath = "/usr/bin/nvidia-ctk"
+)
+
+// CreateNvidiaCTKHook creates a hook which invokes the NVIDIA Container CLI hook subcommand.
+func CreateNvidiaCTKHook(executable string, hookName string, additionalArgs ...string) Hook {
+	return Hook{
+		Lifecycle: cdi.CreateContainerHook,
+		Path:      executable,
+		Args:      append([]string{filepath.Base(executable), "hook", hookName}, additionalArgs...),
+	}
+}
+
+// FindNvidiaCTK locates the nvidia-ctk executable to be used in hooks.
+// If an override is specified, this is used instead.
+func FindNvidiaCTK(logger *logrus.Logger, override string) string {
+	if override != "" {
+		logger.Debugf("Using specified NVIDIA Container Toolkit CLI path %v", override)
+		return override
+	}
+
+	lookup := lookup.NewExecutableLocator(logger, "")
+	hookPath := nvidiaCTKDefaultFilePath
+	targets, err := lookup.Locate(nvidiaCTKExecutable)
+	if err != nil {
+		logger.Warnf("Failed to locate %v: %v", nvidiaCTKExecutable, err)
+	} else if len(targets) == 0 {
+		logger.Warnf("%v not found", nvidiaCTKExecutable)
+	} else {
+		logger.Debugf("Found %v candidates: %v", nvidiaCTKExecutable, targets)
+		hookPath = targets[0]
+	}
+	logger.Debugf("Using NVIDIA Container Toolkit CLI path %v", hookPath)
+
+	return hookPath
+}

internal/discover/ldconfig.go

@@ -22,7 +22,6 @@ import (
 	"strings"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
-	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi"
 	"github.com/sirupsen/logrus"
 )
 
@@ -38,10 +37,6 @@ func NewLDCacheUpdateHook(logger *logrus.Logger, mounts Discover, cfg *Config) (
 	return &d, nil
 }
 
-const (
-	nvidiaCTKDefaultFilePath = "/usr/bin/nvidia-ctk"
-)
-
 type ldconfig struct {
 	None
 	logger        *logrus.Logger
@@ -80,15 +75,6 @@ func CreateLDCacheUpdateHook(executable string, libraries []string) Hook {
 
 }
 
-// CreateNvidiaCTKHook creates a hook which invokes the NVIDIA Container CLI hook subcommand.
-func CreateNvidiaCTKHook(executable string, hookName string, additionalArgs ...string) Hook {
-	return Hook{
-		Lifecycle: cdi.CreateContainerHook,
-		Path:      executable,
-		Args:      append([]string{filepath.Base(executable), "hook", hookName}, additionalArgs...),
-	}
-}
-
 // getLibraryPaths extracts the library dirs from the specified mounts
 func getLibraryPaths(mounts []Mount) []string {
 	var paths []string

internal/discover/symlinks.go

@@ -21,7 +21,6 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
 	"github.com/container-orchestrated-devices/container-device-interface/pkg/cdi"
 	"github.com/sirupsen/logrus"
 )
@@ -29,7 +28,6 @@ import (
 type symlinks struct {
 	None
 	logger        *logrus.Logger
-	lookup        lookup.Locator
 	nvidiaCTKPath string
 	csvFiles      []string
 	mountsFrom    Discover
@@ -39,8 +37,7 @@ type symlinks struct {
 func NewCreateSymlinksHook(logger *logrus.Logger, csvFiles []string, mounts Discover, cfg *Config) (Discover, error) {
 	d := symlinks{
 		logger:        logger,
-		lookup:        lookup.NewExecutableLocator(logger, cfg.Root),
+		nvidiaCTKPath: FindNvidiaCTK(logger, cfg.NvidiaCTKPath),
-		nvidiaCTKPath: cfg.NvidiaCTKPath,
 		csvFiles:      csvFiles,
 		mountsFrom:    mounts,
 	}
@@ -50,19 +47,7 @@ func NewCreateSymlinksHook(logger *logrus.Logger, csvFiles []string, mounts Disc
 
 // Hooks returns a hook to create the symlinks from the required CSV files
 func (d symlinks) Hooks() ([]Hook, error) {
-	hookPath := nvidiaCTKDefaultFilePath
+	args := []string{d.nvidiaCTKPath, "hook", "create-symlinks"}
-	targets, err := d.lookup.Locate(d.nvidiaCTKPath)
-	if err != nil {
-		d.logger.Warnf("Failed to locate %v: %v", d.nvidiaCTKPath, err)
-	} else if len(targets) == 0 {
-		d.logger.Warnf("%v not found", d.nvidiaCTKPath)
-	} else {
-		d.logger.Debugf("Found %v candidates: %v", d.nvidiaCTKPath, targets)
-		hookPath = targets[0]
-	}
-	d.logger.Debugf("Using NVIDIA Container Toolkit CLI path %v", hookPath)
-
-	args := []string{hookPath, "hook", "create-symlinks"}
 	for _, f := range d.csvFiles {
 		args = append(args, "--csv-filename", f)
 	}
@@ -75,7 +60,7 @@ func (d symlinks) Hooks() ([]Hook, error) {
 
 	h := Hook{
 		Lifecycle: cdi.CreateContainerHook,
-		Path:      hookPath,
+		Path:      d.nvidiaCTKPath,
 		Args:      args,
 	}