feat(binder): add NvFractions binding plugin

Signed-off-by: davidLif <davidshani12@gmail.com>

Author: davidLif

.changes/unreleased/added-20260810-110719.yaml

@@ -0,0 +1,3 @@
+kind: Added
+body: |-
+  Bind NvFractions GPU memory devices

pkg/apis/kai/v1/binder/binder.go

@@ -23,6 +23,7 @@ const (
 	DynamicResourcesPluginName = "dynamicresources"
 	GPUSharingPluginName       = "gpusharing"
 	HamiCorePluginName         = "hamicore"
+	NvFractionsPluginName      = "nvfractions"
 
 	BindTimeoutSecondsArgument = "bindTimeoutSeconds"
 	CDIEnabledArgument         = "cdiEnabled"
@@ -35,6 +36,7 @@ var defaultPluginPriorities = map[string]int{
 	VolumeBindingPluginName:    300,
 	DynamicResourcesPluginName: 200,
 	GPUSharingPluginName:       100,
+	NvFractionsPluginName:      90,
 	HamiCorePluginName:         50,
 }
 
@@ -139,13 +141,16 @@ func (b *Binder) setDefaultPlugins() {
 	binderPluginConfig := DefaultPluginsConfig(ptr.Deref(b.VolumeBindingTimeoutSeconds, DefaultBindTimeoutSeconds),
 		ptr.Deref(b.CDIEnabled, DefaultCDIEnabled))
 
-	// When CDIEnabled is unset at the API level, leave the gpusharing cdiEnabled
-	// argument unbaked so the operator can resolve it (auto-detect) without
-	// having to distinguish a defaulted value from a user-supplied one.
+	// When CDIEnabled is unset at the API level, leave the cdiEnabled argument
+	// unbaked on the CDI-aware plugins (gpusharing, nvfractions) so the operator
+	// can resolve it (auto-detect) without having to distinguish a defaulted value
+	// from a user-supplied one.
 	if b.CDIEnabled == nil {
-		gpuSharingDefault := binderPluginConfig[GPUSharingPluginName]
-		delete(gpuSharingDefault.Arguments, CDIEnabledArgument)
-		binderPluginConfig[GPUSharingPluginName] = gpuSharingDefault
+		for _, name := range []string{GPUSharingPluginName, NvFractionsPluginName} {
+			pluginDefault := binderPluginConfig[name]
+			delete(pluginDefault.Arguments, CDIEnabledArgument)
+			binderPluginConfig[name] = pluginDefault
+		}
 	}
 
 	for name, userBinderConfig := range b.Plugins {
@@ -198,6 +203,13 @@ func DefaultPluginsConfig(bindTimeoutSeconds int, cdiEnabled bool) map[string]Pl
 				CDIEnabledArgument: strconv.FormatBool(cdiEnabled),
 			},
 		},
+		NvFractionsPluginName: {
+			Enabled:  ptr.To(false),
+			Priority: ptr.To(defaultPluginPriorities[NvFractionsPluginName]),
+			Arguments: map[string]string{
+				CDIEnabledArgument: strconv.FormatBool(cdiEnabled),
+			},
+		},
 		HamiCorePluginName: {
 			Enabled:  ptr.To(false),
 			Priority: ptr.To(defaultPluginPriorities[HamiCorePluginName]),

pkg/apis/kai/v1/binder/binder_test.go

@@ -32,10 +32,11 @@ var _ = Describe("Binder", func() {
 		Expect(binder.Service.Resources.Requests[v1.ResourceMemory]).To(Equal(resource.MustParse("200Mi")))
 		Expect(binder.Service.Resources.Limits[v1.ResourceCPU]).To(Equal(resource.MustParse("100m")))
 		Expect(binder.Service.Resources.Limits[v1.ResourceMemory]).To(Equal(resource.MustParse("200Mi")))
-		Expect(binder.Plugins).To(HaveLen(4))
+		Expect(binder.Plugins).To(HaveLen(5))
 		Expect(*binder.Plugins[VolumeBindingPluginName].Priority).To(Equal(defaultPluginPriorities[VolumeBindingPluginName]))
 		Expect(*binder.Plugins[DynamicResourcesPluginName].Priority).To(Equal(defaultPluginPriorities[DynamicResourcesPluginName]))
 		Expect(*binder.Plugins[GPUSharingPluginName].Priority).To(Equal(defaultPluginPriorities[GPUSharingPluginName]))
+		Expect(*binder.Plugins[NvFractionsPluginName].Priority).To(Equal(defaultPluginPriorities[NvFractionsPluginName]))
 		Expect(binder.Plugins[VolumeBindingPluginName].Arguments[BindTimeoutSecondsArgument]).
 			To(Equal(strconv.Itoa(DefaultBindTimeoutSeconds)))
 		Expect(binder.Plugins[DynamicResourcesPluginName].Arguments[BindTimeoutSecondsArgument]).

pkg/binder/binding/binder.go

@@ -54,15 +54,25 @@ func (b *Binder) Bind(ctx context.Context, pod *v1.Pod, node *v1.Node, bindReque
 		}
 	}
 	bindingState := &state.BindingState{
+		BindingPodAnnotations: map[string]string{
+			constants.ReceivedResourceType: bindRequest.Spec.ReceivedResourceType,
+		},
 		ReservedGPUIds: reservedGPUIds,
 	}
+	if len(bindRequest.Spec.PredictedNUMAZones) > 0 {
+		placement, err := json.Marshal(bindRequest.Spec.PredictedNUMAZones)
+		if err != nil {
+			return err
+		}
+		bindingState.BindingPodAnnotations[constants.NumaPlacementPredicted] = string(placement)
+	}
 
 	err = b.plugins.PreBind(ctx, pod, node, bindRequest, bindingState)
 	if err != nil {
 		return err
 	}
 
-	err = b.patchBindAnnotations(ctx, pod, bindRequest)
+	err = b.patchPodBindingAnnotations(ctx, pod, bindingState)
 	if err != nil {
 		return fmt.Errorf("failed to patch pod <%s/%s> with bind annotations: %w", pod.Namespace, pod.Name, err)
 	}
@@ -127,22 +137,10 @@ func (b *Binder) reserveGPUs(ctx context.Context, pod *v1.Pod, bindRequest *v1al
 	return gpuIndexes, nil
 }
 
-func (b *Binder) patchBindAnnotations(ctx context.Context, pod *v1.Pod, bindRequest *v1alpha2.BindRequest) error {
-	annotations := map[string]string{
-		constants.ReceivedResourceType: bindRequest.Spec.ReceivedResourceType,
-	}
-
-	if len(bindRequest.Spec.PredictedNUMAZones) > 0 {
-		placement, err := json.Marshal(bindRequest.Spec.PredictedNUMAZones)
-		if err != nil {
-			return err
-		}
-		annotations[constants.NumaPlacementPredicted] = string(placement)
-	}
-
+func (b *Binder) patchPodBindingAnnotations(ctx context.Context, pod *v1.Pod, bindingState *state.BindingState) error {
 	patchBytes, err := json.Marshal(map[string]interface{}{
 		"metadata": map[string]interface{}{
-			"annotations": annotations,
+			"annotations": bindingState.BindingPodAnnotations,
 		},
 	})
 	if err != nil {

pkg/binder/binding/default_binder_test.go

@@ -182,6 +182,9 @@ func TestBindApplyResourceReceivedType(t *testing.T) {
 
 	err := binder.Bind(context.TODO(), pod, &v1.Node{ObjectMeta: metav1.ObjectMeta{
 		Name: "my-node",
+		Labels: map[string]string{
+			constants.NvidiaGpuMemory: "1000",
+		},
 	}}, bindRequest)
 
 	assert.Nil(t, err)

pkg/binder/binding/fraction_binder_test.go

@@ -73,6 +73,9 @@ var happyFlowObjectsBc = []runtime.Object{
 	&v1.Node{
 		ObjectMeta: metav1.ObjectMeta{
 			Name: "my-node",
+			Labels: map[string]string{
+				constants.NvidiaGpuMemory: "1000",
+			},
 		},
 	},
 }
@@ -117,6 +120,9 @@ var happyFlowObjects = []runtime.Object{
 		},
 		ObjectMeta: metav1.ObjectMeta{
 			Name: "my-node",
+			Labels: map[string]string{
+				constants.NvidiaGpuMemory: "1000",
+			},
 		},
 	},
 }
@@ -198,6 +204,9 @@ var _ = Describe("FractionBinder", func() {
 				result := testedBinder.Bind(
 					context.TODO(), pod, &v1.Node{ObjectMeta: metav1.ObjectMeta{
 						Name: "my-node",
+						Labels: map[string]string{
+							constants.NvidiaGpuMemory: "1000",
+						},
 					}}, bindRequest)
 				if testData.expectedErrorContains != "" {
 					Expect(result).NotTo(BeNil())
@@ -260,6 +269,9 @@ var _ = Describe("FractionBinder", func() {
 		result := testedBinder.Bind(
 			context.TODO(), pod, &v1.Node{ObjectMeta: metav1.ObjectMeta{
 				Name: "my-node",
+				Labels: map[string]string{
+					constants.NvidiaGpuMemory: "1000",
+				},
 			}}, bindRequest)
 
 		Expect(result).NotTo(BeNil())

pkg/binder/plugins/config.go

@@ -17,6 +17,7 @@ const (
 	DynamicResourcesPluginName = kaiv1binder.DynamicResourcesPluginName
 	GPUSharingPluginName       = kaiv1binder.GPUSharingPluginName
 	HamiCorePluginName         = kaiv1binder.HamiCorePluginName
+	NvFractionsPluginName      = kaiv1binder.NvFractionsPluginName
 
 	BindTimeoutSecondsArgument = kaiv1binder.BindTimeoutSecondsArgument
 	CDIEnabledArgument         = kaiv1binder.CDIEnabledArgument

pkg/binder/plugins/factory.go

@@ -16,6 +16,7 @@ import (
 	"github.com/kai-scheduler/KAI-scheduler/pkg/binder/plugins/gpusharing"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/binder/plugins/hamicore"
 	k8splugins "github.com/kai-scheduler/KAI-scheduler/pkg/binder/plugins/k8s-plugins"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/binder/plugins/nvfractions"
 )
 
 type PluginBuildContext struct {
@@ -51,6 +52,7 @@ func InitDefaultPlugins() {
 	RegisterPluginBuilder(DynamicResourcesPluginName, newDynamicResourcesPlugin)
 	RegisterPluginBuilder(GPUSharingPluginName, newGPUSharingPlugin)
 	RegisterPluginBuilder(HamiCorePluginName, newHamiCorePlugin)
+	RegisterPluginBuilder(NvFractionsPluginName, newNvFractionsPlugin)
 }
 
 func BuildConfiguredPlugins(buildContext PluginBuildContext, config Config) (*BinderPlugins, error) {
@@ -112,6 +114,14 @@ func newHamiCorePlugin(buildContext PluginBuildContext, _ map[string]string) (Pl
 	return hamicore.New(buildContext.KubeClient), nil
 }
 
+func newNvFractionsPlugin(_ PluginBuildContext, arguments map[string]string) (Plugin, error) {
+	cdiEnabled, err := boolArgument(arguments, CDIEnabledArgument)
+	if err != nil {
+		return nil, err
+	}
+	return nvfractions.New(cdiEnabled), nil
+}
+
 func validateDependentPlugins(config Config) error {
 	hamiCoreCfg, hamiCoreFound := config[HamiCorePluginName]
 	if !hamiCoreFound || (hamiCoreCfg.Enabled != nil && !*hamiCoreCfg.Enabled) {

pkg/binder/plugins/gpusharing/gpu_sharing.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strconv"
 	"strings"
 
 	"golang.org/x/exp/slices"
@@ -15,10 +16,11 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	"github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
-	"github.com/kai-scheduler/KAI-scheduler/pkg/binder/common/gpusharingconfigmap"
-
 	"github.com/kai-scheduler/KAI-scheduler/pkg/binder/common"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/binder/common/gpusharingconfigmap"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/binder/plugins/state"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
 )
 
 const (
@@ -42,7 +44,7 @@ func (p *GPUSharing) Name() string {
 }
 
 func (p *GPUSharing) PreBind(
-	ctx context.Context, pod *v1.Pod, _ *v1.Node, bindRequest *v1alpha2.BindRequest, state *state.BindingState,
+	ctx context.Context, pod *v1.Pod, node *v1.Node, bindRequest *v1alpha2.BindRequest, state *state.BindingState,
 ) error {
 	if !common.IsSharedGPUAllocation(bindRequest) {
 		return nil
@@ -60,6 +62,11 @@ func (p *GPUSharing) PreBind(
 		return fmt.Errorf("failed to get fraction container ref: %w", err)
 	}
 
+	err = addNvFractionsAnnotationIfMissing(pod, node, bindRequest, containerRef, state)
+	if err != nil {
+		return fmt.Errorf("failed to add NvFractions annotation: %w", err)
+	}
+
 	err = p.createCapabilitiesConfigMapIfMissing(ctx, pod, containerRef)
 	if err != nil {
 		return fmt.Errorf("failed to create capabilities configmap: %w", err)
@@ -79,6 +86,44 @@ func (p *GPUSharing) PreBind(
 	return common.SetGPUPortion(ctx, p.kubeClient, pod, containerRef, bindRequest.Spec.ReceivedGPU.Portion)
 }
 
+func addNvFractionsAnnotationIfMissing(pod *v1.Pod, node *v1.Node, bindRequest *v1alpha2.BindRequest,
+	containerRef *gpusharingconfigmap.PodContainerRef, bindingState *state.BindingState) error {
+	annotationKey := resources.CalcGpuFractionAnnotationForContainer(containerRef.Container.Name)
+	if _, found := pod.Annotations[annotationKey]; found {
+		return nil
+	}
+
+	if node == nil || bindRequest == nil || bindRequest.Spec.ReceivedGPU == nil {
+		return fmt.Errorf("missing data for NvFractions annotation calculation")
+	}
+
+	gpuMemoryStr, foundGPUMemory := node.Labels[constants.NvidiaGpuMemory]
+	if !foundGPUMemory {
+		return fmt.Errorf("node does not include %s label - failed to add NvFractions annotation", constants.NvidiaGpuMemory)
+	}
+
+	totalGPUMemoryMiB, err := strconv.ParseFloat(gpuMemoryStr, 64)
+	if err != nil || totalGPUMemoryMiB <= 0 {
+		return fmt.Errorf("invalid %s label value %q - failed to add NvFractions annotation", constants.NvidiaGpuMemory, gpuMemoryStr)
+	}
+
+	gpuPortion, err := strconv.ParseFloat(bindRequest.Spec.ReceivedGPU.Portion, 64)
+	if err != nil || gpuPortion <= 0 {
+		return fmt.Errorf("invalid received gpu portion %q - failed to add NvFractions annotation", bindRequest.Spec.ReceivedGPU.Portion)
+	}
+
+	gpuMemory := uint64(totalGPUMemoryMiB * gpuPortion)
+	if gpuMemory == 0 {
+		return fmt.Errorf("calculated gpu memory request is zero")
+	}
+
+	if bindingState.BindingPodAnnotations == nil {
+		bindingState.BindingPodAnnotations = map[string]string{}
+	}
+	bindingState.BindingPodAnnotations[annotationKey] = resources.GpuMemoryAnnotationToNvFractionsMemoryRequest(gpuMemory).String()
+	return nil
+}
+
 func (p *GPUSharing) createCapabilitiesConfigMapIfMissing(ctx context.Context, pod *v1.Pod,
 	containerRef *gpusharingconfigmap.PodContainerRef) error {
 	capabilitiesConfigMapName, err := gpusharingconfigmap.ExtractCapabilitiesConfigMapName(pod, containerRef)