Skip to content

Commit 2a3f0a6

Browse files
davidLifdavidLif
authored
Support Portional Limit Annotation
Signed-off-by: davidLif <davidshani12@gmail.com>
1 parent 71a630e commit 2a3f0a6

8 files changed

Lines changed: 363 additions & 0 deletions

File tree

‎.changes/unreleased/added-20260813-194544.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+
Support kai.scheduler gpu-memory.portion.limit annotation for per-container GPU memory limits

‎pkg/binder/plugins/nvfractions/nv_fractions.go‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,10 @@ func (p *Plugin) PreBind(
6161
return fmt.Errorf("failed to set NvFractions memory annotation: %w", err)
6262
}
6363

64+
if err := setGpuMemoryPortionLimitAnnotation(pod, node, containerRef.Container.Name, bindingState); err != nil {
65+
return fmt.Errorf("failed to set NvFractions memory limit annotation: %w", err)
66+
}
67+
6468
visibleDevices := bindingState.ReservedGPUIds
6569
if p.gpuDevicePluginUsesCdi {
6670
visibleDevices = make([]string, len(bindingState.ReservedGPUIds))
@@ -113,6 +117,48 @@ func setNvFractionsMemoryAnnotation(pod *v1.Pod, node *v1.Node, bindRequest *v1a
113117
return nil
114118
}
115119

120+
// setGpuMemoryPortionLimitAnnotation translates the kai.scheduler
121+
// gpu-memory.portion.limit annotation into the NvFractions limit form,
122+
// without removing the source annotation.
123+
func setGpuMemoryPortionLimitAnnotation(pod *v1.Pod, node *v1.Node, containerName string, bindingState *state.BindingState) error {
124+
_, rawPortionLimit, found := resources.ExtractGpuMemoryPortionLimitAnnotation(pod)
125+
if !found {
126+
return nil
127+
}
128+
129+
annotationKey := resources.CalcGpuFractionLimitAnnotationForContainer(containerName)
130+
if _, found := pod.Annotations[annotationKey]; found {
131+
return nil
132+
}
133+
134+
if node == nil {
135+
return fmt.Errorf("missing node data for gpu-memory.portion.limit annotation calculation")
136+
}
137+
138+
gpuMemoryStr, found := node.Labels[constants.NvidiaGpuMemory]
139+
if !found {
140+
return fmt.Errorf("node does not include %s label", constants.NvidiaGpuMemory)
141+
}
142+
143+
totalGPUMemoryMiB, err := strconv.ParseFloat(gpuMemoryStr, 64)
144+
if err != nil || totalGPUMemoryMiB <= 0 {
145+
return fmt.Errorf("invalid %s label value %q", constants.NvidiaGpuMemory, gpuMemoryStr)
146+
}
147+
148+
portionLimit, err := strconv.ParseFloat(rawPortionLimit, 64)
149+
if err != nil || portionLimit <= 0 {
150+
return fmt.Errorf("invalid gpu-memory.portion.limit annotation value %q", rawPortionLimit)
151+
}
152+
153+
gpuMemoryLimit := uint64(totalGPUMemoryMiB * portionLimit)
154+
if gpuMemoryLimit == 0 {
155+
return fmt.Errorf("calculated gpu memory limit is zero")
156+
}
157+
158+
bindingState.BindingPodAnnotations[annotationKey] = resources.GpuMemoryAnnotationToNvFractionsMemoryRequest(gpuMemoryLimit).String()
159+
return nil
160+
}
161+
116162
func (p *Plugin) PostBind(
117163
context.Context, *v1.Pod, *v1.Node, *v1alpha2.BindRequest, *state.BindingState,
118164
) {

‎pkg/binder/plugins/nvfractions/nv_fractions_test.go‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99

1010
"github.com/stretchr/testify/assert"
1111
v1 "k8s.io/api/core/v1"
12+
"k8s.io/apimachinery/pkg/api/resource"
1213
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
1314

1415
"github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
@@ -67,6 +68,62 @@ func TestPreBindQualifiesVisibleDevicesAsCdiWhenEnabled(t *testing.T) {
6768
bindingState.BindingPodAnnotations[resources.CalcGpuVisibleDevicesAnnotationForContainer("container-0")])
6869
}
6970

71+
func TestPreBindSetsGpuMemoryPortionLimitAnnotation(t *testing.T) {
72+
pod := &v1.Pod{
73+
ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
74+
constants.GpuFraction: "0.5",
75+
resources.CalcGpuMemoryPortionLimitAnnotationForContainer("container-0"): "0.8",
76+
}},
77+
Spec: v1.PodSpec{Containers: []v1.Container{{Name: "container-0"}}},
78+
}
79+
node := &v1.Node{
80+
ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{constants.NvidiaGpuMemory: "2000"}},
81+
}
82+
bindRequest := &v1alpha2.BindRequest{
83+
Spec: v1alpha2.BindRequestSpec{
84+
ReceivedResourceType: bindercommon.ReceivedTypeFraction,
85+
ReceivedGPU: &v1alpha2.ReceivedGPU{Portion: "0.5"},
86+
},
87+
}
88+
bindingState := &state.BindingState{ReservedGPUIds: []string{"0"}}
89+
90+
err := New(false).PreBind(context.Background(), pod, node, bindRequest, bindingState)
91+
assert.NoError(t, err)
92+
93+
limitKey := resources.CalcGpuFractionLimitAnnotationForContainer("container-0")
94+
limitQuantity := resource.MustParse(bindingState.BindingPodAnnotations[limitKey])
95+
expectedQuantity := resource.MustParse("1600Mi")
96+
assert.Equal(t, expectedQuantity.Value(), limitQuantity.Value())
97+
98+
sourceKey := resources.CalcGpuMemoryPortionLimitAnnotationForContainer("container-0")
99+
assert.Equal(t, "0.8", pod.Annotations[sourceKey])
100+
}
101+
102+
func TestPreBindNoGpuMemoryPortionLimitAnnotationIsNoOp(t *testing.T) {
103+
pod := &v1.Pod{
104+
ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
105+
constants.GpuFraction: "0.5",
106+
}},
107+
Spec: v1.PodSpec{Containers: []v1.Container{{Name: "container-0"}}},
108+
}
109+
node := &v1.Node{
110+
ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{constants.NvidiaGpuMemory: "2000"}},
111+
}
112+
bindRequest := &v1alpha2.BindRequest{
113+
Spec: v1alpha2.BindRequestSpec{
114+
ReceivedResourceType: bindercommon.ReceivedTypeFraction,
115+
ReceivedGPU: &v1alpha2.ReceivedGPU{Portion: "0.5"},
116+
},
117+
}
118+
bindingState := &state.BindingState{ReservedGPUIds: []string{"0"}}
119+
120+
err := New(false).PreBind(context.Background(), pod, node, bindRequest, bindingState)
121+
assert.NoError(t, err)
122+
123+
limitKey := resources.CalcGpuFractionLimitAnnotationForContainer("container-0")
124+
assert.NotContains(t, bindingState.BindingPodAnnotations, limitKey)
125+
}
126+
70127
func TestPreBindNoOpForWholeGpuAllocation(t *testing.T) {
71128
pod := &v1.Pod{Spec: v1.PodSpec{Containers: []v1.Container{{Name: "container-0"}}}}
72129
node := &v1.Node{}

‎pkg/common/constants/constants.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,9 @@ const (
8585
NvFractionsMemoryLimitSuffix = ".gpu-memory.limit"
8686
NvFractionsVisibleDevicesSuffix = ".gpus.devices"
8787

88+
KaiFractionContainerAnnotationPrefix = "kai.scheduler/container."
89+
GpuMemoryPortionLimitSuffix = ".gpu-memory.portion.limit"
90+
8891
// gpu-sharing operator statuses
8992
NvFractionNodeReadyConditionType = "gpu-sharing.nvidia.com/Ready"
9093
)

‎pkg/common/resources/gpu_memory_portion_limit.go‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
// Copyright 2025 NVIDIA CORPORATION
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package resources
5+
6+
import (
7+
"strings"
8+
9+
v1 "k8s.io/api/core/v1"
10+
11+
"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
12+
)
13+
14+
// CalcGpuMemoryPortionLimitAnnotationForContainer returns the kai.scheduler
15+
// per-container GPU memory portion limit annotation key for containerName.
16+
func CalcGpuMemoryPortionLimitAnnotationForContainer(containerName string) string {
17+
return constants.KaiFractionContainerAnnotationPrefix + containerName + constants.GpuMemoryPortionLimitSuffix
18+
}
19+
20+
// ExtractGpuMemoryPortionLimitAnnotation returns the container name and raw
21+
// value of the pod's gpu-memory.portion.limit annotation, if present.
22+
func ExtractGpuMemoryPortionLimitAnnotation(pod *v1.Pod) (containerName string, rawValue string, found bool) {
23+
for annotationKey, annotationValue := range pod.Annotations {
24+
if !isGpuMemoryPortionLimitAnnotation(annotationKey) {
25+
continue
26+
}
27+
return gpuMemoryPortionLimitContainerName(annotationKey), annotationValue, true
28+
}
29+
return "", "", false
30+
}
31+
32+
func isGpuMemoryPortionLimitAnnotation(annotationKey string) bool {
33+
return strings.HasPrefix(annotationKey, constants.KaiFractionContainerAnnotationPrefix) &&
34+
strings.HasSuffix(annotationKey, constants.GpuMemoryPortionLimitSuffix)
35+
}
36+
37+
func gpuMemoryPortionLimitContainerName(annotationKey string) string {
38+
containerName := strings.TrimPrefix(annotationKey, constants.KaiFractionContainerAnnotationPrefix)
39+
return strings.TrimSuffix(containerName, constants.GpuMemoryPortionLimitSuffix)
40+
}

‎pkg/common/resources/gpu_sharing_nvfractions.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,10 @@ func CalcGpuFractionAnnotationForContainer(containerName string) string {
2424
return constants.NvFractionsAnnotationPrefix + containerName + constants.NvFractionsMemoryRequestSuffix
2525
}
2626

27+
func CalcGpuFractionLimitAnnotationForContainer(containerName string) string {
28+
return constants.NvFractionsAnnotationPrefix + containerName + constants.NvFractionsMemoryLimitSuffix
29+
}
30+
2731
func CalcGpuVisibleDevicesAnnotationForContainer(containerName string) string {
2832
return constants.NvFractionsAnnotationPrefix + containerName + constants.NvFractionsVisibleDevicesSuffix
2933
}

‎pkg/common/resources/gpu_sharing_validation.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ package resources
66
import (
77
"fmt"
88
"strconv"
9+
"strings"
910

1011
v1 "k8s.io/api/core/v1"
1112
"k8s.io/apimachinery/pkg/api/resource"
@@ -18,6 +19,10 @@ import (
1819
// the gpu-fraction-container-name annotation, and NvFractions annotations. It is
1920
// configmap-agnostic and shared by the admission plugins.
2021
func ValidateGPUFractionRequest(pod *v1.Pod) error {
22+
if err := validateGpuMemoryPortionLimitAnnotation(pod); err != nil {
23+
return err
24+
}
25+
2126
req, err := ParsePodGPUFractionRequest(pod)
2227
if err != nil {
2328
return err
@@ -90,6 +95,73 @@ func validateGpuMemoryNvFractionsConsistency(pod *v1.Pod) error {
9095
return nil
9196
}
9297

98+
// validateGpuMemoryPortionLimitAnnotation validates the kai.scheduler
99+
// gpu-memory.portion.limit annotation: it may only be used together with
100+
// gpu-fraction, on the same container, as a fraction strictly greater than
101+
// gpu-fraction and strictly smaller than 1.0, with at most 5 decimal digits.
102+
func validateGpuMemoryPortionLimitAnnotation(pod *v1.Pod) error {
103+
containerName, rawValue, found := ExtractGpuMemoryPortionLimitAnnotation(pod)
104+
if !found {
105+
return nil
106+
}
107+
annotationKey := CalcGpuMemoryPortionLimitAnnotationForContainer(containerName)
108+
109+
gpuFractionStr, hasGpuFraction := pod.Annotations[constants.GpuFraction]
110+
if !hasGpuFraction || gpuFractionStr == "" {
111+
return fmt.Errorf("%s annotation can only be used together with the %s annotation",
112+
annotationKey, constants.GpuFraction)
113+
}
114+
115+
if err := validateGpuMemoryPortionLimitContainerName(pod, containerName, annotationKey); err != nil {
116+
return err
117+
}
118+
119+
gpuFraction, err := strconv.ParseFloat(gpuFractionStr, 64)
120+
if err != nil {
121+
return fmt.Errorf("gpu-fraction annotation value must be a positive number smaller than 1.0")
122+
}
123+
124+
if err := validatePortionLimitDecimalPrecision(rawValue, annotationKey); err != nil {
125+
return err
126+
}
127+
128+
portionLimit, err := strconv.ParseFloat(rawValue, 64)
129+
if err != nil || portionLimit <= 0 || portionLimit >= 1 {
130+
return fmt.Errorf("%s annotation value must be a positive number smaller than 1.0", annotationKey)
131+
}
132+
133+
if portionLimit <= gpuFraction {
134+
return fmt.Errorf("%s annotation value (%s) must be greater than %s annotation value (%s)",
135+
annotationKey, rawValue, constants.GpuFraction, gpuFractionStr)
136+
}
137+
138+
return nil
139+
}
140+
141+
func validateGpuMemoryPortionLimitContainerName(pod *v1.Pod, containerName, annotationKey string) error {
142+
if legacyContainerName, hasGpuFractionContainerName := pod.Annotations[constants.GpuFractionContainerName]; hasGpuFractionContainerName {
143+
if legacyContainerName != containerName {
144+
return fmt.Errorf("%s annotation value %s does not match container name %s in %s annotation",
145+
constants.GpuFractionContainerName, legacyContainerName, containerName, annotationKey)
146+
}
147+
return nil
148+
}
149+
150+
if len(pod.Spec.Containers) == 0 || pod.Spec.Containers[0].Name != containerName {
151+
return fmt.Errorf("%s annotation container name %s does not match the gpu-fraction target container",
152+
annotationKey, containerName)
153+
}
154+
return nil
155+
}
156+
157+
func validatePortionLimitDecimalPrecision(rawValue, annotationKey string) error {
158+
_, fraction, hasDecimalPoint := strings.Cut(rawValue, ".")
159+
if hasDecimalPoint && len(fraction) > 5 {
160+
return fmt.Errorf("%s annotation value must have at most 5 digits after the decimal point", annotationKey)
161+
}
162+
return nil
163+
}
164+
93165
func validateNvFractionsAnnotations(hasNvFractionsAnnotation bool, pod *v1.Pod) error {
94166
if !hasNvFractionsAnnotation {
95167
return nil

0 commit comments

Comments
 (0)