Merge pull request #122 from elfinhe/main

Add flags for disk type and machine type; Update default values to im…

gke-disk-image-builder/cli/main.go

@@ -42,7 +42,9 @@ func main() {
 	imageName := flag.String("image-name", "", "name of the image that will be generated")
 	zone := flag.String("zone", "", "zone where the resources will be used to create the image creator resources")
 	gcsPath := flag.String("gcs-path", "", "gcs location to dump the logs")
-	diskSizeGb := flag.Int64("disk-size-gb", 10, "size of a disk that will host the unpacked images")
+	machineType := flag.String("machine-type", "n2-standard-16", "GCE instance machine type to generate the disk image")
+	diskType := flag.String("disk-type", "pd-ssd", "disk type to generate the disk image")
+	diskSizeGb := flag.Int64("disk-size-gb", 60, "disk size to unpack container images")
 	gcpOAuth := flag.String("gcp-oauth", "", "path to GCP service account credential file")
 	imagePullAuth := flag.String("image-pull-auth", "None", "auth mechanism to pull the container image, valid values: [None, ServiceAccountToken].\nNone means that the images are publically available and no authentication is required to pull them.\nServiceAccountToken means the service account oauth token will be used to pull the images.\nFor more information refer to https://cloud.google.com/compute/docs/access/authenticate-workloads#applications")
 	timeout := flag.String("timeout", "20m", "Default timout for each step, defaults to 20m")
@@ -73,6 +75,8 @@ func main() {
 		ProjectName:     *projectName,
 		Zone:            *zone,
 		GCSPath:         *gcsPath,
+		MachineType:     *machineType,
+		DiskType:        *diskType,
 		DiskSizeGB:      *diskSizeGb,
 		GCPOAuth:        *gcpOAuth,
 		ContainerImages: containerImages,

gke-disk-image-builder/imager.go

@@ -51,6 +51,8 @@ type Request struct {
 	ProjectName     string
 	Zone            string
 	GCSPath         string
+	MachineType     string
+	DiskType        string
 	DiskSizeGB      int64
 	GCPOAuth        string
 	ContainerImages []string
@@ -107,7 +109,7 @@ func GenerateDiskImage(ctx context.Context, req Request) error {
 					},
 					Disk: compute.Disk{
 						Name:   fmt.Sprintf("%s-disk", name),
-						Type:   "pd-balanced",
+						Type:   req.DiskType,
 						SizeGb: req.DiskSizeGB,
 					},
 				},
@@ -124,7 +126,8 @@ func GenerateDiskImage(ctx context.Context, req Request) error {
 							StartupScript: "startup.sh",
 						},
 						Instance: compute.Instance{
-							Name: fmt.Sprintf("%s-instance", name),
+							Name:        fmt.Sprintf("%s-instance", name),
+							MachineType: fmt.Sprintf("zones/%s/machineTypes/%s", req.Zone, req.MachineType),
 							Disks: []*compute.AttachedDisk{
 								&compute.AttachedDisk{
 									AutoDelete: true,
@@ -134,7 +137,7 @@ func GenerateDiskImage(ctx context.Context, req Request) error {
 									Mode:       "READ_WRITE",
 									InitializeParams: &compute.AttachedDiskInitializeParams{
 										DiskSizeGb:  req.DiskSizeGB,
-										DiskType:    fmt.Sprintf("projects/%s/zones/%s/diskTypes/pd-balanced", req.ProjectName, req.Zone),
+										DiskType:    fmt.Sprintf("projects/%s/zones/%s/diskTypes/%s", req.ProjectName, req.Zone, req.DiskType),
 										SourceImage: "projects/debian-cloud/global/images/debian-11-bullseye-v20230912",
 									},
 								},