@@ -19,6 +19,7 @@ import (
1919 "sigs.k8s.io/controller-runtime/pkg/client"
2020 "sigs.k8s.io/controller-runtime/pkg/client/fake"
2121
22+ "github.com/kai-scheduler/KAI-scheduler/pkg/podgrouper/podgroup"
2223 "github.com/kai-scheduler/KAI-scheduler/pkg/podgrouper/podgrouper/plugins/constants"
2324 "github.com/kai-scheduler/KAI-scheduler/pkg/podgrouper/podgrouper/plugins/defaultgrouper"
2425 grouperplugin "github.com/kai-scheduler/KAI-scheduler/pkg/podgrouper/podgrouper/plugins/grouper"
@@ -65,7 +66,7 @@ var _ = Describe("SkipTopOwnerGrouper", func() {
6566 supportedTypes = map [metav1.GroupVersionKind ]grouperplugin.Grouper {
6667 {Group : "" , Version : "v1" , Kind : "Pod" }: defaultGrouper ,
6768 }
68- plugin = NewSkipTopOwnerGrouper (client , defaultGrouper , supportedTypes )
69+ plugin = NewSkipTopOwnerGrouper (client , defaultGrouper , resolverFromMap ( supportedTypes ) )
6970 })
7071
7172 Context ("when last owner is a pod" , func () {
@@ -430,6 +431,69 @@ var _ = Describe("SkipTopOwnerGrouper", func() {
430431 })
431432 })
432433
434+ Context ("chained skip-top-owner delegation" , func () {
435+ var (
436+ groveGrouper * recordingGrouper
437+ workloadRunner * unstructured.Unstructured
438+ pod * v1.Pod
439+ otherOwners []* metav1.PartialObjectMetadata
440+ )
441+
442+ BeforeEach (func () {
443+ groveGrouper = & recordingGrouper {name : "Grove Grouper" }
444+ supportedTypes [metav1.GroupVersionKind {
445+ Group : "grove.io" , Version : "v1alpha1" , Kind : "PodCliqueSet" ,
446+ }] = groveGrouper
447+ supportedTypes [metav1.GroupVersionKind {
448+ Group : "nvidia.com" , Version : "*" , Kind : "DynamoGraphDeployment" ,
449+ }] = plugin
450+
451+ // Dynamo 1.2.0+ serves the DGD as v1beta1, matched via the wildcard version.
452+ dgd := newUnstructured ("nvidia.com/v1beta1" , "DynamoGraphDeployment" , "dgd" )
453+ podCliqueSet := newUnstructured ("grove.io/v1alpha1" , "PodCliqueSet" , "dgd-pcs" )
454+ podCliqueSet .SetLabels (map [string ]string {queueLabelKey : queueName })
455+ podClique := newUnstructured ("grove.io/v1alpha1" , "PodClique" , "dgd-pcs-worker" )
456+ workloadRunner = newUnstructured ("run.ai/v1alpha1" , "WorkloadRunner" , "runner" )
457+
458+ pod = examplePod .DeepCopy ()
459+ pod .OwnerReferences = []metav1.OwnerReference {
460+ {Kind : "PodClique" , APIVersion : "grove.io/v1alpha1" , Name : podClique .GetName ()},
461+ }
462+
463+ Expect (client .Create (context .TODO (), dgd )).To (Succeed ())
464+ Expect (client .Create (context .TODO (), podCliqueSet )).To (Succeed ())
465+ Expect (client .Create (context .TODO (), podClique )).To (Succeed ())
466+ Expect (client .Create (context .TODO (), pod )).To (Succeed ())
467+ pod .TypeMeta = metav1.TypeMeta {Kind : examplePod .Kind , APIVersion : examplePod .APIVersion }
468+
469+ otherOwners = []* metav1.PartialObjectMetadata {
470+ objectToPartial (podClique ),
471+ objectToPartial (podCliqueSet ),
472+ objectToPartial (dgd ),
473+ objectToPartial (workloadRunner ),
474+ }
475+ })
476+
477+ It ("skips both WorkloadRunner and DGD and lands on the Grove grouper" , func () {
478+ metadata , err := plugin .GetPodGroupMetadata (workloadRunner , pod , otherOwners ... )
479+
480+ Expect (err ).NotTo (HaveOccurred ())
481+ Expect (metadata ).NotTo (BeNil ())
482+ Expect (metadata .Name ).To (Equal ("Grove Grouper" ))
483+ Expect (groveGrouper .calledWith ).NotTo (BeNil ())
484+ Expect (groveGrouper .calledWith .GetKind ()).To (Equal ("PodCliqueSet" ))
485+ })
486+
487+ It ("propagates the skipped wrappers' labels down to the Grove owner" , func () {
488+ workloadRunner .SetLabels (map [string ]string {"runner-label" : "runner-value" })
489+
490+ _ , err := plugin .GetPodGroupMetadata (workloadRunner , pod , otherOwners ... )
491+
492+ Expect (err ).NotTo (HaveOccurred ())
493+ Expect (groveGrouper .calledWith .GetLabels ()).To (HaveKeyWithValue ("runner-label" , "runner-value" ))
494+ })
495+ })
496+
433497 Context ("default plugin usage" , func () {
434498 It ("uses default plugin when GVK does not have custom plugin" , func () {
435499 // Use StatefulSet which is not in supportedTypes
@@ -492,6 +556,42 @@ var _ = Describe("SkipTopOwnerGrouper", func() {
492556
493557})
494558
559+ type recordingGrouper struct {
560+ name string
561+ calledWith * unstructured.Unstructured
562+ }
563+
564+ func (r * recordingGrouper ) Name () string { return r .name }
565+
566+ func (r * recordingGrouper ) GetPodGroupMetadata (
567+ topOwner * unstructured.Unstructured , _ * v1.Pod , _ ... * metav1.PartialObjectMetadata ,
568+ ) (* podgroup.Metadata , error ) {
569+ r .calledWith = topOwner
570+ return & podgroup.Metadata {Name : r .name }, nil
571+ }
572+
573+ func newUnstructured (apiVersion , kind , name string ) * unstructured.Unstructured {
574+ obj := & unstructured.Unstructured {}
575+ obj .SetAPIVersion (apiVersion )
576+ obj .SetKind (kind )
577+ obj .SetName (name )
578+ obj .SetNamespace ("default" )
579+ return obj
580+ }
581+
582+ func resolverFromMap (table map [metav1.GroupVersionKind ]grouperplugin.Grouper ) GrouperResolver {
583+ return func (gvk metav1.GroupVersionKind ) grouperplugin.Grouper {
584+ if g , found := table [gvk ]; found {
585+ return g
586+ }
587+ gvk .Version = "*"
588+ if g , found := table [gvk ]; found {
589+ return g
590+ }
591+ return nil
592+ }
593+ }
594+
495595func objectToPartial (obj client.Object ) * metav1.PartialObjectMetadata {
496596 objectMeta := metav1.ObjectMeta {
497597 Name : obj .GetName (),
0 commit comments