diff --git a/pkg/cmd/roachprod/vm/aws/aws.go b/pkg/cmd/roachprod/vm/aws/aws.go index dcf5fdb373f8..70290a8edd2d 100644 --- a/pkg/cmd/roachprod/vm/aws/aws.go +++ b/pkg/cmd/roachprod/vm/aws/aws.go @@ -13,6 +13,7 @@ package aws import ( "encoding/json" "fmt" + "io/ioutil" "log" "math/rand" "os" @@ -65,19 +66,115 @@ func init() { vm.Providers[ProviderName] = p } +// ebsDisk represent EBS disk device. +// When marshaled to JSON format, produces JSON specification used +// by AWS sdk to configure attached volumes. +type ebsDisk struct { + VolumeType string `json:"VolumeType"` + VolumeSize int `json:"VolumeSize"` + IOPs int `json:"Iops,omitempty"` + Throughput int `json:"Throughput,omitempty"` + DeleteOnTermination bool `json:"DeleteOnTermination"` +} + +// ebsVolume represents a mounted volume: name + ebsDisk +type ebsVolume struct { + DeviceName string `json:"DeviceName"` + Disk ebsDisk `json:"Ebs"` +} + +const ebsDefaultVolumeSizeGB = 500 + +// Set implements flag Value interface. +func (d *ebsDisk) Set(s string) error { + if err := json.Unmarshal([]byte(s), &d); err != nil { + return err + } + + d.DeleteOnTermination = true + + // Sanity check disk configuration. + // This is not strictly needed since AWS sdk would return error anyway, + // but we can return a nicer error message sooner. + if d.VolumeSize == 0 { + d.VolumeSize = ebsDefaultVolumeSizeGB + } + + switch strings.ToLower(d.VolumeType) { + 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 == 0 { + // 30000 is a base IOPs for gp3. + d.IOPs = 3000 + } + if d.Throughput == 0 { + // 125MB/s is base throughput for gp3. + d.Throughput = 125 + } + case "io1", "io2": + if d.IOPs == 0 { + return errors.AssertionFailedf("Iops required for %s disk", d.VolumeType) + } + default: + return errors.Errorf("Unknown EBS volume type %s", d.VolumeType) + } + return nil +} + +// Type implements flag Value interface. +func (d *ebsDisk) Type() string { + return "JSON" +} + +// String Implements flag Value interface. +func (d *ebsDisk) String() string { + return "EBSDisk" +} + +type ebsVolumeList []*ebsVolume + +func (vl *ebsVolumeList) newVolume() *ebsVolume { + return &ebsVolume{ + DeviceName: fmt.Sprintf("/dev/sd%c", 'd'+len(*vl)), + } +} + +// Set implements flag Value interface. +func (vl *ebsVolumeList) Set(s string) error { + v := vl.newVolume() + if err := v.Disk.Set(s); err != nil { + return err + } + *vl = append(*vl, v) + return nil +} + +// Type implements flag Value interface. +func (vl *ebsVolumeList) Type() string { + return "JSON" +} + +// String Implements flag Value interface. +func (vl *ebsVolumeList) String() string { + return "EBSVolumeList" +} + // providerOpts implements the vm.ProviderFlags interface for aws.Provider. type providerOpts struct { Profile string Config *awsConfig - MachineType string - SSDMachineType string - CPUOptions string - RemoteUserName string - EBSVolumeType string - EBSVolumeSize int - EBSProvisionedIOPs int - UseMultipleDisks bool + MachineType string + SSDMachineType string + CPUOptions string + RemoteUserName string + DefaultEBSVolume ebsVolume + EBSVolumes ebsVolumeList + UseMultipleDisks bool // Use specified ImageAMI when provisioning. // Overrides config.json AMI. @@ -122,7 +219,6 @@ var defaultCreateZones = []string{ // somewhat complicated because different EC2 regions may as well // be parallel universes. func (o *providerOpts) ConfigureCreateFlags(flags *pflag.FlagSet) { - // m5.xlarge is a 4core, 16Gb instance, approximately equal to a GCE n1-standard-4 flags.StringVar(&o.MachineType, ProviderName+"-machine-type", defaultMachineType, "Machine type (see https://aws.amazon.com/ec2/instance-types/)") @@ -139,13 +235,17 @@ func (o *providerOpts) ConfigureCreateFlags(flags *pflag.FlagSet) { flags.StringVar(&o.RemoteUserName, ProviderName+"-user", "ubuntu", "Name of the remote user to SSH as") - flags.StringVar(&o.EBSVolumeType, ProviderName+"-ebs-volume-type", - "gp2", "Type of the EBS volume, only used if local-ssd=false") - flags.IntVar(&o.EBSVolumeSize, ProviderName+"-ebs-volume-size", - 500, "Size in GB of EBS volume, only used if local-ssd=false") - flags.IntVar(&o.EBSProvisionedIOPs, ProviderName+"-ebs-iops", - 1000, "Number of IOPs to provision, only used if "+ProviderName+ - "-ebs-volume-type=io1") + flags.StringVar(&o.DefaultEBSVolume.Disk.VolumeType, ProviderName+"-ebs-volume-type", + "", "Type of the EBS volume, only used if local-ssd=false") + flags.IntVar(&o.DefaultEBSVolume.Disk.VolumeSize, ProviderName+"-ebs-volume-size", + ebsDefaultVolumeSizeGB, "Size in GB of EBS volume, only used if local-ssd=false") + flags.IntVar(&o.DefaultEBSVolume.Disk.IOPs, ProviderName+"-ebs-iops", + 0, "Number of IOPs to provision for supported disk types (io1, io2, gp3)") + flags.IntVar(&o.DefaultEBSVolume.Disk.Throughput, ProviderName+"-ebs-throughput", + 0, "Additional throughput to provision, in MiB/s") + + flags.VarP(&o.EBSVolumes, ProviderName+"-ebs-volume", "", + "Additional EBS disk to attached; specified as JSON: {VolumeType=io2,VolumeSize=213,Iops=321}") flags.StringSliceVar(&o.CreateZones, ProviderName+"-zones", nil, fmt.Sprintf("aws availability zones to use for cluster creation. If zones are formatted\n"+ @@ -723,21 +823,33 @@ func (p *Provider) runInstance(name string, zone string, opts vm.CreateOpts) err // The local NVMe devices are automatically mapped. Otherwise, we need to map an EBS data volume. if !opts.SSDOpts.UseLocalSSD { - var ebsParams string - switch t := p.opts.EBSVolumeType; t { - case "gp2": - ebsParams = fmt.Sprintf("{VolumeSize=%d,VolumeType=%s,DeleteOnTermination=true}", - p.opts.EBSVolumeSize, t) - case "io1", "io2": - ebsParams = fmt.Sprintf("{VolumeSize=%d,VolumeType=%s,Iops=%d,DeleteOnTermination=true}", - p.opts.EBSVolumeSize, t, p.opts.EBSProvisionedIOPs) - default: - return errors.Errorf("Unknown EBS volume type %s", t) + if len(p.opts.EBSVolumes) == 0 && p.opts.DefaultEBSVolume.Disk.VolumeType == "" { + p.opts.DefaultEBSVolume.Disk.VolumeType = "gp2" + } + + if p.opts.DefaultEBSVolume.Disk.VolumeType != "" { + // Add default volume to the list of volumes we'll setup. + v := p.opts.EBSVolumes.newVolume() + v.Disk = p.opts.DefaultEBSVolume.Disk + p.opts.EBSVolumes = append(p.opts.EBSVolumes, v) + } + + mapping, err := json.Marshal(p.opts.EBSVolumes) + if err != nil { + return err + } + + deviceMapping, err := ioutil.TempFile("", "aws-block-device-mapping") + if err != nil { + return err + } + defer deviceMapping.Close() + if _, err := deviceMapping.Write(mapping); err != nil { + return err } args = append(args, "--block-device-mapping", - // Size is measured in GB. gp2 type derives guaranteed iops from size. - "DeviceName=/dev/sdd,Ebs="+ebsParams, + "file://"+deviceMapping.Name(), ) } return p.runJSONCommand(args, &data)