feat: Support Portional Limit Annotation for nvFractions

.changes/unreleased/added-20260813-194544.yaml

@@ -0,0 +1,3 @@
+kind: Added
+body: |-
+  Support kai.scheduler gpu-memory.portion.limit annotation for per-container GPU memory limits

pkg/binder/plugins/nvfractions/nv_fractions.go

@@ -61,6 +61,10 @@ func (p *Plugin) PreBind(
 		return fmt.Errorf("failed to set NvFractions memory annotation: %w", err)
 	}
 
+	if err := setGpuMemoryPortionLimitAnnotation(pod, node, containerRef.Container.Name, bindingState); err != nil {
+		return fmt.Errorf("failed to set NvFractions memory limit annotation: %w", err)
+	}
+
 	visibleDevices := bindingState.ReservedGPUIds
 	if p.gpuDevicePluginUsesCdi {
 		visibleDevices = make([]string, len(bindingState.ReservedGPUIds))
@@ -113,6 +117,48 @@ func setNvFractionsMemoryAnnotation(pod *v1.Pod, node *v1.Node, bindRequest *v1a
 	return nil
 }
 
+// setGpuMemoryPortionLimitAnnotation translates the kai.scheduler
+// gpu-memory.portion.limit annotation into the NvFractions limit form,
+// without removing the source annotation.
+func setGpuMemoryPortionLimitAnnotation(pod *v1.Pod, node *v1.Node, containerName string, bindingState *state.BindingState) error {
+	_, rawPortionLimit, found := resources.ExtractGpuMemoryPortionLimitAnnotation(pod)
+	if !found {
+		return nil
+	}
+
+	annotationKey := resources.CalcGpuFractionLimitAnnotationForContainer(containerName)
+	if _, found := pod.Annotations[annotationKey]; found {
+		return nil
+	}
+
+	if node == nil {
+		return fmt.Errorf("missing node data for gpu-memory.portion.limit annotation calculation")
+	}
+
+	gpuMemoryStr, found := node.Labels[constants.NvidiaGpuMemory]
+	if !found {
+		return fmt.Errorf("node does not include %s label", constants.NvidiaGpuMemory)
+	}
+
+	totalGPUMemoryMiB, err := strconv.ParseFloat(gpuMemoryStr, 64)
+	if err != nil || totalGPUMemoryMiB <= 0 {
+		return fmt.Errorf("invalid %s label value %q", constants.NvidiaGpuMemory, gpuMemoryStr)
+	}
+
+	portionLimit, err := strconv.ParseFloat(rawPortionLimit, 64)
+	if err != nil || portionLimit <= 0 {
+		return fmt.Errorf("invalid gpu-memory.portion.limit annotation value %q", rawPortionLimit)
+	}
+
+	gpuMemoryLimit := uint64(totalGPUMemoryMiB * portionLimit)
+	if gpuMemoryLimit == 0 {
+		return fmt.Errorf("calculated gpu memory limit is zero")
+	}
+
+	bindingState.BindingPodAnnotations[annotationKey] = resources.GpuMemoryAnnotationToNvFractionsMemoryRequest(gpuMemoryLimit).String()
+	return nil
+}
+
 func (p *Plugin) PostBind(
 	context.Context, *v1.Pod, *v1.Node, *v1alpha2.BindRequest, *state.BindingState,
 ) {

pkg/binder/plugins/nvfractions/nv_fractions_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
@@ -67,6 +68,62 @@ func TestPreBindQualifiesVisibleDevicesAsCdiWhenEnabled(t *testing.T) {
 		bindingState.BindingPodAnnotations[resources.CalcGpuVisibleDevicesAnnotationForContainer("container-0")])
 }
 
+func TestPreBindSetsGpuMemoryPortionLimitAnnotation(t *testing.T) {
+	pod := &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
+			constants.GpuFraction: "0.5",
+			resources.CalcGpuMemoryPortionLimitAnnotationForContainer("container-0"): "0.8",
+		}},
+		Spec: v1.PodSpec{Containers: []v1.Container{{Name: "container-0"}}},
+	}
+	node := &v1.Node{
+		ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{constants.NvidiaGpuMemory: "2000"}},
+	}
+	bindRequest := &v1alpha2.BindRequest{
+		Spec: v1alpha2.BindRequestSpec{
+			ReceivedResourceType: bindercommon.ReceivedTypeFraction,
+			ReceivedGPU:          &v1alpha2.ReceivedGPU{Portion: "0.5"},
+		},
+	}
+	bindingState := &state.BindingState{ReservedGPUIds: []string{"0"}}
+
+	err := New(false).PreBind(context.Background(), pod, node, bindRequest, bindingState)
+	assert.NoError(t, err)
+
+	limitKey := resources.CalcGpuFractionLimitAnnotationForContainer("container-0")
+	limitQuantity := resource.MustParse(bindingState.BindingPodAnnotations[limitKey])
+	expectedQuantity := resource.MustParse("1600Mi")
+	assert.Equal(t, expectedQuantity.Value(), limitQuantity.Value())
+
+	sourceKey := resources.CalcGpuMemoryPortionLimitAnnotationForContainer("container-0")
+	assert.Equal(t, "0.8", pod.Annotations[sourceKey])
+}
+
+func TestPreBindNoGpuMemoryPortionLimitAnnotationIsNoOp(t *testing.T) {
+	pod := &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{
+			constants.GpuFraction: "0.5",
+		}},
+		Spec: v1.PodSpec{Containers: []v1.Container{{Name: "container-0"}}},
+	}
+	node := &v1.Node{
+		ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{constants.NvidiaGpuMemory: "2000"}},
+	}
+	bindRequest := &v1alpha2.BindRequest{
+		Spec: v1alpha2.BindRequestSpec{
+			ReceivedResourceType: bindercommon.ReceivedTypeFraction,
+			ReceivedGPU:          &v1alpha2.ReceivedGPU{Portion: "0.5"},
+		},
+	}
+	bindingState := &state.BindingState{ReservedGPUIds: []string{"0"}}
+
+	err := New(false).PreBind(context.Background(), pod, node, bindRequest, bindingState)
+	assert.NoError(t, err)
+
+	limitKey := resources.CalcGpuFractionLimitAnnotationForContainer("container-0")
+	assert.NotContains(t, bindingState.BindingPodAnnotations, limitKey)
+}
+
 func TestPreBindNoOpForWholeGpuAllocation(t *testing.T) {
 	pod := &v1.Pod{Spec: v1.PodSpec{Containers: []v1.Container{{Name: "container-0"}}}}
 	node := &v1.Node{}

pkg/common/constants/constants.go

@@ -85,6 +85,9 @@ const (
 	NvFractionsMemoryLimitSuffix    = ".gpu-memory.limit"
 	NvFractionsVisibleDevicesSuffix = ".gpus.devices"
 
+	KaiFractionContainerAnnotationPrefix = "kai.scheduler/container."
+	GpuMemoryPortionLimitSuffix          = ".gpu-memory.portion.limit"
+
 	// gpu-sharing operator statuses
 	NvFractionNodeReadyConditionType = "gpu-sharing.nvidia.com/Ready"
 )

pkg/common/resources/gpu_memory_portion_limit.go

@@ -0,0 +1,40 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package resources
+
+import (
+	"strings"
+
+	v1 "k8s.io/api/core/v1"
+
+	"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
+)
+
+// CalcGpuMemoryPortionLimitAnnotationForContainer returns the kai.scheduler
+// per-container GPU memory portion limit annotation key for containerName.
+func CalcGpuMemoryPortionLimitAnnotationForContainer(containerName string) string {
+	return constants.KaiFractionContainerAnnotationPrefix + containerName + constants.GpuMemoryPortionLimitSuffix
+}
+
+// ExtractGpuMemoryPortionLimitAnnotation returns the container name and raw
+// value of the pod's gpu-memory.portion.limit annotation, if present.
+func ExtractGpuMemoryPortionLimitAnnotation(pod *v1.Pod) (containerName string, rawValue string, found bool) {
+	for annotationKey, annotationValue := range pod.Annotations {
+		if !isGpuMemoryPortionLimitAnnotation(annotationKey) {
+			continue
+		}
+		return gpuMemoryPortionLimitContainerName(annotationKey), annotationValue, true
+	}
+	return "", "", false
+}
+
+func isGpuMemoryPortionLimitAnnotation(annotationKey string) bool {
+	return strings.HasPrefix(annotationKey, constants.KaiFractionContainerAnnotationPrefix) &&
+		strings.HasSuffix(annotationKey, constants.GpuMemoryPortionLimitSuffix)
+}
+
+func gpuMemoryPortionLimitContainerName(annotationKey string) string {
+	containerName := strings.TrimPrefix(annotationKey, constants.KaiFractionContainerAnnotationPrefix)
+	return strings.TrimSuffix(containerName, constants.GpuMemoryPortionLimitSuffix)
+}

pkg/common/resources/gpu_sharing_nvfractions.go

@@ -24,6 +24,10 @@ func CalcGpuFractionAnnotationForContainer(containerName string) string {
 	return constants.NvFractionsAnnotationPrefix + containerName + constants.NvFractionsMemoryRequestSuffix
 }
 
+func CalcGpuFractionLimitAnnotationForContainer(containerName string) string {
+	return constants.NvFractionsAnnotationPrefix + containerName + constants.NvFractionsMemoryLimitSuffix
+}
+
 func CalcGpuVisibleDevicesAnnotationForContainer(containerName string) string {
 	return constants.NvFractionsAnnotationPrefix + containerName + constants.NvFractionsVisibleDevicesSuffix
 }

pkg/common/resources/gpu_sharing_validation.go

@@ -6,6 +6,7 @@ package resources
 import (
 	"fmt"
 	"strconv"
+	"strings"
 
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -18,6 +19,10 @@ import (
 // the gpu-fraction-container-name annotation, and NvFractions annotations. It is
 // configmap-agnostic and shared by the admission plugins.
 func ValidateGPUFractionRequest(pod *v1.Pod) error {
+	if err := validateGpuMemoryPortionLimitAnnotation(pod); err != nil {
+		return err
+	}
+
 	req, err := ParsePodGPUFractionRequest(pod)
 	if err != nil {
 		return err
@@ -90,6 +95,73 @@ func validateGpuMemoryNvFractionsConsistency(pod *v1.Pod) error {
 	return nil
 }
 
+// validateGpuMemoryPortionLimitAnnotation validates the kai.scheduler
+// gpu-memory.portion.limit annotation: it may only be used together with
+// gpu-fraction, on the same container, as a fraction strictly greater than
+// gpu-fraction and strictly smaller than 1.0, with at most 5 decimal digits.
+func validateGpuMemoryPortionLimitAnnotation(pod *v1.Pod) error {
+	containerName, rawValue, found := ExtractGpuMemoryPortionLimitAnnotation(pod)
+	if !found {
+		return nil
+	}
+	annotationKey := CalcGpuMemoryPortionLimitAnnotationForContainer(containerName)
+
+	gpuFractionStr, hasGpuFraction := pod.Annotations[constants.GpuFraction]
+	if !hasGpuFraction || gpuFractionStr == "" {
+		return fmt.Errorf("%s annotation can only be used together with the %s annotation",
+			annotationKey, constants.GpuFraction)
+	}
+
+	if err := validateGpuMemoryPortionLimitContainerName(pod, containerName, annotationKey); err != nil {
+		return err
+	}
+
+	gpuFraction, err := strconv.ParseFloat(gpuFractionStr, 64)
+	if err != nil {
+		return fmt.Errorf("gpu-fraction annotation value must be a positive number smaller than 1.0")
+	}
+
+	if err := validatePortionLimitDecimalPrecision(rawValue, annotationKey); err != nil {
+		return err
+	}
+
+	portionLimit, err := strconv.ParseFloat(rawValue, 64)
+	if err != nil || portionLimit <= 0 || portionLimit >= 1 {
+		return fmt.Errorf("%s annotation value must be a positive number smaller than 1.0", annotationKey)
+	}
+
+	if portionLimit <= gpuFraction {
+		return fmt.Errorf("%s annotation value (%s) must be greater than %s annotation value (%s)",
+			annotationKey, rawValue, constants.GpuFraction, gpuFractionStr)
+	}
+
+	return nil
+}
+
+func validateGpuMemoryPortionLimitContainerName(pod *v1.Pod, containerName, annotationKey string) error {
+	if legacyContainerName, hasGpuFractionContainerName := pod.Annotations[constants.GpuFractionContainerName]; hasGpuFractionContainerName {
+		if legacyContainerName != containerName {
+			return fmt.Errorf("%s annotation value %s does not match container name %s in %s annotation",
+				constants.GpuFractionContainerName, legacyContainerName, containerName, annotationKey)
+		}
+		return nil
+	}
+
+	if len(pod.Spec.Containers) == 0 || pod.Spec.Containers[0].Name != containerName {
+		return fmt.Errorf("%s annotation container name %s does not match the gpu-fraction target container",
+			annotationKey, containerName)
+	}
+	return nil
+}
+
+func validatePortionLimitDecimalPrecision(rawValue, annotationKey string) error {
+	_, fraction, hasDecimalPoint := strings.Cut(rawValue, ".")
+	if hasDecimalPoint && len(fraction) > 5 {
+		return fmt.Errorf("%s annotation value must have at most 5 digits after the decimal point", annotationKey)
+	}
+	return nil
+}
+
 func validateNvFractionsAnnotations(hasNvFractionsAnnotation bool, pod *v1.Pod) error {
 	if !hasNvFractionsAnnotation {
 		return nil