Update the compute mode annotation according to the updated API with gpu-sharing operator and add validations

Author: davidLif

.changes/unreleased/changed-20260816-152128.yaml

@@ -0,0 +1,3 @@
+kind: Changed
+body: |-
+  Renamed gpu-compute-sharing-mode annotation to per-container nvidia.com/container.<name>.gpu-compute.mode with validation

docs/gpu-sharing/README.md

@@ -314,13 +314,13 @@ kubectl apply -f gpu-sharing-non-default-container.yaml
 
 ### Compute sharing mode within a fractional GPU group
 
-Pods that share a GPU must use compatible compute sharing behavior within the same fractional GPU group. Use `kai.scheduler/gpu-compute-sharing-mode` when the workload needs to declare its compute sharing mode explicitly:
+Pods that share a GPU must use compatible compute sharing behavior within the same fractional GPU group. Use `nvidia.com/container.<container-name>.gpu-compute.mode` when the workload needs to declare its compute sharing mode explicitly. The container name must match the container targeted by the pod's other fractional GPU annotations:
 
 ```yaml
 metadata:
   annotations:
     gpu-fraction: "0.5"
-    kai.scheduler/gpu-compute-sharing-mode: time-slicing
+    nvidia.com/container.gpu-workload.gpu-compute.mode: time-slicing
 ```
 
 KAI keeps pods with incompatible compute sharing modes out of the same fractional GPU group.

docs/gpu-sharing/nv-fraction/README.md

@@ -66,16 +66,18 @@ Setting a `limit` is useful for pods that have occasional "bursts" of gpu memory
 ## Choose a compute-sharing mode
 
 GPU memory and GPU compute mode are selected independently. The
-`kai.scheduler/gpu-compute-sharing-mode` annotation chooses how workloads in
-the same fractional GPU group share compute. KAI supports these values:
+`nvidia.com/container.<container-name>.gpu-compute.mode` annotation chooses how
+workloads in the same fractional GPU group share compute. KAI supports these
+values:
 
 | Mode | How workloads run | Best fit |
 | --- | --- | --- |
 | `time-slicing` | Workloads take turns using the GPU. Unused time is available to other workloads. This is the default. | Bursty development, notebooks, and throughput-oriented or latency-tolerant inference. |
 | `sm-sharing` | Workloads run concurrently and share the GPU streaming multiprocessors (SMs). | Steady or latency-sensitive inference, and coordinated multi-GPU or multi-Pod workloads. Requires MPS. |
 
-Set `kai.scheduler/gpu-compute-sharing-mode` in the NvFractions Pod manifest.
-The [GPU Sharing guide](../README.md#compute-sharing-mode-within-a-fractional-gpu-group)
+Set `nvidia.com/container.<container-name>.gpu-compute.mode` in the NvFractions
+Pod manifest, using the same container name as the pod's other NvFractions
+annotations. The [GPU Sharing guide](../README.md#compute-sharing-mode-within-a-fractional-gpu-group)
 shows the annotation in context.
 
 ![Time-slicing and SM-sharing](compute-mode.png)

pkg/binder/binding/resourcereservation/resource_reservation.go

@@ -578,7 +578,7 @@ func (rsc *service) waitForGPUReservationPodAllocation(
 
 func (rsc *service) createResourceReservationPod(
 	nodeName string, fractionalGpuGroup schedulingv1alpha2.FractionalGpuGroup,
-	podName string, resources v1.ResourceRequirements,
+	podName string, containerResources v1.ResourceRequirements,
 ) (*v1.Pod, error) {
 	fractionalGpuGroup = fractionalGpuGroup.WithDefaults()
 	podSpec := &v1.Pod{
@@ -590,8 +590,8 @@ func (rsc *service) createResourceReservationPod(
 				constants.GPUGroup:     fractionalGpuGroup.ID,
 			},
 			Annotations: map[string]string{
-				karpenterv1.DoNotDisruptAnnotationKey: "true",
-				constants.GpuComputeSharingMode:       string(fractionalGpuGroup.ComputeSharingMode),
+				karpenterv1.DoNotDisruptAnnotationKey:                                          "true",
+				resources.CalcGpuComputeSharingModeAnnotationForContainer(resourceReservation): string(fractionalGpuGroup.ComputeSharingMode),
 			},
 		},
 		Spec: v1.PodSpec{
@@ -609,7 +609,7 @@ func (rsc *service) createResourceReservationPod(
 					Name:            resourceReservation,
 					Image:           rsc.reservationPodImage,
 					ImagePullPolicy: v1.PullIfNotPresent,
-					Resources:       resources,
+					Resources:       containerResources,
 					SecurityContext: rsc.reservationContainerSecurityContext,
 					Env: []v1.EnvVar{
 						{

pkg/binder/binding/resourcereservation/resource_reservation_test.go

@@ -26,6 +26,7 @@ import (
 
 	schedulingv1alpha2 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/common/resources"
 )
 
 const (
@@ -905,7 +906,7 @@ var _ = Describe("ResourceReservationService", func() {
 				kubeClient:          fake.NewClientBuilder().WithScheme(testScheme).Build(),
 			}
 
-			resources := v1.ResourceRequirements{
+			containerResources := v1.ResourceRequirements{
 				Limits: v1.ResourceList{
 					"nvidia.com/gpu": resource.MustParse("1"),
 				},
@@ -918,7 +919,7 @@ var _ = Describe("ResourceReservationService", func() {
 			pod, err := rsc.createResourceReservationPod(nodeName, schedulingv1alpha2.FractionalGpuGroup{
 				ID:                 gpuGroup,
 				ComputeSharingMode: schedulingv1alpha2.GPUComputeSharingModeTimeSlicing,
-			}, podName, resources)
+			}, podName, containerResources)
 			Expect(err).To(BeNil())
 			Expect(pod).NotTo(BeNil())
 
@@ -927,7 +928,7 @@ var _ = Describe("ResourceReservationService", func() {
 			Expect(pod.Namespace).To(Equal("kai-resource-reservation"))
 			Expect(pod.Labels[constants.AppLabelName]).To(Equal("kai-reservation"))
 			Expect(pod.Labels[constants.GPUGroup]).To(Equal(gpuGroup))
-			Expect(pod.Annotations[constants.GpuComputeSharingMode]).To(Equal(
+			Expect(pod.Annotations[resources.CalcGpuComputeSharingModeAnnotationForContainer("resource-reservation")]).To(Equal(
 				string(schedulingv1alpha2.GPUComputeSharingModeTimeSlicing)))
 
 			// PodSpec checks
@@ -941,7 +942,7 @@ var _ = Describe("ResourceReservationService", func() {
 			Expect(container.Name).To(Equal("resource-reservation"))
 			Expect(container.Image).To(Equal("nvidia/kai-reservation:latest"))
 			Expect(container.ImagePullPolicy).To(Equal(v1.PullIfNotPresent))
-			Expect(container.Resources).To(Equal(resources))
+			Expect(container.Resources).To(Equal(containerResources))
 
 			// Check env vars
 			podNameEnv := v1.EnvVar{
@@ -983,7 +984,7 @@ var _ = Describe("ResourceReservationService", func() {
 				v1.ResourceRequirements{},
 			)
 			Expect(err).To(BeNil())
-			Expect(pod.Annotations[constants.GpuComputeSharingMode]).To(Equal(
+			Expect(pod.Annotations[resources.CalcGpuComputeSharingModeAnnotationForContainer("resource-reservation")]).To(Equal(
 				string(schedulingv1alpha2.GPUComputeSharingModeSMSharing)))
 		})
 	})

pkg/common/constants/constants.go

@@ -56,7 +56,6 @@ const (
 	SkipPodGrouperAnnotation      = "kai.scheduler/skip-podgrouper"
 	GpuFraction                   = "gpu-fraction"
 	GpuFractionContainerName      = "gpu-fraction-container-name"
-	GpuComputeSharingMode         = "kai.scheduler/gpu-compute-sharing-mode"
 	GpuMemory                     = "gpu-memory"
 	ReceivedResourceType          = "received-resource-type"
 	GpuFractionsNumDevices        = "gpu-fraction-num-devices"
@@ -84,6 +83,7 @@ const (
 	NvFractionsMemoryRequestSuffix  = ".gpu-memory.request"
 	NvFractionsMemoryLimitSuffix    = ".gpu-memory.limit"
 	NvFractionsVisibleDevicesSuffix = ".gpus.devices"
+	GpuComputeSharingModeSuffix     = ".gpu-compute.mode"
 
 	KaiFractionContainerAnnotationPrefix = "kai.scheduler/container."
 	GpuMemoryPortionLimitSuffix          = ".gpu-memory.portion.limit"

pkg/common/resources/gpu_compute_sharing_mode.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"
+)
+
+// CalcGpuComputeSharingModeAnnotationForContainer returns the per-container
+// GPU compute-sharing-mode annotation key for containerName.
+func CalcGpuComputeSharingModeAnnotationForContainer(containerName string) string {
+	return constants.NvFractionsAnnotationPrefix + containerName + constants.GpuComputeSharingModeSuffix
+}
+
+// ExtractGpuComputeSharingModeAnnotation returns the container name and raw
+// value of the pod's gpu-compute.mode annotation, if present.
+func ExtractGpuComputeSharingModeAnnotation(pod *v1.Pod) (containerName string, rawValue string, found bool) {
+	for annotationKey, annotationValue := range pod.Annotations {
+		if !isGpuComputeSharingModeAnnotation(annotationKey) {
+			continue
+		}
+		return gpuComputeSharingModeContainerName(annotationKey), annotationValue, true
+	}
+	return "", "", false
+}
+
+func isGpuComputeSharingModeAnnotation(annotationKey string) bool {
+	return strings.HasPrefix(annotationKey, constants.NvFractionsAnnotationPrefix) &&
+		strings.HasSuffix(annotationKey, constants.GpuComputeSharingModeSuffix)
+}
+
+func gpuComputeSharingModeContainerName(annotationKey string) string {
+	containerName := strings.TrimPrefix(annotationKey, constants.NvFractionsAnnotationPrefix)
+	return strings.TrimSuffix(containerName, constants.GpuComputeSharingModeSuffix)
+}

pkg/common/resources/gpu_sharing.go

@@ -12,6 +12,7 @@ import (
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 
+	schedulingv1alpha2 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
 )
 
@@ -44,8 +45,9 @@ type PodGPUFractionRequest struct {
 }
 
 type NvFractionsContainerRequest struct {
-	Request *resource.Quantity
-	Limit   *resource.Quantity
+	Request     *resource.Quantity
+	Limit       *resource.Quantity
+	ComputeMode *schedulingv1alpha2.GPUComputeSharingMode
 }
 
 var (
@@ -243,3 +245,8 @@ func GpuMemoryAnnotationToNvFractionsMemoryRequest(gpuMemory uint64) *resource.Q
 	memory := resource.MustParse(fmt.Sprintf("%dMi", gpuMemory))
 	return &memory
 }
+
+func IsValidGPUComputeSharingMode(mode string) bool {
+	return mode == string(schedulingv1alpha2.GPUComputeSharingModeTimeSlicing) ||
+		mode == string(schedulingv1alpha2.GPUComputeSharingModeSMSharing)
+}

pkg/common/resources/gpu_sharing_nvfractions.go

@@ -10,6 +10,7 @@ import (
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 
+	schedulingv1alpha2 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
 )
 
@@ -18,6 +19,7 @@ type nvFractionsAnnotationType int
 const (
 	nvFractionsRequestAnnotation nvFractionsAnnotationType = iota
 	nvFractionsLimitAnnotation
+	nvFractionsComputeModeAnnotation
 )
 
 func CalcGpuFractionAnnotationForContainer(containerName string) string {
@@ -38,18 +40,32 @@ func ExtractNvFractionsData(pod *v1.Pod) (map[string]NvFractionsContainerRequest
 		if !strings.HasPrefix(annotationKey, constants.NvFractionsAnnotationPrefix) {
 			continue
 		}
+		if isNvFractionsDeviceListAnnotation(annotationKey) {
+			continue
+		}
 
 		containerName, annotationType, err := parseNvFractionsAnnotationKey(annotationKey)
 		if err != nil {
 			return nil, err
 		}
 
+		containerData := fractionsData[containerName]
+
+		if annotationType == nvFractionsComputeModeAnnotation {
+			if !IsValidGPUComputeSharingMode(annotationValue) {
+				return nil, fmt.Errorf("invalid NvFractions compute mode: %s", annotationValue)
+			}
+			mode := schedulingv1alpha2.GPUComputeSharingMode(annotationValue)
+			containerData.ComputeMode = &mode
+			fractionsData[containerName] = containerData
+			continue
+		}
+
 		gpuMemory, err := parseNvFractionsAnnotationValue(annotationKey, annotationValue)
 		if err != nil {
 			return nil, err
 		}
 
-		containerData := fractionsData[containerName]
 		if annotationType == nvFractionsRequestAnnotation {
 			containerData.Request = &gpuMemory
 		} else {
@@ -75,6 +91,15 @@ func getNvFractionData(pod *v1.Pod) (*NvFractionsContainerRequest, error) {
 	return nil, nil
 }
 
+// isNvFractionsDeviceListAnnotation reports whether annotationKey is the
+// device-list annotation. It shares the NvFractions prefix but, unlike
+// request/limit/compute-mode, isn't part of the customer's fractional GPU
+// request - it's written by the binder after scheduling - so it must be
+// skipped here rather than treated as an invalid key.
+func isNvFractionsDeviceListAnnotation(annotationKey string) bool {
+	return strings.HasSuffix(annotationKey, constants.NvFractionsVisibleDevicesSuffix)
+}
+
 func parseNvFractionsAnnotationKey(annotationKey string) (string, nvFractionsAnnotationType, error) {
 	containerNameWithSuffix := strings.TrimPrefix(annotationKey, constants.NvFractionsAnnotationPrefix)
 	if strings.HasSuffix(annotationKey, constants.NvFractionsMemoryRequestSuffix) {
@@ -85,6 +110,10 @@ func parseNvFractionsAnnotationKey(annotationKey string) (string, nvFractionsAnn
 		containerName := strings.TrimSuffix(containerNameWithSuffix, constants.NvFractionsMemoryLimitSuffix)
 		return containerName, nvFractionsLimitAnnotation, nil
 	}
+	if strings.HasSuffix(annotationKey, constants.GpuComputeSharingModeSuffix) {
+		containerName := strings.TrimSuffix(containerNameWithSuffix, constants.GpuComputeSharingModeSuffix)
+		return containerName, nvFractionsComputeModeAnnotation, nil
+	}
 	return "", 0, fmt.Errorf("invalid NvFractions annotation key: %s", annotationKey)
 }
 

pkg/common/resources/gpu_sharing_nvfractions_test.go

@@ -7,7 +7,9 @@ import (
 	"testing"
 
 	"k8s.io/apimachinery/pkg/api/resource"
+	"k8s.io/utils/ptr"
 
+	schedulingv1alpha2 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
 )
 
@@ -75,6 +77,31 @@ func TestExtractNvFractionsData(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "compute sharing mode annotation does not break extraction of a sibling container's request",
+			annotations: map[string]string{
+				constants.NvFractionsAnnotationPrefix + "main" + constants.NvFractionsMemoryRequestSuffix: "1Gi",
+				CalcGpuComputeSharingModeAnnotationForContainer("main"):                                   "sm-sharing",
+			},
+			want: map[string]NvFractionsContainerRequest{
+				"main": {
+					Request:     quantityPtr("1Gi"),
+					ComputeMode: ptr.To(schedulingv1alpha2.GPUComputeSharingModeSMSharing),
+				},
+			},
+		},
+		{
+			name: "skips device list annotation",
+			annotations: map[string]string{
+				constants.NvFractionsAnnotationPrefix + "main" + constants.NvFractionsMemoryRequestSuffix: "1Gi",
+				CalcGpuVisibleDevicesAnnotationForContainer("main"):                                       "gpu-0",
+			},
+			want: map[string]NvFractionsContainerRequest{
+				"main": {
+					Request: quantityPtr("1Gi"),
+				},
+			},
+		},
 		{
 			name: "rejects invalid annotation key",
 			annotations: map[string]string{
@@ -130,6 +157,12 @@ func TestParseNvFractionsAnnotationKey(t *testing.T) {
 			wantContainerName: "main",
 			wantType:          nvFractionsLimitAnnotation,
 		},
+		{
+			name:              "compute sharing mode annotation",
+			annotationKey:     constants.NvFractionsAnnotationPrefix + "main" + constants.GpuComputeSharingModeSuffix,
+			wantContainerName: "main",
+			wantType:          nvFractionsComputeModeAnnotation,
+		},
 		{
 			name:              "invalid annotation",
 			annotationKey:     constants.NvFractionsAnnotationPrefix + "main",