fix(podgrouper): skip WorkloadRunner wrapper when selecting the grouping plugin

.changes/unreleased/fixed-20260812-141030.yaml

@@ -0,0 +1,3 @@
+kind: Fixed
+body: |-
+  Podgrouper skips WorkloadRunner wrapper so wrapped workloads keep their gang grouping

deployments/kai-scheduler/templates/rbac/podgrouper.yaml

@@ -264,6 +264,14 @@ rules:
   - create
   - patch
   - update
+- apiGroups:
+  - nvidia.com
+  resources:
+  - dynamographdeployments
+  verbs:
+  - get
+  - list
+  - watch
 - apiGroups:
   - ray.io
   resources:
@@ -294,6 +302,7 @@ rules:
   - kartas
   - runaijobs
   - trainingworkloads
+  - workloadrunners
   verbs:
   - get
   - list

pkg/podgrouper/podgrouper/hub/hub.go

@@ -42,6 +42,7 @@ const (
 	kindDistributedWorkload          = "DistributedWorkload"
 	kindInferenceWorkload            = "InferenceWorkload"
 	kindDistributedInferenceWorkload = "DistributedInferenceWorkload"
+	kindWorkloadRunner               = "WorkloadRunner"
 )
 
 // +kubebuilder:rbac:groups=apps,resources=replicasets;statefulsets,verbs=get;list;watch
@@ -57,7 +58,9 @@ const (
 // +kubebuilder:rbac:groups=tekton.dev,resources=pipelineruns;taskruns,verbs=get;list;watch
 // +kubebuilder:rbac:groups=tekton.dev,resources=pipelineruns/finalizers;taskruns/finalizers,verbs=patch;update;create
 // +kubebuilder:rbac:groups=run.ai,resources=trainingworkloads;interactiveworkloads;distributedworkloads;inferenceworkloads;distributedinferenceworkloads,verbs=get;list;watch
+// +kubebuilder:rbac:groups=run.ai,resources=workloadrunners,verbs=get;list;watch
 // +kubebuilder:rbac:groups=run.ai,resources=kartas,verbs=get;list;watch
+// +kubebuilder:rbac:groups=nvidia.com,resources=dynamographdeployments,verbs=get;list;watch
 // +kubebuilder:rbac:groups=trainer.kubeflow.org,resources=trainjobs,verbs=get;list;watch
 // +kubebuilder:rbac:groups=trainer.kubeflow.org,resources=trainjobs/finalizers,verbs=patch;update;create
 
@@ -298,7 +301,18 @@ func NewDefaultPluginsHub(kubeClient client.Client, searchForLegacyPodGroups,
 		}: groveGrouper,
 	}
 
-	skipTopOwnerGrouper := skiptopowner.NewSkipTopOwnerGrouper(kubeClient, defaultGrouper, table)
+	hub := &DefaultPluginsHub{
+		defaultGroupingHandler: defaultGrouper,
+		customPlugins:          table,
+	}
+
+	// hub is captured by pointer and read only when the closure runs, so the skip
+	// grouper still sees the table entries and Karta fallback assigned below.
+	skipTopOwnerGrouper := skiptopowner.NewSkipTopOwnerGrouper(kubeClient, defaultGrouper,
+		func(gvk metav1.GroupVersionKind) grouper.Grouper {
+			return hub.GetPodGrouperPlugin(gvk)
+		})
+
 	table[metav1.GroupVersionKind{
 		Group:   apiGroupArgo,
 		Version: "v1alpha1",
@@ -337,10 +351,13 @@ func NewDefaultPluginsHub(kubeClient client.Client, searchForLegacyPodGroups,
 		Kind:    "DynamoGraphDeployment",
 	}] = skipTopOwnerGrouper
 
-	hub := &DefaultPluginsHub{
-		defaultGroupingHandler: defaultGrouper,
-		customPlugins:          table,
-	}
+	// WorkloadRunner is a kind-agnostic wrapper around an arbitrary workload template.
+	// Skip it so the wrapped kind's plugin decides the grouping.
+	table[metav1.GroupVersionKind{
+		Group:   apiGroupRunai,
+		Version: "*",
+		Kind:    kindWorkloadRunner,
+	}] = skipTopOwnerGrouper
 
 	if genericKartaFallback {
 		hub.kartaFallbackPlugin = kartaplugin.NewKartaHub(kubeClient, defaultGrouper)

pkg/podgrouper/podgrouper/hub/hub_test.go

@@ -10,7 +10,10 @@ import (
 	. "github.com/onsi/gomega"
 	kartav1alpha1 "github.com/run-ai/karta/pkg/api/runai/v1alpha1"
 
+	appsv1 "k8s.io/api/apps/v1"
+	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/utils/ptr"
@@ -86,6 +89,28 @@ var _ = Describe("SupportedTypes", func() {
 			Expect(hasPlugin).To(BeFalse())
 		})
 
+		It("should return skipTopOwner plugin for WorkloadRunner", func() {
+			gvk := metav1.GroupVersionKind{
+				Group:   "run.ai",
+				Version: "v1alpha1",
+				Kind:    "WorkloadRunner",
+			}
+			plugin := hub.GetPodGrouperPlugin(gvk)
+			Expect(plugin).NotTo(BeNil())
+			Expect(plugin.Name()).To(BeEquivalentTo("SkipTopOwner Grouper"))
+		})
+
+		It("should return skipTopOwner plugin for WorkloadRunner of any served version", func() {
+			gvk := metav1.GroupVersionKind{
+				Group:   "run.ai",
+				Version: "v2beta1",
+				Kind:    "WorkloadRunner",
+			}
+			plugin := hub.GetPodGrouperPlugin(gvk)
+			Expect(plugin).NotTo(BeNil())
+			Expect(plugin.Name()).To(BeEquivalentTo("SkipTopOwner Grouper"))
+		})
+
 		It("should return skipTopOwner plugin for TrainJob", func() {
 			gvk := metav1.GroupVersionKind{
 				Group:   "trainer.kubeflow.org",
@@ -155,6 +180,56 @@ var _ = Describe("SupportedTypes", func() {
 		})
 	})
 
+	Context("Skip Top Owner Resolution Tests", func() {
+		It("should resolve the owner below a skipped WorkloadRunner through the hub", func() {
+			statefulSet := &appsv1.StatefulSet{
+				TypeMeta: metav1.TypeMeta{APIVersion: "apps/v1", Kind: "StatefulSet"},
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "wrapped-sts",
+					Namespace: "default",
+					Labels:    map[string]string{queueLabelKey: "wrapped-queue"},
+				},
+			}
+			kubeClient := fake.NewFakeClient(statefulSet)
+			hub := NewDefaultPluginsHub(
+				kubeClient, false, false, false, queueLabelKey, nodePoolLabelKey, "", "",
+			)
+
+			runner := &unstructured.Unstructured{}
+			runner.SetAPIVersion("run.ai/v1alpha1")
+			runner.SetKind("WorkloadRunner")
+			runner.SetNamespace("default")
+			runner.SetName("runner")
+
+			pod := &v1.Pod{
+				TypeMeta:   metav1.TypeMeta{APIVersion: "v1", Kind: "Pod"},
+				ObjectMeta: metav1.ObjectMeta{Name: "test-pod", Namespace: "default"},
+			}
+			owners := []*metav1.PartialObjectMetadata{
+				{
+					TypeMeta:   metav1.TypeMeta{APIVersion: "apps/v1", Kind: "StatefulSet"},
+					ObjectMeta: metav1.ObjectMeta{Name: "wrapped-sts", Namespace: "default"},
+				},
+				{
+					TypeMeta:   metav1.TypeMeta{APIVersion: "run.ai/v1alpha1", Kind: "WorkloadRunner"},
+					ObjectMeta: metav1.ObjectMeta{Name: "runner", Namespace: "default"},
+				},
+			}
+
+			plugin := hub.GetPodGrouperPlugin(metav1.GroupVersionKind{
+				Group: apiGroupRunai, Version: "v1alpha1", Kind: "WorkloadRunner",
+			})
+			Expect(plugin.Name()).To(BeEquivalentTo("SkipTopOwner Grouper"))
+
+			metadata, err := plugin.GetPodGroupMetadata(runner, pod, owners...)
+
+			Expect(err).NotTo(HaveOccurred())
+			Expect(metadata).NotTo(BeNil())
+			Expect(metadata.Queue).To(Equal("wrapped-queue"))
+			Expect(metadata.Owner.Kind).To(Equal("StatefulSet"))
+		})
+	})
+
 	Context("Wildcard Version Tests", func() {
 		var (
 			kubeClient client.Client

pkg/podgrouper/podgrouper/plugins/skiptopowner/skiptopowner.go

@@ -19,18 +19,21 @@ import (
 	"github.com/kai-scheduler/KAI-scheduler/pkg/podgrouper/podgrouper/plugins/grouper"
 )
 
+// GrouperResolver resolves the grouper handling a GVK, or nil when none matches.
+type GrouperResolver func(gvk metav1.GroupVersionKind) grouper.Grouper
+
 type skipTopOwnerGrouper struct {
-	client        client.Client
-	defaultPlugin *defaultgrouper.DefaultGrouper
-	customPlugins map[metav1.GroupVersionKind]grouper.Grouper
+	client         client.Client
+	defaultPlugin  *defaultgrouper.DefaultGrouper
+	resolveGrouper GrouperResolver
 }
 
 func NewSkipTopOwnerGrouper(client client.Client, defaultGrouper *defaultgrouper.DefaultGrouper,
-	customPlugins map[metav1.GroupVersionKind]grouper.Grouper) *skipTopOwnerGrouper {
+	resolveGrouper GrouperResolver) *skipTopOwnerGrouper {
 	return &skipTopOwnerGrouper{
-		client:        client,
-		defaultPlugin: defaultGrouper,
-		customPlugins: customPlugins,
+		client:         client,
+		defaultPlugin:  defaultGrouper,
+		resolveGrouper: resolveGrouper,
 	}
 }
 
@@ -139,8 +142,10 @@ func (sk *skipTopOwnerGrouper) getSupportedTypePGMetadata(
 	lastOwner *unstructured.Unstructured, pod *v1.Pod, otherOwners ...*metav1.PartialObjectMetadata,
 ) (*podgroup.Metadata, error) {
 	ownerKind := metav1.GroupVersionKind(lastOwner.GroupVersionKind())
-	if grouper, found := sk.customPlugins[ownerKind]; found {
-		return grouper.GetPodGroupMetadata(lastOwner, pod, otherOwners...)
+	if sk.resolveGrouper != nil {
+		if resolved := sk.resolveGrouper(ownerKind); resolved != nil {
+			return resolved.GetPodGroupMetadata(lastOwner, pod, otherOwners...)
+		}
 	}
 	return sk.defaultPlugin.GetPodGroupMetadata(lastOwner, pod, otherOwners...)
 }

pkg/podgrouper/podgrouper/plugins/skiptopowner/skiptopowner_test.go

@@ -19,6 +19,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
+	"github.com/kai-scheduler/KAI-scheduler/pkg/podgrouper/podgroup"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/podgrouper/podgrouper/plugins/constants"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/podgrouper/podgrouper/plugins/defaultgrouper"
 	grouperplugin "github.com/kai-scheduler/KAI-scheduler/pkg/podgrouper/podgrouper/plugins/grouper"
@@ -65,7 +66,7 @@ var _ = Describe("SkipTopOwnerGrouper", func() {
 			supportedTypes = map[metav1.GroupVersionKind]grouperplugin.Grouper{
 				{Group: "", Version: "v1", Kind: "Pod"}: defaultGrouper,
 			}
-			plugin = NewSkipTopOwnerGrouper(client, defaultGrouper, supportedTypes)
+			plugin = NewSkipTopOwnerGrouper(client, defaultGrouper, resolverFromMap(supportedTypes))
 		})
 
 		Context("when last owner is a pod", func() {
@@ -430,6 +431,69 @@ var _ = Describe("SkipTopOwnerGrouper", func() {
 			})
 		})
 
+		Context("chained skip-top-owner delegation", func() {
+			var (
+				groveGrouper   *recordingGrouper
+				workloadRunner *unstructured.Unstructured
+				pod            *v1.Pod
+				otherOwners    []*metav1.PartialObjectMetadata
+			)
+
+			BeforeEach(func() {
+				groveGrouper = &recordingGrouper{name: "Grove Grouper"}
+				supportedTypes[metav1.GroupVersionKind{
+					Group: "grove.io", Version: "v1alpha1", Kind: "PodCliqueSet",
+				}] = groveGrouper
+				supportedTypes[metav1.GroupVersionKind{
+					Group: "nvidia.com", Version: "*", Kind: "DynamoGraphDeployment",
+				}] = plugin
+
+				// Dynamo 1.2.0+ serves the DGD as v1beta1, matched via the wildcard version.
+				dgd := newUnstructured("nvidia.com/v1beta1", "DynamoGraphDeployment", "dgd")
+				podCliqueSet := newUnstructured("grove.io/v1alpha1", "PodCliqueSet", "dgd-pcs")
+				podCliqueSet.SetLabels(map[string]string{queueLabelKey: queueName})
+				podClique := newUnstructured("grove.io/v1alpha1", "PodClique", "dgd-pcs-worker")
+				workloadRunner = newUnstructured("run.ai/v1alpha1", "WorkloadRunner", "runner")
+
+				pod = examplePod.DeepCopy()
+				pod.OwnerReferences = []metav1.OwnerReference{
+					{Kind: "PodClique", APIVersion: "grove.io/v1alpha1", Name: podClique.GetName()},
+				}
+
+				Expect(client.Create(context.TODO(), dgd)).To(Succeed())
+				Expect(client.Create(context.TODO(), podCliqueSet)).To(Succeed())
+				Expect(client.Create(context.TODO(), podClique)).To(Succeed())
+				Expect(client.Create(context.TODO(), pod)).To(Succeed())
+				pod.TypeMeta = metav1.TypeMeta{Kind: examplePod.Kind, APIVersion: examplePod.APIVersion}
+
+				otherOwners = []*metav1.PartialObjectMetadata{
+					objectToPartial(podClique),
+					objectToPartial(podCliqueSet),
+					objectToPartial(dgd),
+					objectToPartial(workloadRunner),
+				}
+			})
+
+			It("skips both WorkloadRunner and DGD and lands on the Grove grouper", func() {
+				metadata, err := plugin.GetPodGroupMetadata(workloadRunner, pod, otherOwners...)
+
+				Expect(err).NotTo(HaveOccurred())
+				Expect(metadata).NotTo(BeNil())
+				Expect(metadata.Name).To(Equal("Grove Grouper"))
+				Expect(groveGrouper.calledWith).NotTo(BeNil())
+				Expect(groveGrouper.calledWith.GetKind()).To(Equal("PodCliqueSet"))
+			})
+
+			It("propagates the skipped wrappers' labels down to the Grove owner", func() {
+				workloadRunner.SetLabels(map[string]string{"runner-label": "runner-value"})
+
+				_, err := plugin.GetPodGroupMetadata(workloadRunner, pod, otherOwners...)
+
+				Expect(err).NotTo(HaveOccurred())
+				Expect(groveGrouper.calledWith.GetLabels()).To(HaveKeyWithValue("runner-label", "runner-value"))
+			})
+		})
+
 		Context("default plugin usage", func() {
 			It("uses default plugin when GVK does not have custom plugin", func() {
 				// Use StatefulSet which is not in supportedTypes
@@ -492,6 +556,42 @@ var _ = Describe("SkipTopOwnerGrouper", func() {
 
 })
 
+type recordingGrouper struct {
+	name       string
+	calledWith *unstructured.Unstructured
+}
+
+func (r *recordingGrouper) Name() string { return r.name }
+
+func (r *recordingGrouper) GetPodGroupMetadata(
+	topOwner *unstructured.Unstructured, _ *v1.Pod, _ ...*metav1.PartialObjectMetadata,
+) (*podgroup.Metadata, error) {
+	r.calledWith = topOwner
+	return &podgroup.Metadata{Name: r.name}, nil
+}
+
+func newUnstructured(apiVersion, kind, name string) *unstructured.Unstructured {
+	obj := &unstructured.Unstructured{}
+	obj.SetAPIVersion(apiVersion)
+	obj.SetKind(kind)
+	obj.SetName(name)
+	obj.SetNamespace("default")
+	return obj
+}
+
+func resolverFromMap(table map[metav1.GroupVersionKind]grouperplugin.Grouper) GrouperResolver {
+	return func(gvk metav1.GroupVersionKind) grouperplugin.Grouper {
+		if g, found := table[gvk]; found {
+			return g
+		}
+		gvk.Version = "*"
+		if g, found := table[gvk]; found {
+			return g
+		}
+		return nil
+	}
+}
+
 func objectToPartial(obj client.Object) *metav1.PartialObjectMetadata {
 	objectMeta := metav1.ObjectMeta{
 		Name:            obj.GetName(),