roachprod: add GCE hyperdisk-balanced support and misc improvements

Add support for GCE hyperdisk-balanced disk type with new flags:
- --gce-pd-volume-provisioned-iops: Required for hyperdisk-balanced,
  optional for pd-extreme
- --gce-pd-volume-provisioned-throughput: Required for hyperdisk-balanced

Additional improvements:
- Move google_compute_engine to front of SSH key list to avoid
  passphrase prompts when connecting to GCE VMs (the key is typically
  not passphrase-protected)
- Increase AWS gp3 max IOPS validation from 16000 to 80000
- Add AWS m8g/m8gd ARM instance type detection

Release note: None
Epic: None

Author: petermattis

pkg/roachprod/config/config.go

@@ -190,13 +190,16 @@ var localClusterRegex = regexp.MustCompile(`^local(|-[a-zA-Z0-9\-]+)$`)
 // establishing a remote session.
 // See https://github.com/openssh/openssh-portable/blob/86bdd385/ssh_config.5#L1123-L1130
 var DefaultPubKeyNames = []string{
+	// google_compute_engine is listed first because it's typically not
+	// passphrase-protected, avoiding passphrase prompts when connecting to GCE
+	// VMs if the ssh-agent is slow to respond.
+	"google_compute_engine",
 	"id_rsa",
 	"id_ecdsa",
 	"id_ecdsa_sk",
 	"id_ed25519",
 	"id_ed25519_sk",
 	"id_dsa",
-	"google_compute_engine",
 }
 
 // SSHPublicKeyPath returns the path to the default public key expected by

pkg/roachprod/vm/aws/aws.go

@@ -166,8 +166,8 @@ func (d *ebsDisk) Set(s string) error {
 	case "gp2":
 		// Nothing -- size checked above.
 	case "gp3":
-		if d.IOPs > 16000 {
-			return errors.AssertionFailedf("Iops required for gp3 disk: [3000, 16000]")
+		if d.IOPs > 80000 {
+			return errors.AssertionFailedf("Iops required for gp3 disk: [3000, 80000]")
 		}
 		if d.IOPs == 0 {
 			// 3000 is a base IOPs for gp3.
@@ -1378,7 +1378,8 @@ func (p *Provider) runInstance(
 	}
 	imageID := withFlagOverride(az.Region.AMI_X86_64, &providerOpts.ImageAMI)
 	useArmAMI := strings.Index(machineType, "6g.") == 1 || strings.Index(machineType, "6gd.") == 1 ||
-		strings.Index(machineType, "7g.") == 1 || strings.Index(machineType, "7gd.") == 1
+		strings.Index(machineType, "7g.") == 1 || strings.Index(machineType, "7gd.") == 1 ||
+		strings.Index(machineType, "8g.") == 1 || strings.Index(machineType, "8gd.") == 1
 	if useArmAMI && (opts.Arch != "" && opts.Arch != string(vm.ArchARM64)) {
 		return nil, errors.Errorf("machine type %s is arm64, but requested arch is %s", machineType, opts.Arch)
 	}

pkg/roachprod/vm/gce/gcloud.go

@@ -347,16 +347,18 @@ type ProviderOpts struct {
 	// projects represent the GCE projects to operate on. Accessed through
 	// GetProject() or GetProjects() depending on whether the command accepts
 	// multiple projects or a single one.
-	MachineType      string
-	MinCPUPlatform   string
-	BootDiskType     string
-	Zones            []string
-	Image            string
-	SSDCount         int
-	PDVolumeType     string
-	PDVolumeSize     int
-	PDVolumeCount    int
-	UseMultipleDisks bool
+	MachineType                   string
+	MinCPUPlatform                string
+	BootDiskType                  string
+	Zones                         []string
+	Image                         string
+	SSDCount                      int
+	PDVolumeType                  string
+	PDVolumeSize                  int
+	PDVolumeCount                 int
+	PDVolumeProvisionedIOPS       int
+	PDVolumeProvisionedThroughput int
+	UseMultipleDisks              bool
 	// use spot instances (i.e., latest version of preemptibles which can run > 24 hours)
 	UseSpot bool
 	// Use an instance template and a managed instance group to create VMs. This
@@ -740,12 +742,12 @@ func (p *Provider) CreateVolume(
 	}
 
 	switch vco.Type {
-	case "local-ssd", "pd-balanced", "pd-extreme", "pd-ssd", "pd-standard":
+	case "local-ssd", "pd-balanced", "pd-extreme", "pd-ssd", "pd-standard", "hyperdisk-balanced":
 		args = append(args, "--type", vco.Type)
 	case "":
 	// use the default
 	default:
-		return vol, errors.Newf("Expected type to be one of local-ssd, pd-balanced, pd-extreme, pd-ssd, pd-standard got %s\n", vco.Type)
+		return vol, errors.Newf("Expected type to be one of local-ssd, pd-balanced, pd-extreme, pd-ssd, pd-standard, hyperdisk-balanced got %s\n", vco.Type)
 	}
 
 	var commandResponse []describeVolumeCommandResponse
@@ -1140,11 +1142,16 @@ func (o *ProviderOpts) ConfigureCreateFlags(flags *pflag.FlagSet) {
 	flags.IntVar(&o.SSDCount, ProviderName+"-local-ssd-count", 1,
 		"Number of local SSDs to create, only used if local-ssd=true")
 	flags.StringVar(&o.PDVolumeType, ProviderName+"-pd-volume-type", "pd-ssd",
-		"Type of the persistent disk volume, only used if local-ssd=false")
+		"Type of the persistent disk volume, only used if local-ssd=false "+
+			"(pd-ssd, pd-balanced, pd-extreme, pd-standard, hyperdisk-balanced)")
 	flags.IntVar(&o.PDVolumeSize, ProviderName+"-pd-volume-size", 500,
 		"Size in GB of persistent disk volume, only used if local-ssd=false")
 	flags.IntVar(&o.PDVolumeCount, ProviderName+"-pd-volume-count", 1,
 		"Number of persistent disk volumes, only used if local-ssd=false")
+	flags.IntVar(&o.PDVolumeProvisionedIOPS, ProviderName+"-pd-volume-provisioned-iops", 0,
+		"Provisioned IOPS for the disk volume (required for hyperdisk-balanced, optional for pd-extreme)")
+	flags.IntVar(&o.PDVolumeProvisionedThroughput, ProviderName+"-pd-volume-provisioned-throughput", 0,
+		"Provisioned throughput in MiB/s for the disk volume (required for hyperdisk-balanced)")
 	flags.BoolVar(&o.UseMultipleDisks, ProviderName+"-enable-multiple-stores",
 		false, "Enable the use of multiple stores by creating one store directory per disk. "+
 			"Default is to raid0 stripe all disks.")
@@ -1503,9 +1510,14 @@ func (p *Provider) computeInstanceArgs(
 					fmt.Sprintf("size=%dGB", providerOpts.PDVolumeSize),
 					"auto-delete=yes",
 				}
-				// TODO(pavelkalinnikov): support disk types with "provisioned-throughput"
-				// option, such as Hyperdisk Throughput:
-				// https://cloud.google.com/compute/docs/disks/add-hyperdisk#hyperdisk-throughput.
+				// Add provisioned IOPS if specified (required for hyperdisk-balanced, optional for pd-extreme).
+				if providerOpts.PDVolumeProvisionedIOPS > 0 {
+					pdProps = append(pdProps, fmt.Sprintf("provisioned-iops=%d", providerOpts.PDVolumeProvisionedIOPS))
+				}
+				// Add provisioned throughput if specified (required for hyperdisk-balanced).
+				if providerOpts.PDVolumeProvisionedThroughput > 0 {
+					pdProps = append(pdProps, fmt.Sprintf("provisioned-throughput=%d", providerOpts.PDVolumeProvisionedThroughput))
+				}
 				args = append(args, "--create-disk", strings.Join(pdProps, ","))
 			}
 			// Enable DISCARD commands for persistent disks, as is advised in: