Add support for Ultra Disks as Persistent Volumes

api/v1alpha3/azuremachine_conversion.go

@@ -54,6 +54,10 @@ func (src *AzureMachine) ConvertTo(dstRaw conversion.Hub) error {
 		dst.Spec.Image.ComputeGallery = restored.Spec.Image.ComputeGallery
 	}
 
+	if restored.Spec.AdditionalCapabilities != nil {
+		dst.Spec.AdditionalCapabilities = restored.Spec.AdditionalCapabilities
+	}
+
 	dst.Spec.SubnetName = restored.Spec.SubnetName
 
 	dst.Status.LongRunningOperationStates = restored.Status.LongRunningOperationStates

api/v1alpha3/azuremachinetemplate_conversion.go

@@ -54,6 +54,10 @@ func (src *AzureMachineTemplate) ConvertTo(dstRaw conversion.Hub) error {
 		dst.Spec.Template.Spec.Image.ComputeGallery = restored.Spec.Template.Spec.Image.ComputeGallery
 	}
 
+	if restored.Spec.Template.Spec.AdditionalCapabilities != nil {
+		dst.Spec.Template.Spec.AdditionalCapabilities = restored.Spec.Template.Spec.AdditionalCapabilities
+	}
+
 	dst.Spec.Template.Spec.SubnetName = restored.Spec.Template.Spec.SubnetName
 	dst.Spec.Template.ObjectMeta = restored.Spec.Template.ObjectMeta
 

api/v1alpha4/azuremachine_conversion.go

@@ -40,6 +40,10 @@ func (src *AzureMachine) ConvertTo(dstRaw conversion.Hub) error {
 		dst.Spec.Image.ComputeGallery = restored.Spec.Image.ComputeGallery
 	}
 
+	if restored.Spec.AdditionalCapabilities != nil {
+		dst.Spec.AdditionalCapabilities = restored.Spec.AdditionalCapabilities
+	}
+
 	return nil
 }
 
@@ -85,3 +89,8 @@ func Convert_v1alpha4_AzureMarketplaceImage_To_v1beta1_AzureMarketplaceImage(in
 func Convert_v1beta1_Image_To_v1alpha4_Image(in *v1beta1.Image, out *Image, s apiconversion.Scope) error {
 	return autoConvert_v1beta1_Image_To_v1alpha4_Image(in, out, s)
 }
+
+// Convert_v1beta1_AzureMachineSpec_To_v1alpha4_AzureMachineSpec is an autogenerated conversion function.
+func Convert_v1beta1_AzureMachineSpec_To_v1alpha4_AzureMachineSpec(in *v1beta1.AzureMachineSpec, out *AzureMachineSpec, s apiconversion.Scope) error {
+	return autoConvert_v1beta1_AzureMachineSpec_To_v1alpha4_AzureMachineSpec(in, out, s)
+}

api/v1alpha4/azuremachinetemplate_conversion.go

@@ -40,6 +40,10 @@ func (src *AzureMachineTemplate) ConvertTo(dstRaw conversion.Hub) error {
 		dst.Spec.Template.Spec.Image.ComputeGallery = restored.Spec.Template.Spec.Image.ComputeGallery
 	}
 
+	if restored.Spec.Template.Spec.AdditionalCapabilities != nil {
+		dst.Spec.Template.Spec.AdditionalCapabilities = restored.Spec.Template.Spec.AdditionalCapabilities
+	}
+
 	dst.Spec.Template.ObjectMeta = restored.Spec.Template.ObjectMeta
 
 	return nil

api/v1beta1/azuremachine_types.go

@@ -86,6 +86,10 @@ type AzureMachineSpec struct {
 	// +optional
 	AdditionalTags Tags `json:"additionalTags,omitempty"`
 
+	// AdditionalCapabilities specifies additional capabilities enabled or disabled on the virtual machine.
+	// +optional
+	AdditionalCapabilities *AdditionalCapabilities `json:"additionalCapabilities,omitempty"`
+
 	// AllocatePublicIP allows the ability to create dynamic public ips for machines where this value is true.
 	// +optional
 	AllocatePublicIP bool `json:"allocatePublicIP,omitempty"`
@@ -185,6 +189,15 @@ type AzureMachineStatus struct {
 	LongRunningOperationStates Futures `json:"longRunningOperationStates,omitempty"`
 }
 
+// AdditionalCapabilities enables or disables a capability on the virtual machine.
+type AdditionalCapabilities struct {
+	// UltraSSDEnabled enables or disables Azure UltraSSD capability for the virtual machine.
+	// Defaults to true if Ultra SSD data disks are specified,
+	// otherwise it doesn't set the capability on the VM.
+	// +optional
+	UltraSSDEnabled *bool `json:"ultraSSDEnabled,omitempty"`
+}
+
 // +kubebuilder:object:root=true
 // +kubebuilder:printcolumn:name="Ready",type="string",JSONPath=".status.conditions[?(@.type=='Ready')].status"
 // +kubebuilder:printcolumn:name="Reason",type="string",JSONPath=".status.conditions[?(@.type=='Ready')].reason"

azure/scope/machine.go

@@ -164,6 +164,7 @@ func (m *MachineScope) VMSpec() azure.ResourceSpecGetter {
 		SpotVMOptions:          m.AzureMachine.Spec.SpotVMOptions,
 		SecurityProfile:        m.AzureMachine.Spec.SecurityProfile,
 		AdditionalTags:         m.AdditionalTags(),
+		AdditionalCapabilities: m.AzureMachine.Spec.AdditionalCapabilities,
 		ProviderID:             m.ProviderID(),
 	}
 	if m.cache != nil {

azure/services/scalesets/scalesets.go

@@ -335,17 +335,30 @@ func (s *Service) validateSpec(ctx context.Context) error {
 		return azure.WithTerminalError(errors.Errorf("encryption at host is not supported for VM type %s", spec.Size))
 	}
 
-	// check the support for ultra disks based on location and vm size
-	for _, disks := range spec.DataDisks {
-		location := s.Scope.Location()
-		zones, err := s.resourceSKUCache.GetZones(ctx, location)
-		if err != nil {
-			return azure.WithTerminalError(errors.Wrapf(err, "failed to get the zones for location %s", location))
-		}
+	// Fetch location and zone to check for their support of ultra disks.
+	location := s.Scope.Location()
+	zones, err := s.resourceSKUCache.GetZones(ctx, location)
+	if err != nil {
+		return azure.WithTerminalError(errors.Wrapf(err, "failed to get the zones for location %s", location))
+	}
+
+	for _, zone := range zones {
+		hasLocationCapability := sku.HasLocationCapability(resourceskus.UltraSSDAvailable, location, zone)
+		err := fmt.Errorf("vm size %s does not support ultra disks in location %s. select a different vm size or disable ultra disks", spec.Size, location)
 
-		for _, zone := range zones {
-			if disks.ManagedDisk != nil && disks.ManagedDisk.StorageAccountType == string(compute.StorageAccountTypesUltraSSDLRS) && !sku.HasLocationCapability(resourceskus.UltraSSDAvailable, location, zone) {
-				return azure.WithTerminalError(fmt.Errorf("vm size %s does not support ultra disks in location %s. select a different vm size or disable ultra disks", spec.Size, location))
+		// Check support for ultra disks as data disks.
+		for _, disks := range spec.DataDisks {
+			if disks.ManagedDisk != nil &&
+				disks.ManagedDisk.StorageAccountType == string(compute.StorageAccountTypesUltraSSDLRS) &&
+				!hasLocationCapability {
+				return azure.WithTerminalError(err)
+			}
+		}
+		// Check support for ultra disks as persistent volumes.
+		if spec.AdditionalCapabilities != nil && spec.AdditionalCapabilities.UltraSSDEnabled != nil {
+			if *spec.AdditionalCapabilities.UltraSSDEnabled &&
+				!hasLocationCapability {
+				return azure.WithTerminalError(err)
 			}
 		}
 	}
@@ -491,6 +504,8 @@ func (s *Service) buildVMSSFromSpec(ctx context.Context, vmssSpec azure.ScaleSet
 		}
 	}
 
+	// Provisionally detect whether there is any Data Disk defined which uses UltraSSDs.
+	// If that's the case, enable the UltraSSD capability.
 	for _, dataDisk := range vmssSpec.DataDisks {
 		if dataDisk.ManagedDisk != nil && dataDisk.ManagedDisk.StorageAccountType == string(compute.StorageAccountTypesUltraSSDLRS) {
 			vmss.VirtualMachineScaleSetProperties.AdditionalCapabilities = &compute.AdditionalCapabilities{
@@ -499,6 +514,14 @@ func (s *Service) buildVMSSFromSpec(ctx context.Context, vmssSpec azure.ScaleSet
 		}
 	}
 
+	// Set Additional Capabilities if any is present on the spec.
+	if vmssSpec.AdditionalCapabilities != nil {
+		// Set UltraSSDEnabled if a specific value is set on the spec for it.
+		if vmssSpec.AdditionalCapabilities.UltraSSDEnabled != nil {
+			vmss.AdditionalCapabilities.UltraSSDEnabled = vmssSpec.AdditionalCapabilities.UltraSSDEnabled
+		}
+	}
+
 	if vmssSpec.TerminateNotificationTimeout != nil {
 		vmss.VirtualMachineScaleSetProperties.VirtualMachineProfile.ScheduledEventsProfile = &compute.ScheduledEventsProfile{
 			TerminateNotificationProfile: &compute.TerminateNotificationProfile{

azure/services/scalesets/scalesets_test.go

@@ -500,7 +500,7 @@ func TestReconcileVMSS(t *testing.T) {
 			},
 		},
 		{
-			name:          "fail to create a vm with ultra disk enabled",
+			name:          "fail to create a vm with ultra disk implicitly enabled by data disk, when location not supported",
 			expectedError: "reconcile error that cannot be recovered occurred: vm size VM_SIZE_USSD does not support ultra disks in location test-location. select a different vm size or disable ultra disks. Object will not be requeued",
 			expect: func(g *WithT, s *mock_scalesets.MockScaleSetScopeMockRecorder, m *mock_scalesets.MockClientMockRecorder) {
 				s.ScaleSetSpec().Return(azure.ScaleSetSpec{
@@ -519,6 +519,45 @@ func TestReconcileVMSS(t *testing.T) {
 				s.Location().AnyTimes().Return("test-location")
 			},
 		},
+		{
+			name:          "fail to create a vm with ultra disk explicitly enabled via additional capabilities, when location not supported",
+			expectedError: "reconcile error that cannot be recovered occurred: vm size VM_SIZE_USSD does not support ultra disks in location test-location. select a different vm size or disable ultra disks. Object will not be requeued",
+			expect: func(g *WithT, s *mock_scalesets.MockScaleSetScopeMockRecorder, m *mock_scalesets.MockClientMockRecorder) {
+				s.ScaleSetSpec().Return(azure.ScaleSetSpec{
+					Name:       defaultVMSSName,
+					Size:       "VM_SIZE_USSD",
+					Capacity:   2,
+					SSHKeyData: "ZmFrZXNzaGtleQo=",
+					AdditionalCapabilities: &infrav1.AdditionalCapabilities{
+						UltraSSDEnabled: to.BoolPtr(true),
+					},
+				})
+				s.Location().AnyTimes().Return("test-location")
+			},
+		},
+		{
+			name:          "fail to create a vm with ultra disk explicitly enabled via additional capabilities, when location not supported",
+			expectedError: "reconcile error that cannot be recovered occurred: vm size VM_SIZE_USSD does not support ultra disks in location test-location. select a different vm size or disable ultra disks. Object will not be requeued",
+			expect: func(g *WithT, s *mock_scalesets.MockScaleSetScopeMockRecorder, m *mock_scalesets.MockClientMockRecorder) {
+				s.ScaleSetSpec().Return(azure.ScaleSetSpec{
+					Name:       defaultVMSSName,
+					Size:       "VM_SIZE_USSD",
+					Capacity:   2,
+					SSHKeyData: "ZmFrZXNzaGtleQo=",
+					DataDisks: []infrav1.DataDisk{
+						{
+							ManagedDisk: &infrav1.ManagedDiskParameters{
+								StorageAccountType: "UltraSSD_LRS",
+							},
+						},
+					},
+					AdditionalCapabilities: &infrav1.AdditionalCapabilities{
+						UltraSSDEnabled: to.BoolPtr(false),
+					},
+				})
+				s.Location().AnyTimes().Return("test-location")
+			},
+		},
 	}
 
 	for _, tc := range testcases {

azure/services/virtualmachines/spec.go

@@ -49,6 +49,7 @@ type VMSpec struct {
 	SpotVMOptions          *infrav1.SpotVMOptions
 	SecurityProfile        *infrav1.SecurityProfile
 	AdditionalTags         infrav1.Tags
+	AdditionalCapabilities *infrav1.AdditionalCapabilities
 	SKU                    resourceskus.SKU
 	Image                  *infrav1.Image
 	BootstrapData          string
@@ -306,6 +307,9 @@ func (s *VMSpec) generateNICRefs() *[]compute.NetworkInterfaceReference {
 
 func (s *VMSpec) generateAdditionalCapabilities() *compute.AdditionalCapabilities {
 	var capabilities *compute.AdditionalCapabilities
+
+	// Provisionally detect whether there is any Data Disk defined which uses UltraSSDs.
+	// If that's the case, enable the UltraSSD capability.
 	for _, dataDisk := range s.DataDisks {
 		if dataDisk.ManagedDisk != nil && dataDisk.ManagedDisk.StorageAccountType == string(compute.StorageAccountTypesUltraSSDLRS) {
 			capabilities = &compute.AdditionalCapabilities{
@@ -314,6 +318,18 @@ func (s *VMSpec) generateAdditionalCapabilities() *compute.AdditionalCapabilitie
 			break
 		}
 	}
+
+	// Set Additional Capabilities if any is present on the spec.
+	if s.AdditionalCapabilities != nil {
+		if capabilities == nil {
+			capabilities = &compute.AdditionalCapabilities{}
+		}
+		// Set UltraSSDEnabled if a specific value is set on the spec for it.
+		if s.AdditionalCapabilities.UltraSSDEnabled != nil {
+			capabilities.UltraSSDEnabled = s.AdditionalCapabilities.UltraSSDEnabled
+		}
+	}
+
 	return capabilities
 }