Skip to content

Commit 23afd8f

Browse files
davidLifdavidLif
committed
feat(scheduler): account for NvFractions GPU memory
Signed-off-by: davidLif <davidshani12@gmail.com>
1 parent be5c501 commit 23afd8f

11 files changed

Lines changed: 177 additions & 98 deletions

File tree

‎.changes/unreleased/added-20260810-110236.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
kind: Added
2+
body: |-
3+
Schedule NvFractions GPU memory requests

‎pkg/nodescaleadjuster/scale_adjuster/calculator.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,11 +77,11 @@ func (c *calculator) getGPUFraction(pod *v1.Pod) (float64, error) {
7777
}
7878
return gpuFraction, nil
7979
}
80-
if pod.Annotations[constants.GpuMemory] != "" {
81-
_, err := resources.GetGPUMemory(pod)
82-
if err != nil {
83-
return 0, err
84-
}
80+
gpuMemory, err := resources.GetGPUMemory(pod)
81+
if err != nil {
82+
return 0, err
83+
}
84+
if gpuMemory > 0 {
8585
return c.gpuMemoryToFractionRatio, nil
8686
}
8787
return 0, fmt.Errorf("pod %v/%v does not have GPU fraction or memory annotation", pod.Namespace, pod.Name)

‎pkg/nodescaleadjuster/scale_adjuster/scale_adjuster.go‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ import (
1414
corev1 "k8s.io/api/core/v1"
1515
"sigs.k8s.io/controller-runtime/pkg/client"
1616

17-
"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
1817
"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
1918
"github.com/kai-scheduler/KAI-scheduler/pkg/nodescaleadjuster/scaler"
2019
)
@@ -160,7 +159,7 @@ func (sa *ScaleAdjuster) getUnschedulablePods() ([]*corev1.Pod, error) {
160159
if pod.Spec.SchedulerName != sa.schedulerName {
161160
continue
162161
}
163-
if !requestFractionalGPU(&pod) {
162+
if !resources.RequestsGPUFraction(&pod) {
164163
continue
165164
}
166165
if !isPodAlive(&pod) {
@@ -175,10 +174,6 @@ func (sa *ScaleAdjuster) getUnschedulablePods() ([]*corev1.Pod, error) {
175174
return pods, nil
176175
}
177176

178-
func requestFractionalGPU(pod *corev1.Pod) bool {
179-
return pod.Annotations[constants.GpuFraction] != "" || pod.Annotations[constants.GpuMemory] != ""
180-
}
181-
182177
func isPodAlive(pod *corev1.Pod) bool {
183178
return !slices.Contains([]corev1.PodPhase{corev1.PodSucceeded, corev1.PodFailed}, pod.Status.Phase)
184179
}

‎pkg/podgroupcontroller/controllers/resources/fraction.go‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"sigs.k8s.io/controller-runtime/pkg/client"
1717

1818
"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
19+
"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
1920
)
2021

2122
const (
@@ -31,28 +32,23 @@ var (
3132
func calculateAllocatedFraction(
3233
ctx context.Context, pod *v1.Pod, kubeClient client.Client,
3334
) (resource.Quantity, error) {
34-
gpuFractionStr, hasFractionAnnotation := pod.Annotations[constants.GpuFraction]
35-
if hasFractionAnnotation {
36-
return resource.MustParse(gpuFractionStr), nil
35+
req, err := resources.ParsePodGPUFractionRequest(pod)
36+
if err != nil {
37+
return resource.Quantity{}, fmt.Errorf("failed to parse GPU fraction for pod %s/%s: %s", pod.Namespace, pod.Name, err)
3738
}
38-
39-
gpuMemoryStr, hasMemoryAnnotation := pod.Annotations[constants.GpuMemory]
40-
if !hasMemoryAnnotation {
41-
return resource.Quantity{}, fmt.Errorf(
42-
"cannot calculate fraction because the pod doesn't a fraction or memory annotation")
39+
if req == nil {
40+
return resource.Quantity{}, fmt.Errorf("cannot calculate fraction because the pod doesn't have a fraction or memory annotation")
4341
}
4442

45-
return getFractionFromMemoryRequest(ctx, gpuMemoryStr, pod.Spec.NodeName, kubeClient)
43+
if req.Portion > 0 {
44+
return resource.MustParse(fmt.Sprintf("%g", req.Portion)), nil
45+
}
46+
return getFractionFromMemoryRequest(ctx, req.Memory.Value()/resources.BytesInMiB, pod.Spec.NodeName, kubeClient)
4647
}
4748

4849
func getFractionFromMemoryRequest(
49-
ctx context.Context, gpuMemoryStr string, nodeName string, kubeClient client.Client,
50+
ctx context.Context, gpuMemory int64, nodeName string, kubeClient client.Client,
5051
) (resource.Quantity, error) {
51-
gpuMemory, err := strconv.ParseInt(gpuMemoryStr, 10, 64)
52-
if err != nil {
53-
return resource.Quantity{}, fmt.Errorf("failed to parse %s annotation to int: %w", constants.GpuMemory, err)
54-
}
55-
5652
nodeGpuMemory, err := getNodeSingleGpuMemory(ctx, nodeName, kubeClient)
5753
if err != nil {
5854
return resource.Quantity{}, fmt.Errorf("failed extract node gpu memory for node %s : %w",

‎pkg/podgroupcontroller/controllers/resources/fraction_test.go‎

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"sigs.k8s.io/controller-runtime/pkg/client/fake"
1616

1717
"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
18+
"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
1819
)
1920

2021
func Test_getReceivedFraction(t *testing.T) {
@@ -53,6 +54,23 @@ func Test_getReceivedFraction(t *testing.T) {
5354
resource.MustParse("0.5"),
5455
false,
5556
},
57+
{
58+
"NvFractions request + Nvidia node",
59+
&v1.Pod{
60+
ObjectMeta: metav1.ObjectMeta{
61+
Annotations: map[string]string{resources.CalcGpuFractionAnnotationForContainer(""): "2000Mi"},
62+
},
63+
Spec: v1.PodSpec{NodeName: "n1"},
64+
},
65+
&v1.Node{
66+
ObjectMeta: metav1.ObjectMeta{
67+
Name: "n1",
68+
Labels: map[string]string{constants.NvidiaGpuMemory: "4000"},
69+
},
70+
},
71+
resource.MustParse("0.5"),
72+
false,
73+
},
5674
{
5775
"Memory request + Amd node",
5876
&v1.Pod{
@@ -143,8 +161,8 @@ func Test_getReceivedFraction(t *testing.T) {
143161

144162
func Test_getFractionFromMemoryRequest(t *testing.T) {
145163
type args struct {
146-
gpuMemoryStr string
147-
nodeName string
164+
gpuMemory int64
165+
nodeName string
148166
}
149167
tests := []struct {
150168
name string
@@ -156,8 +174,8 @@ func Test_getFractionFromMemoryRequest(t *testing.T) {
156174
{
157175
"Node with Nvidia memory label",
158176
args{
159-
gpuMemoryStr: "2000",
160-
nodeName: "n1",
177+
gpuMemory: 2000,
178+
nodeName: "n1",
161179
},
162180
&v1.Node{
163181
ObjectMeta: metav1.ObjectMeta{
@@ -171,8 +189,8 @@ func Test_getFractionFromMemoryRequest(t *testing.T) {
171189
{
172190
"Node with Amd memory label",
173191
args{
174-
gpuMemoryStr: "4000",
175-
nodeName: "n1",
192+
gpuMemory: 4000,
193+
nodeName: "n1",
176194
},
177195
&v1.Node{
178196
ObjectMeta: metav1.ObjectMeta{
@@ -193,7 +211,7 @@ func Test_getFractionFromMemoryRequest(t *testing.T) {
193211
}
194212
kubeClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(tt.node).Build()
195213

196-
got, err := getFractionFromMemoryRequest(context.TODO(), tt.args.gpuMemoryStr, tt.args.nodeName, kubeClient)
214+
got, err := getFractionFromMemoryRequest(context.TODO(), tt.args.gpuMemory, tt.args.nodeName, kubeClient)
197215
if (err != nil) != tt.wantErr {
198216
t.Errorf("getFractionFromMemoryRequest() error = %v, wantErr %v", err, tt.wantErr)
199217
return

‎pkg/podgroupcontroller/controllers/resources/requested.go‎

Lines changed: 20 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"k8s.io/apimachinery/pkg/api/resource"
1111

1212
"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
13+
common_resources "github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
1314
)
1415

1516
const (
@@ -19,56 +20,36 @@ const (
1920
func ExtractGPUSharingRequestedResources(pod *v1.Pod) (v1.ResourceList, error) {
2021
resources := v1.ResourceList{}
2122

22-
fractionsCount := int64(1)
23-
gpuFractionsCountStr, hasAnnotation := pod.Annotations[constants.GpuFractionsNumDevices]
24-
if hasAnnotation {
25-
quantity, err := resource.ParseQuantity(gpuFractionsCountStr)
26-
if err != nil {
27-
return v1.ResourceList{},
28-
fmt.Errorf("failed to parse gpu fraction count annotation value <%s>, error: %s",
29-
gpuFractionsCountStr, err.Error())
30-
}
31-
var successfulIntExtraction bool
32-
fractionsCount, successfulIntExtraction = quantity.AsInt64()
33-
if !successfulIntExtraction {
34-
return v1.ResourceList{},
35-
fmt.Errorf("failed to extract int value from gpu fraction count annotation. value <%s>",
36-
gpuFractionsCountStr)
37-
}
23+
req, err := common_resources.ParsePodGPUFractionRequest(pod)
24+
if err != nil {
25+
return v1.ResourceList{},
26+
fmt.Errorf("failed to parse GPU fraction for pod %s/%s: %s", pod.Namespace, pod.Name, err)
27+
}
28+
if req == nil {
29+
return resources, nil
3830
}
3931

40-
gpuFractionStr, hasAnnotation := pod.Annotations[constants.GpuFraction]
41-
if hasAnnotation {
42-
quantity, err := resource.ParseQuantity(gpuFractionStr)
32+
if req.Portion > 0 {
33+
fractionStr := fmt.Sprintf("%g", req.Portion)
34+
quantity, err := resource.ParseQuantity(fractionStr)
4335
if err != nil {
4436
return v1.ResourceList{},
45-
fmt.Errorf("failed to parse gpu fraction annotation value <%s>, error: %s",
46-
gpuFractionStr, err.Error())
37+
fmt.Errorf("failed to parse gpu fraction value <%s>: %s", fractionStr, err)
4738
}
48-
successfulMulti := quantity.Mul(fractionsCount)
49-
if !successfulMulti {
39+
if ok := quantity.Mul(req.NumDevices); !ok {
5040
return v1.ResourceList{},
51-
fmt.Errorf("failed to multiple gpu fraction by the fraction count. "+
52-
"Please check resource.Quantity restrictions. fraction <%s>, count: %d",
53-
gpuFractionStr, fractionsCount)
41+
fmt.Errorf("failed to multiply gpu fraction by device count. fraction <%s>, count: %d",
42+
fractionStr, req.NumDevices)
5443
}
5544
resources[v1.ResourceName(constants.NvidiaGpuResource)] = quantity
5645
}
5746

58-
gpuMemoryStr, hasAnnotation := pod.Annotations[constants.GpuMemory]
59-
if hasAnnotation {
60-
quantity, err := resource.ParseQuantity(gpuMemoryStr)
61-
if err != nil {
62-
return v1.ResourceList{},
63-
fmt.Errorf("failed to parse gpu memory annotation value <%s>, error: %s",
64-
gpuMemoryStr, err.Error())
65-
}
66-
successfulMulti := quantity.Mul(fractionsCount)
67-
if !successfulMulti {
47+
if req.Memory != nil {
48+
quantity := req.Memory.DeepCopy()
49+
if ok := quantity.Mul(req.NumDevices); !ok {
6850
return v1.ResourceList{},
69-
fmt.Errorf("failed to multiple gpu memory by the fraction count. "+
70-
"Please check resource.Quantity restrictions. fraction <%s>, count: %d",
71-
gpuMemoryStr, fractionsCount)
51+
fmt.Errorf("failed to multiply gpu memory by device count. memory <%s>, count: %d",
52+
req.Memory.String(), req.NumDevices)
7253
}
7354
resources[v1.ResourceName(gpuMemoryResourceName)] = quantity
7455
}

‎pkg/podgroupcontroller/controllers/resources/requested_test.go‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
1313

1414
"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
15+
"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
1516
)
1617

1718
func Test_extractRequestedResources(t *testing.T) {
@@ -48,7 +49,16 @@ func Test_extractRequestedResources(t *testing.T) {
4849
Annotations: map[string]string{constants.GpuMemory: "2000"},
4950
},
5051
},
51-
v1.ResourceList{gpuMemoryResourceName: resource.MustParse("2000")},
52+
v1.ResourceList{gpuMemoryResourceName: resource.MustParse("2000Mi")},
53+
},
54+
{
55+
"Pod with NvFractions request",
56+
&v1.Pod{
57+
ObjectMeta: metav1.ObjectMeta{
58+
Annotations: map[string]string{resources.CalcGpuFractionAnnotationForContainer(""): "2000Mi"},
59+
},
60+
},
61+
v1.ResourceList{gpuMemoryResourceName: resource.MustParse("2000Mi")},
5262
},
5363
{
5464
"Regular cpu pod",
@@ -92,15 +102,15 @@ func compareResourceLists(rightList, leftList v1.ResourceList) error {
92102
" right-list: %v, left-list: %v", rightList, leftList)
93103
}
94104

95-
for resourceName, rightListResourceValue := range rightList {
96-
leftListResourceValue, ok := leftList[resourceName]
105+
for resourceName, rightListResource := range rightList {
106+
leftListResource, ok := leftList[resourceName]
97107
if !ok {
98108
return fmt.Errorf("the resource %s exists in the rightList but doesn't exist in the leftList",
99109
resourceName)
100110
}
101-
if !rightListResourceValue.Equal(leftListResourceValue) {
111+
if rightListResource.Value() != leftListResource.Value() {
102112
return fmt.Errorf("for the resource %s, the values differ between the list."+
103-
" right-list: %v, left-list: %v", resourceName, rightListResourceValue, leftListResourceValue)
113+
" right-list: %v, left-list: %v", resourceName, rightListResource, leftListResource)
104114
}
105115
}
106116

‎pkg/podgroupcontroller/controllers/status_updater_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -720,7 +720,7 @@ func Test_handlePodGroupStatus(t *testing.T) {
720720
Requested: map[v1.ResourceName]resource.Quantity{
721721
v1.ResourceCPU: resource.MustParse("500m"),
722722
v1.ResourceMemory: resource.MustParse("1G"),
723-
"run.ai/gpu.memory": resource.MustParse("1k"),
723+
"run.ai/gpu.memory": resource.MustParse("1000Mi"),
724724
},
725725
},
726726
},

‎pkg/scheduler/api/node_info/node_info.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -670,7 +670,7 @@ func getNodeGpuMemory(node *v1.Node) (int64, bool) {
670670
gpuMemoryLabelValue = convertBytesToMib(gpuMemoryLabelValue)
671671
}
672672

673-
return gpuMemoryLabelValue - (gpuMemoryLabelValue % 100), true // Floor the memory count to make sure its divided by 100 so there will not be 2 jobs that get same bytes
673+
return gpuMemoryLabelValue, true
674674
}
675675

676676
func checkGpuMemoryIsInMib(gpuMemoryValue int64) bool {

0 commit comments

Comments
 (0)