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`.
BUG=b/229285915
TEST=run a GCB workflow with run-script, script content:
`sudo cos-extensions install gpu` and verify the output
image have the driver installed
RELEASE_NOTE=None
Change-Id: Iad8ff1e1cf30b8804a243fa0b375b07156253aa5
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/49530
Tested-by: He Gao <hegao@google.com>
Reviewed-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
diff --git a/src/cmd/cos_customizer/README.md b/src/cmd/cos_customizer/README.md
index 79bc794..9247669 100644
--- a/src/cmd/cos_customizer/README.md
+++ b/src/cmd/cos_customizer/README.md
@@ -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
+`finish-image-build`.
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(&f.zone, "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(&f.network, "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 = f.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.")
f.StringVar(
- &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'.")
f.BoolVar(
&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.")
