nvidia-container-toolkit/tools/container/toolkit/installer/executables.go

/**
# 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 installer

import (
	"bytes"
	"fmt"
	"html/template"
	"io"
	"path/filepath"
	"strings"

	log "github.com/sirupsen/logrus"

	"github.com/NVIDIA/nvidia-container-toolkit/tools/container/operator"
)

type executable struct {
	path    string
	symlink string
	args    []string
	env     map[string]string
}

func (t *toolkitInstaller) collectExecutables(destDir string) ([]Installer, error) {
	configHome := filepath.Join(destDir, ".config")
	configDir := filepath.Join(configHome, "nvidia-container-runtime")
	configPath := filepath.Join(configDir, "config.toml")

	executables := []executable{
		{
			path: "nvidia-ctk",
		},
		{
			path: "nvidia-cdi-hook",
		},
	}
	for _, runtime := range operator.GetRuntimes() {
		e := executable{
			path: runtime.Path,
			env: map[string]string{
				"XDG_CONFIG_HOME": configHome,
			},
		}
		executables = append(executables, e)
	}
	executables = append(executables,
		executable{
			path: "nvidia-container-cli",
			env:  map[string]string{"LD_LIBRARY_PATH": destDir + ":$LD_LIBRARY_PATH"},
		},
	)

	executables = append(executables,
		executable{
			path:    "nvidia-container-runtime-hook",
			symlink: "nvidia-container-toolkit",
			args:    []string{fmt.Sprintf("-config %s", configPath)},
		},
	)

	var installers []Installer
	for _, executable := range executables {
		executablePath, err := t.artifactRoot.findExecutable(executable.path)
		if err != nil {
			if t.ignoreErrors {
				log.Errorf("Ignoring error: %v", err)
				continue
			}
			return nil, err
		}

		wrappedExecutableFilename := filepath.Base(executablePath)
		dotRealFilename := wrappedExecutableFilename + ".real"

		w := &wrapper{
			Source:            executablePath,
			WrappedExecutable: dotRealFilename,
			Args:              executable.args,
			Envvars: map[string]string{
				"PATH": strings.Join([]string{destDir, "$PATH"}, ":"),
			},
		}
		for k, v := range executable.env {
			w.Envvars[k] = v
		}

		installers = append(installers, w)

		if executable.symlink == "" {
			continue
		}
		link := symlink{
			linkname: executable.symlink,
			target:   filepath.Base(executablePath),
		}
		installers = append(installers, link)
	}

	return installers, nil

}

type wrapper struct {
	Source            string
	Envvars           map[string]string
	WrappedExecutable string
	Args              []string
}

func (w *wrapper) Install(destDir string) error {
	// Copy the executable with a .real extension.
	mode, err := installFile(w.Source, filepath.Join(destDir, w.WrappedExecutable))
	if err != nil {
		return err
	}

	// Create a wrapper file.
	content, err := w.render()
	if err != nil {
		return nil
	}
	wrapperFile := filepath.Join(destDir, filepath.Base(w.Source))
	return installContent(content, wrapperFile, mode|0111)
}

func (w *wrapper) render() (io.Reader, error) {
	wrapperTemplate := `#! /bin/sh

cat /proc/modules | grep -e \"^nvidia \" >/dev/null 2>&1
if [ "${?}" != "0" ]; then
	echo "nvidia driver modules are not yet loaded, invoking runc directly"
	exec runc "$@"
fi
{{- range $key, $value := .Envvars }}
{{$key}}={{$value}} \
{{- end }}
	{{ .WrappedExecutable }} \
{{- range $arg := .Args }}
		{{$arg}} \
{{- end }}
		"$@"
`

	var content bytes.Buffer
	tmpl, err := template.New("wrapper").Parse(wrapperTemplate)
	if err != nil {
		return nil, err
	}
	if err := tmpl.Execute(&content, w); err != nil {
		return nil, err
	}

	return &content, nil
}