@@ -34,16 +34,18 @@ type Interface interface {
3434 Sync (ctx context.Context ) error
3535 SyncForNode (ctx context.Context , nodeName string ) error
3636 SyncForGpuGroup (ctx context.Context , gpuGroup string ) error
37- ReserveGpuDevice (ctx context.Context , pod * v1.Pod , nodeName string , gpuGroup string ) (string , error )
37+ ReserveGpuDevice (
38+ ctx context.Context , pod * v1.Pod , nodeName string ,
39+ fractionalGpuGroup schedulingv1alpha2.FractionalGpuGroup ,
40+ ) (string , error )
3841 RemovePodGpuGroupsConnection (ctx context.Context , pod * v1.Pod ) error
3942}
4043
4144const (
42- resourceReservation = "resource-reservation"
43- gpuReservationPodPrefix = "gpu-reservation"
44- gpuIndexAnnotationName = "run.ai/reserve_for_gpu_index"
45- numberOfGPUsToReserve = 1
46- unknownGpuIndicator = "-1"
45+ resourceReservation = "resource-reservation"
46+ gpuIndexAnnotationName = "run.ai/reserve_for_gpu_index"
47+ numberOfGPUsToReserve = 1
48+ unknownGpuIndicator = "-1"
4749)
4850
4951type service struct {
@@ -236,7 +238,7 @@ func (rsc *service) hasActiveBindRequestsForGpuGroup(ctx context.Context, gpuGro
236238 }
237239
238240 for _ , br := range bindRequestList .Items {
239- if ! slices .Contains (br .Spec .SelectedGPUGroups , gpuGroup ) {
241+ if ! slices .Contains (br .Spec .SelectedFractionalGpuGroupIDs () , gpuGroup ) {
240242 continue
241243 }
242244
@@ -278,21 +280,25 @@ func (rsc *service) hasLivePodForBindRequest(ctx context.Context, bindRequest *s
278280 return false , nil
279281 }
280282
281- for _ , gpuGroup := range bindRequest .Spec .SelectedGPUGroups {
283+ for _ , gpuGroup := range bindRequest .Spec .SelectedFractionalGpuGroupIDs () {
282284 if slices .Contains (resources .GetGpuGroups (pod ), gpuGroup ) {
283285 return true , nil
284286 }
285287 }
286288
287289 return true , nil
288290}
289- func (rsc * service ) ReserveGpuDevice (ctx context.Context , pod * v1.Pod , nodeName string , gpuGroup string ) (string , error ) {
291+ func (rsc * service ) ReserveGpuDevice (
292+ ctx context.Context , pod * v1.Pod , nodeName string , fractionalGpuGroup schedulingv1alpha2.FractionalGpuGroup ,
293+ ) (string , error ) {
290294 logger := log .FromContext (ctx )
295+ fractionalGpuGroup = fractionalGpuGroup .WithDefaults ()
296+ gpuGroup := fractionalGpuGroup .ID
291297
292298 rsc .gpuGroupMutex .LockMutexForGroup (gpuGroup )
293299 defer rsc .gpuGroupMutex .ReleaseMutex (gpuGroup )
294300
295- gpuIndex , err := rsc .acquireGPUIndexByGroup (ctx , pod , nodeName , gpuGroup )
301+ gpuIndex , err := rsc .acquireGPUIndexByGroup (ctx , nodeName , fractionalGpuGroup )
296302 if err != nil {
297303 return unknownGpuIndicator , err
298304 }
@@ -373,16 +379,16 @@ func escapeJSONPointer(s string) string {
373379}
374380
375381func (rsc * service ) acquireGPUIndexByGroup (
376- ctx context.Context , sourcePod * v1. Pod , nodeName , gpuGroup string ,
382+ ctx context.Context , nodeName string , fractionalGpuGroup schedulingv1alpha2. FractionalGpuGroup ,
377383) (string , error ) {
378- gpuIndex , err := rsc .findGPUIndexByGroup (gpuGroup )
384+ gpuIndex , err := rsc .findGPUIndexByGroup (fractionalGpuGroup . ID )
379385 if err != nil {
380386 return "" , err
381387 }
382388 if gpuIndex != "" {
383389 return gpuIndex , err
384390 }
385- return rsc .createGPUReservationPodAndGetIndex (ctx , sourcePod , nodeName , gpuGroup )
391+ return rsc .createGPUReservationPodAndGetIndex (ctx , nodeName , fractionalGpuGroup )
386392}
387393
388394func (rsc * service ) findGPUIndexByGroup (gpuGroup string ) (
@@ -409,11 +415,11 @@ func (rsc *service) findGPUIndexByGroup(gpuGroup string) (
409415}
410416
411417func (rsc * service ) createGPUReservationPodAndGetIndex (
412- ctx context.Context , sourcePod * v1. Pod , nodeName , gpuGroup string ,
418+ ctx context.Context , nodeName string , fractionalGpuGroup schedulingv1alpha2. FractionalGpuGroup ,
413419) (
414420 gpuIndex string , err error ) {
415421 logger := log .FromContext (ctx )
416- pod , err := rsc .createGPUReservationPod (ctx , sourcePod , nodeName , gpuGroup )
422+ pod , err := rsc .createGPUReservationPod (ctx , nodeName , fractionalGpuGroup )
417423 if err != nil {
418424 return unknownGpuIndicator , err
419425 }
@@ -466,14 +472,14 @@ func (rsc *service) deleteReservationPod(ctx context.Context, pod *v1.Pod) error
466472}
467473
468474func (rsc * service ) createGPUReservationPod (
469- ctx context.Context , sourcePod * v1. Pod , nodeName , gpuGroup string ,
475+ ctx context.Context , nodeName string , fractionalGpuGroup schedulingv1alpha2. FractionalGpuGroup ,
470476) (* v1.Pod , error ) {
471477 logger := log .FromContext (ctx )
472478 if rsc .isScalingUp (ctx ) {
473479 return nil , fmt .Errorf ("cluster is scaling up, could not create reservation pod" )
474480 }
475481
476- podName := reservationPodName (nodeName , gpuGroup )
482+ podName := reservationPodName (nodeName , fractionalGpuGroup . ID )
477483
478484 // Build resource requirements starting with GPU resources
479485 resources := v1.ResourceRequirements {
@@ -496,15 +502,15 @@ func (rsc *service) createGPUReservationPod(
496502 }
497503 }
498504
499- pod , err := rsc .createResourceReservationPod (sourcePod , nodeName , gpuGroup , podName , resources )
505+ pod , err := rsc .createResourceReservationPod (nodeName , fractionalGpuGroup , podName , resources )
500506 if err != nil {
501507 // The reservation pod name is deterministic per (node, gpu-group). AlreadyExists
502508 // means another actor (a concurrent bind, a retry, or another binder replica)
503509 // already created the reservation pod for this gpu-group, so reuse it rather than
504510 // creating a duplicate on a different physical GPU.
505511 if apierrors .IsAlreadyExists (err ) {
506512 logger .Info ("GPU reservation pod already exists for gpu group, reusing" ,
507- "nodeName" , nodeName , "namespace" , rsc .namespace , "name" , podName , "gpuGroup" , gpuGroup )
513+ "nodeName" , nodeName , "namespace" , rsc .namespace , "name" , podName , "gpuGroup" , fractionalGpuGroup . ID )
508514 return pod , nil
509515 }
510516 logger .Error (err , "Failed to create GPU reservation pod on node" ,
@@ -571,29 +577,25 @@ func (rsc *service) waitForGPUReservationPodAllocation(
571577}
572578
573579func (rsc * service ) createResourceReservationPod (
574- sourcePod * v1.Pod , nodeName , gpuGroup , podName string , resources v1.ResourceRequirements ,
580+ nodeName string , fractionalGpuGroup schedulingv1alpha2.FractionalGpuGroup ,
581+ podName string , resources v1.ResourceRequirements ,
575582) (* v1.Pod , error ) {
576- var tolerations []v1.Toleration
577- if sourcePod != nil {
578- sourcePodSpec := sourcePod .Spec .DeepCopy ()
579- tolerations = sourcePodSpec .Tolerations
580- }
581-
583+ fractionalGpuGroup = fractionalGpuGroup .WithDefaults ()
582584 podSpec := & v1.Pod {
583585 ObjectMeta : metav1.ObjectMeta {
584586 Name : podName ,
585587 Namespace : rsc .namespace ,
586588 Labels : map [string ]string {
587589 constants .AppLabelName : rsc .appLabelValue ,
588- constants .GPUGroup : gpuGroup ,
590+ constants .GPUGroup : fractionalGpuGroup . ID ,
589591 },
590592 Annotations : map [string ]string {
591593 karpenterv1 .DoNotDisruptAnnotationKey : "true" ,
594+ constants .GpuComputeSharingMode : string (fractionalGpuGroup .ComputeSharingMode ),
592595 },
593596 },
594597 Spec : v1.PodSpec {
595- NodeName : nodeName ,
596- Tolerations : tolerations ,
598+ NodeName : nodeName ,
597599 RuntimeClassName : func () * string {
598600 if len (rsc .runtimeClassName ) == 0 {
599601 return nil
@@ -673,7 +675,7 @@ func (rsc *service) isScalingUp(ctx context.Context) bool {
673675}
674676
675677func IsGPUReservationPod (pod * v1.Pod ) bool {
676- return strings .HasPrefix (pod .Name , gpuReservationPodPrefix )
678+ return strings .HasPrefix (pod .Name , constants . GPUReservationPodPrefix )
677679}
678680
679681// reservationPodName derives a deterministic reservation pod name from the node and
@@ -682,5 +684,5 @@ func IsGPUReservationPod(pod *v1.Pod) bool {
682684// pods on different physical GPUs.
683685func reservationPodName (nodeName , gpuGroup string ) string {
684686 hash := sha256 .Sum256 ([]byte (nodeName + "/" + gpuGroup ))
685- return fmt .Sprintf ("%s-%s" , gpuReservationPodPrefix , hex .EncodeToString (hash [:8 ]))
687+ return fmt .Sprintf ("%s-%s" , constants . GPUReservationPodPrefix , hex .EncodeToString (hash [:8 ]))
686688}
0 commit comments