Auto-enable pod integration for pod-dependent frameworks (#6736)

IrvingMg · web-flow · commit 0d6a4aee95cb · 2025-10-03T01:42:59.000-07:00
diff --git a/pkg/controller/jobframework/integrationmanager.go b/pkg/controller/jobframework/integrationmanager.go
@@ -81,7 +81,10 @@ type IntegrationCallbacks struct {
 	// The job's MultiKueue adapter (optional)
 	MultiKueueAdapter MultiKueueAdapter
 	// The list of integration that need to be enabled along with the current one.
+	// Deprecated: Use ImplicitlyEnabledFrameworkNames instead.
 	DependencyList []string
+	// The list of integrations implicitly enabled as dependencies of the integration.
+	ImplicitlyEnabledFrameworkNames []string
 }
 
 func (i *IntegrationCallbacks) getGVK() schema.GroupVersionKind {
@@ -100,11 +103,13 @@ func (i *IntegrationCallbacks) matchingOwnerReference(ownerRef *metav1.OwnerRefe
 }
 
 type integrationManager struct {
-	names                []string
-	integrations         map[string]IntegrationCallbacks
-	enabledIntegrations  set.Set[string]
-	externalIntegrations map[string]runtime.Object
-	mu                   sync.RWMutex
+	names                         []string
+	integrations                  map[string]IntegrationCallbacks
+	enabledIntegrations           set.Set[string]
+	externalIntegrations          map[string]runtime.Object
+	implicitlyEnabledIntegrations sets.Set[string]
+	gvkToName                     map[schema.GroupVersionKind]string
+	mu                            sync.RWMutex
 }
 
 var manager integrationManager
@@ -132,6 +137,11 @@ func (m *integrationManager) register(name string, cb IntegrationCallbacks) erro
 	m.integrations[name] = cb
 	m.names = append(m.names, name)
 
+	if m.gvkToName == nil {
+		m.gvkToName = make(map[schema.GroupVersionKind]string)
+	}
+	m.gvkToName[cb.getGVK()] = name
+
 	return nil
 }
 
@@ -323,6 +333,15 @@ func GetIntegrationByGVK(gvk schema.GroupVersionKind) (IntegrationCallbacks, boo
 	return IntegrationCallbacks{}, false
 }
 
+// HasImplicitlyEnabledFramework returns true if the given GVK maps to an implicitly enabled framework.
+func HasImplicitlyEnabledFramework(gvk schema.GroupVersionKind) bool {
+	name, found := manager.gvkToName[gvk]
+	if !found {
+		return false
+	}
+	return manager.implicitlyEnabledIntegrations.Has(name)
+}
+
 func ownerReferenceMatchingGVK(ownerRef *metav1.OwnerReference, gvk schema.GroupVersionKind) bool {
 	apiVersion, kind := gvk.ToAPIVersionAndKind()
 	return ownerRef.APIVersion == apiVersion && ownerRef.Kind == kind
@@ -370,3 +389,24 @@ func GetMultiKueueAdapters(enabledIntegrations sets.Set[string]) (map[string]Mul
 	}
 	return ret, nil
 }
+
+func (m *integrationManager) collectImplicitlyEnabledIntegrations(enabledFrameworks sets.Set[string]) sets.Set[string] {
+	result := sets.New[string]()
+	for frameworkName := range enabledFrameworks {
+		callbacks, found := m.get(frameworkName)
+		if !found {
+			continue
+		}
+
+		for _, implicitFramework := range callbacks.ImplicitlyEnabledFrameworkNames {
+			if !enabledFrameworks.Has(implicitFramework) {
+				result.Insert(implicitFramework)
+			}
+		}
+	}
+	return result
+}
+
+func (m *integrationManager) setImplicitlyEnabledIntegrations(implicitlyIntegrations sets.Set[string]) {
+	m.implicitlyEnabledIntegrations = implicitlyIntegrations
+}
diff --git a/pkg/controller/jobframework/integrationmanager_test.go b/pkg/controller/jobframework/integrationmanager_test.go
@@ -458,14 +458,14 @@ func TestEnabledIntegrationsDependencies(t *testing.T) {
 			enabled:   []string{"i1"},
 			wantError: errIntegrationNotFound,
 		},
-		"dependecncy not enabled": {
+		"dependency not enabled": {
 			integrationsDependencies: map[string][]string{
 				"i1": {"i2"},
 			},
 			enabled:   []string{"i1"},
 			wantError: errDependencyIntegrationNotEnabled,
 		},
-		"dependecncy not found": {
+		"dependency not found": {
 			integrationsDependencies: map[string][]string{
 				"i1": {"i2"},
 			},
@@ -486,8 +486,8 @@ func TestEnabledIntegrationsDependencies(t *testing.T) {
 			manager := integrationManager{
 				integrations: map[string]IntegrationCallbacks{},
 			}
-			for inegration, deps := range tc.integrationsDependencies {
-				manager.integrations[inegration] = IntegrationCallbacks{
+			for integration, deps := range tc.integrationsDependencies {
+				manager.integrations[integration] = IntegrationCallbacks{
 					DependencyList: deps,
 				}
 			}
@@ -498,3 +498,59 @@ func TestEnabledIntegrationsDependencies(t *testing.T) {
 		})
 	}
 }
+
+func TestImplicitlyEnabledIntegrations(t *testing.T) {
+	cases := map[string]struct {
+		integrationsImplicit map[string][]string
+		enabled              []string
+		wantImplicit         []string
+	}{
+		"empty": {},
+		"no implicit dependencies": {
+			integrationsImplicit: map[string][]string{
+				"i1": nil,
+			},
+			enabled:      []string{"i1"},
+			wantImplicit: []string{},
+		},
+		"single implicit dependency": {
+			integrationsImplicit: map[string][]string{
+				"i1": {"i2"},
+			},
+			enabled:      []string{"i1"},
+			wantImplicit: []string{"i2"},
+		},
+		"multiple implicit dependencies": {
+			integrationsImplicit: map[string][]string{
+				"i1": {"i2", "i3"},
+			},
+			enabled:      []string{"i1"},
+			wantImplicit: []string{"i2", "i3"},
+		},
+		"nested implicit dependencies": {
+			integrationsImplicit: map[string][]string{
+				"i1": {"i2", "i3"},
+				"i2": {"i3"},
+				"i3": nil,
+			},
+			enabled:      []string{"i1", "i2"},
+			wantImplicit: []string{"i3"},
+		},
+	}
+	for tcName, tc := range cases {
+		t.Run(tcName, func(t *testing.T) {
+			mgr := integrationManager{
+				integrations: map[string]IntegrationCallbacks{},
+			}
+			for integration, implicitDeps := range tc.integrationsImplicit {
+				mgr.integrations[integration] = IntegrationCallbacks{
+					ImplicitlyEnabledFrameworkNames: implicitDeps,
+				}
+			}
+			gotImplicit := mgr.collectImplicitlyEnabledIntegrations(sets.New(tc.enabled...))
+			if diff := cmp.Diff(sets.New(tc.wantImplicit...), gotImplicit); diff != "" {
+				t.Errorf("Unexpected implicitly enabled integrations (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
diff --git a/pkg/controller/jobframework/setup.go b/pkg/controller/jobframework/setup.go
@@ -56,7 +56,11 @@ func SetupControllers(ctx context.Context, mgr ctrl.Manager, log logr.Logger, op
 func (m *integrationManager) setupControllers(ctx context.Context, mgr ctrl.Manager, log logr.Logger, opts ...Option) error {
 	options := ProcessOptions(opts...)
 
-	if err := m.checkEnabledListDependencies(options.EnabledFrameworks); err != nil {
+	implicitlyEnabledIntegrations := m.collectImplicitlyEnabledIntegrations(options.EnabledFrameworks)
+	m.setImplicitlyEnabledIntegrations(implicitlyEnabledIntegrations)
+	allEnabledIntegrations := options.EnabledFrameworks.Union(implicitlyEnabledIntegrations)
+
+	if err := m.checkEnabledListDependencies(allEnabledIntegrations); err != nil {
 		return fmt.Errorf("check enabled frameworks list: %w", err)
 	}
 
@@ -69,7 +73,7 @@ func (m *integrationManager) setupControllers(ctx context.Context, mgr ctrl.Mana
 		logger := log.WithValues("jobFrameworkName", name)
 		fwkNamePrefix := fmt.Sprintf("jobFrameworkName %q", name)
 
-		if options.EnabledFrameworks.Has(name) {
+		if allEnabledIntegrations.Has(name) {
 			if cb.CanSupportIntegration != nil {
 				if canSupport, err := cb.CanSupportIntegration(opts...); !canSupport || err != nil {
 					return fmt.Errorf("failed to configure reconcilers: %w", err)
@@ -178,8 +182,10 @@ func restMappingExists(mgr ctrl.Manager, gvk schema.GroupVersionKind) error {
 // Note that the second argument, "indexer" needs to be the fieldIndexer obtained from the Manager.
 func SetupIndexes(ctx context.Context, indexer client.FieldIndexer, opts ...Option) error {
 	options := ProcessOptions(opts...)
+
+	allEnabledIntegrations := options.EnabledFrameworks.Union(manager.collectImplicitlyEnabledIntegrations(options.EnabledFrameworks))
 	return ForEachIntegration(func(name string, cb IntegrationCallbacks) error {
-		if options.EnabledFrameworks.Has(name) {
+		if allEnabledIntegrations.Has(name) {
 			if err := cb.SetupIndexes(ctx, indexer); err != nil {
 				return fmt.Errorf("jobFrameworkName %q: %w", name, err)
 			}
diff --git a/pkg/controller/jobs/deployment/deployment_controller.go b/pkg/controller/jobs/deployment/deployment_controller.go
@@ -38,13 +38,13 @@ const (
 
 func init() {
 	utilruntime.Must(jobframework.RegisterIntegration(FrameworkName, jobframework.IntegrationCallbacks{
-		SetupIndexes:   SetupIndexes,
-		NewReconciler:  jobframework.NewNoopReconcilerFactory(gvk),
-		GVK:            gvk,
-		SetupWebhook:   SetupWebhook,
-		JobType:        &appsv1.Deployment{},
-		AddToScheme:    appsv1.AddToScheme,
-		DependencyList: []string{"pod"},
+		SetupIndexes:                    SetupIndexes,
+		NewReconciler:                   jobframework.NewNoopReconcilerFactory(gvk),
+		GVK:                             gvk,
+		SetupWebhook:                    SetupWebhook,
+		JobType:                         &appsv1.Deployment{},
+		AddToScheme:                     appsv1.AddToScheme,
+		ImplicitlyEnabledFrameworkNames: []string{"pod"},
 	}))
 }
 
diff --git a/pkg/controller/jobs/leaderworkerset/leaderworkerset_controller.go b/pkg/controller/jobs/leaderworkerset/leaderworkerset_controller.go
@@ -38,14 +38,14 @@ const (
 
 func init() {
 	utilruntime.Must(jobframework.RegisterIntegration(FrameworkName, jobframework.IntegrationCallbacks{
-		SetupIndexes:             SetupIndexes,
-		NewReconciler:            NewReconciler,
-		NewAdditionalReconcilers: []jobframework.ReconcilerFactory{NewPodReconciler},
-		SetupWebhook:             SetupWebhook,
-		JobType:                  &leaderworkersetv1.LeaderWorkerSet{},
-		AddToScheme:              leaderworkersetv1.AddToScheme,
-		DependencyList:           []string{"pod"},
-		GVK:                      gvk,
+		SetupIndexes:                    SetupIndexes,
+		NewReconciler:                   NewReconciler,
+		NewAdditionalReconcilers:        []jobframework.ReconcilerFactory{NewPodReconciler},
+		SetupWebhook:                    SetupWebhook,
+		JobType:                         &leaderworkersetv1.LeaderWorkerSet{},
+		AddToScheme:                     leaderworkersetv1.AddToScheme,
+		ImplicitlyEnabledFrameworkNames: []string{"pod"},
+		GVK:                             gvk,
 	}))
 }
 
diff --git a/pkg/controller/jobs/pod/pod_controller.go b/pkg/controller/jobs/pod/pod_controller.go
@@ -565,6 +565,11 @@ func (p *Pod) Skip() bool {
 	if v, ok := p.pod.GetLabels()[constants.ManagedByKueueLabelKey]; p.isFound && (!ok || v != constants.ManagedByKueueLabelValue) {
 		return true
 	}
+	if jobframework.HasImplicitlyEnabledFramework(p.pod.GroupVersionKind()) &&
+		p.pod.GetAnnotations()[podconstants.SuspendedByParentAnnotation] == "" {
+		ctrl.Log.V(3).Info("Pod Integration was implicitly enabled but object lacks parent annotation, skipping")
+		return true
+	}
 	return false
 }
 
diff --git a/pkg/controller/jobs/statefulset/statefulset_controller.go b/pkg/controller/jobs/statefulset/statefulset_controller.go
@@ -38,13 +38,13 @@ const (
 
 func init() {
 	utilruntime.Must(jobframework.RegisterIntegration(FrameworkName, jobframework.IntegrationCallbacks{
-		SetupIndexes:   SetupIndexes,
-		NewReconciler:  NewReconciler,
-		SetupWebhook:   SetupWebhook,
-		JobType:        &appsv1.StatefulSet{},
-		AddToScheme:    appsv1.AddToScheme,
-		DependencyList: []string{"pod"},
-		GVK:            gvk,
+		SetupIndexes:                    SetupIndexes,
+		NewReconciler:                   NewReconciler,
+		SetupWebhook:                    SetupWebhook,
+		JobType:                         &appsv1.StatefulSet{},
+		AddToScheme:                     appsv1.AddToScheme,
+		ImplicitlyEnabledFrameworkNames: []string{"pod"},
+		GVK:                             gvk,
 	}))
 }
 
diff --git a/test/e2e/customconfigs/podintegrationautoenablement_test.go b/test/e2e/customconfigs/podintegrationautoenablement_test.go

Original file line number	Diff line number	Diff line change
`@@ -565,6 +565,11 @@ func (p *Pod) Skip() bool {`
`565`	`565`	`if v, ok := p.pod.GetLabels()[constants.ManagedByKueueLabelKey]; p.isFound && (!ok \|\| v != constants.ManagedByKueueLabelValue) {`
`566`	`566`	`return true`
`567`	`567`	`}`
	`568`	`+ if jobframework.HasImplicitlyEnabledFramework(p.pod.GroupVersionKind()) &&`
	`569`	`+ p.pod.GetAnnotations()[podconstants.SuspendedByParentAnnotation] == "" {`
	`570`	`+ ctrl.Log.V(3).Info("Pod Integration was implicitly enabled but object lacks parent annotation, skipping")`
	`571`	`+ return true`
	`572`	`+ }`
`568`	`573`	`return false`
`569`	`574`	`}`
`570`	`575`