scheduler: support designated vf

Signed-off-by: 乔普 <wangjianyu.wjy@alibaba-inc.com>

Author: 乔普

pkg/scheduler/plugins/deviceshare/device_allocator.go

@@ -57,6 +57,8 @@ type requestContext struct {
 	preferred                 map[schedulingv1alpha1.DeviceType]sets.Int
 	allocationScorer          *resourceAllocationScorer
 	nodeDevice                *nodeDevice
+
+	designatedVF map[schedulingv1alpha1.DeviceType]map[int32]sets.Set[string]
 }
 
 type AutopilotAllocator struct {
@@ -101,6 +103,7 @@ func (a *AutopilotAllocator) Allocate(
 		return nil, status
 	}
 
+	designatedVF := constructDesignatedVF(a.state.designatedAllocation)
 	nodeDevice := a.filterNodeDevice(requiredDeviceResources, preemptibleDeviceResources)
 	requestCtx := &requestContext{
 		pod:                       a.pod,
@@ -114,6 +117,7 @@ func (a *AutopilotAllocator) Allocate(
 		required:                  required,
 		preferred:                 preferred,
 		nodeDevice:                a.nodeDevice,
+		designatedVF:              designatedVF,
 	}
 	var deviceAllocations apiext.DeviceAllocations
 	var status *framework.Status
@@ -137,6 +141,30 @@ func (a *AutopilotAllocator) Allocate(
 	return deviceAllocations, nil
 }
 
+func constructDesignatedVF(designatedAllocation apiext.DeviceAllocations) map[schedulingv1alpha1.DeviceType]map[int32]sets.Set[string] {
+	if designatedAllocation == nil {
+		return nil
+	}
+	designatedVF := make(map[schedulingv1alpha1.DeviceType]map[int32]sets.Set[string])
+	for deviceType, allocations := range designatedAllocation {
+		vfOfType := make(map[int32]sets.Set[string])
+		for _, allocation := range allocations {
+			if allocation.Extension != nil && len(allocation.Extension.VirtualFunctions) != 0 {
+				minor := allocation.Minor
+				vfOfMinor := sets.New[string]()
+				for _, function := range allocation.Extension.VirtualFunctions {
+					vfOfMinor.Insert(function.BusID)
+				}
+				vfOfType[minor] = vfOfMinor
+			}
+		}
+		if len(vfOfType) != 0 {
+			designatedVF[deviceType] = vfOfType
+		}
+	}
+	return designatedVF
+}
+
 func (a *AutopilotAllocator) filterNodeDevice(
 	requiredDeviceResources, preemptibleDeviceResources map[schedulingv1alpha1.DeviceType]deviceResources,
 ) *nodeDevice {
@@ -370,6 +398,7 @@ func defaultAllocateDevices(
 		minor := ptr.Deref[int32](v.Minor, 0)
 		deviceInfos[int(minor)] = v
 	}
+	designatedVFOfType := requestCtx.designatedVF[deviceType]
 
 	var allocations []*apiext.DeviceAllocation
 	resourceMinorPairs := scoreDevices(podRequestPerInstance, nodeDeviceTotal, freeDevices, requestCtx.allocationScorer)
@@ -396,9 +425,15 @@ func defaultAllocateDevices(
 			Minor:     int32(resourceMinorPair.minor),
 			Resources: podRequestPerInstance,
 		}
+
+		var designatedVFOfMinor sets.Set[string]
+		if designatedVFOfType != nil {
+			designatedVFOfMinor = designatedVFOfType[r.Minor]
+		}
+
 		if mustAllocateVF(hint) {
 			// TODO Device allocation logic hotspots discovered through flame graphs
-			vf := allocateVF(vfAllocation, deviceInfos, resourceMinorPair.minor, vfSelector)
+			vf := allocateVF(vfAllocation, deviceInfos, resourceMinorPair.minor, vfSelector, designatedVFOfMinor)
 			if vf == nil {
 				continue
 			}
@@ -426,7 +461,7 @@ func defaultAllocateDevices(
 	return allocations, nil
 }
 
-func allocateVF(vfAllocation *VFAllocation, deviceInfos map[int]*schedulingv1alpha1.DeviceInfo, minor int, vfSelector labels.Selector) *schedulingv1alpha1.VirtualFunction {
+func allocateVF(vfAllocation *VFAllocation, deviceInfos map[int]*schedulingv1alpha1.DeviceInfo, minor int, vfSelector labels.Selector, designatedVFOfMinor sets.Set[string]) *schedulingv1alpha1.VirtualFunction {
 	deviceInfo := deviceInfos[minor]
 	if deviceInfo == nil {
 		return nil
@@ -440,9 +475,13 @@ func allocateVF(vfAllocation *VFAllocation, deviceInfos map[int]*schedulingv1alp
 	for _, vfGroup := range deviceInfo.VFGroups {
 		if vfSelector == nil || vfSelector.Matches(labels.Set(vfGroup.Labels)) {
 			for _, vf := range vfGroup.VFs {
-				if !allocated.Has(vf.BusID) {
-					remainingVFs = append(remainingVFs, vf)
+				if allocated.Has(vf.BusID) {
+					continue
+				}
+				if designatedVFOfMinor != nil && !designatedVFOfMinor.Has(vf.BusID) {
+					continue
 				}
+				remainingVFs = append(remainingVFs, vf)
 			}
 		}
 	}