Skip to content

Commit be5c501

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

8 files changed

Lines changed: 68 additions & 42 deletions

File tree

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

Lines changed: 0 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -183,21 +183,6 @@ func Test_getFractionFromMemoryRequest(t *testing.T) {
183183
resource.MustParse("0.25"),
184184
false,
185185
},
186-
{
187-
"invalid gpu memory value",
188-
args{
189-
gpuMemoryStr: "abc",
190-
nodeName: "n1",
191-
},
192-
&v1.Node{
193-
ObjectMeta: metav1.ObjectMeta{
194-
Name: "n1",
195-
Labels: map[string]string{constants.NvidiaGpuMemory: "4000"},
196-
},
197-
},
198-
resource.Quantity{},
199-
true,
200-
},
201186
}
202187
for _, tt := range tests {
203188
t.Run(tt.name, func(t *testing.T) {

‎pkg/scheduler/actions/common/solvers/accumulated_scenario_filters/idle_gpus/idle_gpus_test.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
1515

1616
commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
17+
"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
1718
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/actions/common/solvers/scenario"
1819
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api"
1920
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
@@ -1303,8 +1304,8 @@ func TestAccumulatedIdleGpus_Filter(t *testing.T) {
13031304
Name: "pv1",
13041305
Namespace: "n2",
13051306
Annotations: map[string]string{
1306-
commonconstants.GpuMemory: "20",
1307-
commonconstants.PodGroupAnnotationForPod: "potential_victims",
1307+
resources.CalcGpuFractionAnnotationForContainer("main"): "20Mi",
1308+
commonconstants.PodGroupAnnotationForPod: "potential_victims",
13081309
},
13091310
},
13101311
Spec: v1.PodSpec{

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"k8s.io/apimachinery/pkg/types"
1414

1515
commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
16+
"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
1617
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
1718
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/node_info"
1819
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_affinity"
@@ -300,14 +301,15 @@ func createTaskGPUMemory() *pod_info.PodInfo {
300301
Name: "gpu-memory-task",
301302
Namespace: "default",
302303
Annotations: map[string]string{
303-
commonconstants.GpuMemory: "40000",
304-
commonconstants.GpuFraction: "0.5",
304+
resources.CalcGpuFractionAnnotationForContainer("main"): "40000Mi",
305+
commonconstants.GpuFraction: "0.5",
305306
},
306307
},
307308
Spec: v1.PodSpec{
308309
NodeName: "",
309310
Containers: []v1.Container{
310311
{
312+
Name: "main",
311313
Resources: v1.ResourceRequirements{
312314
Requests: common_info.BuildResourceList("1000m", "4G"),
313315
},

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import (
3636
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
3737

3838
commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
39+
"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
3940
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
4041
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_affinity"
4142
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
@@ -992,8 +993,8 @@ func TestNodeInfo_isTaskAllocatableOnNonAllocatedResources(t *testing.T) {
992993
Name: "p1",
993994
Namespace: "n1",
994995
Annotations: map[string]string{
995-
commonconstants.PodGroupAnnotationForPod: "pg1",
996-
commonconstants.GpuMemory: "1500",
996+
commonconstants.PodGroupAnnotationForPod: "pg1",
997+
resources.CalcGpuFractionAnnotationForContainer("c1"): "1500Mi",
997998
},
998999
},
9991000
Spec: v1.PodSpec{
@@ -1038,8 +1039,8 @@ func TestNodeInfo_isTaskAllocatableOnNonAllocatedResources(t *testing.T) {
10381039
Name: "p1",
10391040
Namespace: "n1",
10401041
Annotations: map[string]string{
1041-
commonconstants.PodGroupAnnotationForPod: "pg1",
1042-
commonconstants.GpuMemory: "1000",
1042+
commonconstants.PodGroupAnnotationForPod: "pg1",
1043+
resources.CalcGpuFractionAnnotationForContainer("c1"): "1000Mi",
10431044
},
10441045
},
10451046
Spec: v1.PodSpec{

‎pkg/scheduler/api/pod_info/pod_info.go‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ const (
4545
GPUGroup = "runai-gpu-group"
4646
ReceivedResourceTypeAnnotationName = "received-resource-type"
4747
WholeGpuIndicator = "-2"
48+
bytesPerMiB = int64(1024 * 1024)
4849
)
4950

5051
type ResourceRequestType string
@@ -518,11 +519,7 @@ func (pi *PodInfo) updatePodAdditionalFields(bindRequest *bindrequest_info.BindR
518519
}
519520
}
520521

521-
gpuMemory, err := strconv.ParseInt(pi.Pod.Annotations[commonconstants.GpuMemory], 10, 64)
522-
if err == nil && gpuMemory > 0 {
523-
pi.GpuRequirement = *resource_info.NewGpuResourceRequirementWithGpus(0, gpuMemory)
524-
pi.ResourceRequestType = RequestTypeGpuMemory
525-
}
522+
gpuMemoryRequest := pi.updateGpuMemoryRequest()
526523

527524
gpuFractionString := pi.Pod.Annotations[commonconstants.GpuFraction]
528525
gpuFraction, GPUFractionErr := strconv.ParseFloat(gpuFractionString, 64)
@@ -537,7 +534,7 @@ func (pi *PodInfo) updatePodAdditionalFields(bindRequest *bindrequest_info.BindR
537534
numFractionDevices, numFractionDevicesErr := strconv.ParseInt(numFractionDevicesStr, 10, 64)
538535
if numFractionDevicesErr == nil {
539536
pi.GpuRequirement = *resource_info.NewGpuResourceRequirementWithMultiFraction(
540-
numFractionDevices, gpuFraction, gpuMemory)
537+
numFractionDevices, gpuFraction, gpuMemoryRequest)
541538
}
542539
}
543540
}
@@ -554,6 +551,18 @@ func (pi *PodInfo) updatePodAdditionalFields(bindRequest *bindrequest_info.BindR
554551
pi.rebuildResReqVector()
555552
}
556553

554+
func (pi *PodInfo) updateGpuMemoryRequest() int64 {
555+
gpuMemoryRequest, err := resources.GetGPUMemory(pi.Pod)
556+
if err != nil {
557+
log.InfraLogger.Errorf("Failed to get GPU memory for pod %s/%s. Error: %s", pi.Pod.Namespace, pi.Pod.Name, err.Error())
558+
}
559+
if gpuMemoryRequest > 0 {
560+
pi.GpuRequirement = *resource_info.NewGpuResourceRequirementWithGpus(0, gpuMemoryRequest)
561+
pi.ResourceRequestType = RequestTypeGpuMemory
562+
}
563+
return gpuMemoryRequest
564+
}
565+
557566
// updateLegacyMigResourceRequestFromAnnotations updates the mig resource request of legacy MIG pods
558567
func (pi *PodInfo) updateLegacyMigResourceRequestFromAnnotations() {
559568
for annotationName, annotationValue := range pi.Pod.Annotations {

‎pkg/scheduler/api/pod_info/pod_info_test.go‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import (
3333

3434
schedulingv1alpha2 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
3535
commonconstants "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
36+
"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
3637
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/bindrequest_info"
3738
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
3839
"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_status"
@@ -414,7 +415,7 @@ func TestPodInfo_updatePodAdditionalFields(t *testing.T) {
414415
nil,
415416
map[string]string{},
416417
map[string]string{
417-
commonconstants.GpuMemory: "1024",
418+
resources.CalcGpuFractionAnnotationForContainer("main"): "1Gi",
418419
}),
419420
},
420421
expected{
@@ -438,8 +439,8 @@ func TestPodInfo_updatePodAdditionalFields(t *testing.T) {
438439
nil,
439440
map[string]string{},
440441
map[string]string{
441-
commonconstants.GpuMemory: "1024",
442-
commonconstants.GpuFractionsNumDevices: "2",
442+
resources.CalcGpuFractionAnnotationForContainer("main"): "1Gi",
443+
commonconstants.GpuFractionsNumDevices: "2",
443444
}),
444445
},
445446
expected{
@@ -556,6 +557,30 @@ func TestPodInfo_updatePodAdditionalFields(t *testing.T) {
556557
GPUGroups: nil,
557558
},
558559
},
560+
{
561+
"Legacy gpu-memory annotation fallback (pod predates NvFractions webhook)",
562+
podFields{
563+
Job: common_info.FakePogGroupId,
564+
Name: "p1",
565+
Namespace: "ns1",
566+
Status: pod_status.Pending,
567+
Pod: common_info.BuildPod("ns1", "p1", "node1", v1.PodPending,
568+
common_info.BuildResourceList("2000m", "2G"),
569+
nil,
570+
map[string]string{},
571+
map[string]string{
572+
commonconstants.GpuMemory: "2048",
573+
}),
574+
},
575+
expected{
576+
GpuRequirement: *resource_info.NewGpuResourceRequirementWithGpus(0, 2048),
577+
ResourceRequestType: "GpuMemory",
578+
GPUGroups: nil,
579+
SelectedMigProfile: "",
580+
IsBound: false,
581+
IsChiefPod: true,
582+
},
583+
},
559584
}
560585
for _, tt := range tests {
561586
t.Run(tt.name, func(t *testing.T) {

‎pkg/scheduler/test_utils/jobs_fake/jobs.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -205,12 +205,12 @@ func generateTasks(
205205
for taskIndex, task := range job.Tasks {
206206
gpuGroups := tasks_fake.GetTestTaskGPUIndex(task)
207207

208-
podResourceList, gpuMemory, gpuFraction, gpuGroups :=
208+
podResourceList, gpuMemoryMiB, gpuFraction, gpuGroups :=
209209
CalcJobAndPodResources(job, jobAllocatedResource, task, gpuGroups,
210210
usedSharedGPUs)
211211

212212
podOfTask := createPodOfTask(job, taskIndex, task, podResourceList, gpuFraction,
213-
gpuMemory, gpuGroups)
213+
gpuMemoryMiB, gpuGroups)
214214
if job.QOSClass != "" {
215215
podOfTask.Status.QOSClass = job.QOSClass
216216
}
@@ -278,10 +278,10 @@ func getDraClaimsForPod(task *tasks_fake.TestTaskBasic, draClaimsMap map[string]
278278

279279
func CalcJobAndPodResources(job *TestJobBasic, jobAllocatedResource *resource_info.Resource,
280280
task *tasks_fake.TestTaskBasic, gpuGroups []string,
281-
usedSharedGPUs map[string]map[string]bool) (*v1.ResourceList, string, string, []string) {
281+
usedSharedGPUs map[string]map[string]bool) (*v1.ResourceList, uint64, string, []string) {
282282
var podResourceList *v1.ResourceList
283283
var gpuFraction string
284-
var gpuMemory string
284+
var gpuMemoryMiB uint64
285285
if job.IsBestEffortJob {
286286
podResourceList =
287287
resources_fake.BuildResourceList(nil, nil, nil, nil)
@@ -292,7 +292,7 @@ func CalcJobAndPodResources(job *TestJobBasic, jobAllocatedResource *resource_in
292292
requiredMemoryInput = CalcRequiredMemory(job, requiredMemoryAsString, requiredMemoryInput)
293293

294294
// whole GPU job
295-
gpuMemory = strconv.FormatUint(job.RequiredGpuMemory, 10)
295+
gpuMemoryMiB = job.RequiredGpuMemory
296296
if float64(int(job.RequiredGPUsPerTask)) == job.RequiredGPUsPerTask {
297297
requiredGPUsAsString = strconv.Itoa(int(job.RequiredGPUsPerTask))
298298
} else {
@@ -305,7 +305,7 @@ func CalcJobAndPodResources(job *TestJobBasic, jobAllocatedResource *resource_in
305305

306306
(*podResourceList)[v1.ResourcePods] = resource.MustParse("1")
307307

308-
return podResourceList, gpuMemory, gpuFraction, gpuGroups
308+
return podResourceList, gpuMemoryMiB, gpuFraction, gpuGroups
309309
}
310310

311311
func CalcRequiredMemory(job *TestJobBasic, requiredMemoryAsString string, requiredMemoryInput *string) *string {
@@ -352,9 +352,9 @@ func resourceFractionCalc(job *TestJobBasic, jobAllocatedResource *resource_info
352352

353353
func createPodOfTask(job *TestJobBasic, taskIndex int,
354354
task *tasks_fake.TestTaskBasic, podResourceList *v1.ResourceList,
355-
gpuFraction string, gpuMemory string, gpuGroups []string) *v1.Pod {
355+
gpuFraction string, gpuMemoryMiB uint64, gpuGroups []string) *v1.Pod {
356356
podName := fmt.Sprintf("%s-%d", job.Name, taskIndex)
357-
podOfTask := tasks_fake.BuildPod(podName, job.Namespace, task, v1.PodPending, *podResourceList, gpuFraction, gpuMemory,
357+
podOfTask := tasks_fake.BuildPod(podName, job.Namespace, task, v1.PodPending, *podResourceList, gpuFraction, gpuMemoryMiB,
358358
gpuGroups, job.Name)
359359
addPersistentVolumeClaimVolumes(podOfTask, task.PersistentVolumeClaimNames)
360360

‎pkg/scheduler/test_utils/tasks_fake/tasks.go‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ func BuildPod(
4646
name, namespace string,
4747
task *TestTaskBasic,
4848
phase v1.PodPhase, req v1.ResourceList,
49-
gpuFraction, gpuMemory string, gpuGroups []string, jobName string,
49+
gpuFraction string, gpuMemoryMiB uint64, gpuGroups []string, jobName string,
5050
) *v1.Pod {
5151
controllerBool := true
5252
pod := &v1.Pod{
@@ -68,7 +68,6 @@ func BuildPod(
6868
}(),
6969
Annotations: map[string]string{
7070
commonconstants.GpuFraction: gpuFraction,
71-
commonconstants.GpuMemory: gpuMemory,
7271
commonconstants.PodGroupAnnotationForPod: jobName,
7372
},
7473
},
@@ -79,6 +78,7 @@ func BuildPod(
7978
NodeName: task.NodeName,
8079
Containers: []v1.Container{
8180
{
81+
Name: "main",
8282
Resources: v1.ResourceRequirements{
8383
Requests: req,
8484
},
@@ -87,6 +87,9 @@ func BuildPod(
8787
SchedulerName: "kai-scheduler",
8888
},
8989
}
90+
if gpuMemoryMiB > 0 {
91+
pod.Annotations[resources.CalcGpuFractionAnnotationForContainer("main")] = resources.GpuMemoryAnnotationToNvFractionsMemoryRequest(gpuMemoryMiB).String()
92+
}
9093
if len(gpuGroups) > 1 {
9194
for _, gpuGroup := range gpuGroups {
9295
multiGroupKey, multiGroupValue := resources.GetMultiFractionGpuGroupLabel(gpuGroup)

0 commit comments

Comments
 (0)