DRA: Allow AllocationMode: All from multi-node resource pools

staging/src/k8s.io/dynamic-resource-allocation/structured/internal/allocatortesting/allocator_testing.go

@@ -5090,6 +5090,40 @@ func TestAllocator(t *testing.T,
 				),
 			},
 		},
+		"allocation-mode-all-with-multi-host-resource-pool": {
+			claimsToAllocate: objects(claimWithRequests(claim0, nil, resourceapi.DeviceRequest{
+				Name: req0,
+				Exactly: &resourceapi.ExactDeviceRequest{
+					AllocationMode:  resourceapi.DeviceAllocationModeAll,
+					DeviceClassName: classA,
+				},
+			})),
+			classes: objects(class(classA, driverA)),
+			slices: unwrap(
+				func() wrapResourceSlice {
+					s := slice(slice1, node1, pool1, driverA,
+						device(device1, nil, nil),
+					)
+					s.Spec.Pool.ResourceSliceCount = 2
+					return s
+				}(),
+				func() wrapResourceSlice {
+					s := slice(slice2, node2, pool1, driverA,
+						device(device2, nil, nil),
+					)
+					s.Spec.Pool.ResourceSliceCount = 2
+					return s
+				}(),
+			),
+			node: node(node1, region1),
+
+			expectResults: []any{
+				allocationResult(
+					localNodeSelector(node1),
+					deviceAllocationResult(req0, driverA, pool1, device1, false),
+				),
+			},
+		},
 	}
 
 	for name, tc := range testcases {

staging/src/k8s.io/dynamic-resource-allocation/structured/internal/experimental/pools_experimental.go

@@ -45,6 +45,10 @@ func nodeMatches(node *v1.Node, nodeNameToMatch string, allNodesMatch bool, node
 	return false, nil
 }
 
+type poolIdentifier struct {
+	driver, pool string
+}
+
 // GatherPools collects information about all resource pools which provide
 // devices that are accessible from the given node.
 //
@@ -52,9 +56,46 @@ func nodeMatches(node *v1.Node, nodeNameToMatch string, allNodesMatch bool, node
 // required slices available) or invalid (for example, device names not unique).
 // Both is recorded in the result.
 func GatherPools(ctx context.Context, slices []*resourceapi.ResourceSlice, node *v1.Node, features Features) ([]*Pool, error) {
+	slicesByPool := make(map[poolIdentifier][]*resourceapi.ResourceSlice)
+	for _, slice := range slices {
+		poolID := poolIdentifier{
+			driver: slice.Spec.Driver,
+			pool:   slice.Spec.Pool.Name,
+		}
+		slicesByPool[poolID] = append(slicesByPool[poolID], slice)
+	}
+
+	// We need to check whether a pool is complete while we have all
+	// the slices. Once we discard slices that don't target the node, we
+	// no longer have the information needed to find out.
+	incompletePools := sets.New[poolIdentifier]()
+	for poolID, slices := range slicesByPool {
+		complete := true
+		sliceCount := len(slices)
+		generation := slices[0].Spec.Pool.Generation
+		for _, slice := range slices {
+			// If the number of slices in the pool specified in any of the slices
+			// doesn't match what we found, the pool is most likely being updated
+			// by the controller.
+			if slice.Spec.Pool.ResourceSliceCount != int64(sliceCount) {
+				complete = false
+			}
+			// If the generation of the pool isn't the same across all slices,
+			// the pool is most likely being updated by the controller. We can't
+			// allocate devices from it.
+			if slice.Spec.Pool.Generation != generation {
+				complete = false
+			}
+		}
+		// We still need to keep incomplete pools, since we need to make sure
+		// all devices available on a node is considered for allocationMode All.
+		if !complete {
+			incompletePools.Insert(poolID)
+		}
+	}
+
 	pools := make(map[PoolID]*Pool)
 	var slicesWithBindingConditions []*resourceapi.ResourceSlice
-
 	for _, slice := range slices {
 		if !features.PartitionableDevices && slice.Spec.PerDeviceNodeSelection != nil {
 			continue
@@ -121,7 +162,7 @@ func GatherPools(ctx context.Context, slices []*resourceapi.ResourceSlice, node
 	result := make([]*Pool, 0, len(pools))
 	var resultWithBindingConditions []*Pool
 	for _, pool := range pools {
-		pool.IsIncomplete = int64(len(pool.Slices)) != pool.Slices[0].Spec.Pool.ResourceSliceCount
+		pool.IsIncomplete = incompletePools.Has(poolIdentifier{driver: pool.Driver.String(), pool: pool.Pool.String()})
 		pool.IsInvalid, pool.InvalidReason = poolIsInvalid(pool)
 		// if pool has binding conditions, add the pool to the end of the result
 		if poolHasBindingConditions(*pool) {

staging/src/k8s.io/dynamic-resource-allocation/structured/internal/incubating/pools_incubating.go

@@ -45,6 +45,10 @@ func nodeMatches(node *v1.Node, nodeNameToMatch string, allNodesMatch bool, node
 	return false, nil
 }
 
+type poolIdentifier struct {
+	driver, pool string
+}
+
 // GatherPools collects information about all resource pools which provide
 // devices that are accessible from the given node.
 //
@@ -52,8 +56,45 @@ func nodeMatches(node *v1.Node, nodeNameToMatch string, allNodesMatch bool, node
 // required slices available) or invalid (for example, device names not unique).
 // Both is recorded in the result.
 func GatherPools(ctx context.Context, slices []*resourceapi.ResourceSlice, node *v1.Node, features Features) ([]*Pool, error) {
-	pools := make(map[PoolID]*Pool)
+	slicesByPool := make(map[poolIdentifier][]*resourceapi.ResourceSlice)
+	for _, slice := range slices {
+		poolID := poolIdentifier{
+			driver: slice.Spec.Driver,
+			pool:   slice.Spec.Pool.Name,
+		}
+		slicesByPool[poolID] = append(slicesByPool[poolID], slice)
+	}
 
+	// We need to check whether a pool is complete while we have all
+	// the slices. Once we discard slices that don't target the node, we
+	// no longer have the information needed to find out.
+	incompletePools := sets.New[poolIdentifier]()
+	for poolID, slices := range slicesByPool {
+		complete := true
+		sliceCount := len(slices)
+		generation := slices[0].Spec.Pool.Generation
+		for _, slice := range slices {
+			// If the number of slices in the pool specified in any of the slices
+			// doesn't match what we found, the pool is most likely being updated
+			// by the controller.
+			if slice.Spec.Pool.ResourceSliceCount != int64(sliceCount) {
+				complete = false
+			}
+			// If the generation of the pool isn't the same across all slices,
+			// the pool is most likely being updated by the controller. We can't
+			// allocate devices from it.
+			if slice.Spec.Pool.Generation != generation {
+				complete = false
+			}
+		}
+		// We still need to keep incomplete pools, since we need to make sure
+		// all devices available on a node is considered for allocationMode All.
+		if !complete {
+			incompletePools.Insert(poolID)
+		}
+	}
+
+	pools := make(map[PoolID]*Pool)
 	for _, slice := range slices {
 		if !features.PartitionableDevices && slice.Spec.PerDeviceNodeSelection != nil {
 			continue
@@ -99,7 +140,7 @@ func GatherPools(ctx context.Context, slices []*resourceapi.ResourceSlice, node
 	// Find incomplete pools and flatten into a single slice.
 	result := make([]*Pool, 0, len(pools))
 	for _, pool := range pools {
-		pool.IsIncomplete = int64(len(pool.Slices)) != pool.Slices[0].Spec.Pool.ResourceSliceCount
+		pool.IsIncomplete = incompletePools.Has(poolIdentifier{driver: pool.Driver.String(), pool: pool.Pool.String()})
 		pool.IsInvalid, pool.InvalidReason = poolIsInvalid(pool)
 		result = append(result, pool)
 	}

staging/src/k8s.io/dynamic-resource-allocation/structured/internal/stable/pools_stable.go

@@ -45,6 +45,10 @@ func nodeMatches(node *v1.Node, nodeNameToMatch string, allNodesMatch bool, node
 	return false, nil
 }
 
+type poolIdentifier struct {
+	driver, pool string
+}
+
 // GatherPools collects information about all resource pools which provide
 // devices that are accessible from the given node.
 //
@@ -52,8 +56,45 @@ func nodeMatches(node *v1.Node, nodeNameToMatch string, allNodesMatch bool, node
 // required slices available) or invalid (for example, device names not unique).
 // Both is recorded in the result.
 func GatherPools(ctx context.Context, slices []*resourceapi.ResourceSlice, node *v1.Node, features Features) ([]*Pool, error) {
-	pools := make(map[PoolID]*Pool)
+	slicesByPool := make(map[poolIdentifier][]*resourceapi.ResourceSlice)
+	for _, slice := range slices {
+		poolID := poolIdentifier{
+			driver: slice.Spec.Driver,
+			pool:   slice.Spec.Pool.Name,
+		}
+		slicesByPool[poolID] = append(slicesByPool[poolID], slice)
+	}
 
+	// We need to check whether a pool is complete while we have all
+	// the slices. Once we discard slices that don't target the node, we
+	// no longer have the information needed to find out.
+	incompletePools := sets.New[poolIdentifier]()
+	for poolID, slices := range slicesByPool {
+		complete := true
+		sliceCount := len(slices)
+		generation := slices[0].Spec.Pool.Generation
+		for _, slice := range slices {
+			// If the number of slices in the pool specified in any of the slices
+			// doesn't match what we found, the pool is most likely being updated
+			// by the controller.
+			if slice.Spec.Pool.ResourceSliceCount != int64(sliceCount) {
+				complete = false
+			}
+			// If the generation of the pool isn't the same across all slices,
+			// the pool is most likely being updated by the controller. We can't
+			// allocate devices from it.
+			if slice.Spec.Pool.Generation != generation {
+				complete = false
+			}
+		}
+		// We still need to keep incomplete pools, since we need to make sure
+		// all devices available on a node is considered for allocationMode All.
+		if !complete {
+			incompletePools.Insert(poolID)
+		}
+	}
+
+	pools := make(map[PoolID]*Pool)
 	for _, slice := range slices {
 		if !features.PartitionableDevices && slice.Spec.PerDeviceNodeSelection != nil {
 			continue
@@ -99,7 +140,7 @@ func GatherPools(ctx context.Context, slices []*resourceapi.ResourceSlice, node
 	// Find incomplete pools and flatten into a single slice.
 	result := make([]*Pool, 0, len(pools))
 	for _, pool := range pools {
-		pool.IsIncomplete = int64(len(pool.Slices)) != pool.Slices[0].Spec.Pool.ResourceSliceCount
+		pool.IsIncomplete = incompletePools.Has(poolIdentifier{driver: pool.Driver.String(), pool: pool.Pool.String()})
 		pool.IsInvalid, pool.InvalidReason = poolIsInvalid(pool)
 		result = append(result, pool)
 	}