src/cmd/cos_gpu_driver_builder/internal/builder/builder.go

package builder

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"cloud.google.com/go/storage"
	"cos.googlesource.com/cos/tools.git/src/pkg/cos"
	"cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig"
	"cos.googlesource.com/cos/tools.git/src/pkg/utils"
)

const (
	linkerLocation        = "/bin/ld"
	patchDir              = "/build/patches"
	installDirTemplate    = "/install/%s"
	defaultFilePermission = 0755
	kernelSrcTemplate     = "usr/src/linux-headers-*"
)

func kernelSrcDirectory(dirName string) string {
	files, err := filepath.Glob(filepath.Join(dirName, kernelSrcTemplate))
	if err != nil || len(files) != 1 {
		return ""
	}
	return files[0]
}

func nvidiaInstallerBuildCommand(dirName, runfile string, config gpuconfig.GPUPrecompilationConfig) *exec.Cmd {
	cmd := exec.Command(filepath.Join(dirName, runfile), "--kernel-source-path="+kernelSrcDirectory(dirName), "--add-this-kernel", "--no-install-compat32-libs", "--silent", "--accept-license")
	cmd.Dir = dirName
	return cmd
}

func nvidiaInstallerPatchCommand(dirName, runfile, patchDir, patchfile string) *exec.Cmd {
	cmd := exec.Command(filepath.Join(dirName, runfile), "--apply-patch", filepath.Join(patchDir, patchfile))
	cmd.Dir = dirName
	return cmd
}

func BuildPrecompiledDriver(ctx context.Context, client *storage.Client, config gpuconfig.GPUPrecompilationConfig) (string, string, error) {
	var err error
	dirName := fmt.Sprintf(installDirTemplate, config.Version)
	if err = os.MkdirAll(dirName, defaultFilePermission); err != nil {
		return "", "", fmt.Errorf("failed to create installation dir: %v", err)
	}
	downloader := gpuconfig.NewGPUArtifactsDownloader(client, config)
	// download NVIDIA runfile
	var nvidiaInstaller string
	if nvidiaInstaller, err = downloader.DownloadNVIDIARunfile(ctx, dirName); err != nil {
		return "", "", fmt.Errorf("failed to download NVIDIA runfile: %v", err)
	}
	// install kernel headers and toolchain
	// sets SYSROOT and PATH env vars
	if err = cos.InstallCrossToolchain(ctx, downloader, dirName); err != nil {
		return "", "", fmt.Errorf("failed to install toolchain: %v", err)
	}
	// set CC CXX env vars from toolchain_env
	if err = cos.SetCompilationEnv(ctx, downloader); err != nil {
		return "", "", fmt.Errorf("failed to set compilation env vars: %v", err)
	}
	// create symlink to ld - required by NVIDIA driver package
	if err = cos.ForceSymlinkLinker(filepath.Join(dirName, linkerLocation)); err != nil {
		return "", "", fmt.Errorf("failed to create symlink to COS linker: %v", err)
	}
	cc := os.Getenv("CC")
	if cc == "" {
		return "", "", fmt.Errorf("failed to find CC in env")
	} else {
		// create a wrapper removing -Werror=strict-prototypes from the CC command line.
		if err = cos.AddCCWrapperToPath(dirName, dirName, cc); err != nil {
			return "", "", fmt.Errorf("failed to create CC wrapper: %v", err)
		}
	}
	// run NVIDIA driver package
	if err = os.Chmod(filepath.Join(dirName, nvidiaInstaller), defaultFilePermission); err != nil {
		return "", "", err
	}

	// TODO: add driver patches arg to pb.COSGPUBuildRequest
	// hardcoded list for known patches for now.
	patches := patchesForDriverBuild(config)
	for _, patch := range patches {
		cmd := nvidiaInstallerPatchCommand(dirName, nvidiaInstaller, patchDir, patch)
		if err := utils.RunCommandAndLogOutput(cmd, false); err != nil {
			return "", "", fmt.Errorf("error applying provided patch to NVIDIA driver package: %v", err)
		}
		patchedInstaller := strings.Split(nvidiaInstaller, ".run")[0] + "-custom.run"
		if err := os.Rename(filepath.Join(dirName, patchedInstaller), filepath.Join(dirName, nvidiaInstaller)); err != nil {
			return "", "", fmt.Errorf("error moving patched NVIDIA driver package: %v", err)
		}
	}

	cmd := nvidiaInstallerBuildCommand(dirName, nvidiaInstaller, config)
	if err = utils.RunCommandAndLogOutput(cmd, false); err != nil {
		return "", "", fmt.Errorf("error running NVIDIA driver installation package: %v", err)
	}

	outputFileName := strings.Split(nvidiaInstaller, ".run")[0] + "-custom.run"
	return dirName, outputFileName, nil
}

func patchesForDriverBuild(config gpuconfig.GPUPrecompilationConfig) []string {
	pfnAffectedDriverVersions := []string{"470.199.02", "470.223.02", "525.125.06", "525.147.05", "535.104.05", "535.104.12", "535.129.03", "535.154.05"}
	pfnAffectedMilestones := []string{"97", "101", "105"}
	var patches []string
	if utils.StringSliceContains(pfnAffectedDriverVersions, config.DriverVersion) && utils.StringSliceContains(pfnAffectedMilestones, config.Milestone) {
		var kernelVersionDot string
		if config.Milestone == "97" {
			kernelVersionDot = "5.10"
		} else {
			kernelVersionDot = "5.15"
		}
		majorDriverVersion := strings.Split(config.DriverVersion, ".")[0]
		patchFile := fmt.Sprintf("nvidia-drivers-%s-%s-gpl-pfn_valid.patch", majorDriverVersion, kernelVersionDot)
		patches = append(patches, patchFile)
	}
	return patches
}