cos_gpu_installer: add --prepare-build-tools flag which populates toolchain cache

This flag can be used to populate the build tools cache. This includes
downloading and installing the toolchain and the kernel headers which
are a pre-requisite to the nvidia driver installer. Using this flag will NOT
install the driver and this will not depend on existence of a GPU
attached to the instance.

BUG=b/175236375
TEST=manual invokation of installer with --prepare-build-tools

Change-Id: I7414efdb8ee9ff380358fe3e6cd86f3c232778ab
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/30822
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: He Gao <hegao@google.com>
Tested-by: Arnav Kansal <rnv@google.com>
diff --git a/src/cmd/cos_gpu_installer/internal/commands/install.go b/src/cmd/cos_gpu_installer/internal/commands/install.go
index 5a4a50d..543ba72 100644
--- a/src/cmd/cos_gpu_installer/internal/commands/install.go
+++ b/src/cmd/cos_gpu_installer/internal/commands/install.go
@@ -43,6 +43,7 @@
 	signatureURL       string
 	debug              bool
 	test               bool
+	prepareBuildTools  bool
 }
 
 // Name implements subcommands.Command.Name.
@@ -84,6 +85,8 @@
 	f.BoolVar(&c.test, "test", false,
 		"Enable test mode. "+
 			"In test mode, `-nvidia-installer-url` can be used without `-allow-unsigned-driver`.")
+	f.BoolVar(&c.prepareBuildTools, "prepare-build-tools", false, "Whether to populate the build tools cache, i.e. to download and install the toolchain and the kernel headers. Drivers are NOT installed when this flag is set and running with this flag does not require GPU attached to the instance.")
+
 }
 
 func (c *InstallCommand) validateFlags() error {
@@ -124,15 +127,17 @@
 		return subcommands.ExitFailure
 	}
 
-	var isGpuConfigured bool
-	if isGpuConfigured, err = c.isGpuConfigured(); err != nil {
-		c.logError(errors.Wrapf(err, "failed to check if GPU is configured"))
-		return subcommands.ExitFailure
-	}
+	if !c.prepareBuildTools {
+		var isGpuConfigured bool
+		if isGpuConfigured, err = c.isGpuConfigured(); err != nil {
+			c.logError(errors.Wrapf(err, "failed to check if GPU is configured"))
+			return subcommands.ExitFailure
+		}
 
-	if !isGpuConfigured {
-		c.logError(fmt.Errorf("Please have GPU device configured"))
-		return subcommands.ExitFailure
+		if !isGpuConfigured {
+			c.logError(fmt.Errorf("Please have GPU device configured"))
+			return subcommands.ExitFailure
+		}
 	}
 
 	downloader := cos.NewGCSDownloader(envReader, c.gcsDownloadBucket, c.gcsDownloadPrefix)
@@ -271,6 +276,11 @@
 		return errors.Wrap(err, "failed to install toolchain")
 	}
 
+	// Skip driver installation if we are only populating build tools cache
+	if c.prepareBuildTools {
+		return nil
+	}
+
 	if err := installer.RunDriverInstaller(toolchainPkgDir, installerFile, !c.unsignedDriver, false); err != nil {
 		if errors.Is(err, installer.ErrDriverLoad) {
 			// Drivers were linked, but couldn't load; try again with legacy linking