cos-customizer: add flag -gpu-type to finish-image-build

To allow users to use `cos-extensions install gpu` for gpu
driver installation via `run-script`, we need to add a flag
for attaching GPU to the preloading VM.

Deprecated the same flag in step `install-gpu`.

TEST=run a GCB workflow with run-script, script content:
`sudo cos-extensions install gpu` and verify the output
image have the driver installed

Change-Id: Iad8ff1e1cf30b8804a243fa0b375b07156253aa5
Tested-by: He Gao <>
Reviewed-by: Robert Kolchmeyer <>
Cloud-Build: GCB Service account <>
diff --git a/src/cmd/cos_customizer/ b/src/cmd/cos_customizer/
index 79bc794..9247669 100644
--- a/src/cmd/cos_customizer/
+++ b/src/cmd/cos_customizer/
@@ -272,6 +272,11 @@
 separate from the Cloud Build machine type option, which sets the machine type
 of the Cloud Build VM, which is different from the COS Customizer preload VM.
+`-gpu-type`: The GPU type attached to the COS Customizer preload VM e.g.
+`nvidia-tesla-t4`. Useful for installing GPU drivers via subcommand
+`run-script`. If this flag is set, it will overwrite the same flag set in
+step `install-gpu`.
 `-network`: The network/VPC to use for the COS Customizer preload VM.
 The network must have access to Google Cloud Storage. Defaults to
 default network `global/networks/default`.  If -subnet is also specified subnet 
@@ -424,12 +429,8 @@
 choose a directory that will persist across reboots; for the most part, this
 means a subdirectory of `/var` or `/home`.
-`-gpu-type`: The type of GPU to use to verify correct installation of GPU
-drivers. The valid values here are nvidia-tesla-k80, nvidia-tesla-p100, and
-nvidia-tesla-v100. This value has no impact on the drivers that are installed on
-the image; it is only used when verifying that the driver installation
-succeeded. Make sure that the zone you are running the image build in has quota
-for a GPU of this type.
+`-gpu-type`: This flag is deprecated, please use the same flag in step
 An example `install-gpu` step looks like the following:
diff --git a/src/cmd/cos_customizer/finish_image_build.go b/src/cmd/cos_customizer/finish_image_build.go
index 0a1a8a0..1fb5044 100644
--- a/src/cmd/cos_customizer/finish_image_build.go
+++ b/src/cmd/cos_customizer/finish_image_build.go
@@ -48,6 +48,7 @@
 	zone           string
 	project        string
 	machineType    string
+	gpuType        string
 	imageName      string
 	imageSuffix    string
 	imageFamily    string
@@ -100,6 +101,7 @@
 	flags.StringVar(&, "zone", "", "Zone to make GCE resources in.")
 	flags.StringVar(&f.project, "project", "", "Project to make GCE resources in.")
 	flags.StringVar(&f.machineType, "machine-type", "n1-standard-1", "Machine type to use during the build.")
+	flags.StringVar(&f.gpuType, "gpu-type", "", "GPU type e.g. 'nvidia-tesla-t4' used for GPU driver installation via subcommand 'run-script'.")
 	flags.StringVar(&, "network", "", "Network to use"+
 		" during the build. The network must have access to Google Cloud Storage."+
 		" If not specified, the network named default is used."+
@@ -187,6 +189,9 @@
 	buildConfig.Network =
 	buildConfig.Subnet = f.subnet
 	buildConfig.Timeout = f.timeout.String()
+	if f.gpuType != "" {
+		buildConfig.GPUType = f.gpuType
+	}
 	provConfig := &provisioner.Config{}
 	if err := config.LoadFromFile(files.ProvConfig, provConfig); err != nil {
 		return nil, nil, nil, nil, err
diff --git a/src/cmd/cos_customizer/install_gpu.go b/src/cmd/cos_customizer/install_gpu.go
index 317d932..6874500 100644
--- a/src/cmd/cos_customizer/install_gpu.go
+++ b/src/cmd/cos_customizer/install_gpu.go
@@ -83,8 +83,7 @@
 	f.StringVar(&i.NvidiaInstallDirHost, "install-dir", "/var/lib/nvidia",
 		"Location to install drivers on the image.")
-		&i.gpuType, "gpu-type", "nvidia-tesla-p100",
-		fmt.Sprintf("The type of GPU to verify drivers for. Must be one of: %v", validGPUs))
+		&i.gpuType, "gpu-type", "nvidia-tesla-p100", "This flag is deprecated, please use the same flag in step 'finish-image-build'.")
 		&i.getValidDrivers, "get-valid-drivers", false,
 		"Print the list of supported GPU driver versions. If this flag is given, no other actions will be taken.")