roachtest: support boot disk only workload VMs

Previously, workload machines automatically inherited the same disk setup as the
rest of the cluster. To optimize for cost, we now allow workload machines to
attach only boot disks.

The default behavior is to attach a standard data disk. Hence, a new explicit
flag "BootDiskOnly" has been added to the provider options to prevent the
default behavior, specifically for workload machines.

Epic: None
Release note: None

Author: herkolategan

pkg/cmd/roachtest/spec/cluster_spec.go

@@ -104,9 +104,10 @@ type ClusterSpec struct {
 	// treated as workload node. Defaults to a VM with 4 CPUs if not specified
 	// by WorkloadNodeCPUs.
 	// TODO(GouravKumar): remove use of WorkloadNode, use WorkloadNodeCount instead
-	WorkloadNode      bool
-	WorkloadNodeCount int
-	WorkloadNodeCPUs  int
+	WorkloadNode         bool
+	WorkloadNodeCount    int
+	WorkloadNodeCPUs     int
+	WorkloadRequiresDisk bool
 	// CPUs is the number of CPUs per node.
 	CPUs                 int
 	Mem                  MemPerCPU
@@ -226,7 +227,12 @@ func awsMachineSupportsSSD(machineType string) bool {
 }
 
 func getAWSOpts(
-	machineType string, volumeSize, ebsThroughput int, ebsIOPS int, localSSD bool, useSpotVMs bool,
+	machineType string,
+	volumeSize, ebsThroughput int,
+	ebsIOPS int,
+	localSSD bool,
+	useSpotVMs bool,
+	bootDiskOnly bool,
 ) vm.ProviderOpts {
 	opts := aws.DefaultProviderOpts()
 	if volumeSize != 0 {
@@ -244,6 +250,7 @@ func getAWSOpts(
 		opts.MachineType = machineType
 	}
 	opts.UseSpot = useSpotVMs
+	opts.BootDiskOnly = bootDiskOnly
 	return opts
 }
 
@@ -258,6 +265,7 @@ func getGCEOpts(
 	volumeType string,
 	volumeCount int,
 	useSpot bool,
+	bootDiskOnly bool,
 ) vm.ProviderOpts {
 	opts := gce.DefaultProviderOpts()
 	opts.MachineType = machineType
@@ -286,16 +294,17 @@ func getGCEOpts(
 	if volumeType != "" {
 		opts.PDVolumeType = volumeType
 	}
-
+	opts.BootDiskOnly = bootDiskOnly
 	return opts
 }
 
-func getAzureOpts(machineType string, volumeSize int) vm.ProviderOpts {
+func getAzureOpts(machineType string, volumeSize int, bootDiskOnly bool) vm.ProviderOpts {
 	opts := azure.DefaultProviderOpts()
 	opts.MachineType = machineType
 	if volumeSize != 0 {
 		opts.NetworkDiskSize = int32(volumeSize)
 	}
+	opts.BootDiskOnly = bootDiskOnly
 	return opts
 }
 
@@ -307,6 +316,7 @@ func getIBMOpts(
 	volumeIOPS int,
 	extraVolumeCount int,
 	RAID0 bool,
+	bootDiskOnly bool,
 ) vm.ProviderOpts {
 	opts := ibm.DefaultProviderOpts()
 	opts.MachineType = machineType
@@ -334,6 +344,7 @@ func getIBMOpts(
 		}
 		opts.UseMultipleDisks = !RAID0
 	}
+	opts.BootDiskOnly = bootDiskOnly
 
 	return opts
 }
@@ -510,7 +521,7 @@ func (s *ClusterSpec) RoachprodOpts(
 	var err error
 	switch cloud {
 	case AWS:
-		workloadMachineType, _, err = SelectAWSMachineType(s.WorkloadNodeCPUs, s.Mem, preferLocalSSD && s.VolumeSize == 0, selectedArch)
+		workloadMachineType, _, err = SelectAWSMachineType(s.WorkloadNodeCPUs, s.Mem, false, selectedArch)
 	case GCE:
 		workloadMachineType, _ = SelectGCEMachineType(s.WorkloadNodeCPUs, s.Mem, selectedArch)
 	case Azure:
@@ -532,27 +543,29 @@ func (s *ClusterSpec) RoachprodOpts(
 	switch cloud {
 	case AWS:
 		providerOpts = getAWSOpts(machineType, s.VolumeSize, s.AWS.VolumeThroughput, s.AWS.VolumeIOPS,
-			createVMOpts.SSDOpts.UseLocalSSD, s.UseSpotVMs)
-		workloadProviderOpts = getAWSOpts(workloadMachineType, s.VolumeSize, s.AWS.VolumeThroughput,
-			s.AWS.VolumeIOPS, createVMOpts.SSDOpts.UseLocalSSD, s.UseSpotVMs)
+			createVMOpts.SSDOpts.UseLocalSSD, s.UseSpotVMs, false)
+		workloadProviderOpts = getAWSOpts(workloadMachineType, s.VolumeSize, s.AWS.VolumeThroughput, s.AWS.VolumeIOPS,
+			createVMOpts.SSDOpts.UseLocalSSD, s.UseSpotVMs, !s.WorkloadRequiresDisk)
 	case GCE:
 		providerOpts = getGCEOpts(machineType, s.VolumeSize, ssdCount,
 			createVMOpts.SSDOpts.UseLocalSSD, s.RAID0, s.TerminateOnMigration,
-			s.GCE.MinCPUPlatform, vm.ParseArch(createVMOpts.Arch), s.GCE.VolumeType, s.GCE.VolumeCount, s.UseSpotVMs,
+			s.GCE.MinCPUPlatform, vm.ParseArch(createVMOpts.Arch), s.GCE.VolumeType,
+			s.GCE.VolumeCount, s.UseSpotVMs, false,
 		)
 		workloadProviderOpts = getGCEOpts(workloadMachineType, s.VolumeSize, ssdCount,
 			createVMOpts.SSDOpts.UseLocalSSD, s.RAID0, s.TerminateOnMigration,
-			s.GCE.MinCPUPlatform, vm.ParseArch(createVMOpts.Arch), s.GCE.VolumeType, s.GCE.VolumeCount, s.UseSpotVMs,
+			s.GCE.MinCPUPlatform, vm.ParseArch(createVMOpts.Arch), s.GCE.VolumeType,
+			s.GCE.VolumeCount, s.UseSpotVMs, !s.WorkloadRequiresDisk,
 		)
 	case Azure:
-		providerOpts = getAzureOpts(machineType, s.VolumeSize)
-		workloadProviderOpts = getAzureOpts(workloadMachineType, s.VolumeSize)
+		providerOpts = getAzureOpts(machineType, s.VolumeSize, false)
+		workloadProviderOpts = getAzureOpts(workloadMachineType, s.VolumeSize, true)
 	case IBM:
 		providerOpts = getIBMOpts(machineType, s.TerminateOnMigration, s.VolumeSize,
-			s.IBM.VolumeType, s.IBM.VolumeIOPS, s.IBM.VolumeCount, s.RAID0,
+			s.IBM.VolumeType, s.IBM.VolumeIOPS, s.IBM.VolumeCount, s.RAID0, false,
 		)
 		workloadProviderOpts = getIBMOpts(workloadMachineType, s.TerminateOnMigration, s.VolumeSize,
-			s.IBM.VolumeType, s.IBM.VolumeIOPS, s.IBM.VolumeCount, s.RAID0,
+			s.IBM.VolumeType, s.IBM.VolumeIOPS, s.IBM.VolumeCount, s.RAID0, true,
 		)
 	}
 

pkg/cmd/roachtest/spec/option.go

@@ -57,6 +57,15 @@ func WorkloadNodeCPU(n int) Option {
 	}
 }
 
+// WorkloadRequiresDisk should be used if the workload nodes should have the
+// exact same disk configuration as the rest of the cluster. Otherwise, all
+// workload nodes only have a boot disk.
+func WorkloadRequiresDisk() Option {
+	return func(spec *ClusterSpec) {
+		spec.WorkloadRequiresDisk = true
+	}
+}
+
 // Mem requests nodes with low/standard/high ratio of memory per CPU.
 func Mem(level MemPerCPU) Option {
 	return func(spec *ClusterSpec) {

pkg/roachprod/vm/aws/aws.go

@@ -279,6 +279,9 @@ type ProviderOpts struct {
 	// use spot vms, spot vms are significantly cheaper, but can be preempted AWS.
 	// see https://aws.amazon.com/ec2/spot/ for more details.
 	UseSpot bool
+	// BootDiskOnly ensures that no additional disks will be attached, other than
+	// the boot disk.
+	BootDiskOnly bool
 }
 
 // Provider implements the vm.Provider interface for AWS.
@@ -1348,13 +1351,15 @@ func (p *Provider) runInstance(
 			extraMountOpts = "nobarrier"
 		}
 	}
+
 	filename, err := writeStartupScript(
 		name,
 		extraMountOpts,
 		opts.SSDOpts.FileSystem,
 		providerOpts.UseMultipleDisks,
 		opts.Arch == string(vm.ArchFIPS),
 		providerOpts.RemoteUserName,
+		providerOpts.BootDiskOnly,
 	)
 	if err != nil {
 		return nil, errors.Wrapf(err, "could not write AWS startup script to temp file")
@@ -1605,7 +1610,7 @@ func assignEBSVolumes(opts *vm.CreateOpts, providerOpts *ProviderOpts) ebsVolume
 	// Make a local copy of providerOpts.EBSVolumes to prevent data races
 	ebsVolumes := providerOpts.EBSVolumes
 	// The local NVMe devices are automatically mapped.  Otherwise, we need to map an EBS data volume.
-	if !opts.SSDOpts.UseLocalSSD {
+	if !opts.SSDOpts.UseLocalSSD && !providerOpts.BootDiskOnly {
 		if len(ebsVolumes) == 0 && providerOpts.DefaultEBSVolume.Disk.VolumeType == "" {
 			providerOpts.DefaultEBSVolume.Disk.VolumeType = defaultEBSVolumeType
 			providerOpts.DefaultEBSVolume.Disk.DeleteOnTermination = true

pkg/roachprod/vm/aws/support.go

@@ -30,6 +30,11 @@ const awsStartupScriptTemplate = `#!/usr/bin/env bash
 # Script for setting up a AWS machine for roachprod use.
 
 function setup_disks() {
+{{ if .BootDiskOnly }}
+	echo "VM has no disk attached other than the boot disk."
+	return 0
+{{ end }}
+
 {{ if not .Zfs }}
 	mount_opts="defaults,nofail"
 	{{if .ExtraMountOpts}}mount_opts="${mount_opts},{{.ExtraMountOpts}}"{{end}}
@@ -181,11 +186,13 @@ func writeStartupScript(
 	useMultiple bool,
 	enableFips bool,
 	remoteUser string,
+	bootDiskOnly bool,
 ) (string, error) {
 	type tmplParams struct {
 		vm.StartupArgs
 		ExtraMountOpts   string
 		UseMultipleDisks bool
+		BootDiskOnly     bool
 	}
 
 	args := tmplParams{
@@ -198,6 +205,7 @@ func writeStartupScript(
 		),
 		ExtraMountOpts:   extraMountOpts,
 		UseMultipleDisks: useMultiple,
+		BootDiskOnly:     bootDiskOnly,
 	}
 
 	tmpfile, err := os.CreateTemp("", "aws-startup-script")

pkg/roachprod/vm/aws/support_test.go

@@ -17,7 +17,8 @@ import (
 
 // TestWriteStartupScriptTemplate mainly tests the startup script tpl compiles.
 func TestWriteStartupScriptTemplate(t *testing.T) {
-	file, err := writeStartupScript("vm_name", "", vm.Zfs, false, false, "ubuntu")
+	file, err := writeStartupScript("vm_name", "", vm.Zfs, false,
+		false, "ubuntu", false)
 	require.NoError(t, err)
 
 	f, err := os.ReadFile(file)

pkg/roachprod/vm/aws/testdata/startup_script

@@ -7,6 +7,8 @@ echo
 function setup_disks() {
 
 
+
+
 	use_multiple_disks=''
 
 	mount_prefix="/mnt/data"

pkg/roachprod/vm/azure/azure.go

@@ -931,6 +931,9 @@ func (p *Provider) createVM(
 		startupArgs.AttachedDiskLun = &lun
 	}
 
+	// Check if we only require a boot disk (workload only machines).
+	startupArgs.BootDiskOnly = providerOpts.BootDiskOnly
+
 	startupScript, err := evalStartupTemplate(startupArgs)
 	if err != nil {
 		return machine, err
@@ -1058,7 +1061,7 @@ func (p *Provider) createVM(
 		machine.VirtualMachineProperties.StorageProfile.DiskControllerType = compute.NVMe
 	}
 
-	if !opts.SSDOpts.UseLocalSSD {
+	if !opts.SSDOpts.UseLocalSSD && !providerOpts.BootDiskOnly {
 		caching := compute.CachingTypesNone
 
 		switch providerOpts.DiskCaching {

pkg/roachprod/vm/azure/flags.go

@@ -26,6 +26,7 @@ type ProviderOpts struct {
 	NetworkDiskSize int32
 	UltraDiskIOPS   int64
 	DiskCaching     string
+	BootDiskOnly    bool
 }
 
 // These default locations support availability zones. At the time of

pkg/roachprod/vm/azure/testdata/startup_script

@@ -6,7 +6,9 @@ echo
 # Script for setting up a Azure machine for roachprod use.
 
 function setup_disks() {
-	mount_opts="defaults"
+
+
+mount_opts="defaults"
 
 	devices=()
 

pkg/roachprod/vm/azure/utils.go

@@ -23,14 +23,20 @@ type azureStartupArgs struct {
 	vm.StartupArgs
 	AttachedDiskLun    *int // Use attached disk, with specified LUN; Use local ssd if nil.
 	DiskControllerNVMe bool // Interface data disk via NVMe
+	BootDiskOnly       bool
 }
 
 const azureStartupTemplate = `#!/bin/bash
 
 # Script for setting up a Azure machine for roachprod use.
 
 function setup_disks() {
-	mount_opts="defaults"
+{{ if .BootDiskOnly }}
+	echo "VM has no disk attached other than the boot disk."
+	return 0
+{{ end }}
+
+mount_opts="defaults"
 
 	devices=()
 {{if .DiskControllerNVMe}}