Merge cockroachdb#130862

130862: roachprod: use NVMe when provisioning clusters on Azure. r=renatolabs,srosenberg,herkolategan a=shailendra-patel

 Currently, when provisioning clusters on Azure, we interface with the
 data disk via SCSI (the default). For improved performance and
 uniformity with other clouds use NVMe.

 Azure support DiskControllerType NVMe for `E Series v5` VM family.
 OS disk and network disks are interfaced with NVMe.
 Local storage are interfaced via SCSI, NVMe is not supported for
 local disk.

 Changes as part of this PR:
  - For `E Series v5` VM family enabled NVMe.

Epic: none
Fixes: cockroachdb#129009

Release note: None

Co-authored-by: Shailendra Patel <[email protected]>

pkg/cmd/roachtest/spec/machine_type.go

@@ -214,6 +214,7 @@ func SelectGCEMachineType(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.
 // N.B. cpus is expected to be an even number; validation is deferred to a specific cloud provider.
 //
 // See ExampleSelectAzureMachineType for an exhaustive list of selected machine types.
+// TODO: Add Ebsv5 machine type to leverage NVMe
 func SelectAzureMachineType(cpus int, mem MemPerCPU, arch vm.CPUArch) (string, vm.CPUArch, error) {
 	series := "Ddsv5" // 4 GB RAM per CPU
 	selectedArch := vm.ArchAMD64

pkg/roachprod/vm/azure/azure.go

@@ -746,7 +746,11 @@ func (p *Provider) createVM(
 		OSInitializedFile:    vm.OSInitializedFile,
 		StartupLogs:          vm.StartupLogs,
 	}
-	if !opts.SSDOpts.UseLocalSSD {
+	useNVMe := MachineSupportsNVMe(providerOpts.MachineType)
+	if useNVMe {
+		startupArgs.DiskControllerNVMe = true
+	}
+	if !opts.SSDOpts.UseLocalSSD && !useNVMe {
 		// We define lun42 explicitly in the data disk request below.
 		lun := 42
 		startupArgs.AttachedDiskLun = &lun
@@ -864,6 +868,10 @@ func (p *Provider) createVM(
 			},
 		},
 	}
+	if useNVMe {
+		machine.VirtualMachineProperties.StorageProfile.DiskControllerType = compute.NVMe
+	}
+
 	if !opts.SSDOpts.UseLocalSSD {
 		caching := compute.CachingTypesNone
 
@@ -1690,3 +1698,25 @@ func MachineFamilyVersionFromMachineType(machineType string) int {
 	}
 	return -1
 }
+
+// MachineSupportsNVMe Azure supports Nvme for E series v5 machine family.
+// OS disk and network disk support nvme. Local storage do not support nvme.
+func MachineSupportsNVMe(machineType string) bool {
+	version := MachineFamilyVersionFromMachineType(machineType)
+	if version == 5 {
+		matches := azureMachineTypes.FindStringSubmatch(machineType)
+		if len(matches) >= 4 {
+			family, features := matches[1], matches[3]
+			// additive features of azure vm are represented by lower case letter
+			// b = Block Storage performance
+			// d = diskful (that is, a local temp disk is present);
+			// s = Premium Storage capable, including possible use of Ultra SSD
+			// https://learn.microsoft.com/en-us/azure/virtual-machines/vm-naming-conventions
+			// example of supported machine types Standard_E2bds_v5, Standard_E2bs_v5
+			if family == "Standard_E" && (features == "bs" || features == "bds") {
+				return true
+			}
+		}
+	}
+	return false
+}

pkg/roachprod/vm/azure/utils.go

@@ -27,6 +27,7 @@ type azureStartupArgs struct {
 	DisksInitializedFile string // File to touch when disks are initialized.
 	OSInitializedFile    string // File to touch when OS is initialized.
 	StartupLogs          string // File to redirect startup script output logs.
+	DiskControllerNVMe   bool   // Interface data disk via NVMe
 }
 
 const azureStartupTemplate = `#!/bin/bash
@@ -42,12 +43,13 @@ exec &> >(tee -a {{ .StartupLogs }})
 echo "startup script starting: $(date -u)"
 mount_opts="defaults"
 
-{{if .AttachedDiskLun}}
+devices=()
+{{if .DiskControllerNVMe}}
+# Setup nvme network storage, need to remove nvme OS disk from the device list.
+devices=($(realpath -qe /dev/disk/by-id/nvme-* | grep -v "nvme0n1" | sort -u))
+{{else if .AttachedDiskLun}}
 # Setup network attached storage
 devices=("/dev/disk/azure/scsi1/lun{{.AttachedDiskLun}}")
-{{else}}
-# Setup local storage.
-devices=($(realpath -qe /dev/disk/by-id/nvme-* | sort -u))
 {{end}}
 
 if (( ${#devices[@]} == 0 ));