cos-customizer: Add -disk-type flag
This new flag allows users to select the disk type of the preload VM.
BUG=b/229419550
TEST=./run_tests.sh
Change-Id: Ifd1908aa794a0474bd3682cc35e344cae9578be9
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/32087
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Tested-by: Robert Kolchmeyer <rkolchmeyer@google.com>
diff --git a/src/cmd/cos_customizer/README.md b/src/cmd/cos_customizer/README.md
index e9d772e..80d83d3 100644
--- a/src/cmd/cos_customizer/README.md
+++ b/src/cmd/cos_customizer/README.md
@@ -223,6 +223,9 @@
image labels present on the source image. The labels specified by the `-labels`
flag take precedence over labels assigned with this flag.
+`-disk-type`: The disk type for the preload VM to use when creating the image.
+Defaults to `pd-standard`.
+
`-disk-size-gb`: The disk size in GB to use when creating the image.
This value should never be smaller than 10 (the default size of a COS image).
If `-oem-size` is set, the lower limit of `-disk-size-gb` is as shown in the
diff --git a/src/cmd/cos_customizer/finish_image_build.go b/src/cmd/cos_customizer/finish_image_build.go
index 51232f5..2c46f7f 100644
--- a/src/cmd/cos_customizer/finish_image_build.go
+++ b/src/cmd/cos_customizer/finish_image_build.go
@@ -54,6 +54,7 @@
inheritLabels bool
oemSize string
oemFSSize4K uint64
+ diskType string
diskSize int
timeout time.Duration
}
@@ -114,6 +115,7 @@
flags.StringVar(&f.oemSize, "oem-size", "", "Size of the new OEM partition, "+
"can be a number with unit like 10G, 10M, 10K or 10B, "+
"or without unit indicating the number of 512B sectors.")
+ flags.StringVar(&f.diskType, "disk-type", "pd-standard", "The disk type to use when creating the image.")
flags.IntVar(&f.diskSize, "disk-size-gb", 0, "The disk size to use when creating the image in GB. Value of '0' "+
"indicates the default size.")
flags.DurationVar(&f.timeout, "timeout", time.Hour, "Timeout value of the image build process. Must be formatted "+
@@ -166,6 +168,7 @@
buildConfig.Project = f.project
buildConfig.Zone = f.zone
buildConfig.MachineType = f.machineType
+ buildConfig.DiskType = f.diskType
buildConfig.DiskSize = f.diskSize
buildConfig.Network = f.network
buildConfig.Subnet = f.subnet
diff --git a/src/data/build_image.wf.json b/src/data/build_image.wf.json
index 4c96e3d..1ac7a67 100644
--- a/src/data/build_image.wf.json
+++ b/src/data/build_image.wf.json
@@ -6,6 +6,7 @@
"output_image_family": {"Value": "", "Description": "Family of output image."},
"output_image_project": {"Required": true, "Description": "Project of output image."},
"cidata_img": {"Required": true, "Description": "Path to CIDATA vfat image containing cloud-init user-data and the provisioner program. Must be in .tar.gz format."},
+ "disk_type": {"Value": "pd-standard", "Description": "The disk type to use for preloading."},
"disk_size_gb": {"Value": "10", "Description": "The disk size to use for preloading."},
"host_maintenance": {"Value": "MIGRATE", "Description": "VM behavior when there is maintenance."},
"machine_type": {"Required": true, "Description": "Machine type of the preload VM."},
@@ -40,11 +41,13 @@
{
"Name": "boot-disk",
"SourceImage": "${source_image}",
+ "Type": "${disk_type}",
"SizeGb": "${disk_size_gb}"
},
{
"Name": "cidata-disk",
- "SourceImage": "cidata"
+ "SourceImage": "cidata",
+ "Type": "${disk_type}"
}
]
},
diff --git a/src/pkg/config/config.go b/src/pkg/config/config.go
index 57debde..f60752b 100644
--- a/src/pkg/config/config.go
+++ b/src/pkg/config/config.go
@@ -67,6 +67,7 @@
Project string
Zone string
MachineType string
+ DiskType string
DiskSize int
GPUType string
Timeout string
diff --git a/src/pkg/preloader/BUILD.bazel b/src/pkg/preloader/BUILD.bazel
index c7e110a..b6b088a 100644
--- a/src/pkg/preloader/BUILD.bazel
+++ b/src/pkg/preloader/BUILD.bazel
@@ -62,6 +62,7 @@
go_test(
name = "preloader_test",
+ size = "large",
srcs = [
"gcs_test.go",
"preload_test.go",
diff --git a/src/pkg/preloader/preload.go b/src/pkg/preloader/preload.go
index 1e8973f..cae3ffe 100644
--- a/src/pkg/preloader/preload.go
+++ b/src/pkg/preloader/preload.go
@@ -308,6 +308,8 @@
ciDataFile,
"-var:machine_type",
buildSpec.MachineType,
+ "-var:disk_type",
+ buildSpec.DiskType,
"-var:host_maintenance",
hostMaintenance,
"-var:network",
diff --git a/src/pkg/preloader/preload_test.go b/src/pkg/preloader/preload_test.go
index 2f7f1ee..7de309a 100644
--- a/src/pkg/preloader/preload_test.go
+++ b/src/pkg/preloader/preload_test.go
@@ -316,6 +316,13 @@
want: []string{"-var:cidata_img"},
},
{
+ testName: "DiskType",
+ inputImage: config.NewImage("", ""),
+ outputImage: config.NewImage("", ""),
+ buildConfig: &config.Build{DiskType: "pd-ssd", GCSBucket: "bucket", GCSDir: "dir"},
+ want: []string{"-var:disk_type", "pd-ssd"},
+ },
+ {
testName: "DiskSize",
inputImage: config.NewImage("", ""),
outputImage: config.NewImage("", ""),