src/pkg/preloader/gpu.go - cos/tools - Git at Google

 // Copyright 2025 Google LLC
 //
 // 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 preloader

 import (
 	"context"
 	"crypto/md5"
 	"fmt"
 	"io"
 	"io/fs"
 	"log"
 	"net/http"
 	"os"
 	"os/exec"
 	"path/filepath"

 	"cos.googlesource.com/cos/tools.git/src/pkg/config"
 	"cos.googlesource.com/cos/tools.git/src/pkg/cos"
 	cosfs "cos.googlesource.com/cos/tools.git/src/pkg/fs"
 	"cos.googlesource.com/cos/tools.git/src/pkg/gpubuild"
 	"cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig"
 	"cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig/pb"
 	"cos.googlesource.com/cos/tools.git/src/pkg/utils"

 	"cloud.google.com/go/storage"
 	"google.golang.org/protobuf/proto"
 )

 func extractTarball(srcPath, destDir string) error {
 	cmd := exec.Command("tar", "xf", srcPath, "-C", destDir)
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 	return cmd.Run()
 }

 type envReaderCreator func(context.Context) (*cos.EnvReader, error)

 // cachedEnvReader returns a closure that creates an EnvReader from a VMImage.
 // Since creating an EnvReader is very expensive in this context (it boots a VM
 // from the image to read data from it), we use a closure to cache the resulting
 // envReader.
 //
 // We could create the EnvReader once and pass it around everywhere. However,
 // under specific conditions, the EnvReader won't need to be created at all,
 // so this closure allows us to lazily create the EnvReader when it is needed.
 func cachedEnvReader(input *config.Image, gcs *gcsManager, files *cosfs.Files, buildSpec *config.Build) envReaderCreator {
 	var envReader *cos.EnvReader
 	return func(ctx context.Context) (*cos.EnvReader, error) {
 		var err error
 		if envReader != nil {
 			return envReader, nil
 		}
 		envReader, err = cos.NewEnvReaderFromVMImage(input.Name, input.Project, cos.EnvReaderVMImageConfig{
 			DaisyBin:       files.DaisyBin,
 			DaisyWorkflow:  "/data/read_cos_env.wf.json",
 			Project:        buildSpec.Project,
 			Zone:           buildSpec.Zone,
 			GCSDir:         gcs.managedDirURL(),
 			DiskType:       buildSpec.DiskType,
 			MachineType:    buildSpec.MachineType,
 			Network:        buildSpec.Network,
 			SubNet:         buildSpec.Subnet,
 			ServiceAccount: buildSpec.ServiceAccount,
 		})
 		if err != nil {
 			return nil, fmt.Errorf("failed to read env from VM image: %v", err)
 		}
 		return envReader, nil
 	}
 }

 func installToolchain(ctx context.Context, envReaderFunc envReaderCreator, gcsClient *storage.Client, buildSpec *config.Build, toolchainDir string) error {
 	if buildSpec.Toolchain != nil {
 		log.Println("Extracting toolchain package")
 		if err := extractTarball(buildSpec.Toolchain.ToolchainTarPath, toolchainDir); err != nil {
 			return err
 		}
 		if err := extractTarball(buildSpec.Toolchain.KernelHeadersPath, toolchainDir); err != nil {
 			return err
 		}
 		return nil
 	}
 	log.Println("No toolchain package provided; getting toolchain from input image")
 	envReader, err := envReaderFunc(ctx)
 	if err != nil {
 		return err
 	}
 	downloader := cos.NewGCSDownloader(gcsClient, envReader, "", "", "", "")
 	if err := cos.InstallCrossToolchain(ctx, downloader, toolchainDir); err != nil {
 		return err
 	}
 	// InstallCrossToolchain has an undesirable side effect of setting the SYSROOT
 	// environment variable. It's unclear if removing this behavior from
 	// InstallCrossToolchain is safe for all usages.
 	if err := os.Unsetenv("SYSROOT"); err != nil {
 		log.Printf("WARNING: unable to unset SYSROOT: %v", err)
 	}
 	return nil
 }

 func getArchFromToolchain(toolchainDir string) (string, error) {
 	x86Matches, err := filepath.Glob(filepath.Join(toolchainDir, "bin/x86_64-cros-linux-gnu-*"))
 	if err != nil {
 		return "", err
 	}
 	if len(x86Matches) != 0 {
 		return "x86_64", nil
 	}
 	aarch64Matches, err := filepath.Glob(filepath.Join(toolchainDir, "bin/aarch64-cros-linux-gnu-*"))
 	if err != nil {
 		return "", err
 	}
 	if len(aarch64Matches) != 0 {
 		return "aarch64", nil
 	}
 	return "", fmt.Errorf("could not identify arch from toolchain at %q; are programs missing from the toolchain?")
 }

 func downloadRunfile(version, arch, md5Sum, outDir string) (_ string, err error) {
 	outName := fmt.Sprintf("NVIDIA-Linux-%s-%s.run", arch, version)
 	url := fmt.Sprintf("https://us.download.nvidia.com/tesla/%s/%s", version, outName)
 	log.Printf("Downloading runfile from %q", url)
 	outPath := filepath.Join(outDir, outName)
 	outFile, err := os.Create(outPath)
 	if err != nil {
 		return "", err
 	}
 	defer utils.CheckClose(outFile, "error closing downloaded runfile", &err)
 	resp, err := http.Get(url)
 	if err != nil {
 		return "", err
 	}
 	defer utils.CheckClose(resp.Body, "error closing runfile HTTP body", &err)
 	if resp.StatusCode != http.StatusOK {
 		return "", fmt.Errorf("unhealthy status downloading Nvidia runfile: %v", resp.Status)
 	}
 	if md5Sum != "" {
 		h := md5.New()
 		mw := io.MultiWriter(outFile, h)
 		if _, err := io.Copy(mw, resp.Body); err != nil {
 			return "", err
 		}
 		got := fmt.Sprintf("%x", h.Sum(nil))
 		if got != md5Sum {
 			return "", fmt.Errorf("md5sum mismatch when downloading runfile from %q; got %q, want %q", url, got, md5Sum)
 		}
 	} else {
 		if _, err := io.Copy(outFile, resp.Body); err != nil {
 			return "", err
 		}
 	}
 	return outPath, nil
 }

 func listModules(dir string) ([]string, error) {
 	var result []string
 	if err := filepath.WalkDir(dir, func(path string, _ fs.DirEntry, err error) error {
 		if err != nil {
 			return err
 		}
 		if filepath.Ext(path) == ".ko" {
 			result = append(result, filepath.Base(path))
 		}
 		return nil
 	}); err != nil {
 		return nil, err
 	}
 	return result, nil
 }

 func setDiff(a, b []string) []string {
 	aMap := map[string]bool{}
 	for _, e := range a {
 		aMap[e] = true
 	}
 	for _, e := range b {
 		if _, ok := aMap[e]; ok {
 			delete(aMap, e)
 		}
 	}
 	var result []string
 	for k, _ := range aMap {
 		result = append(result, k)
 	}
 	return result
 }

 func mergeDriverPackages(generatedDriver *gpubuild.DriverPackage, preBuiltPath, outDir string) (*gpubuild.DriverPackage, error) {
 	log.Println("Merging any missing precompiled modules from precompiled package into custom package")
 	nvidiaGenerated, err := os.MkdirTemp("", "nvidia-generated")
 	if err != nil {
 		return nil, err
 	}
 	defer os.RemoveAll(nvidiaGenerated)
 	if err := extractTarball(generatedDriver.Path, nvidiaGenerated); err != nil {
 		return nil, err
 	}
 	nvidiaPreBuilt, err := os.MkdirTemp("", "prebuilt-generated")
 	if err != nil {
 		return nil, err
 	}
 	defer os.RemoveAll(nvidiaPreBuilt)
 	if err := extractTarball(preBuiltPath, nvidiaPreBuilt); err != nil {
 		return nil, err
 	}
 	generatedModules, err := listModules(nvidiaGenerated)
 	if err != nil {
 		return nil, err
 	}
 	log.Printf("Generated package includes: %v", generatedModules)
 	preBuiltModules, err := listModules(nvidiaPreBuilt)
 	if err != nil {
 		return nil, err
 	}
 	log.Printf("Prebuilt package includes: %v", preBuiltModules)
 	notGeneratedModules := setDiff(preBuiltModules, generatedModules)
 	log.Printf("Adding the following to generated package: %v", notGeneratedModules)
 	for _, module := range notGeneratedModules {
 		src := filepath.Join(nvidiaPreBuilt, "drivers", module)
 		dst := filepath.Join(nvidiaGenerated, "drivers", module)
 		if err := utils.CopyFile(src, dst); err != nil {
 			return nil, err
 		}
 	}
 	result := &gpubuild.DriverPackage{
 		Path:    filepath.Join(outDir, fmt.Sprintf("nvidia-drivers-%s.tgz", generatedDriver.Version)),
 		Version: generatedDriver.Version,
 	}
 	if err := cosfs.TarDir(nvidiaGenerated, result.Path); err != nil {
 		return nil, err
 	}
 	return result, nil
 }

 // installNonRunfileModules searches for kernel modules in any existing
 // precompiled packages that are not present in the runfile we just compiled. If
 // it finds anything, it adds it to the generated package we just compiled. This
 // helps with supporting gdrcopy.
 func installNonRunfileModules(ctx context.Context, envReaderFunc envReaderCreator, gcsClient *storage.Client, gpuConfig *config.GPUConfig, generatedDriver *gpubuild.DriverPackage, outDir string) (*gpubuild.DriverPackage, error) {
 	if gpuConfig.PrePackagedDriversPath != "" {
 		return mergeDriverPackages(generatedDriver, gpuConfig.PrePackagedDriversPath, outDir)
 	}
 	log.Println("No precompiled driver package provided, looking for an existing driver package for the given VM image")
 	envReader, err := envReaderFunc(ctx)
 	if err != nil {
 		return nil, err
 	}
 	downloader := cos.NewGCSDownloader(gcsClient, envReader, "", "", "", "")
 	precompiled := fmt.Sprintf("nvidia-drivers-%s.tgz", generatedDriver.Version)
 	if exists, err := downloader.ArtifactExists(ctx, precompiled); err != nil || !exists {
 		if err != nil {
 			log.Printf("WARNING: error searching for non-runfile modules for version %q: %v", generatedDriver.Version, err)
 		}
 		log.Printf("WARNING: could not find non-runfile modules for version %q, will not install them", generatedDriver.Version)
 		return generatedDriver, nil
 	}
 	tmpDir, err := os.MkdirTemp("", "non-runfile-modules")
 	if err != nil {
 		return nil, err
 	}
 	defer os.RemoveAll(tmpDir)
 	if err := downloader.DownloadArtifact(ctx, tmpDir, precompiled); err != nil {
 		return nil, err
 	}
 	precompiledPath := filepath.Join(tmpDir, precompiled)
 	return mergeDriverPackages(generatedDriver, precompiledPath, outDir)
 }

 func downloadImexDriver(ctx context.Context, envReaderFunc envReaderCreator, gcsClient *storage.Client, generatedDriver *gpubuild.DriverPackage, arch, outDir string) (string, error) {
 	if arch != "aarch64" {
 		log.Printf("Arch is %s, not searching for IMEX drivers", arch)
 		return "", nil
 	}
 	log.Println("Searching for IMEX driver to install")
 	envReader, err := envReaderFunc(ctx)
 	if err != nil {
 		return "", err
 	}
 	downloader := cos.NewGCSDownloader(gcsClient, envReader, "", "", "", "")
 	result, err := downloader.DownloadImexDriver(ctx, outDir, generatedDriver.Version)
 	if err != nil {
 		log.Printf("WARNING: could not download IMEX driver for %s", generatedDriver.Version)
 		return "", nil
 	}
 	return result, nil
 }

 func updateVersionInfo(versionsProto []byte, driverVersion, outDir string) (string, error) {
 	driverVersionInfo, err := gpuconfig.ParseGPUDriverVersionInfoList(versionsProto)
 	if err != nil {
 		return "", err
 	}
 	for _, gpuVersionInfo := range driverVersionInfo.GetGpuDriverVersionInfo() {
 		found := false
 		for _, v := range gpuVersionInfo.SupportedDriverVersions {
 			if v.Version == driverVersion {
 				found = true
 				break
 			}
 		}
 		if !found {
 			gpuVersionInfo.SupportedDriverVersions = append(gpuVersionInfo.SupportedDriverVersions, &pb.DriverVersion{Version: driverVersion})
 		}
 	}
 	resultBytes, err := proto.Marshal(driverVersionInfo)
 	if err != nil {
 		return "", err
 	}
 	result := filepath.Join(outDir, "gpu_driver_versions.bin")
 	if err := os.WriteFile(result, resultBytes, 0644); err != nil {
 		return "", err
 	}
 	return result, nil
 }

 // installDriverVersionsProto adds the version we just compiled to the
 // gpu_driver_versions.bin associated with the input image. For each GPU type,
 // add the version to it, if the version is not present at all for that GPU
 // type.
 //
 // Some older COS images do not have a gpu_driver_versions.bin, and those
 // versions do not require one. Do not install anything in this case and return
 // an empty string.
 func installDriverVersionsProto(ctx context.Context, envReaderFunc envReaderCreator, gcsClient *storage.Client, gpuConfig *config.GPUConfig, driverVersion, outDir string) (string, error) {
 	if gpuConfig.VersionsProtoPath != "" {
 		log.Println("Using provided gpu_driver_versions.bin")
 		versionsProto, err := os.ReadFile(gpuConfig.VersionsProtoPath)
 		if err != nil {
 			return "", err
 		}
 		return updateVersionInfo(versionsProto, driverVersion, outDir)
 	}
 	log.Println("No provided gpu_driver_versions.bin, looking for the one associated with the input image")
 	envReader, err := envReaderFunc(ctx)
 	if err != nil {
 		return "", err
 	}
 	downloader := cos.NewGCSDownloader(gcsClient, envReader, "", "", "", "")
 	if exists, err := downloader.ArtifactExists(ctx, "gpu_driver_versions.bin"); err != nil || !exists {
 		if err != nil {
 			log.Printf("WARNING: error searching for gpu_driver_versions.bin for input COS image: %v", err)
 		}
 		log.Println("WARNING: could not find gpu_driver_versions.bin for input image, will not provide one to cos-gpu-installer")
 		return "", nil
 	}
 	tmpDir, err := os.MkdirTemp("", "image-driver-versions")
 	if err != nil {
 		return "", err
 	}
 	defer os.RemoveAll(tmpDir)
 	if err := downloader.DownloadArtifact(ctx, tmpDir, "gpu_driver_versions.bin"); err != nil {
 		return "", err
 	}
 	versionsProto, err := os.ReadFile(filepath.Join(tmpDir, "gpu_driver_versions.bin"))
 	if err != nil {
 		return "", err
 	}
 	return updateVersionInfo(versionsProto, driverVersion, outDir)
 }

 func buildGPUPackage(ctx context.Context, input *config.Image, gcs *gcsManager, files *cosfs.Files, buildSpec *config.Build, workDir string) error {
 	log.Println("Building GPU package")
 	envReaderFunc := cachedEnvReader(input, gcs, files, buildSpec)
 	toolchainDir := filepath.Join(workDir, "toolchain")
 	if err := os.Mkdir(toolchainDir, 0755); err != nil {
 		return err
 	}
 	if err := installToolchain(ctx, envReaderFunc, gcs.gcsClient, buildSpec, toolchainDir); err != nil {
 		return fmt.Errorf("unable to install toolchain: %v", err)
 	}
 	arch, err := getArchFromToolchain(toolchainDir)
 	if err != nil {
 		return err
 	}
 	var runfile string
 	if filepath.Ext(buildSpec.GPUConfig.Version) == ".run" {
 		runfile = buildSpec.GPUConfig.Version
 	} else {
 		runfile, err = downloadRunfile(buildSpec.GPUConfig.Version, arch, buildSpec.GPUConfig.MD5Sum, workDir)
 		if err != nil {
 			return err
 		}
 	}
 	if err := os.Chmod(runfile, 0755); err != nil {
 		return err
 	}
 	driverPackage, err := gpubuild.Compile(runfile, toolchainDir, workDir)
 	if err != nil {
 		return err
 	}
 	mergedDir := filepath.Join(workDir, "merged")
 	if err := os.Mkdir(mergedDir, 0755); err != nil {
 		return err
 	}
 	mergedPackage, err := installNonRunfileModules(ctx, envReaderFunc, gcs.gcsClient, buildSpec.GPUConfig, driverPackage, mergedDir)
 	if err != nil {
 		return err
 	}
 	imexPath, err := downloadImexDriver(ctx, envReaderFunc, gcs.gcsClient, driverPackage, arch, mergedDir)
 	if err != nil {
 		return err
 	}
 	generatedProtoPath, err := installDriverVersionsProto(ctx, envReaderFunc, gcs.gcsClient, buildSpec.GPUConfig, driverPackage.Version, mergedDir)
 	if err != nil {
 		return err
 	}
 	buildSpec.GCSFiles = append(buildSpec.GCSFiles, runfile, mergedPackage.Path)
 	if imexPath != "" {
 		buildSpec.GCSFiles = append(buildSpec.GCSFiles, imexPath)
 	}
 	if generatedProtoPath != "" {
 		buildSpec.GCSFiles = append(buildSpec.GCSFiles, generatedProtoPath)
 	}
 	buildSpec.GPUConfig.Version = driverPackage.Version
 	return nil
 }
	// Copyright 2025 Google LLC
	//
	// 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 preloader

	import (
	"context"
	"crypto/md5"
	"fmt"
	"io"
	"io/fs"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"

	"cos.googlesource.com/cos/tools.git/src/pkg/config"
	"cos.googlesource.com/cos/tools.git/src/pkg/cos"
	cosfs "cos.googlesource.com/cos/tools.git/src/pkg/fs"
	"cos.googlesource.com/cos/tools.git/src/pkg/gpubuild"
	"cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig"
	"cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig/pb"
	"cos.googlesource.com/cos/tools.git/src/pkg/utils"

	"cloud.google.com/go/storage"
	"google.golang.org/protobuf/proto"
	)

	func extractTarball(srcPath, destDir string) error {
	cmd := exec.Command("tar", "xf", srcPath, "-C", destDir)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
	}

	type envReaderCreator func(context.Context) (*cos.EnvReader, error)

	// cachedEnvReader returns a closure that creates an EnvReader from a VMImage.
	// Since creating an EnvReader is very expensive in this context (it boots a VM
	// from the image to read data from it), we use a closure to cache the resulting
	// envReader.
	//
	// We could create the EnvReader once and pass it around everywhere. However,
	// under specific conditions, the EnvReader won't need to be created at all,
	// so this closure allows us to lazily create the EnvReader when it is needed.
	func cachedEnvReader(input config.Image, gcs gcsManager, files cosfs.Files, buildSpec config.Build) envReaderCreator {
	var envReader *cos.EnvReader
	return func(ctx context.Context) (*cos.EnvReader, error) {
	var err error
	if envReader != nil {
	return envReader, nil
	}
	envReader, err = cos.NewEnvReaderFromVMImage(input.Name, input.Project, cos.EnvReaderVMImageConfig{
	DaisyBin: files.DaisyBin,
	DaisyWorkflow: "/data/read_cos_env.wf.json",
	Project: buildSpec.Project,
	Zone: buildSpec.Zone,
	GCSDir: gcs.managedDirURL(),
	DiskType: buildSpec.DiskType,
	MachineType: buildSpec.MachineType,
	Network: buildSpec.Network,
	SubNet: buildSpec.Subnet,
	ServiceAccount: buildSpec.ServiceAccount,
	})
	if err != nil {
	return nil, fmt.Errorf("failed to read env from VM image: %v", err)
	}
	return envReader, nil
	}
	}

	func installToolchain(ctx context.Context, envReaderFunc envReaderCreator, gcsClient storage.Client, buildSpec config.Build, toolchainDir string) error {
	if buildSpec.Toolchain != nil {
	log.Println("Extracting toolchain package")
	if err := extractTarball(buildSpec.Toolchain.ToolchainTarPath, toolchainDir); err != nil {
	return err
	}
	if err := extractTarball(buildSpec.Toolchain.KernelHeadersPath, toolchainDir); err != nil {
	return err
	}
	return nil
	}
	log.Println("No toolchain package provided; getting toolchain from input image")
	envReader, err := envReaderFunc(ctx)
	if err != nil {
	return err
	}
	downloader := cos.NewGCSDownloader(gcsClient, envReader, "", "", "", "")
	if err := cos.InstallCrossToolchain(ctx, downloader, toolchainDir); err != nil {
	return err
	}
	// InstallCrossToolchain has an undesirable side effect of setting the SYSROOT
	// environment variable. It's unclear if removing this behavior from
	// InstallCrossToolchain is safe for all usages.
	if err := os.Unsetenv("SYSROOT"); err != nil {
	log.Printf("WARNING: unable to unset SYSROOT: %v", err)
	}
	return nil
	}

	func getArchFromToolchain(toolchainDir string) (string, error) {
	x86Matches, err := filepath.Glob(filepath.Join(toolchainDir, "bin/x86_64-cros-linux-gnu-*"))
	if err != nil {
	return "", err
	}
	if len(x86Matches) != 0 {
	return "x86_64", nil
	}
	aarch64Matches, err := filepath.Glob(filepath.Join(toolchainDir, "bin/aarch64-cros-linux-gnu-*"))
	if err != nil {
	return "", err
	}
	if len(aarch64Matches) != 0 {
	return "aarch64", nil
	}
	return "", fmt.Errorf("could not identify arch from toolchain at %q; are programs missing from the toolchain?")
	}

	func downloadRunfile(version, arch, md5Sum, outDir string) (_ string, err error) {
	outName := fmt.Sprintf("NVIDIA-Linux-%s-%s.run", arch, version)
	url := fmt.Sprintf("https://us.download.nvidia.com/tesla/%s/%s", version, outName)
	log.Printf("Downloading runfile from %q", url)
	outPath := filepath.Join(outDir, outName)
	outFile, err := os.Create(outPath)
	if err != nil {
	return "", err
	}
	defer utils.CheckClose(outFile, "error closing downloaded runfile", &err)
	resp, err := http.Get(url)
	if err != nil {
	return "", err
	}
	defer utils.CheckClose(resp.Body, "error closing runfile HTTP body", &err)
	if resp.StatusCode != http.StatusOK {
	return "", fmt.Errorf("unhealthy status downloading Nvidia runfile: %v", resp.Status)
	}
	if md5Sum != "" {
	h := md5.New()
	mw := io.MultiWriter(outFile, h)
	if _, err := io.Copy(mw, resp.Body); err != nil {
	return "", err
	}
	got := fmt.Sprintf("%x", h.Sum(nil))
	if got != md5Sum {
	return "", fmt.Errorf("md5sum mismatch when downloading runfile from %q; got %q, want %q", url, got, md5Sum)
	}
	} else {
	if _, err := io.Copy(outFile, resp.Body); err != nil {
	return "", err
	}
	}
	return outPath, nil
	}

	func listModules(dir string) ([]string, error) {
	var result []string
	if err := filepath.WalkDir(dir, func(path string, _ fs.DirEntry, err error) error {
	if err != nil {
	return err
	}
	if filepath.Ext(path) == ".ko" {
	result = append(result, filepath.Base(path))
	}
	return nil
	}); err != nil {
	return nil, err
	}
	return result, nil
	}

	func setDiff(a, b []string) []string {
	aMap := map[string]bool{}
	for _, e := range a {
	aMap[e] = true
	}
	for _, e := range b {
	if _, ok := aMap[e]; ok {
	delete(aMap, e)
	}
	}
	var result []string
	for k, _ := range aMap {
	result = append(result, k)
	}
	return result
	}

	func mergeDriverPackages(generatedDriver gpubuild.DriverPackage, preBuiltPath, outDir string) (gpubuild.DriverPackage, error) {
	log.Println("Merging any missing precompiled modules from precompiled package into custom package")
	nvidiaGenerated, err := os.MkdirTemp("", "nvidia-generated")
	if err != nil {
	return nil, err
	}
	defer os.RemoveAll(nvidiaGenerated)
	if err := extractTarball(generatedDriver.Path, nvidiaGenerated); err != nil {
	return nil, err
	}
	nvidiaPreBuilt, err := os.MkdirTemp("", "prebuilt-generated")
	if err != nil {
	return nil, err
	}
	defer os.RemoveAll(nvidiaPreBuilt)
	if err := extractTarball(preBuiltPath, nvidiaPreBuilt); err != nil {
	return nil, err
	}
	generatedModules, err := listModules(nvidiaGenerated)
	if err != nil {
	return nil, err
	}
	log.Printf("Generated package includes: %v", generatedModules)
	preBuiltModules, err := listModules(nvidiaPreBuilt)
	if err != nil {
	return nil, err
	}
	log.Printf("Prebuilt package includes: %v", preBuiltModules)
	notGeneratedModules := setDiff(preBuiltModules, generatedModules)
	log.Printf("Adding the following to generated package: %v", notGeneratedModules)
	for _, module := range notGeneratedModules {
	src := filepath.Join(nvidiaPreBuilt, "drivers", module)
	dst := filepath.Join(nvidiaGenerated, "drivers", module)
	if err := utils.CopyFile(src, dst); err != nil {
	return nil, err
	}
	}
	result := &gpubuild.DriverPackage{
	Path: filepath.Join(outDir, fmt.Sprintf("nvidia-drivers-%s.tgz", generatedDriver.Version)),
	Version: generatedDriver.Version,
	}
	if err := cosfs.TarDir(nvidiaGenerated, result.Path); err != nil {
	return nil, err
	}
	return result, nil
	}

	// installNonRunfileModules searches for kernel modules in any existing
	// precompiled packages that are not present in the runfile we just compiled. If
	// it finds anything, it adds it to the generated package we just compiled. This
	// helps with supporting gdrcopy.
	func installNonRunfileModules(ctx context.Context, envReaderFunc envReaderCreator, gcsClient storage.Client, gpuConfig config.GPUConfig, generatedDriver gpubuild.DriverPackage, outDir string) (gpubuild.DriverPackage, error) {
	if gpuConfig.PrePackagedDriversPath != "" {
	return mergeDriverPackages(generatedDriver, gpuConfig.PrePackagedDriversPath, outDir)
	}
	log.Println("No precompiled driver package provided, looking for an existing driver package for the given VM image")
	envReader, err := envReaderFunc(ctx)
	if err != nil {
	return nil, err
	}
	downloader := cos.NewGCSDownloader(gcsClient, envReader, "", "", "", "")
	precompiled := fmt.Sprintf("nvidia-drivers-%s.tgz", generatedDriver.Version)
	if exists, err := downloader.ArtifactExists(ctx, precompiled); err != nil \|\| !exists {
	if err != nil {
	log.Printf("WARNING: error searching for non-runfile modules for version %q: %v", generatedDriver.Version, err)
	}
	log.Printf("WARNING: could not find non-runfile modules for version %q, will not install them", generatedDriver.Version)
	return generatedDriver, nil
	}
	tmpDir, err := os.MkdirTemp("", "non-runfile-modules")
	if err != nil {
	return nil, err
	}
	defer os.RemoveAll(tmpDir)
	if err := downloader.DownloadArtifact(ctx, tmpDir, precompiled); err != nil {
	return nil, err
	}
	precompiledPath := filepath.Join(tmpDir, precompiled)
	return mergeDriverPackages(generatedDriver, precompiledPath, outDir)
	}

	func downloadImexDriver(ctx context.Context, envReaderFunc envReaderCreator, gcsClient storage.Client, generatedDriver gpubuild.DriverPackage, arch, outDir string) (string, error) {
	if arch != "aarch64" {
	log.Printf("Arch is %s, not searching for IMEX drivers", arch)
	return "", nil
	}
	log.Println("Searching for IMEX driver to install")
	envReader, err := envReaderFunc(ctx)
	if err != nil {
	return "", err
	}
	downloader := cos.NewGCSDownloader(gcsClient, envReader, "", "", "", "")
	result, err := downloader.DownloadImexDriver(ctx, outDir, generatedDriver.Version)
	if err != nil {
	log.Printf("WARNING: could not download IMEX driver for %s", generatedDriver.Version)
	return "", nil
	}
	return result, nil
	}

	func updateVersionInfo(versionsProto []byte, driverVersion, outDir string) (string, error) {
	driverVersionInfo, err := gpuconfig.ParseGPUDriverVersionInfoList(versionsProto)
	if err != nil {
	return "", err
	}
	for _, gpuVersionInfo := range driverVersionInfo.GetGpuDriverVersionInfo() {
	found := false
	for _, v := range gpuVersionInfo.SupportedDriverVersions {
	if v.Version == driverVersion {
	found = true
	break
	}
	}
	if !found {
	gpuVersionInfo.SupportedDriverVersions = append(gpuVersionInfo.SupportedDriverVersions, &pb.DriverVersion{Version: driverVersion})
	}
	}
	resultBytes, err := proto.Marshal(driverVersionInfo)
	if err != nil {
	return "", err
	}
	result := filepath.Join(outDir, "gpu_driver_versions.bin")
	if err := os.WriteFile(result, resultBytes, 0644); err != nil {
	return "", err
	}
	return result, nil
	}

	// installDriverVersionsProto adds the version we just compiled to the
	// gpu_driver_versions.bin associated with the input image. For each GPU type,
	// add the version to it, if the version is not present at all for that GPU
	// type.
	//
	// Some older COS images do not have a gpu_driver_versions.bin, and those
	// versions do not require one. Do not install anything in this case and return
	// an empty string.
	func installDriverVersionsProto(ctx context.Context, envReaderFunc envReaderCreator, gcsClient storage.Client, gpuConfig config.GPUConfig, driverVersion, outDir string) (string, error) {
	if gpuConfig.VersionsProtoPath != "" {
	log.Println("Using provided gpu_driver_versions.bin")
	versionsProto, err := os.ReadFile(gpuConfig.VersionsProtoPath)
	if err != nil {
	return "", err
	}
	return updateVersionInfo(versionsProto, driverVersion, outDir)
	}
	log.Println("No provided gpu_driver_versions.bin, looking for the one associated with the input image")
	envReader, err := envReaderFunc(ctx)
	if err != nil {
	return "", err
	}
	downloader := cos.NewGCSDownloader(gcsClient, envReader, "", "", "", "")
	if exists, err := downloader.ArtifactExists(ctx, "gpu_driver_versions.bin"); err != nil \|\| !exists {
	if err != nil {
	log.Printf("WARNING: error searching for gpu_driver_versions.bin for input COS image: %v", err)
	}
	log.Println("WARNING: could not find gpu_driver_versions.bin for input image, will not provide one to cos-gpu-installer")
	return "", nil
	}
	tmpDir, err := os.MkdirTemp("", "image-driver-versions")
	if err != nil {
	return "", err
	}
	defer os.RemoveAll(tmpDir)
	if err := downloader.DownloadArtifact(ctx, tmpDir, "gpu_driver_versions.bin"); err != nil {
	return "", err
	}
	versionsProto, err := os.ReadFile(filepath.Join(tmpDir, "gpu_driver_versions.bin"))
	if err != nil {
	return "", err
	}
	return updateVersionInfo(versionsProto, driverVersion, outDir)
	}

	func buildGPUPackage(ctx context.Context, input config.Image, gcs gcsManager, files cosfs.Files, buildSpec config.Build, workDir string) error {
	log.Println("Building GPU package")
	envReaderFunc := cachedEnvReader(input, gcs, files, buildSpec)
	toolchainDir := filepath.Join(workDir, "toolchain")
	if err := os.Mkdir(toolchainDir, 0755); err != nil {
	return err
	}
	if err := installToolchain(ctx, envReaderFunc, gcs.gcsClient, buildSpec, toolchainDir); err != nil {
	return fmt.Errorf("unable to install toolchain: %v", err)
	}
	arch, err := getArchFromToolchain(toolchainDir)
	if err != nil {
	return err
	}
	var runfile string
	if filepath.Ext(buildSpec.GPUConfig.Version) == ".run" {
	runfile = buildSpec.GPUConfig.Version
	} else {
	runfile, err = downloadRunfile(buildSpec.GPUConfig.Version, arch, buildSpec.GPUConfig.MD5Sum, workDir)
	if err != nil {
	return err
	}
	}
	if err := os.Chmod(runfile, 0755); err != nil {
	return err
	}
	driverPackage, err := gpubuild.Compile(runfile, toolchainDir, workDir)
	if err != nil {
	return err
	}
	mergedDir := filepath.Join(workDir, "merged")
	if err := os.Mkdir(mergedDir, 0755); err != nil {
	return err
	}
	mergedPackage, err := installNonRunfileModules(ctx, envReaderFunc, gcs.gcsClient, buildSpec.GPUConfig, driverPackage, mergedDir)
	if err != nil {
	return err
	}
	imexPath, err := downloadImexDriver(ctx, envReaderFunc, gcs.gcsClient, driverPackage, arch, mergedDir)
	if err != nil {
	return err
	}
	generatedProtoPath, err := installDriverVersionsProto(ctx, envReaderFunc, gcs.gcsClient, buildSpec.GPUConfig, driverPackage.Version, mergedDir)
	if err != nil {
	return err
	}
	buildSpec.GCSFiles = append(buildSpec.GCSFiles, runfile, mergedPackage.Path)
	if imexPath != "" {
	buildSpec.GCSFiles = append(buildSpec.GCSFiles, imexPath)
	}
	if generatedProtoPath != "" {
	buildSpec.GCSFiles = append(buildSpec.GCSFiles, generatedProtoPath)
	}
	buildSpec.GPUConfig.Version = driverPackage.Version
	return nil
	}