fix: only count topology domains the pod can actually reach

pkg/controllers/provisioning/scheduling/topology.go

@@ -113,6 +113,10 @@ func buildDomainGroups(nodePools []*v1.NodePool, instanceTypes map[string][]*clo
 	domainGroups := map[string]TopologyDomainGroup{}
 	for npName, its := range instanceTypes {
 		np := nodePoolIndex[npName]
+		// Requirements carried by every node this NodePool launches, regardless of instance type.
+		// Domains are attributed to them so that a pod which can't select the NodePool doesn't have
+		// its topology spread computed against domains only this NodePool can supply.
+		nodePoolRequirements := nodePoolDomainRequirements(np)
 		for _, it := range its {
 			// We need to intersect the instance type requirements with the current nodePool requirements.  This
 			// ensures that something like zones from an instance type don't expand the universe of valid domains.
@@ -125,7 +129,7 @@ func buildDomainGroups(nodePools []*v1.NodePool, instanceTypes map[string][]*clo
 					domainGroups[topologyKey] = NewTopologyDomainGroup()
 				}
 				for _, domain := range requirement.Values() {
-					domainGroups[topologyKey].Insert(domain, np.Spec.Template.Spec.Taints...)
+					domainGroups[topologyKey].Insert(domain, npName, np.Spec.Template.Spec.Taints, nodePoolRequirements)
 				}
 			}
 		}
@@ -138,14 +142,29 @@ func buildDomainGroups(nodePools []*v1.NodePool, instanceTypes map[string][]*clo
 					domainGroups[key] = NewTopologyDomainGroup()
 				}
 				for _, value := range requirement.Values() {
-					domainGroups[key].Insert(value, np.Spec.Template.Spec.Taints...)
+					domainGroups[key].Insert(value, npName, np.Spec.Template.Spec.Taints, nodePoolRequirements)
 				}
 			}
 		}
 	}
 	return domainGroups
 }
 
+// nodePoolDomainRequirements returns the requirements shared by every node the NodePool can launch:
+// its template requirements and labels, plus the labels Karpenter stamps onto each of its NodeClaims.
+// Instance type requirements are deliberately left out, since a domain can be offered by only some of
+// the pool's instance types while this set has to hold for any node supplying the domain.
+func nodePoolDomainRequirements(np *v1.NodePool) scheduling.Requirements {
+	requirements := scheduling.NewNodeSelectorRequirementsWithMinValues(np.Spec.Template.Spec.Requirements...)
+	requirements.Add(scheduling.NewLabelRequirements(np.Spec.Template.Labels).Values()...)
+	nodeClaimLabels := map[string]string{v1.NodePoolLabelKey: np.Name}
+	if ref := np.Spec.Template.Spec.NodeClassRef; ref != nil {
+		nodeClaimLabels[v1.NodeClassLabelKey(ref.GroupKind())] = ref.Name
+	}
+	requirements.Add(scheduling.NewLabelRequirements(nodeClaimLabels).Values()...)
+	return requirements
+}
+
 // topologyError allows lazily generating the error string in the topology error.  If a pod fails to schedule, most often
 // we are only interested in the fact that it failed to schedule and not why.
 type topologyError struct {

pkg/controllers/provisioning/scheduling/topology_test.go

@@ -170,6 +170,31 @@ var _ = Describe("Topology", func() {
 			// should spread the two pods evenly across the only valid zones in our universe (the two zones from our single nodePool)
 			ExpectSkew(ctx, env.Client, "default", &topology[0]).To(ConsistOf(2, 2))
 		})
+		It("should not count zones that are only offered by NodePools the pod can't select", func() {
+			// Zones only reachable through another NodePool hold no pods, so counting them would peg the
+			// spread's minimum at zero and leave the pod's own zones permanently outside maxSkew.
+			nodePool.Spec.Template.Spec.Requirements = []v1.NodeSelectorRequirementWithMinValues{
+				{Key: corev1.LabelTopologyZone, Operator: corev1.NodeSelectorOpIn, Values: []string{"test-zone-1", "test-zone-2"}}}
+			otherNodePool := test.NodePool(v1.NodePool{Spec: v1.NodePoolSpec{Template: v1.NodeClaimTemplate{Spec: v1.NodeClaimTemplateSpec{
+				Requirements: []v1.NodeSelectorRequirementWithMinValues{
+					{Key: corev1.LabelTopologyZone, Operator: corev1.NodeSelectorOpIn, Values: []string{"test-zone-3"}}},
+			}}}})
+			topology := []corev1.TopologySpreadConstraint{{
+				TopologyKey:       corev1.LabelTopologyZone,
+				WhenUnsatisfiable: corev1.DoNotSchedule,
+				LabelSelector:     &metav1.LabelSelector{MatchLabels: labels},
+				MaxSkew:           1,
+			}}
+			ExpectApplied(ctx, env.Client, nodePool, otherNodePool)
+			ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, prov,
+				test.UnschedulablePods(test.PodOptions{
+					ObjectMeta:                metav1.ObjectMeta{Labels: labels},
+					NodeSelector:              map[string]string{v1.NodePoolLabelKey: nodePool.Name},
+					TopologySpreadConstraints: topology,
+				}, 4)...,
+			)
+			ExpectSkew(ctx, env.Client, "default", &topology[0]).To(ConsistOf(2, 2))
+		})
 		It("should respect NodePool zonal constraints (subset) with labels", func() {
 			nodePool.Spec.Template.Labels = lo.Assign(nodePool.Spec.Template.Labels, map[string]string{corev1.LabelTopologyZone: "test-zone-1"})
 			topology := []corev1.TopologySpreadConstraint{{

pkg/controllers/provisioning/scheduling/topologydomaingroup.go

@@ -22,51 +22,60 @@ import (
 	"sigs.k8s.io/karpenter/pkg/scheduling"
 )
 
-// TopologyDomainGroup tracks the domains for a single topology. Additionally, it tracks the taints associated with
-// each of these domains. This enables us to determine which domains should be considered by a pod if its
-// NodeTaintPolicy is honor.
-type TopologyDomainGroup map[string][][]v1.Taint
+// TopologyDomainSource is a NodePool that can supply a domain, along with the taints and
+// requirements a node launched from that NodePool would carry. Taints answer whether a pod with
+// NodeTaintsPolicy honor counts the domain; requirements answer the same for NodeAffinityPolicy.
+type TopologyDomainSource struct {
+	Taints       []v1.Taint
+	Requirements scheduling.Requirements
+}
+
+// TopologyDomainGroup tracks the domains for a single topology, keyed by domain and then by the
+// name of each NodePool that can supply it.
+type TopologyDomainGroup map[string]map[string]TopologyDomainSource
 
 func NewTopologyDomainGroup() TopologyDomainGroup {
-	return map[string][][]v1.Taint{}
+	return TopologyDomainGroup{}
 }
 
-// Insert either adds a new domain to the TopologyDomainGroup or updates an existing domain.
-func (t TopologyDomainGroup) Insert(domain string, taints ...v1.Taint) {
-	// If the domain is not currently tracked, insert it with the associated taints. Additionally, if there are no taints
-	// provided, override the taints associated with the domain. Generally, we could remove any sets of taints for which
-	// the provided set is a proper subset. This is because if a pod tolerates the supersets, it will also tolerate the
-	// proper subset, and removing the superset reduces the number of taint sets we need to traverse. For now we only
-	// implement the simplest case, the empty set, but we could do additional performance testing to determine if
-	// implementing the general case is worth the precomputation cost.
-	if _, ok := t[domain]; !ok || len(taints) == 0 {
-		t[domain] = [][]v1.Taint{taints}
-		return
+// Insert records that nodePool can supply domain, on nodes carrying the given taints and
+// requirements.
+func (t TopologyDomainGroup) Insert(domain string, nodePool string, taints []v1.Taint, requirements scheduling.Requirements) {
+	sources, ok := t[domain]
+	if !ok {
+		sources = map[string]TopologyDomainSource{}
+		t[domain] = sources
 	}
-	if len(t[domain][0]) == 0 {
-		// This is the base case, where we're already tracking the empty set of taints for the domain. Pods will always
-		// be eligible for NodeClaims with this domain (based on taints), so there is no need to track additional taints.
-		return
-	}
-	t[domain] = append(t[domain], taints)
+	sources[nodePool] = TopologyDomainSource{Taints: taints, Requirements: requirements}
 }
 
-// ForEachDomain calls f on each domain tracked by the topology group. If the taintHonorPolicy is honor, only domains
-// available on nodes tolerated by the provided pod will be included.
-func (t TopologyDomainGroup) ForEachDomain(pod *v1.Pod, taintHonorPolicy v1.NodeInclusionPolicy, f func(domain string)) {
-	for domain, taintGroups := range t {
-		if taintHonorPolicy == v1.NodeInclusionPolicyIgnore {
-			f(domain)
-			continue
-		}
-		// Since the taint policy is honor, we should only call f if there is a set of taints associated with the domain which
-		// the pod tolerates.
-		// Perf Note: We could consider hashing the pod's tolerations and using that to look up a set of tolerated domains.
-		for _, taints := range taintGroups {
-			if err := scheduling.Taints(taints).ToleratesPod(pod); err == nil {
-				f(domain)
-				break
+// ForEachDomain calls f on each domain tracked by the topology group that the pod could actually
+// land in, given at least one NodePool supplying that domain. If the taint policy is honor, the pod
+// must tolerate that NodePool's taints; if the affinity policy is honor, the pod's node selector and
+// required node affinity must not conflict with the NodePool's requirements. A domain is only
+// dropped on an outright conflict, so a pod selecting a label that the NodePool leaves to its
+// instance types keeps the domain.
+//
+// Honoring affinity here is what keeps a spread's global minimum meaningful in a cluster whose
+// NodePools do not all offer the same domains. A pod pinned to one NodePool would otherwise count
+// every domain reachable only through the other pools (for example the zones of a pool spanning
+// another region), each with a pod count of zero, pinning the global minimum at zero and leaving no
+// domain within maxSkew of it, so a DoNotSchedule spread could never be satisfied.
+func (t TopologyDomainGroup) ForEachDomain(pod *v1.Pod, nodeFilter TopologyNodeFilter, f func(domain string)) {
+	for domain, sources := range t {
+		for _, source := range sources {
+			if nodeFilter.TaintPolicy != v1.NodeInclusionPolicyIgnore {
+				// Perf Note: We could consider hashing the pod's tolerations and using that to look up a set of
+				// tolerated domains.
+				if err := scheduling.Taints(source.Taints).ToleratesPod(pod); err != nil {
+					continue
+				}
+			}
+			if nodeFilter.AffinityPolicy == v1.NodeInclusionPolicyHonor && nodeFilter.ConflictsWithRequirements(source.Requirements) {
+				continue
 			}
+			f(domain)
+			break
 		}
 	}
 }