fix: respect NodePool zone requirement (if any) (Azure#47)

* fix: respect zone requirements * test: zone-aware provisioning
dtzar · Nov 28, 2023 · 01acc1c · 01acc1c
1 parent 61a5891
commit 01acc1c
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/pkg/providers/instance/instance.go b/pkg/providers/instance/instance.go
@@ -553,8 +553,11 @@ func (p *Provider) pickSkuSizePriorityAndZone(ctx context.Context, nodeClaim *co
 	logging.FromContext(ctx).Infof("Selected instance type %s", instanceType.Name)
 	// Priority - Provisioner defaults to Regular, so pick Spot if it is explicitly included in requirements (and is offered in at least one zone)
 	priority := p.getPriorityForInstanceType(nodeClaim, instanceType)
-	// Zone - ideally random/spread from zones that support given Priority
-	priorityOfferings := lo.Filter(instanceType.Offerings.Available(), func(o corecloudprovider.Offering, _ int) bool { return o.CapacityType == priority })
+	// Zone - ideally random/spread from requested zones that support given Priority
+	requestedZones := scheduling.NewNodeSelectorRequirements(nodeClaim.Spec.Requirements...).Get(v1.LabelTopologyZone)
+	priorityOfferings := lo.Filter(instanceType.Offerings.Available(), func(o corecloudprovider.Offering, _ int) bool {
+		return o.CapacityType == priority && requestedZones.Has(o.Zone)
+	})
 	zonesWithPriority := lo.Map(priorityOfferings, func(o corecloudprovider.Offering, _ int) string { return o.Zone })
 	if zone, ok := sets.New(zonesWithPriority...).PopAny(); ok {
 		// Zones in Offerings have <region>-<number> format; the zone returned from here will be used for VM instantiation,

diff --git a/pkg/providers/instancetype/suite_test.go b/pkg/providers/instancetype/suite_test.go
@@ -880,6 +880,25 @@ var _ = Describe("InstanceType Provider", func() {
 		})
 	})
 
+	Context("Zone aware provisioning", func() {
+		It("should launch in the NodePool-requested zone", func() {
+			zone, vmZone := "eastus-3", "3"
+			nodePool.Spec.Template.Spec.Requirements = []v1.NodeSelectorRequirement{
+				{Key: corev1beta1.CapacityTypeLabelKey, Operator: v1.NodeSelectorOpIn, Values: []string{corev1beta1.CapacityTypeSpot, corev1beta1.CapacityTypeOnDemand}},
+				{Key: v1.LabelTopologyZone, Operator: v1.NodeSelectorOpIn, Values: []string{zone}},
+			}
+			ExpectApplied(ctx, env.Client, nodePool, nodeClass)
+			pod := coretest.UnschedulablePod()
+			ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, coreProvisioner, pod)
+			node := ExpectScheduled(ctx, env.Client, pod)
+			Expect(node.Labels).To(HaveKeyWithValue(v1alpha2.AlternativeLabelTopologyZone, zone))
+
+			vm := azureEnv.VirtualMachinesAPI.VirtualMachineCreateOrUpdateBehavior.CalledWithInput.Pop().VM
+			Expect(vm).NotTo(BeNil())
+			Expect(vm.Zones).To(ConsistOf(&vmZone))
+		})
+	})
+
 })
 
 var _ = Describe("Tax Calculator", func() {