Skip to content

Commit 7f194bc

Browse files
authored
feat(binder): persist fractional GPU compute mode
Signed-off-by: davidLif <davidshani12@gmail.com>
1 parent b901e51 commit 7f194bc

6 files changed

Lines changed: 103 additions & 68 deletions

File tree

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
kind: Added
2+
body: |-
3+
Persist fractional GPU compute mode

pkg/binder/binding/binder.go

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -119,18 +119,22 @@ func (b *Binder) Rollback(ctx context.Context, pod *v1.Pod, node *v1.Node, bindR
119119
}
120120

121121
func (b *Binder) reserveGPUs(ctx context.Context, pod *v1.Pod, bindRequest *v1alpha2.BindRequest) ([]string, error) {
122-
if len(bindRequest.Spec.SelectedGPUGroups) == 0 {
122+
fractionalGpuGroups := bindRequest.Spec.SelectedFractionalGpuGroupsOrDefault()
123+
if len(fractionalGpuGroups) == 0 {
123124
// Old bindingRequest bad conversion. delete the binding request.
124-
return nil, fmt.Errorf("no SelectedGPUGroups for fractional pod: %w", InvalidCrdWarning)
125+
return nil, fmt.Errorf("no selected GPU groups for fractional pod: %w", InvalidCrdWarning)
125126
}
126127

127128
var gpuIndexes []string
128-
for _, gpuGroup := range bindRequest.Spec.SelectedGPUGroups {
129-
gpuIndex, err := b.resourceReservationService.ReserveGpuDevice(ctx, pod, bindRequest.Spec.SelectedNode, gpuGroup)
129+
for _, fractionalGpuGroup := range fractionalGpuGroups {
130+
gpuIndex, err := b.resourceReservationService.ReserveGpuDevice(
131+
ctx, pod, bindRequest.Spec.SelectedNode, fractionalGpuGroup)
130132
if err != nil {
131133
// Cleanup will be handled by the rollback function
132134

133-
return nil, fmt.Errorf("failed to reserve GPUs for pod <%s/%s> in gpu group <%s>: %w", pod.Namespace, pod.Name, gpuGroup, err)
135+
return nil, fmt.Errorf(
136+
"failed to reserve GPUs for pod <%s/%s> in gpu group <%s>: %w",
137+
pod.Namespace, pod.Name, fractionalGpuGroup.ID, err)
134138
}
135139
gpuIndexes = append(gpuIndexes, gpuIndex)
136140
}

pkg/binder/binding/resourcereservation/mock/resource_reservation_mock.go

Lines changed: 5 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pkg/binder/binding/resourcereservation/resource_reservation.go

Lines changed: 33 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -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

4144
const (
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

4951
type 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

375381
func (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

388394
func (rsc *service) findGPUIndexByGroup(gpuGroup string) (
@@ -409,11 +415,11 @@ func (rsc *service) findGPUIndexByGroup(gpuGroup string) (
409415
}
410416

411417
func (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

468474
func (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

573579
func (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

675677
func 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.
683685
func 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

Comments
 (0)