diff --git a/cmd/binder/app/app.go b/cmd/binder/app/app.go index da35c0c50..ebf9b3c8c 100644 --- a/cmd/binder/app/app.go +++ b/cmd/binder/app/app.go @@ -114,6 +114,10 @@ func New(options *Options, config *rest.Config) (*App, error) { options.ResourceReservationPodResources.Value, options.ResourceReservationPodSecurityContext.Value, options.ResourceReservationContainerSecurityContext.Value) + if options.ReservationSchedulerName != "" { + rrs = rrs.WithExternalReservation(options.ReservationSchedulerName, + options.ReservationGpuResourceName, options.ReservationPodAnnotations) + } reconcilerParams := &controllers.ReconcilerParams{ MaxConcurrentReconciles: options.MaxConcurrentReconciles, @@ -174,6 +178,27 @@ func (app *App) Run(ctx context.Context) error { setupLog.Error(err, "unable to create controller", "controller", "BindRequest") return err } + + if err = (&controllers.ReserveAheadReconciler{ + Client: app.manager.GetClient(), + Scheme: app.manager.GetScheme(), + ResourceReservation: app.rrs, + SchedulerName: app.Options.SchedulerName, + Enabled: app.Options.ReservationSchedulerName != "", + }).SetupWithManager(app.manager, app.reconcilerParams); err != nil { + setupLog.Error(err, "unable to create controller", "controller", "ReserveAhead") + return err + } + + if err = (&controllers.ReservationCascadeReconciler{ + Client: app.manager.GetClient(), + Scheme: app.manager.GetScheme(), + ReservationNamespace: app.Options.ResourceReservationNamespace, + Enabled: app.Options.ReservationSchedulerName != "", + }).SetupWithManager(app.manager, app.reconcilerParams); err != nil { + setupLog.Error(err, "unable to create controller", "controller", "ReservationCascade") + return err + } // +kubebuilder:scaffold:builder setupLog.Info("starting manager") diff --git a/cmd/binder/app/options.go b/cmd/binder/app/options.go index b3450248c..0defb48cb 100644 --- a/cmd/binder/app/options.go +++ b/cmd/binder/app/options.go @@ -36,6 +36,9 @@ type Options struct { FakeGPUNodes bool Plugins flags.JSONFlag[binderplugins.Config] RuntimeClassName string + ReservationSchedulerName string + ReservationGpuResourceName string + ReservationPodAnnotations map[string]string } func InitOptions(fs *pflag.FlagSet) *Options { @@ -109,6 +112,19 @@ func InitOptions(fs *pflag.FlagSet) *Options { fs.StringVar(&options.RuntimeClassName, "runtime-class-name", "", "Runtime class for GPU reservation pods. Defaults to empty (no runtime class).") + fs.StringVar(&options.ReservationSchedulerName, + "reservation-scheduler-name", "", + "When set, reservation pods are created unbound with this spec.schedulerName so an "+ + "external scheduler (e.g. slurm-bridge-scheduler) places them, instead of being "+ + "pre-bound to the scheduler-selected node. Empty keeps the default pinned behavior.") + fs.StringVar(&options.ReservationGpuResourceName, + "reservation-gpu-resource-name", "", + "Resource key the reservation pod requests for a whole GPU in external mode. "+ + "Empty defaults to nvidia.com/gpu; e.g. deviceclass.resource.kubernetes.io/gpu.nvidia.com.") + fs.StringToStringVar(&options.ReservationPodAnnotations, + "reservation-pod-annotations", nil, + "Extra annotations to stamp on externally-scheduled reservation pods (k=v,k2=v2), "+ + "e.g. the slurmjob.slinky.slurm.net/{qos,exclusive} annotations slurm-bridge requires.") utilfeature.DefaultMutableFeatureGate.AddFlag(fs) diff --git a/deployments/kai-scheduler/crds/kai.scheduler_configs.yaml b/deployments/kai-scheduler/crds/kai.scheduler_configs.yaml index d336bb37b..64273430e 100644 --- a/deployments/kai-scheduler/crds/kai.scheduler_configs.yaml +++ b/deployments/kai-scheduler/crds/kai.scheduler_configs.yaml @@ -4048,6 +4048,30 @@ spec: If enabled, this prevents pods of each microservice from being scheduled on the same node. If another podAntiAffinity term is defined (either globally or locally for a specific microservice), this will be ignored. type: boolean + reservationGpuResourceName: + description: |- + ReservationGpuResourceName overrides the resource key the reservation pod requests for a + whole GPU (only used when ReservationSchedulerName is set). Empty defaults to + "nvidia.com/gpu"; set it to "deviceclass.resource.kubernetes.io/gpu.nvidia.com" to request + a DRA-backed extended resource that slurm-bridge accepts and Slurm accounts. + type: string + reservationPodAnnotations: + additionalProperties: + type: string + description: |- + ReservationPodAnnotations are extra annotations stamped onto externally-scheduled + reservation pods (only used when ReservationSchedulerName is set), e.g. the + slurmjob.slinky.slurm.net/{qos,exclusive} annotations slurm-bridge requires to build a + Slurm job. Ignored in the default pinned mode. + type: object + reservationSchedulerName: + description: |- + ReservationSchedulerName, when set, makes the binder create GPU-sharing reservation + pods with this `spec.schedulerName` (instead of pre-binding them via `spec.nodeName`), + so an external scheduler (e.g. slurm-bridge-scheduler) places the whole-GPU reservation + pod and KAI only sub-divides fractions onto it. Empty/nil preserves the default behavior + where the binder pins the reservation pod to the scheduler-selected node. + type: string schedulerName: description: |- SchedulerName specifies the name of the KAI scheduler. Pods must set this value diff --git a/deployments/kai-scheduler/templates/_helpers.tpl b/deployments/kai-scheduler/templates/_helpers.tpl index 0090e6bac..778e4a881 100644 --- a/deployments/kai-scheduler/templates/_helpers.tpl +++ b/deployments/kai-scheduler/templates/_helpers.tpl @@ -72,6 +72,16 @@ spec: {{- if .Values.global.jsonLog }} jsonLog: true {{- end }} + {{- if .Values.global.reservationSchedulerName }} + reservationSchedulerName: {{ .Values.global.reservationSchedulerName | quote }} + {{- end }} + {{- if .Values.global.reservationGpuResourceName }} + reservationGpuResourceName: {{ .Values.global.reservationGpuResourceName | quote }} + {{- end }} + {{- if .Values.global.reservationPodAnnotations }} + reservationPodAnnotations: + {{- toYaml .Values.global.reservationPodAnnotations | nindent 6 }} + {{- end }} {{- if .Values.global.affinity }} affinity: {{- toYaml .Values.global.affinity | nindent 6 }} diff --git a/pkg/apis/kai/v1/global.go b/pkg/apis/kai/v1/global.go index 2200a5329..d583981c0 100644 --- a/pkg/apis/kai/v1/global.go +++ b/pkg/apis/kai/v1/global.go @@ -74,6 +74,28 @@ type GlobalConfig struct { // +kubebuilder:validation:Optional SchedulerName *string `json:"schedulerName,omitempty"` + // ReservationSchedulerName, when set, makes the binder create GPU-sharing reservation + // pods with this `spec.schedulerName` (instead of pre-binding them via `spec.nodeName`), + // so an external scheduler (e.g. slurm-bridge-scheduler) places the whole-GPU reservation + // pod and KAI only sub-divides fractions onto it. Empty/nil preserves the default behavior + // where the binder pins the reservation pod to the scheduler-selected node. + // +kubebuilder:validation:Optional + ReservationSchedulerName *string `json:"reservationSchedulerName,omitempty"` + + // ReservationPodAnnotations are extra annotations stamped onto externally-scheduled + // reservation pods (only used when ReservationSchedulerName is set), e.g. the + // slurmjob.slinky.slurm.net/{qos,exclusive} annotations slurm-bridge requires to build a + // Slurm job. Ignored in the default pinned mode. + // +kubebuilder:validation:Optional + ReservationPodAnnotations map[string]string `json:"reservationPodAnnotations,omitempty"` + + // ReservationGpuResourceName overrides the resource key the reservation pod requests for a + // whole GPU (only used when ReservationSchedulerName is set). Empty defaults to + // "nvidia.com/gpu"; set it to "deviceclass.resource.kubernetes.io/gpu.nvidia.com" to request + // a DRA-backed extended resource that slurm-bridge accepts and Slurm accounts. + // +kubebuilder:validation:Optional + ReservationGpuResourceName *string `json:"reservationGpuResourceName,omitempty"` + // NodePoolLabelKey is the label name by with to filter nodes, pods and other resources that the scheduler is watching // +kubebuilder:validation:Optional NodePoolLabelKey *string `json:"nodePoolLabelKey,omitempty"` diff --git a/pkg/apis/kai/v1/zz_generated.deepcopy.go b/pkg/apis/kai/v1/zz_generated.deepcopy.go index 3766fc416..b87774086 100644 --- a/pkg/apis/kai/v1/zz_generated.deepcopy.go +++ b/pkg/apis/kai/v1/zz_generated.deepcopy.go @@ -275,6 +275,23 @@ func (in *GlobalConfig) DeepCopyInto(out *GlobalConfig) { *out = new(string) **out = **in } + if in.ReservationSchedulerName != nil { + in, out := &in.ReservationSchedulerName, &out.ReservationSchedulerName + *out = new(string) + **out = **in + } + if in.ReservationPodAnnotations != nil { + in, out := &in.ReservationPodAnnotations, &out.ReservationPodAnnotations + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } + if in.ReservationGpuResourceName != nil { + in, out := &in.ReservationGpuResourceName, &out.ReservationGpuResourceName + *out = new(string) + **out = **in + } if in.NodePoolLabelKey != nil { in, out := &in.NodePoolLabelKey, &out.NodePoolLabelKey *out = new(string) diff --git a/pkg/binder/binding/resourcereservation/mock/resource_reservation_mock.go b/pkg/binder/binding/resourcereservation/mock/resource_reservation_mock.go index 79e9f46e0..0f60d46b6 100644 --- a/pkg/binder/binding/resourcereservation/mock/resource_reservation_mock.go +++ b/pkg/binder/binding/resourcereservation/mock/resource_reservation_mock.go @@ -55,6 +55,37 @@ func (mr *MockInterfaceMockRecorder) RemovePodGpuGroupsConnection(ctx, pod any) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemovePodGpuGroupsConnection", reflect.TypeOf((*MockInterface)(nil).RemovePodGpuGroupsConnection), ctx, pod) } +// EnsureReservation mocks base method. +func (m *MockInterface) EnsureReservation(ctx context.Context, sourcePod *v1.Pod, gpuGroup string) (*v1.Pod, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "EnsureReservation", ctx, sourcePod, gpuGroup) + ret0, _ := ret[0].(*v1.Pod) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// EnsureReservation indicates an expected call of EnsureReservation. +func (mr *MockInterfaceMockRecorder) EnsureReservation(ctx, sourcePod, gpuGroup any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EnsureReservation", reflect.TypeOf((*MockInterface)(nil).EnsureReservation), ctx, sourcePod, gpuGroup) +} + +// ReservationPlacement mocks base method. +func (m *MockInterface) ReservationPlacement(ctx context.Context, gpuGroup string) (bool, string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReservationPlacement", ctx, gpuGroup) + ret0, _ := ret[0].(bool) + ret1, _ := ret[1].(string) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// ReservationPlacement indicates an expected call of ReservationPlacement. +func (mr *MockInterfaceMockRecorder) ReservationPlacement(ctx, gpuGroup any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReservationPlacement", reflect.TypeOf((*MockInterface)(nil).ReservationPlacement), ctx, gpuGroup) +} + // ReserveGpuDevice mocks base method. func (m *MockInterface) ReserveGpuDevice(ctx context.Context, pod *v1.Pod, nodeName, gpuGroup string) (string, error) { m.ctrl.T.Helper() diff --git a/pkg/binder/binding/resourcereservation/resource_reservation.go b/pkg/binder/binding/resourcereservation/resource_reservation.go index 1944827e3..32255a876 100644 --- a/pkg/binder/binding/resourcereservation/resource_reservation.go +++ b/pkg/binder/binding/resourcereservation/resource_reservation.go @@ -8,6 +8,7 @@ import ( "crypto/sha256" "encoding/hex" "encoding/json" + "errors" "fmt" "maps" "strings" @@ -36,6 +37,16 @@ type Interface interface { SyncForGpuGroup(ctx context.Context, gpuGroup string) error ReserveGpuDevice(ctx context.Context, pod *v1.Pod, nodeName string, gpuGroup string) (string, error) RemovePodGpuGroupsConnection(ctx context.Context, pod *v1.Pod) error + // EnsureReservation creates (idempotently, by gpuGroup) a whole-GPU reservation pod for the + // given gpuGroup, WITHOUT binding it to a node — the external scheduler (slurm-bridge) places + // it. Used by the reserve-ahead controller to enqueue a reservation for a Pending fraction + // that cannot yet pack, so Slurm can make room (preempt) while the fraction waits. External + // reservation mode only. + EnsureReservation(ctx context.Context, sourcePod *v1.Pod, gpuGroup string) (*v1.Pod, error) + // ReservationPlacement reports whether the reservation pod for gpuGroup has been placed and is + // ready to back a fraction (Running, scheduled to a node, and has self-reported its GPU index), + // along with the node it landed on. + ReservationPlacement(ctx context.Context, gpuGroup string) (ready bool, nodeName string, err error) } const ( @@ -44,8 +55,30 @@ const ( gpuIndexAnnotationName = "run.ai/reserve_for_gpu_index" numberOfGPUsToReserve = 1 unknownGpuIndicator = "-1" + // pendingGpuIndicator is returned by waitForGPUReservationPodAllocation when the + // reservation pod did not allocate within the timeout but is still Pending (e.g. an + // external scheduler such as slurm-bridge has enqueued its Slurm job and is waiting + // for a whole GPU to free). Distinct from unknownGpuIndicator, which means the + // reservation genuinely failed and should be reaped. + pendingGpuIndicator = "-2" + // ReservationCascadeFinalizer is placed on externally-scheduled reservation pods so that + // deleting a reservation cascades to the GPU-fraction workloads it backs: a controller + // deletes those fractions and only removes this finalizer once they are gone. Without it, + // deleting the reservation frees the whole GPU in Slurm (slurm-bridge cancels the GRES job + // on pod termination) while the fractions keep running on it — a silent double-booking. + ReservationCascadeFinalizer = "kai.scheduler/reservation-cascade" ) +// ErrReservationPending signals that a GPU reservation pod was created and is still +// Pending (parked, e.g. waiting in an external scheduler's queue) rather than failed. +// The BindRequest reconciler treats it as "retry later" — it must NOT roll back the +// fraction pod's GPU-group binding or mark the BindRequest Failed, so the parked +// reservation keeps its place in the queue and is reused on the next reconcile. The +// parked reservation is protected from GC while its BindRequest stays non-terminal +// (see hasActiveBindRequestsForGpuGroup) and is torn down by the reservation sync once +// the fraction binds elsewhere (packs onto a freed GPU) and its BindRequest goes terminal. +var ErrReservationPending = errors.New("gpu reservation pod is still pending allocation") + type service struct { fakeGPuNodes bool kubeClient client.WithWatch @@ -60,6 +93,9 @@ type service struct { podResources *v1.ResourceRequirements reservationPodSecurityContext *v1.PodSecurityContext reservationContainerSecurityContext *v1.SecurityContext + reservationSchedulerName string + reservationGpuResourceName string + reservationPodAnnotations map[string]string } func NewService( @@ -93,6 +129,39 @@ func NewService( } } +// WithExternalReservation switches the service into "external reservation" mode: reservation +// pods are created unbound with the given spec.schedulerName (so an external scheduler places +// them) instead of being pre-bound via spec.nodeName. They request gpuResourceName for the whole +// GPU (defaulting to nvidia.com/gpu when empty) and carry the given extra annotations. Passing an +// empty schedulerName leaves the service in the default pinned mode. +func (rsc *service) WithExternalReservation(schedulerName, gpuResourceName string, podAnnotations map[string]string) *service { + rsc.reservationSchedulerName = schedulerName + rsc.reservationGpuResourceName = gpuResourceName + rsc.reservationPodAnnotations = podAnnotations + return rsc +} + +func (rsc *service) externalReservationMode() bool { + return rsc.reservationSchedulerName != "" +} + +func (rsc *service) gpuResourceName() v1.ResourceName { + if rsc.reservationGpuResourceName != "" { + return v1.ResourceName(rsc.reservationGpuResourceName) + } + return constants.NvidiaGpuResource +} + +// reservationPodAnnotationsMap returns the annotations to stamp on a reservation pod: always the +// karpenter do-not-disrupt annotation, plus any configured external-mode annotations. +func (rsc *service) reservationPodAnnotationsMap() map[string]string { + annotations := map[string]string{ + karpenterv1.DoNotDisruptAnnotationKey: "true", + } + maps.Copy(annotations, rsc.reservationPodAnnotations) + return annotations +} + func (rsc *service) Sync(ctx context.Context) error { podsList := &v1.PodList{} err := rsc.kubeClient.List(ctx, podsList, @@ -324,6 +393,38 @@ func (rsc *service) RemovePodGpuGroupsConnection(ctx context.Context, pod *v1.Po return nil } +func (rsc *service) EnsureReservation(ctx context.Context, sourcePod *v1.Pod, gpuGroup string) (*v1.Pod, error) { + if !rsc.externalReservationMode() { + return nil, fmt.Errorf("reserve-ahead requires external reservation mode (reservationSchedulerName)") + } + // nodeName "" — the external scheduler (slurm-bridge) picks the node and places the pod. + return rsc.createGPUReservationPod(ctx, sourcePod, "", gpuGroup) +} + +func (rsc *service) ReservationPlacement(ctx context.Context, gpuGroup string) (bool, string, error) { + gpuIndex, err := rsc.findGPUIndexByGroup(gpuGroup) + if err != nil { + // The reservation pod exists but has not self-reported its GPU index yet — not ready. + return false, "", nil + } + if gpuIndex == "" { + // No reservation pod for this group. + return false, "", nil + } + pods := &v1.PodList{} + if listErr := rsc.kubeClient.List(ctx, pods, + client.InNamespace(rsc.namespace), + client.MatchingLabels{constants.GPUGroup: gpuGroup}); listErr != nil { + return false, "", listErr + } + if len(pods.Items) == 0 { + return false, "", nil + } + pod := pods.Items[0] + ready := pod.Status.Phase == v1.PodRunning && pod.Spec.NodeName != "" + return ready, pod.Spec.NodeName, nil +} + // escapeJSONPointer escapes a string for use in a JSON Pointer path (RFC 6901). // ~ must be escaped as ~0, and / must be escaped as ~1. func escapeJSONPointer(s string) string { @@ -379,6 +480,15 @@ func (rsc *service) createGPUReservationPodAndGetIndex( } gpuIndex = rsc.waitForGPUReservationPodAllocation(ctx, nodeName, pod.Name) + if gpuIndex == pendingGpuIndicator { + // The reservation is still Pending (parked): do NOT delete it. Keep it in place + // so it holds its spot in the external scheduler's queue and is reused on the next + // reconcile. The caller (BindRequest reconciler) requeues without rollback. + logger.Info("GPU reservation pod is still pending, parking it for retry", + "name", pod.Name, "gpuGroup", gpuGroup) + return unknownGpuIndicator, fmt.Errorf( + "gpu reservation pod %v/%v: %w", rsc.namespace, pod.Name, ErrReservationPending) + } if gpuIndex == unknownGpuIndicator { deleteErr := rsc.deleteReservationPod(ctx, pod) if deleteErr != nil { @@ -388,6 +498,30 @@ func (rsc *service) createGPUReservationPodAndGetIndex( "failed waiting for GPU reservation pod to allocate: %v/%v", rsc.namespace, pod.Name) } + if rsc.externalReservationMode() && nodeName != "" { + // The reservation pod is running now (it self-reported its GPU index), so the external + // scheduler has set its nodeName. In v1 (static single-GPU-node pool) the fraction pod is + // bound to the scheduler-selected node, so the reservation must have landed there too; + // otherwise the fraction would share a GPU that lives on a different node. Fail loudly + // instead of binding to the wrong node. True multi-node "follow the reservation" is a + // Phase-2 change to the bind path. + placed := &v1.Pod{} + if getErr := rsc.kubeClient.Get(ctx, + types.NamespacedName{Namespace: rsc.namespace, Name: pod.Name}, placed); getErr != nil { + return unknownGpuIndicator, fmt.Errorf( + "failed to read externally-scheduled reservation pod %v/%v: %w", rsc.namespace, pod.Name, getErr) + } + if placed.Spec.NodeName != "" && placed.Spec.NodeName != nodeName { + if deleteErr := rsc.deleteReservationPod(ctx, placed); deleteErr != nil { + logger.Error(deleteErr, "failed to delete diverged reservation pod", "name", pod.Name) + } + return unknownGpuIndicator, fmt.Errorf( + "external reservation pod %v/%v was placed on node %q but the scheduler selected %q; "+ + "multi-node external placement is not supported in v1 (static single-GPU-node pool)", + rsc.namespace, pod.Name, placed.Spec.NodeName, nodeName) + } + } + return gpuIndex, err } @@ -411,6 +545,25 @@ func (rsc *service) deleteNonReservedPods(ctx context.Context, gpuGroup string, func (rsc *service) deleteReservationPod(ctx context.Context, pod *v1.Pod) error { logger := log.FromContext(ctx) + // This is the INTERNAL reservation deleter (reap-on-allocation-timeout, node-divergence + // guard, and orphan cleanup in syncForPods). Strip our cascade finalizer first so this + // system-initiated deletion does NOT cascade to the backed fraction workloads — otherwise a + // transient/racy internal reservation delete would kill healthy fractions. Only an EXTERNAL + // delete (which leaves the finalizer intact) should trigger the cascade. + if slices.Contains(pod.Finalizers, ReservationCascadeFinalizer) { + original := pod.DeepCopy() + kept := make([]string, 0, len(pod.Finalizers)) + for _, f := range pod.Finalizers { + if f != ReservationCascadeFinalizer { + kept = append(kept, f) + } + } + pod.Finalizers = kept + if patchErr := rsc.kubeClient.Patch(ctx, pod, client.MergeFrom(original)); patchErr != nil && + !apierrors.IsNotFound(patchErr) { + return fmt.Errorf("failed to strip cascade finalizer from reservation pod %s: %w", pod.Name, patchErr) + } + } logger.Info("Deleting reservation pod", "name", pod.Name) err := rsc.kubeClient.Delete(ctx, pod, @@ -433,25 +586,35 @@ func (rsc *service) createGPUReservationPod( return nil, fmt.Errorf("cluster is scaling up, could not create reservation pod") } - podName := reservationPodName(nodeName, gpuGroup) + // In external-reservation mode the pod is not pre-bound, so its name must not depend on a + // node (there isn't one yet) — key it on the gpu-group only so concurrent/retried creates + // still collide on a single object. + nameNode := nodeName + if rsc.externalReservationMode() { + nameNode = "" + } + podName := reservationPodName(nameNode, gpuGroup) - // Build resource requirements starting with GPU resources + // Build resource requirements starting with the whole-GPU request. The resource key is + // configurable (e.g. deviceclass.resource.kubernetes.io/gpu.nvidia.com for a DRA-backed + // extended resource under slurm-bridge) and defaults to nvidia.com/gpu. + gpuResource := rsc.gpuResourceName() resources := v1.ResourceRequirements{ Limits: v1.ResourceList{ - constants.NvidiaGpuResource: *resource.NewQuantity(numberOfGPUsToReserve, resource.DecimalSI), + gpuResource: *resource.NewQuantity(numberOfGPUsToReserve, resource.DecimalSI), }, Requests: v1.ResourceList{ - constants.NvidiaGpuResource: *resource.NewQuantity(numberOfGPUsToReserve, resource.DecimalSI), + gpuResource: *resource.NewQuantity(numberOfGPUsToReserve, resource.DecimalSI), }, } if rsc.podResources != nil { if rsc.podResources.Limits != nil { - delete(rsc.podResources.Limits, constants.NvidiaGpuResource) + delete(rsc.podResources.Limits, gpuResource) maps.Copy(resources.Limits, rsc.podResources.Limits) } if rsc.podResources.Requests != nil { - delete(rsc.podResources.Requests, constants.NvidiaGpuResource) + delete(rsc.podResources.Requests, gpuResource) maps.Copy(resources.Requests, rsc.podResources.Requests) } } @@ -506,7 +669,10 @@ func (rsc *service) waitForGPUReservationPodAllocation( logger.Error(fmt.Errorf("timeout"), "Reached timeout while waiting for GPU reservation pod to be allocated", "nodeName", nodeName, "name", gpuReservationPodName) - return unknownGpuIndicator + // A reservation that is still Pending after the timeout has not failed — an + // external scheduler (slurm-bridge) may have enqueued it and be waiting for a + // whole GPU. Signal "parked" so the caller keeps it instead of reaping it. + return rsc.timeoutAllocationIndicator(ctx, gpuReservationPodName) case event, ok := <-watcher.ResultChan(): if !ok { logger.Error(nil, @@ -530,6 +696,23 @@ func (rsc *service) waitForGPUReservationPodAllocation( } } +// timeoutAllocationIndicator classifies a reservation pod that did not report its GPU +// index within the allocation timeout: pendingGpuIndicator if it still exists and is +// Pending (parked, e.g. waiting in slurm-bridge's queue) so the caller keeps it, or +// unknownGpuIndicator otherwise (gone / failed) so the caller reaps it. +func (rsc *service) timeoutAllocationIndicator(ctx context.Context, gpuReservationPodName string) string { + pod := &v1.Pod{} + if err := rsc.kubeClient.Get(ctx, + types.NamespacedName{Namespace: rsc.namespace, Name: gpuReservationPodName}, pod); err != nil { + // Cannot confirm it is parked (e.g. not found): treat as failed so it is reaped. + return unknownGpuIndicator + } + if pod.Status.Phase == v1.PodPending { + return pendingGpuIndicator + } + return unknownGpuIndicator +} + func (rsc *service) createResourceReservationPod( sourcePod *v1.Pod, nodeName, gpuGroup, podName string, resources v1.ResourceRequirements, ) (*v1.Pod, error) { @@ -547,12 +730,9 @@ func (rsc *service) createResourceReservationPod( constants.AppLabelName: rsc.appLabelValue, constants.GPUGroup: gpuGroup, }, - Annotations: map[string]string{ - karpenterv1.DoNotDisruptAnnotationKey: "true", - }, + Annotations: rsc.reservationPodAnnotationsMap(), }, Spec: v1.PodSpec{ - NodeName: nodeName, Tolerations: tolerations, RuntimeClassName: func() *string { if len(rsc.runtimeClassName) == 0 { @@ -590,14 +770,26 @@ func (rsc *service) createResourceReservationPod( }, }, }, - Status: v1.PodStatus{ + } + + if rsc.externalReservationMode() { + // Leave the pod unbound (no nodeName, no faked PodScheduled) so the external scheduler + // places it; the whole-GPU request keeps it from double-booking. + podSpec.Spec.SchedulerName = rsc.reservationSchedulerName + // Cascade-delete the backed fractions if this reservation is deleted (see finalizer doc). + podSpec.Finalizers = append(podSpec.Finalizers, ReservationCascadeFinalizer) + } else { + // Default behavior: pre-bind the reservation pod to the selected node and mark it + // PodScheduled so the kubelet/device-plugin admits it without any scheduling pass. + podSpec.Spec.NodeName = nodeName + podSpec.Status = v1.PodStatus{ Conditions: []v1.PodCondition{ { Type: v1.PodScheduled, Status: v1.ConditionTrue, }, }, - }, + } } if rsc.fakeGPuNodes { diff --git a/pkg/binder/binding/resourcereservation/resource_reservation_test.go b/pkg/binder/binding/resourcereservation/resource_reservation_test.go index 0749cdef9..d1d306b50 100644 --- a/pkg/binder/binding/resourcereservation/resource_reservation_test.go +++ b/pkg/binder/binding/resourcereservation/resource_reservation_test.go @@ -1640,3 +1640,148 @@ var _ = Describe("Reservation pod duplicate gpu-group race", func() { func ptrBool(b bool) *bool { return &b } func ptrInt64(i int64) *int64 { return &i } + +var _ = Describe("External reservation mode", func() { + const ( + extNode = "kai-node" + extGroup = "ext-gpu-group" + externalSched = "slurm-bridge-scheduler" + draGpuResource = "deviceclass.resource.kubernetes.io/gpu.nvidia.com" + ) + + buildClient := func() runtimeClient.WithWatch { + return fake.NewClientBuilder().WithScheme(testScheme). + WithIndex(&v1.Pod{}, "spec.nodeName", nodeNameIndexer).Build() + } + + It("creates an unbound reservation pod placed by the external scheduler with the DRA-backed GPU resource and configured annotations", func() { + rsc := initializeTestService(buildClient()).WithExternalReservation( + externalSched, draGpuResource, + map[string]string{ + "slurmjob.slinky.slurm.net/qos": "inference", + "slurmjob.slinky.slurm.net/exclusive": "false", + }, + ) + + pod, err := rsc.createGPUReservationPod(context.TODO(), nil, extNode, extGroup) + Expect(err).To(Succeed()) + + // Unbound: an external scheduler (e.g. slurm-bridge) must place it, so no nodeName and no + // faked PodScheduled=True condition (which is what makes the default mode bypass scheduling). + Expect(pod.Spec.SchedulerName).To(Equal(externalSched)) + Expect(pod.Spec.NodeName).To(BeEmpty()) + Expect(pod.Status.Conditions).To(BeEmpty()) + + // Whole GPU is requested as the DRA-backed extended resource, not the classic device-plugin + // resource, so slurm-bridge accepts it and Slurm accounts it. + Expect(pod.Spec.Containers[0].Resources.Requests).To(HaveKey(v1.ResourceName(draGpuResource))) + Expect(pod.Spec.Containers[0].Resources.Limits).To(HaveKey(v1.ResourceName(draGpuResource))) + Expect(pod.Spec.Containers[0].Resources.Requests).ToNot(HaveKey(v1.ResourceName(constants.NvidiaGpuResource))) + + // Configured annotations are stamped so slurm-bridge can build the Slurm job. + Expect(pod.Annotations).To(HaveKeyWithValue("slurmjob.slinky.slurm.net/qos", "inference")) + Expect(pod.Annotations).To(HaveKeyWithValue("slurmjob.slinky.slurm.net/exclusive", "false")) + + // Name is keyed on the gpu-group only (no node at creation time), distinct from the pinned + // node-keyed name. + Expect(pod.Name).To(Equal(reservationPodName("", extGroup))) + Expect(pod.Name).ToNot(Equal(reservationPodName(extNode, extGroup))) + }) + + It("keeps the default pinned behavior when no external scheduler is configured", func() { + rsc := initializeTestService(buildClient()) + + pod, err := rsc.createGPUReservationPod(context.TODO(), nil, extNode, extGroup) + Expect(err).To(Succeed()) + + Expect(pod.Spec.SchedulerName).To(BeEmpty()) + Expect(pod.Spec.NodeName).To(Equal(extNode)) + Expect(pod.Status.Conditions).ToNot(BeEmpty()) + Expect(pod.Spec.Containers[0].Resources.Requests).To(HaveKey(v1.ResourceName(constants.NvidiaGpuResource))) + Expect(pod.Name).To(Equal(reservationPodName(extNode, extGroup))) + }) +}) + +func TestTimeoutAllocationIndicator(t *testing.T) { + const ns = "kai-resource-reservation" + const podName = "gpu-reservation-abc123" + + newPod := func(phase v1.PodPhase) *v1.Pod { + return &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: podName, Namespace: ns}, + Status: v1.PodStatus{Phase: phase}, + } + } + + tests := []struct { + name string + objects []runtimeClient.Object + want string + }{ + { + name: "still pending -> parked", + objects: []runtimeClient.Object{newPod(v1.PodPending)}, + want: pendingGpuIndicator, + }, + { + name: "failed -> reap", + objects: []runtimeClient.Object{newPod(v1.PodFailed)}, + want: unknownGpuIndicator, + }, + { + name: "running-but-no-index -> reap", + objects: []runtimeClient.Object{newPod(v1.PodRunning)}, + want: unknownGpuIndicator, + }, + { + name: "not found -> reap", + objects: nil, + want: unknownGpuIndicator, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rsc := &service{ + namespace: ns, + kubeClient: fake.NewClientBuilder().WithScheme(testScheme).WithObjects(tt.objects...).Build(), + } + if got := rsc.timeoutAllocationIndicator(context.Background(), podName); got != tt.want { + t.Errorf("timeoutAllocationIndicator = %q, want %q", got, tt.want) + } + }) + } +} + +func TestDeleteReservationPodStripsCascadeFinalizer(t *testing.T) { + const ns = "kai-resource-reservation" + resv := &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "gpu-reservation-strip", + Namespace: ns, + Finalizers: []string{ReservationCascadeFinalizer, "scheduler.slinky.slurm.net/finalizer"}, + }, + } + c := fake.NewClientBuilder().WithScheme(testScheme).WithObjects(resv).Build() + rsc := &service{namespace: ns, kubeClient: c} + + if err := rsc.deleteReservationPod(context.Background(), resv.DeepCopy()); err != nil { + t.Fatalf("deleteReservationPod: %v", err) + } + + // slurm-bridge's finalizer still holds the object, so it should still exist — but OUR cascade + // finalizer must be gone, so the cascade controller will NOT treat this internal delete as a + // signal to delete the backed fractions. + got := &v1.Pod{} + if err := c.Get(context.Background(), runtimeClient.ObjectKey{Namespace: ns, Name: "gpu-reservation-strip"}, got); err != nil { + t.Fatalf("reservation should still exist while slurm-bridge finalizer holds it: %v", err) + } + for _, f := range got.Finalizers { + if f == ReservationCascadeFinalizer { + t.Errorf("cascade finalizer should have been stripped, got finalizers=%v", got.Finalizers) + } + } + if got.DeletionTimestamp == nil { + t.Errorf("reservation should be terminating (deletionTimestamp set) after internal delete") + } +} diff --git a/pkg/binder/controllers/bindrequest_controller.go b/pkg/binder/controllers/bindrequest_controller.go index d94026852..9202e69d9 100644 --- a/pkg/binder/controllers/bindrequest_controller.go +++ b/pkg/binder/controllers/bindrequest_controller.go @@ -36,6 +36,11 @@ import ( const ( podBoundCondition = "PodBound" + // reservationPendingRequeue is how often a BindRequest is retried while its GPU + // reservation pod is parked (still Pending, e.g. queued in slurm-bridge waiting for a + // whole GPU). Short enough that the fraction binds promptly once the reservation is + // placed, without a terminal-failure backoff. + reservationPendingRequeue = 10 * time.Second ) // BindRequestReconciler reconciles a BindRequest object @@ -93,6 +98,9 @@ func (r *BindRequestReconciler) Reconcile(ctx context.Context, req ctrl.Request) var bindRequest = &schedulingv1alpha2.BindRequest{} var pod *v1.Pod + // Set when Bind reports the GPU reservation is parked (still Pending). The fraction is + // neither bound nor failed: keep the BindRequest non-terminal and requeue. + var reservationPending bool // Fetch the BindRequest instance if err = r.Client.Get(ctx, req.NamespacedName, bindRequest); err != nil { @@ -117,6 +125,16 @@ func (r *BindRequestReconciler) Reconcile(ctx context.Context, req ctrl.Request) err = fmt.Errorf("Internal Error: %v", r) } + if reservationPending && finalError == nil { + // Parked reservation: keep the BindRequest non-terminal (no phase change, no + // failed-attempt increment) so it is protected from GC and the fraction stays + // Pending, and requeue to poll for the reservation being placed. Skip the + // pod-condition update so we don't falsely report the pod as Bound. + result.RequeueAfter = reservationPendingRequeue + err = nil + return + } + result, err = r.UpdateStatus(ctx, bindRequest, result, err) if pod != nil { r.updatePodCondition(ctx, bindRequest, pod, result, err) @@ -161,6 +179,16 @@ func (r *BindRequestReconciler) Reconcile(ctx context.Context, req ctrl.Request) "namespace", pod.Namespace) err = r.Client.Delete(ctx, bindRequest) } + if errors.Is(err, resourcereservation.ErrReservationPending) { + // The GPU reservation is parked (still Pending, e.g. queued in slurm-bridge waiting + // for a whole GPU). Do NOT roll back the fraction pod's GPU-group binding: keep the + // parked reservation in place so it holds its queue spot and is reused next reconcile. + // The deferred handler requeues without marking the BindRequest Failed. + reservationPending = true + logger.Info("GPU reservation still pending; keeping it parked and requeueing", + "pod", pod.Name, "namespace", pod.Namespace, "node", node.Name) + return result, err + } if err != nil { logger.Error(err, "Failed to bind pod to node", "pod", pod.Name, "namespace", pod.Namespace, "node", node.Name) diff --git a/pkg/binder/controllers/reservationcascade_controller.go b/pkg/binder/controllers/reservationcascade_controller.go new file mode 100644 index 000000000..1f86f3142 --- /dev/null +++ b/pkg/binder/controllers/reservationcascade_controller.go @@ -0,0 +1,160 @@ +// Copyright 2025 NVIDIA CORPORATION +// SPDX-License-Identifier: Apache-2.0 + +package controllers + +import ( + "context" + "time" + + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/client-go/util/workqueue" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + "sigs.k8s.io/controller-runtime/pkg/event" + "sigs.k8s.io/controller-runtime/pkg/handler" + "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/reconcile" + + "github.com/kai-scheduler/KAI-scheduler/pkg/binder/binding/resourcereservation" + "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants" +) + +// reservationCascadePollInterval is how often we re-check that the fractions backed by a +// deleting reservation have terminated before releasing the reservation's finalizer. +const reservationCascadePollInterval = 5 * time.Second + +// ReservationCascadeReconciler makes deleting a reservation ("capacity") pod cascade to the +// GPU-fraction workloads it backs. slurm-bridge frees the whole GPU (cancels the GRES job) the +// moment the reservation pod terminates, so a reservation must never outlive-in-reverse its +// fractions: if it is deleted, the fractions on that GPU have to go too, or they run on a GPU +// Slurm now considers free (silent double-booking). +// +// The reservation pod carries ReservationCascadeFinalizer. On deletion, this controller deletes +// the fraction pods sharing its GPU group and holds the finalizer until they are actually gone — +// so the reservation stays Terminating until the workloads have drained, then is released. +type ReservationCascadeReconciler struct { + Client client.Client + Scheme *runtime.Scheme + ReservationNamespace string + // Enabled only in external-reservation mode (slurm-bridge), where the binder stamps the + // finalizer; otherwise this controller would strand pods with a finalizer nothing removes. + Enabled bool +} + +func (r *ReservationCascadeReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { + if !r.Enabled { + return ctrl.Result{}, nil + } + logger := log.FromContext(ctx) + + pod := &corev1.Pod{} + if err := r.Client.Get(ctx, req.NamespacedName, pod); err != nil { + return ctrl.Result{}, client.IgnoreNotFound(err) + } + if pod.Namespace != r.ReservationNamespace { + return ctrl.Result{}, nil + } + group := pod.Labels[constants.GPUGroup] + + if pod.DeletionTimestamp.IsZero() { + // Not being deleted — nothing to do. We deliberately do NOT (re-)add the finalizer here: + // the binder stamps it at creation, and re-adding it would race-undo the strip that the + // internal reservation deleter does to opt a system-initiated deletion out of cascading. + return ctrl.Result{}, nil + } + + // Reservation is being deleted. + if !controllerutil.ContainsFinalizer(pod, resourcereservation.ReservationCascadeFinalizer) { + return ctrl.Result{}, nil + } + if group == "" { + // Nothing to cascade — release immediately. + return ctrl.Result{}, r.releaseFinalizer(ctx, pod) + } + + remaining, err := r.deleteFractionsForGroup(ctx, group) + if err != nil { + return ctrl.Result{}, err + } + if remaining > 0 { + logger.Info("Reservation deleted; deleting backed fractions and waiting for them to terminate", + "gpuGroup", group, "remaining", remaining) + return ctrl.Result{RequeueAfter: reservationCascadePollInterval}, nil + } + + logger.Info("All fractions backed by the reservation have terminated; releasing reservation", + "gpuGroup", group) + return ctrl.Result{}, r.releaseFinalizer(ctx, pod) +} + +// deleteFractionsForGroup deletes all GPU-fraction pods (outside the reservation namespace) that +// share the given GPU group and returns how many still exist (including those still terminating). +func (r *ReservationCascadeReconciler) deleteFractionsForGroup(ctx context.Context, group string) (int, error) { + pods := &corev1.PodList{} + if err := r.Client.List(ctx, pods, client.MatchingLabels{constants.GPUGroup: group}); err != nil { + return 0, err + } + remaining := 0 + for i := range pods.Items { + p := &pods.Items[i] + if p.Namespace == r.ReservationNamespace { + continue // skip reservation pods themselves + } + remaining++ + if p.DeletionTimestamp.IsZero() { + if err := r.Client.Delete(ctx, p); err != nil && !apierrors.IsNotFound(err) { + return remaining, err + } + } + } + return remaining, nil +} + +func (r *ReservationCascadeReconciler) releaseFinalizer(ctx context.Context, pod *corev1.Pod) error { + if controllerutil.RemoveFinalizer(pod, resourcereservation.ReservationCascadeFinalizer) { + return r.Client.Update(ctx, pod) + } + return nil +} + +// SetupWithManager sets up the controller with the Manager. +func (r *ReservationCascadeReconciler) SetupWithManager(mgr ctrl.Manager, params *ReconcilerParams) error { + return ctrl.NewControllerManagedBy(mgr). + Named("reservationcascade"). + For(&corev1.Pod{}). + Watches(&corev1.Pod{}, r.eventHandlers()). + WithOptions(controller.Options{ + MaxConcurrentReconciles: params.MaxConcurrentReconciles, + RateLimiter: workqueue.NewTypedItemExponentialFailureRateLimiter[ctrl.Request]( + time.Duration(params.RateLimiterBaseDelaySeconds)*time.Second, + time.Duration(params.RateLimiterMaxDelaySeconds)*time.Second, + ), + SkipNameValidation: &[]bool{true}[0], + }). + Complete(r) +} + +func (r *ReservationCascadeReconciler) eventHandlers() handler.Funcs { + enqueue := func(_ context.Context, obj client.Object, q workqueue.TypedRateLimitingInterface[reconcile.Request]) { + if obj.GetNamespace() != r.ReservationNamespace { + return + } + q.Add(reconcile.Request{NamespacedName: client.ObjectKeyFromObject(obj)}) + } + return handler.Funcs{ + CreateFunc: func(ctx context.Context, e event.CreateEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) { + enqueue(ctx, e.Object, q) + }, + UpdateFunc: func(ctx context.Context, e event.UpdateEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) { + enqueue(ctx, e.ObjectNew, q) + }, + DeleteFunc: func(ctx context.Context, e event.DeleteEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) { + enqueue(ctx, e.Object, q) + }, + } +} diff --git a/pkg/binder/controllers/reservationcascade_controller_test.go b/pkg/binder/controllers/reservationcascade_controller_test.go new file mode 100644 index 000000000..f9664fe98 --- /dev/null +++ b/pkg/binder/controllers/reservationcascade_controller_test.go @@ -0,0 +1,128 @@ +// Copyright 2025 NVIDIA CORPORATION +// SPDX-License-Identifier: Apache-2.0 + +package controllers + +import ( + "context" + "testing" + + v1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + "github.com/kai-scheduler/KAI-scheduler/pkg/binder/binding/resourcereservation" + "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants" +) + +const rcResNs = "kai-resource-reservation" + +func rcScheme(t *testing.T) *runtime.Scheme { + t.Helper() + s := runtime.NewScheme() + if err := v1.AddToScheme(s); err != nil { + t.Fatalf("add corev1: %v", err) + } + return s +} + +func rcReservation(group string) *v1.Pod { + return &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "gpu-reservation-x", + Namespace: rcResNs, + Labels: map[string]string{constants.GPUGroup: group}, + Finalizers: []string{resourcereservation.ReservationCascadeFinalizer}, + }, + } +} + +func rcFraction(name, group string) *v1.Pod { + return &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: name, + Namespace: "team", + Labels: map[string]string{constants.GPUGroup: group}, + }, + } +} + +func rcReq() ctrl.Request { + return ctrl.Request{NamespacedName: types.NamespacedName{Namespace: rcResNs, Name: "gpu-reservation-x"}} +} + +func TestReservationCascadeDeletesFractionsAndWaits(t *testing.T) { + group := "g-1" + resv := rcReservation(group) + f1 := rcFraction("frac-1", group) + f2 := rcFraction("frac-2", group) + c := fake.NewClientBuilder().WithScheme(rcScheme(t)).WithObjects(resv, f1, f2).Build() + r := &ReservationCascadeReconciler{Client: c, ReservationNamespace: rcResNs, Enabled: true} + ctx := context.Background() + + // Delete the reservation -> finalizer keeps it Terminating. + if err := c.Delete(ctx, resv); err != nil { + t.Fatalf("delete reservation: %v", err) + } + + // First reconcile: fractions deleted, finalizer retained (requeue while they drain). + res, err := r.Reconcile(ctx, rcReq()) + if err != nil { + t.Fatalf("reconcile 1: %v", err) + } + if res.RequeueAfter != reservationCascadePollInterval { + t.Errorf("expected requeue %v, got %v", reservationCascadePollInterval, res.RequeueAfter) + } + for _, n := range []string{"frac-1", "frac-2"} { + if err := c.Get(ctx, types.NamespacedName{Namespace: "team", Name: n}, &v1.Pod{}); !apierrors.IsNotFound(err) { + t.Errorf("fraction %s should have been deleted, err=%v", n, err) + } + } + if err := c.Get(ctx, rcReq().NamespacedName, &v1.Pod{}); err != nil { + t.Errorf("reservation should still exist (finalizer held) while fractions drain, err=%v", err) + } + + // Second reconcile: no fractions remain -> finalizer removed -> reservation GC'd. + if _, err := r.Reconcile(ctx, rcReq()); err != nil { + t.Fatalf("reconcile 2: %v", err) + } + if err := c.Get(ctx, rcReq().NamespacedName, &v1.Pod{}); !apierrors.IsNotFound(err) { + t.Errorf("reservation should be released once fractions are gone, err=%v", err) + } +} + +func TestReservationCascadeNoFractionsReleasesImmediately(t *testing.T) { + resv := rcReservation("g-empty") + c := fake.NewClientBuilder().WithScheme(rcScheme(t)).WithObjects(resv).Build() + r := &ReservationCascadeReconciler{Client: c, ReservationNamespace: rcResNs, Enabled: true} + ctx := context.Background() + if err := c.Delete(ctx, resv); err != nil { + t.Fatalf("delete: %v", err) + } + if _, err := r.Reconcile(ctx, rcReq()); err != nil { + t.Fatalf("reconcile: %v", err) + } + if err := c.Get(ctx, rcReq().NamespacedName, &v1.Pod{}); !apierrors.IsNotFound(err) { + t.Errorf("reservation with no backed fractions should be released immediately, err=%v", err) + } +} + +func TestReservationCascadeDisabledIsNoop(t *testing.T) { + resv := rcReservation("g-1") + f1 := rcFraction("frac-1", "g-1") + c := fake.NewClientBuilder().WithScheme(rcScheme(t)).WithObjects(resv, f1).Build() + r := &ReservationCascadeReconciler{Client: c, ReservationNamespace: rcResNs, Enabled: false} + ctx := context.Background() + _ = c.Delete(ctx, resv) + if _, err := r.Reconcile(ctx, rcReq()); err != nil { + t.Fatalf("reconcile: %v", err) + } + // Fraction must be untouched when disabled. + if err := c.Get(ctx, types.NamespacedName{Namespace: "team", Name: "frac-1"}, &v1.Pod{}); err != nil { + t.Errorf("fraction should be untouched when disabled, err=%v", err) + } +} diff --git a/pkg/binder/controllers/reserveahead_controller.go b/pkg/binder/controllers/reserveahead_controller.go new file mode 100644 index 000000000..a9e4ef358 --- /dev/null +++ b/pkg/binder/controllers/reserveahead_controller.go @@ -0,0 +1,298 @@ +// Copyright 2025 NVIDIA CORPORATION +// SPDX-License-Identifier: Apache-2.0 + +package controllers + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "time" + + corev1 "k8s.io/api/core/v1" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/client-go/util/workqueue" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller" + "sigs.k8s.io/controller-runtime/pkg/event" + "sigs.k8s.io/controller-runtime/pkg/handler" + "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/reconcile" + + schedulingv1alpha2 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2" + "github.com/kai-scheduler/KAI-scheduler/pkg/binder/binding/resourcereservation" + bindercommon "github.com/kai-scheduler/KAI-scheduler/pkg/binder/common" + "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants" + "github.com/kai-scheduler/KAI-scheduler/pkg/common/resources" +) + +const ( + // defaultReserveAheadDelay is how long a fraction must stay Pending+unschedulable before we + // reserve a whole GPU ahead for it. It gives the scheduler a chance to pack it onto a freed + // GPU first (the common, cheaper outcome) before we enqueue a whole-GPU reservation to Slurm. + defaultReserveAheadDelay = 15 * time.Second + // reserveAheadPollInterval is how often we re-check a placed-but-not-yet-running reservation. + reserveAheadPollInterval = 10 * time.Second +) + +// ReserveAheadReconciler enqueues a whole-GPU reservation for a GPU-fraction pod that has been +// Pending and unschedulable for a while (it could not pack onto an existing reserved GPU). The +// reservation is placed by the external scheduler (slurm-bridge), which may preempt lower-QoS +// Slurm jobs to make room — KAI does not decide whether the cluster is "full". +// +// The fraction is left Pending so the KAI scheduler keeps trying to pack it onto a GPU that frees +// up; whichever happens first wins the race: +// - scheduler packs it onto a freed reserved GPU -> its GPU-group label flips, the reserve-ahead +// reservation is orphaned and torn down by the reservation sync; +// - the reserve-ahead reservation gets placed -> we create the BindRequest to bind the fraction +// onto it. +// +// The fraction is labeled with the reserve-ahead GPU group so the existing reservation sync keeps +// the parked reservation alive while the fraction references it and reaps it once it does not. +type ReserveAheadReconciler struct { + Client client.Client + Scheme *runtime.Scheme + ResourceReservation resourcereservation.Interface + // SchedulerName is the fraction pods' schedulerName (the KAI scheduler). + SchedulerName string + // Enabled is true only in external-reservation mode (reservationSchedulerName set); otherwise + // there is nothing that can place a node-less reservation, so the reconciler no-ops. + Enabled bool + // Delay is how long a fraction must be Pending+unschedulable before reserving ahead. + Delay time.Duration +} + +func (r *ReserveAheadReconciler) delay() time.Duration { + if r.Delay <= 0 { + return defaultReserveAheadDelay + } + return r.Delay +} + +// reserveAheadGroup derives a stable GPU-group id for a fraction from its identity, so repeated +// reconciles (and reconcile-on-delete, which only has the namespaced name) address the same +// reservation. Hashed to stay within label-value constraints. +func reserveAheadGroup(namespace, name string) string { + sum := sha256.Sum256([]byte(namespace + "/" + name)) + return "reserve-ahead-" + hex.EncodeToString(sum[:8]) +} + +func (r *ReserveAheadReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { + if !r.Enabled { + return ctrl.Result{}, nil + } + logger := log.FromContext(ctx) + gpuGroup := reserveAheadGroup(req.Namespace, req.Name) + + pod := &corev1.Pod{} + if err := r.Client.Get(ctx, req.NamespacedName, pod); err != nil { + if apierrors.IsNotFound(err) { + // Fraction gone: tear down any reserve-ahead reservation we created for it. + if syncErr := r.ResourceReservation.SyncForGpuGroup(ctx, gpuGroup); syncErr != nil { + logger.Error(syncErr, "failed to sync reserve-ahead reservation after fraction deletion", + "gpuGroup", gpuGroup) + } + return ctrl.Result{}, nil + } + return ctrl.Result{}, err + } + + // Only act on GPU-fraction pods managed by our scheduler. + if pod.Spec.SchedulerName != r.SchedulerName || !resources.RequestsGPUFraction(pod) { + return ctrl.Result{}, nil + } + + // Once a BindRequest exists, the fraction's placement is decided — the reconciler MUST stop + // touching it (re-ensuring or re-labeling here would race the binder and corrupt accounting). + br := &schedulingv1alpha2.BindRequest{} + brErr := r.Client.Get(ctx, req.NamespacedName, br) + if brErr == nil { + if bindRequestUsesGroup(br, gpuGroup) { + // The reserve-ahead reservation won the race and is being bound onto — keep it. + return ctrl.Result{}, nil + } + // The scheduler packed the fraction onto a different GPU that freed up (pack won the + // race). Point the pod's group label at the real bound group so accounting is correct, + // then tear down our now-orphaned reserve-ahead reservation. + if podHasGpuGroup(pod, gpuGroup) && len(br.Spec.SelectedGPUGroups) > 0 { + if err := r.labelFractionGpuGroup(ctx, pod, br.Spec.SelectedGPUGroups[0]); err != nil { + return ctrl.Result{}, err + } + } + if err := r.ResourceReservation.SyncForGpuGroup(ctx, gpuGroup); err != nil { + return ctrl.Result{}, fmt.Errorf("failed to tear down reserve-ahead reservation for <%s/%s>: %w", + pod.Namespace, pod.Name, err) + } + return ctrl.Result{}, nil + } else if !apierrors.IsNotFound(brErr) { + return ctrl.Result{}, brErr + } + + // No BindRequest. If the pod is already bound/terminal/terminating (without a BindRequest, e.g. + // completed), reap any reserve-ahead reservation it no longer references. + if pod.Spec.NodeName != "" || pod.Status.Phase != corev1.PodPending || pod.DeletionTimestamp != nil { + if !podHasGpuGroup(pod, gpuGroup) { + if syncErr := r.ResourceReservation.SyncForGpuGroup(ctx, gpuGroup); syncErr != nil { + logger.Error(syncErr, "failed to sync reserve-ahead reservation for settled fraction", + "gpuGroup", gpuGroup) + } + } + return ctrl.Result{}, nil + } + + // Give the scheduler time to pack onto a freed GPU before reserving ahead. + if wait := r.delay() - time.Since(pod.CreationTimestamp.Time); wait > 0 { + return ctrl.Result{RequeueAfter: wait}, nil + } + + // If the reservation is already placed and running, bind the fraction onto it. + ready, nodeName, err := r.ResourceReservation.ReservationPlacement(ctx, gpuGroup) + if err != nil { + return ctrl.Result{}, err + } + if ready { + return ctrl.Result{}, r.bindFractionToReservation(ctx, pod, nodeName, gpuGroup) + } + + // Ensure exactly one parked reservation exists for this fraction, and tie the fraction to its + // group so the reservation sync keeps it alive (and reaps it once the fraction moves on). + if _, err := r.ResourceReservation.EnsureReservation(ctx, pod, gpuGroup); err != nil { + return ctrl.Result{}, fmt.Errorf("failed to ensure reserve-ahead reservation for <%s/%s>: %w", + pod.Namespace, pod.Name, err) + } + if err := r.labelFractionGpuGroup(ctx, pod, gpuGroup); err != nil { + return ctrl.Result{}, err + } + + logger.Info("Reserved a whole GPU ahead for pending fraction; waiting for external scheduler to place it", + "pod", pod.Name, "namespace", pod.Namespace, "gpuGroup", gpuGroup) + return ctrl.Result{RequeueAfter: reserveAheadPollInterval}, nil +} + +func (r *ReserveAheadReconciler) bindFractionToReservation( + ctx context.Context, pod *corev1.Pod, nodeName, gpuGroup string, +) error { + logger := log.FromContext(ctx) + portion, err := resources.GetGPUFraction(pod) + if err != nil { + return fmt.Errorf("failed to read gpu-fraction for <%s/%s>: %w", pod.Namespace, pod.Name, err) + } + bindRequest := &schedulingv1alpha2.BindRequest{ + ObjectMeta: metav1.ObjectMeta{ + Name: pod.Name, + Namespace: pod.Namespace, + OwnerReferences: []metav1.OwnerReference{{ + APIVersion: "v1", + Kind: "Pod", + Name: pod.Name, + UID: pod.UID, + }}, + }, + Spec: schedulingv1alpha2.BindRequestSpec{ + PodName: pod.Name, + SelectedNode: nodeName, + SelectedGPUGroups: []string{gpuGroup}, + ReceivedResourceType: bindercommon.ReceivedTypeFraction, + ReceivedGPU: &schedulingv1alpha2.ReceivedGPU{ + Count: 1, + Portion: fmt.Sprintf("%.2f", portion), + }, + BackoffLimit: ptrInt32(defaultReserveAheadBackoffLimit), + }, + } + if err := r.Client.Create(ctx, bindRequest); err != nil { + if apierrors.IsAlreadyExists(err) { + // The scheduler (or a previous reconcile) already created a BindRequest for this pod — + // e.g. it packed the fraction onto a GPU that freed up first. Let that one proceed; our + // now-orphaned reservation is reaped by the sync once the fraction settles elsewhere. + logger.Info("BindRequest already exists for fraction; not overriding", + "pod", pod.Name, "namespace", pod.Namespace) + return nil + } + return fmt.Errorf("failed to create BindRequest for reserve-ahead fraction <%s/%s>: %w", + pod.Namespace, pod.Name, err) + } + logger.Info("Bound pending fraction onto its reserve-ahead reservation", + "pod", pod.Name, "namespace", pod.Namespace, "node", nodeName, "gpuGroup", gpuGroup) + return nil +} + +func (r *ReserveAheadReconciler) labelFractionGpuGroup(ctx context.Context, pod *corev1.Pod, gpuGroup string) error { + if podHasGpuGroup(pod, gpuGroup) { + return nil + } + updated := pod.DeepCopy() + if updated.Labels == nil { + updated.Labels = map[string]string{} + } + updated.Labels[constants.GPUGroup] = gpuGroup + if err := r.Client.Patch(ctx, updated, client.MergeFrom(pod)); err != nil { + return fmt.Errorf("failed to label fraction <%s/%s> with reserve-ahead group: %w", + pod.Namespace, pod.Name, err) + } + return nil +} + +func podHasGpuGroup(pod *corev1.Pod, gpuGroup string) bool { + for _, g := range resources.GetGpuGroups(pod) { + if g == gpuGroup { + return true + } + } + return false +} + +func bindRequestUsesGroup(br *schedulingv1alpha2.BindRequest, gpuGroup string) bool { + for _, g := range br.Spec.SelectedGPUGroups { + if g == gpuGroup { + return true + } + } + return false +} + +// SetupWithManager sets up the controller with the Manager. +func (r *ReserveAheadReconciler) SetupWithManager(mgr ctrl.Manager, params *ReconcilerParams) error { + return ctrl.NewControllerManagedBy(mgr). + Named("reserveahead"). + For(&corev1.Pod{}). + Watches(&corev1.Pod{}, r.eventHandlers()). + WithOptions(controller.Options{ + MaxConcurrentReconciles: params.MaxConcurrentReconciles, + RateLimiter: workqueue.NewTypedItemExponentialFailureRateLimiter[ctrl.Request]( + time.Duration(params.RateLimiterBaseDelaySeconds)*time.Second, + time.Duration(params.RateLimiterMaxDelaySeconds)*time.Second, + ), + SkipNameValidation: &[]bool{true}[0], + }). + Complete(r) +} + +func (r *ReserveAheadReconciler) eventHandlers() handler.Funcs { + enqueue := func(ctx context.Context, obj client.Object, q workqueue.TypedRateLimitingInterface[reconcile.Request]) { + pod, ok := obj.(*corev1.Pod) + if !ok || pod.Spec.SchedulerName != r.SchedulerName || !resources.RequestsGPUFraction(pod) { + return + } + q.Add(reconcile.Request{NamespacedName: client.ObjectKeyFromObject(pod)}) + } + return handler.Funcs{ + CreateFunc: func(ctx context.Context, e event.CreateEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) { + enqueue(ctx, e.Object, q) + }, + UpdateFunc: func(ctx context.Context, e event.UpdateEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) { + enqueue(ctx, e.ObjectNew, q) + }, + DeleteFunc: func(ctx context.Context, e event.DeleteEvent, q workqueue.TypedRateLimitingInterface[reconcile.Request]) { + enqueue(ctx, e.Object, q) + }, + } +} + +const defaultReserveAheadBackoffLimit int32 = 100 + +func ptrInt32(v int32) *int32 { return &v } diff --git a/pkg/binder/controllers/reserveahead_controller_test.go b/pkg/binder/controllers/reserveahead_controller_test.go new file mode 100644 index 000000000..0ea183261 --- /dev/null +++ b/pkg/binder/controllers/reserveahead_controller_test.go @@ -0,0 +1,207 @@ +// Copyright 2025 NVIDIA CORPORATION +// SPDX-License-Identifier: Apache-2.0 + +package controllers + +import ( + "context" + "testing" + "time" + + "go.uber.org/mock/gomock" + v1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + kubeaischedulerscheme "github.com/kai-scheduler/KAI-scheduler/pkg/apis/client/clientset/versioned/scheme" + schedulingv1alpha2 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/scheduling/v1alpha2" + mock_resourcereservation "github.com/kai-scheduler/KAI-scheduler/pkg/binder/binding/resourcereservation/mock" + bindercommon "github.com/kai-scheduler/KAI-scheduler/pkg/binder/common" + "github.com/kai-scheduler/KAI-scheduler/pkg/common/constants" +) + +const raSchedulerName = "kai-scheduler" + +func raTestScheme(t *testing.T) *runtime.Scheme { + t.Helper() + s := runtime.NewScheme() + if err := v1.AddToScheme(s); err != nil { + t.Fatalf("add corev1 to scheme: %v", err) + } + if err := kubeaischedulerscheme.AddToScheme(s); err != nil { + t.Fatalf("add scheduling scheme: %v", err) + } + return s +} + +func raFraction(name string, age time.Duration) *v1.Pod { + return &v1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: name, + Namespace: "team", + UID: types.UID(name + "-uid"), + CreationTimestamp: metav1.NewTime(time.Now().Add(-age)), + Annotations: map[string]string{constants.GpuFraction: "0.5"}, + }, + Spec: v1.PodSpec{SchedulerName: raSchedulerName}, + Status: v1.PodStatus{Phase: v1.PodPending}, + } +} + +func TestReserveAheadGroupIsStableAndDeterministic(t *testing.T) { + a := reserveAheadGroup("team", "frac-1") + b := reserveAheadGroup("team", "frac-1") + c := reserveAheadGroup("team", "frac-2") + if a != b { + t.Errorf("group not deterministic: %q vs %q", a, b) + } + if a == c { + t.Errorf("distinct fractions collided on group %q", a) + } + if len(a) > 63 { + t.Errorf("group %q exceeds label-value length", a) + } +} + +func TestReserveAheadDisabledIsNoop(t *testing.T) { + ctrlm := gomock.NewController(t) + rrs := mock_resourcereservation.NewMockInterface(ctrlm) + // No mock expectations set -> any call fails the test. + pod := raFraction("frac-1", time.Hour) + c := fake.NewClientBuilder().WithScheme(raTestScheme(t)).WithObjects(pod).Build() + r := &ReserveAheadReconciler{Client: c, ResourceReservation: rrs, SchedulerName: raSchedulerName, Enabled: false} + + res, err := r.Reconcile(context.Background(), ctrl.Request{NamespacedName: types.NamespacedName{Namespace: "team", Name: "frac-1"}}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if res.RequeueAfter != 0 { + t.Errorf("disabled reconciler should not requeue, got %v", res.RequeueAfter) + } +} + +func TestReserveAheadWaitsBeforeReserving(t *testing.T) { + ctrlm := gomock.NewController(t) + rrs := mock_resourcereservation.NewMockInterface(ctrlm) + // Fresh pod (age 0) -> must requeue without touching the reservation service. + pod := raFraction("frac-1", 0) + c := fake.NewClientBuilder().WithScheme(raTestScheme(t)).WithObjects(pod).Build() + r := &ReserveAheadReconciler{Client: c, ResourceReservation: rrs, SchedulerName: raSchedulerName, Enabled: true, Delay: 15 * time.Second} + + res, err := r.Reconcile(context.Background(), ctrl.Request{NamespacedName: types.NamespacedName{Namespace: "team", Name: "frac-1"}}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if res.RequeueAfter <= 0 { + t.Errorf("expected a requeue while waiting, got %v", res.RequeueAfter) + } +} + +func TestReserveAheadEnsuresReservationAndLabelsFraction(t *testing.T) { + ctrlm := gomock.NewController(t) + rrs := mock_resourcereservation.NewMockInterface(ctrlm) + pod := raFraction("frac-1", time.Hour) + group := reserveAheadGroup("team", "frac-1") + c := fake.NewClientBuilder().WithScheme(raTestScheme(t)).WithObjects(pod).Build() + r := &ReserveAheadReconciler{Client: c, ResourceReservation: rrs, SchedulerName: raSchedulerName, Enabled: true, Delay: 15 * time.Second} + + rrs.EXPECT().ReservationPlacement(gomock.Any(), group).Return(false, "", nil) + rrs.EXPECT().EnsureReservation(gomock.Any(), gomock.Any(), group).Return(pod, nil) + + res, err := r.Reconcile(context.Background(), ctrl.Request{NamespacedName: types.NamespacedName{Namespace: "team", Name: "frac-1"}}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if res.RequeueAfter != reserveAheadPollInterval { + t.Errorf("expected poll requeue %v, got %v", reserveAheadPollInterval, res.RequeueAfter) + } + + got := &v1.Pod{} + if err := c.Get(context.Background(), types.NamespacedName{Namespace: "team", Name: "frac-1"}, got); err != nil { + t.Fatalf("get pod: %v", err) + } + if got.Labels[constants.GPUGroup] != group { + t.Errorf("fraction not labeled with reserve-ahead group: got %q want %q", got.Labels[constants.GPUGroup], group) + } +} + +func TestReserveAheadPackWonTearsDownReservation(t *testing.T) { + ctrlm := gomock.NewController(t) + rrs := mock_resourcereservation.NewMockInterface(ctrlm) + group := reserveAheadGroup("team", "frac-1") + // Fraction was reserve-ahead labeled, but the scheduler packed it onto a DIFFERENT freed group. + pod := raFraction("frac-1", time.Hour) + pod.Labels = map[string]string{constants.GPUGroup: group} + br := &schedulingv1alpha2.BindRequest{ + ObjectMeta: metav1.ObjectMeta{Name: "frac-1", Namespace: "team"}, + Spec: schedulingv1alpha2.BindRequestSpec{PodName: "frac-1", SelectedNode: "node-y", SelectedGPUGroups: []string{"real-group-abc"}}, + } + c := fake.NewClientBuilder().WithScheme(raTestScheme(t)).WithObjects(pod, br).Build() + r := &ReserveAheadReconciler{Client: c, ResourceReservation: rrs, SchedulerName: raSchedulerName, Enabled: true, Delay: 15 * time.Second} + + // Teardown of the orphaned reserve-ahead reservation; NO EnsureReservation. + rrs.EXPECT().SyncForGpuGroup(gomock.Any(), group).Return(nil) + + if _, err := r.Reconcile(context.Background(), ctrl.Request{NamespacedName: types.NamespacedName{Namespace: "team", Name: "frac-1"}}); err != nil { + t.Fatalf("unexpected error: %v", err) + } + got := &v1.Pod{} + if err := c.Get(context.Background(), types.NamespacedName{Namespace: "team", Name: "frac-1"}, got); err != nil { + t.Fatalf("get pod: %v", err) + } + if got.Labels[constants.GPUGroup] != "real-group-abc" { + t.Errorf("pod group label = %q, want the real bound group real-group-abc", got.Labels[constants.GPUGroup]) + } +} + +func TestReserveAheadReservationWonIsKept(t *testing.T) { + ctrlm := gomock.NewController(t) + rrs := mock_resourcereservation.NewMockInterface(ctrlm) + group := reserveAheadGroup("team", "frac-1") + pod := raFraction("frac-1", time.Hour) + // BindRequest binds onto OUR reserve-ahead group (the reservation won the race). + br := &schedulingv1alpha2.BindRequest{ + ObjectMeta: metav1.ObjectMeta{Name: "frac-1", Namespace: "team"}, + Spec: schedulingv1alpha2.BindRequestSpec{PodName: "frac-1", SelectedNode: "node-x", SelectedGPUGroups: []string{group}}, + } + c := fake.NewClientBuilder().WithScheme(raTestScheme(t)).WithObjects(pod, br).Build() + r := &ReserveAheadReconciler{Client: c, ResourceReservation: rrs, SchedulerName: raSchedulerName, Enabled: true, Delay: 15 * time.Second} + + // No reservation-service calls at all — the reservation is in use, nothing to do. + if _, err := r.Reconcile(context.Background(), ctrl.Request{NamespacedName: types.NamespacedName{Namespace: "team", Name: "frac-1"}}); err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestReserveAheadBindsWhenReservationPlaced(t *testing.T) { + ctrlm := gomock.NewController(t) + rrs := mock_resourcereservation.NewMockInterface(ctrlm) + pod := raFraction("frac-1", time.Hour) + group := reserveAheadGroup("team", "frac-1") + c := fake.NewClientBuilder().WithScheme(raTestScheme(t)).WithObjects(pod).Build() + r := &ReserveAheadReconciler{Client: c, ResourceReservation: rrs, SchedulerName: raSchedulerName, Enabled: true, Delay: 15 * time.Second} + + rrs.EXPECT().ReservationPlacement(gomock.Any(), group).Return(true, "node-x", nil) + // EnsureReservation must NOT be called on the ready path. + + if _, err := r.Reconcile(context.Background(), ctrl.Request{NamespacedName: types.NamespacedName{Namespace: "team", Name: "frac-1"}}); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + br := &schedulingv1alpha2.BindRequest{} + if err := c.Get(context.Background(), types.NamespacedName{Namespace: "team", Name: "frac-1"}, br); err != nil { + t.Fatalf("expected a BindRequest to be created: %v", err) + } + if br.Spec.SelectedNode != "node-x" { + t.Errorf("BindRequest node = %q, want node-x", br.Spec.SelectedNode) + } + if len(br.Spec.SelectedGPUGroups) != 1 || br.Spec.SelectedGPUGroups[0] != group { + t.Errorf("BindRequest groups = %v, want [%s]", br.Spec.SelectedGPUGroups, group) + } + if br.Spec.ReceivedResourceType != bindercommon.ReceivedTypeFraction { + t.Errorf("BindRequest resource type = %q, want %q", br.Spec.ReceivedResourceType, bindercommon.ReceivedTypeFraction) + } +} diff --git a/pkg/operator/operands/binder/resources.go b/pkg/operator/operands/binder/resources.go index 028d506df..b675cf119 100644 --- a/pkg/operator/operands/binder/resources.go +++ b/pkg/operator/operands/binder/resources.go @@ -7,6 +7,7 @@ import ( "context" "encoding/json" "fmt" + "sort" "strconv" "golang.org/x/mod/semver" @@ -234,6 +235,27 @@ func buildArgsList(kaiConfig *kaiv1.Config, config *kaiv1binder.Binder, fakeGPU *config.ResourceReservation.AllocationTimeout)) } + // External-reservation mode: thread the global config into the binder so it creates + // reservation pods for an external scheduler (e.g. slurm-bridge-scheduler) instead of + // pinning them to the KAI-selected node. Empty/nil leaves the default pinned behavior + // (no args added — so existing deployments are unaffected). + if name := kaiConfig.Spec.Global.ReservationSchedulerName; name != nil && *name != "" { + args = append(args, "--reservation-scheduler-name", *name) + } + if res := kaiConfig.Spec.Global.ReservationGpuResourceName; res != nil && *res != "" { + args = append(args, "--reservation-gpu-resource-name", *res) + } + if ann := kaiConfig.Spec.Global.ReservationPodAnnotations; len(ann) > 0 { + keys := make([]string, 0, len(ann)) + for k := range ann { + keys = append(keys, k) + } + sort.Strings(keys) // stable arg order so the binder Deployment doesn't churn on reconcile + for _, k := range keys { + args = append(args, "--reservation-pod-annotations", fmt.Sprintf("%s=%s", k, ann[k])) + } + } + pluginsConfig := binderplugins.FromAPIConfig(config.Plugins) pluginsJSON, err := json.Marshal(pluginsConfig) if err != nil { diff --git a/pkg/scheduler/api/node_info/node_info.go b/pkg/scheduler/api/node_info/node_info.go index 582e6414f..1d1f071c9 100644 --- a/pkg/scheduler/api/node_info/node_info.go +++ b/pkg/scheduler/api/node_info/node_info.go @@ -90,6 +90,16 @@ type NodeInfo struct { // HasDRAGPUs indicates GPUs were added via DRA ResourceSlices. Temporary fix - remove when device-plugin pods are supported on DRA nodes. HasDRAGPUs bool + // DevicePluginGPUs is the node's GPU count as advertised by the device plugin + // (nvidia.com/gpu in Status.Allocatable), captured at construction BEFORE any + // DRA ResourceSlice GPUs are merged into AllocatableVector. It is the signal for + // whether KAI's fractional-sharing accounting can work on the node: a value > 0 + // means the device plugin also supplies the nvidia.com/gpu.memory GFD label the + // accounting anchors on. On a DRA-only node this is 0 even though AllocatableVector + // GPU > 0 (DRA GPUs are added there), so it — not AllocatableVector — must gate the + // "fractional not supported on DRA-only nodes" check. + DevicePluginGPUs float64 + NodeResourceTopology *nrtv1alpha2.NodeResourceTopology NumaTopology *NumaTopology @@ -103,6 +113,9 @@ func NewNodeInfo(node *v1.Node, podAffinityInfo pod_affinity.NodePodAffinityInfo gpuMemory, exists := getNodeGpuMemory(node) allocatableVector := resource_info.ResourceFromResourceList(node.Status.Allocatable).ToVector(vectorMap) + // Capture the device-plugin GPU count before any DRA ResourceSlice GPUs are merged + // in via AddDRAGPUs (see DevicePluginGPUs field). + devicePluginGPUs := allocatableVector.Get(resource_info.GPUIndex) idleVector := allocatableVector.Clone() usedVector := resource_info.NewResourceVector(vectorMap) releasingVector := resource_info.NewResourceVector(vectorMap) @@ -122,6 +135,7 @@ func NewNodeInfo(node *v1.Node, podAffinityInfo pod_affinity.NodePodAffinityInfo PodInfos: make(map[common_info.PodID]*pod_info.PodInfo), MemoryOfEveryGpuOnNode: gpuMemory, GpuMemorySynced: exists, + DevicePluginGPUs: devicePluginGPUs, LegacyMIGTasks: map[common_info.PodID]string{}, GpuSharingNodeInfo: *newGpuSharingNodeInfo(), @@ -314,7 +328,15 @@ func (ni *NodeInfo) PredicateByNodeResourcesType(task *pod_info.PodInfo) error { return nil } - if ni.HasDRAGPUs && task.IsSharedGPURequest() { + // Reject fractional/shared GPU pods only on DRA-ONLY nodes (matching this message). A node + // that ALSO advertises device-plugin GPUs (DevicePluginGPUs > 0, i.e. it carries the + // nvidia.com/gpu.memory GFD label) has the signals KAI's sharing accounting needs, so + // fractional sharing works there even though the whole-GPU reservation is held via DRA + // (meshy fork: the reservation pod goes through slurm-bridge as a DRA-backed extended + // resource). Note we must test DevicePluginGPUs, not AllocatableVector[GPU]: AddDRAGPUs + // merges DRA ResourceSlice GPUs into AllocatableVector, so on a DRA-only node that vector + // is > 0 and would wrongly bypass this guard. + if ni.HasDRAGPUs && ni.DevicePluginGPUs <= 0 && task.IsSharedGPURequest() { return common_info.NewFitError(task.Name, task.Namespace, ni.Name, "fractional/shared GPU pods are not yet supported on DRA-only nodes") } diff --git a/pkg/scheduler/api/node_info/node_info_test.go b/pkg/scheduler/api/node_info/node_info_test.go index 55ea9b972..9780fd0d0 100644 --- a/pkg/scheduler/api/node_info/node_info_test.go +++ b/pkg/scheduler/api/node_info/node_info_test.go @@ -1637,6 +1637,16 @@ func TestPredicateByNodeResourcesType_SharedGPU_DRANode(t *testing.T) { podResources: common_info.BuildResourceListWithGPU("1000m", "1G", "500m"), wantErr: false, }, + { + // meshy fork: a node exposing GPUs through BOTH the device plugin and DRA + // has the device-plugin accounting signals, so fractional pods are accepted + // even though HasDRAGPUs is true (the reservation is held via DRA). + name: "fraction pod accepted on both device-plugin + DRA node", + node: dpNode, + draGPUs: 4, + podResources: common_info.BuildResourceListWithGPU("1000m", "1G", "500m"), + wantErr: false, + }, } for _, tt := range tests { diff --git a/pkg/scheduler/cache/cluster_info/cluster_info.go b/pkg/scheduler/cache/cluster_info/cluster_info.go index 9609c9885..ac9c873d8 100644 --- a/pkg/scheduler/cache/cluster_info/cluster_info.go +++ b/pkg/scheduler/cache/cluster_info/cluster_info.go @@ -347,10 +347,19 @@ func (c *ClusterInfo) populateDRAGPUs(nodes map[string]*node_info.NodeInfo) { if draGPUCount > 0 { log.InfraLogger.V(6).Infof("Node %s has %d DRA GPUs from ResourceSlices", nodeName, draGPUCount) - if nodeInfo.AllocatableVector.Get(resource_info.GPUIndex) > 0 { - log.InfraLogger.Warningf("Node %s has both device-plugin GPUs and DRA GPUs", nodeName) + // A node may expose GPUs through BOTH the device plugin and DRA (meshy fork: + // the whole-GPU reservation is held via DRA while the device plugin advertises + // the same physical GPUs so KAI's fractional-sharing accounting has its + // nvidia.com/gpu + nvidia.com/gpu.memory signals). The two describe the SAME + // hardware, so the node's true GPU count is max(device-plugin, DRA), not the + // sum. Adding both would double-count the physical GPUs (e.g. one T4 -> 2), + // leaving a phantom idle GPU that the scheduler keeps trying to fill with new + // reservation pods that can never be placed. Only add the DRA devices not + // already covered by the device-plugin allocatable. + devicePluginGPUs := int64(nodeInfo.AllocatableVector.Get(resource_info.GPUIndex)) + if extraDRAGPUs := draGPUCount - devicePluginGPUs; extraDRAGPUs > 0 { + nodeInfo.AddDRAGPUs(float64(extraDRAGPUs)) } - nodeInfo.AddDRAGPUs(float64(draGPUCount)) nodeInfo.HasDRAGPUs = true } } diff --git a/pkg/scheduler/cache/cluster_info/cluster_info_test.go b/pkg/scheduler/cache/cluster_info/cluster_info_test.go index c4a142a7f..fc1090f7b 100644 --- a/pkg/scheduler/cache/cluster_info/cluster_info_test.go +++ b/pkg/scheduler/cache/cluster_info/cluster_info_test.go @@ -2569,6 +2569,38 @@ func TestSnapshotNodesWithDRAGPUs(t *testing.T) { expectedDRAGPUs: map[string]float64{"node-1": 6}, hasDRAGPUs: map[string]bool{"node-1": true}, }, + // meshy fork: a node that exposes the SAME physical GPUs through both the + // device plugin and DRA must not double-count them. Total = max(device-plugin, DRA). + "Both device-plugin and DRA for the same physical GPUs": { + nodes: []*corev1.Node{ + { + ObjectMeta: metav1.ObjectMeta{Name: "node-1"}, + Status: corev1.NodeStatus{Allocatable: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("1"), + }}, + }, + }, + resourceSlices: []*resourceapi.ResourceSlice{ + createTestResourceSlice("slice-1", "node-1", "nvidia.com/gpu", 1), + }, + expectedDRAGPUs: map[string]float64{"node-1": 1}, // not 2 + hasDRAGPUs: map[string]bool{"node-1": true}, + }, + "More DRA GPUs than device-plugin GPUs (add only the excess)": { + nodes: []*corev1.Node{ + { + ObjectMeta: metav1.ObjectMeta{Name: "node-1"}, + Status: corev1.NodeStatus{Allocatable: corev1.ResourceList{ + "nvidia.com/gpu": resource.MustParse("2"), + }}, + }, + }, + resourceSlices: []*resourceapi.ResourceSlice{ + createTestResourceSlice("slice-1", "node-1", "nvidia.com/gpu", 4), + }, + expectedDRAGPUs: map[string]float64{"node-1": 4}, // max(2,4) + hasDRAGPUs: map[string]bool{"node-1": true}, + }, } for name, test := range tests {