gpu: setting default gpu partition policy (#2245)

Signed-off-by: wangjianyu.wjy <[email protected]> Co-authored-by: wangjianyu.wjy <[email protected]>
koordinator-sh · Nov 4, 2024 · 4555e24 · 4555e24
1 parent a5f0fbd
commit 4555e24
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 39 deletions.
diff --git a/apis/extension/device_share.go b/apis/extension/device_share.go
@@ -323,11 +323,11 @@ func GetGPUPartitionTable(device *schedulingv1alpha1.Device) (GPUPartitionTable,
 	return nil, nil
 }
 
-func GetGPUPartitionPolicy(device *schedulingv1alpha1.Device) GPUPartitionPolicy {
-	if device == nil {
+func GetGPUPartitionPolicy(nodeOrDevice metav1.Object) GPUPartitionPolicy {
+	if nodeOrDevice == nil {
 		return GPUPartitionPolicyPrefer
 	}
-	if allocatePolicy := device.Labels[LabelGPUPartitionPolicy]; GPUPartitionPolicy(allocatePolicy) == GPUPartitionPolicyHonor {
+	if allocatePolicy := nodeOrDevice.GetLabels()[LabelGPUPartitionPolicy]; GPUPartitionPolicy(allocatePolicy) == GPUPartitionPolicyHonor {
 		return GPUPartitionPolicyHonor
 	}
 	return GPUPartitionPolicyPrefer

diff --git a/apis/extension/device_share_test.go b/apis/extension/device_share_test.go
@@ -289,11 +289,6 @@ func TestGetNodeLevelGPUAllocatePolicy(t *testing.T) {
 		node     *schedulingv1alpha1.Device
 		expected GPUPartitionPolicy
 	}{
-		{
-			name:     "Nil node",
-			node:     nil,
-			expected: GPUPartitionPolicyPrefer,
-		},
 		{
 			name: "Node with Honor policy",
 			node: &schedulingv1alpha1.Device{

diff --git a/pkg/scheduler/plugins/deviceshare/allocator_gpu.go b/pkg/scheduler/plugins/deviceshare/allocator_gpu.go
@@ -76,10 +76,7 @@ func (a *GPUAllocator) Allocate(requestCtx *requestContext, nodeDevice *nodeDevi
 	nodeHonorPartition := nodeDevice.nodeHonorGPUPartition
 	gpuPartitionIndexer := nodeDevice.gpuPartitionIndexer
 	if gpuPartitionIndexer == nil {
-		gpuPartitionIndexer = GetDesignatedGPUPartitionIndexer(requestCtx.node)
-		if gpuPartitionIndexer != nil {
-			nodeHonorPartition = true
-		}
+		gpuPartitionIndexer, nodeHonorPartition = GetDesignatedGPUPartitionIndexer(requestCtx.node)
 	}
 	honorGPUPartition := gpuRequirements.honorGPUPartition || nodeHonorPartition
 	realUsed := getRealUsed(requestCtx.nodeDevice.deviceUsed[schedulingv1alpha1.GPU], nodeDevice.deviceTotal[schedulingv1alpha1.GPU], nodeDevice.deviceUsed[schedulingv1alpha1.GPU])

diff --git a/pkg/scheduler/plugins/deviceshare/allocator_gpu_helper.go b/pkg/scheduler/plugins/deviceshare/allocator_gpu_helper.go
@@ -147,13 +147,15 @@ var (
 		},
 	}
 
-	GetDesignatedGPUPartitionIndexer = func(node *corev1.Node) GPUPartitionIndexer {
+	GetDesignatedGPUPartitionIndexer = func(node *corev1.Node) (GPUPartitionIndexer, bool) {
+		var partitionIndexer GPUPartitionIndexer
+		partitionPolicy := apiext.GetGPUPartitionPolicy(node)
 		model := node.Labels[apiext.LabelGPUModel]
 		switch model {
 		case "H100", "H800", "H20":
-			return gpuPartitionIndexOfNVIDIAHopper
+			partitionIndexer = gpuPartitionIndexOfNVIDIAHopper
 		}
-		return nil
+		return partitionIndexer, partitionPolicy == apiext.GPUPartitionPolicyHonor
 	}
 )
 

diff --git a/pkg/scheduler/plugins/deviceshare/allocator_gpu_test.go b/pkg/scheduler/plugins/deviceshare/allocator_gpu_test.go
@@ -51,14 +51,15 @@ var (
 
 func TestAllocateByPartition(t *testing.T) {
 	tests := []struct {
-		name            string
-		deviceCR        *schedulingv1alpha1.Device
-		gpuWanted       int
-		hostNetwork     bool
-		assignedDevices apiext.DeviceAllocations
-		modelSeries     string
-		want            apiext.DeviceAllocations
-		wantErr         bool
+		name               string
+		deviceCR           *schedulingv1alpha1.Device
+		gpuPartitionPolicy apiext.GPUPartitionPolicy
+		gpuWanted          int
+		hostNetwork        bool
+		assignedDevices    apiext.DeviceAllocations
+		modelSeries        string
+		want               apiext.DeviceAllocations
+		wantErr            bool
 	}{
 		{
 			name:        "allocate 0 GPU and 1 VF",
@@ -636,9 +637,9 @@ func TestAllocateByPartition(t *testing.T) {
 			},
 		},
 		{
-			name:        "allocate 2 GPU and 2 VF with assigned devices; use partition table from device",
-			deviceCR:    fakeH800DeviceCR,
-			modelSeries: "fakeModelSeries",
+			name:        "allocate 2 GPU and 2 VF with assigned devices; H100",
+			deviceCR:    fakeH100DeviceCR,
+			modelSeries: "H100",
 			gpuWanted:   2,
 			assignedDevices: apiext.DeviceAllocations{
 				schedulingv1alpha1.GPU: []*apiext.DeviceAllocation{
@@ -727,10 +728,11 @@ func TestAllocateByPartition(t *testing.T) {
 			},
 		},
 		{
-			name:        "allocate 2 GPU and 2 VF with assigned devices; H100",
-			deviceCR:    fakeH100DeviceCR,
-			modelSeries: "H100",
-			gpuWanted:   2,
+			name:               "allocate 2 GPU and 2 VF with assigned devices; H100; gpuPartitionPolicyPrefer",
+			deviceCR:           fakeH100DeviceCR,
+			modelSeries:        "H100",
+			gpuWanted:          3,
+			gpuPartitionPolicy: apiext.GPUPartitionPolicyPrefer,
 			assignedDevices: apiext.DeviceAllocations{
 				schedulingv1alpha1.GPU: []*apiext.DeviceAllocation{
 					{
@@ -777,39 +779,57 @@ func TestAllocateByPartition(t *testing.T) {
 			want: apiext.DeviceAllocations{
 				schedulingv1alpha1.GPU: []*apiext.DeviceAllocation{
 					{
-						Minor:     0,
+						Minor:     4,
 						Resources: gpuResourceList,
 					},
 					{
-						Minor:     1,
+						Minor:     5,
+						Resources: gpuResourceList,
+					},
+					{
+						Minor:     6,
 						Resources: gpuResourceList,
 					},
 				},
 				schedulingv1alpha1.RDMA: []*apiext.DeviceAllocation{
 					{
-						Minor: 1,
+						Minor: 5,
 						Resources: corev1.ResourceList{
 							apiext.ResourceRDMA: *resource.NewQuantity(1, resource.DecimalSI),
 						},
 						Extension: &apiext.DeviceAllocationExtension{
 							VirtualFunctions: []apiext.VirtualFunction{
 								{
-									BusID: "0000:09:00.4",
-									Minor: 2,
+									BusID: "0001:08:00.3",
+									Minor: 1,
 								},
 							},
 						},
 					},
 					{
-						Minor: 2,
+						Minor: 6,
 						Resources: corev1.ResourceList{
 							apiext.ResourceRDMA: *resource.NewQuantity(1, resource.DecimalSI),
 						},
 						Extension: &apiext.DeviceAllocationExtension{
 							VirtualFunctions: []apiext.VirtualFunction{
 								{
-									BusID: "0000:7f:00.4",
-									Minor: 2,
+									BusID: "0001:7e:00.3",
+									Minor: 1,
+								},
+							},
+						},
+					},
+					{
+						Minor: 7,
+						Resources: corev1.ResourceList{
+							apiext.ResourceRDMA: *resource.NewQuantity(1, resource.DecimalSI),
+						},
+						Extension: &apiext.DeviceAllocationExtension{
+							VirtualFunctions: []apiext.VirtualFunction{
+								{
+									BusID: "0001:a2:00.3",
+									Minor: 1,
 								},
 							},
 						},
@@ -831,10 +851,14 @@ func TestAllocateByPartition(t *testing.T) {
 				ObjectMeta: metav1.ObjectMeta{
 					Name: "test-node-1",
 					Labels: map[string]string{
-						apiext.LabelGPUModel: tt.modelSeries,
+						apiext.LabelGPUModel:           tt.modelSeries,
+						apiext.LabelGPUPartitionPolicy: string(apiext.GPUPartitionPolicyHonor),
 					},
 				},
 			}
+			if tt.gpuPartitionPolicy != "" {
+				node.Labels[apiext.LabelGPUPartitionPolicy] = string(tt.gpuPartitionPolicy)
+			}
 			kubeFakeClient := kubefake.NewSimpleClientset(node)
 			sharedInformerFactory := informers.NewSharedInformerFactory(kubeFakeClient, 0)