[jobframework] Only check if the owner is managed for enabled integrations.

pkg/controller/jobframework/integrationmanager.go

@@ -26,6 +26,7 @@ import (
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/client-go/tools/record"
+	"k8s.io/utils/set"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
@@ -72,6 +73,7 @@ type IntegrationCallbacks struct {
 type integrationManager struct {
 	names                []string
 	integrations         map[string]IntegrationCallbacks
+	enabledIntegrations  set.Set[string]
 	externalIntegrations map[string]runtime.Object
 }
 
@@ -144,6 +146,14 @@ func (m *integrationManager) getExternal(kindArg string) (runtime.Object, bool)
 	return jt, f
 }
 
+func (m *integrationManager) enableIntegration(name string) {
+	if m.enabledIntegrations == nil {
+		m.enabledIntegrations = set.New(name)
+	} else {
+		m.enabledIntegrations.Insert(name)
+	}
+}
+
 func (m *integrationManager) getList() []string {
 	ret := make([]string, len(m.names))
 	copy(ret, m.names)
@@ -152,8 +162,9 @@ func (m *integrationManager) getList() []string {
 }
 
 func (m *integrationManager) getJobTypeForOwner(ownerRef *metav1.OwnerReference) runtime.Object {
-	for _, cbs := range m.integrations {
-		if cbs.IsManagingObjectsOwner != nil && cbs.IsManagingObjectsOwner(ownerRef) {
+	for jobKey := range m.enabledIntegrations {
+		cbs, found := m.integrations[jobKey]
+		if found && cbs.IsManagingObjectsOwner != nil && cbs.IsManagingObjectsOwner(ownerRef) {
 			return cbs.JobType
 		}
 	}
@@ -186,6 +197,23 @@ func ForEachIntegration(f func(name string, cb IntegrationCallbacks) error) erro
 	return manager.forEach(f)
 }
 
+// EnableIntegration marks the integration identified by name as enabled.
+func EnableIntegration(name string) {
+	manager.enableIntegration(name)
+}
+
+// EnableIntegrationsForTest - should be used only in tests
+// Mark the frameworks identified by names and return a revert function.
+func EnableIntegrationsForTest(names ...string) func() {
+	old := manager.enabledIntegrations.Clone()
+	for _, name := range names {
+		manager.enableIntegration(name)
+	}
+	return func() {
+		manager.enabledIntegrations = old
+	}
+}
+
 // GetIntegration looks-up the framework identified by name in the currently registered
 // list of frameworks returning its callbacks and true if found.
 func GetIntegration(name string) (IntegrationCallbacks, bool) {

pkg/controller/jobframework/integrationmanager_test.go

@@ -363,18 +363,28 @@ func TestGetJobTypeForOwner(t *testing.T) {
 	externalK3 := func() runtime.Object {
 		return &metav1.PartialObjectMetadata{TypeMeta: metav1.TypeMeta{Kind: "K3"}}
 	}()
+	manageK4 := func() IntegrationCallbacks {
+		ret := dontManage
+		ret.IsManagingObjectsOwner = func(owner *metav1.OwnerReference) bool { return owner.Kind == "K4" }
+		ret.JobType = &metav1.PartialObjectMetadata{TypeMeta: metav1.TypeMeta{Kind: "K4"}}
+		return ret
+	}()
 
 	mgr := integrationManager{
-		names: []string{"manageK1", "dontManage", "manageK2"},
+		names: []string{"manageK1", "dontManage", "manageK2", "manageK4"},
 		integrations: map[string]IntegrationCallbacks{
 			"dontManage": dontManage,
 			"manageK1":   manageK1,
 			"manageK2":   manageK2,
+			"manageK4":   manageK4,
 		},
 		externalIntegrations: map[string]runtime.Object{
 			"externalK3": externalK3,
 		},
 	}
+	mgr.enableIntegration("dontManage")
+	mgr.enableIntegration("manageK1")
+	mgr.enableIntegration("manageK2")
 
 	cases := map[string]struct {
 		owner       *metav1.OwnerReference
@@ -396,6 +406,10 @@ func TestGetJobTypeForOwner(t *testing.T) {
 			owner:       &metav1.OwnerReference{Kind: "K4"},
 			wantJobType: nil,
 		},
+		"K5": {
+			owner:       &metav1.OwnerReference{Kind: "K5"},
+			wantJobType: nil,
+		},
 	}
 
 	for tcName, tc := range cases {

pkg/controller/jobframework/setup.go

@@ -83,6 +83,7 @@ func SetupControllers(mgr ctrl.Manager, log logr.Logger, opts ...Option) error {
 				if err = cb.SetupWebhook(mgr, opts...); err != nil {
 					return fmt.Errorf("%s: unable to create webhook: %w", fwkNamePrefix, err)
 				}
+				EnableIntegration(name)
 				logger.Info("Set up controller and webhook for job framework")
 				return nil
 			}

pkg/controller/jobs/job/job_controller_test.go

@@ -396,6 +396,7 @@ var (
 )
 
 func TestReconciler(t *testing.T) {
+	defer jobframework.EnableIntegrationsForTest(FrameworkName)()
 	baseJobWrapper := utiltestingjob.MakeJob("job", "ns").
 		Suspend(true).
 		Queue("foo").

pkg/controller/jobs/pod/pod_controller_test.go

@@ -4879,6 +4879,7 @@ func TestReconciler_ErrorFinalizingPod(t *testing.T) {
 }
 
 func TestIsPodOwnerManagedByQueue(t *testing.T) {
+	defer jobframework.EnableIntegrationsForTest("batch/job", "ray.io/raycluster")()
 	testCases := map[string]struct {
 		ownerReference metav1.OwnerReference
 		wantRes        bool

pkg/controller/jobs/pod/pod_webhook_test.go

@@ -31,6 +31,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 
 	configapi "sigs.k8s.io/kueue/apis/config/v1beta1"
+	"sigs.k8s.io/kueue/pkg/controller/jobframework"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/mpijob"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster"
@@ -64,6 +65,7 @@ func TestDefault(t *testing.T) {
 		manageJobsWithoutQueueName bool
 		namespaceSelector          *metav1.LabelSelector
 		podSelector                *metav1.LabelSelector
+		enableIntegrations         []string
 		want                       *corev1.Pod
 	}{
 		"pod with queue nil ns selector": {
@@ -102,6 +104,22 @@ func TestDefault(t *testing.T) {
 				KueueFinalizer().
 				Obj(),
 		},
+		"pod with owner managed by kueue (Job) while not enabled": {
+			initObjects:       []client.Object{defaultNamespace},
+			podSelector:       &metav1.LabelSelector{},
+			namespaceSelector: defaultNamespaceSelector,
+			pod: testingpod.MakePod("test-pod", defaultNamespace.Name).
+				Queue("test-queue").
+				OwnerReference("parent-job", batchv1.SchemeGroupVersion.WithKind("Job")).
+				Obj(),
+			want: testingpod.MakePod("test-pod", defaultNamespace.Name).
+				Queue("test-queue").
+				Label("kueue.x-k8s.io/managed", "true").
+				KueueSchedulingGate().
+				KueueFinalizer().
+				OwnerReference("parent-job", batchv1.SchemeGroupVersion.WithKind("Job")).
+				Obj(),
+		},
 		"pod with owner managed by kueue (Job)": {
 			initObjects:       []client.Object{defaultNamespace},
 			podSelector:       &metav1.LabelSelector{},
@@ -110,6 +128,7 @@ func TestDefault(t *testing.T) {
 				Queue("test-queue").
 				OwnerReference("parent-job", batchv1.SchemeGroupVersion.WithKind("Job")).
 				Obj(),
+			enableIntegrations: []string{"batch/job"},
 			want: testingpod.MakePod("test-pod", defaultNamespace.Name).
 				Queue("test-queue").
 				OwnerReference("parent-job", batchv1.SchemeGroupVersion.WithKind("Job")).
@@ -123,6 +142,7 @@ func TestDefault(t *testing.T) {
 				Queue("test-queue").
 				OwnerReference("parent-ray-cluster", rayv1.GroupVersion.WithKind("RayCluster")).
 				Obj(),
+			enableIntegrations: []string{"ray.io/raycluster"},
 			want: testingpod.MakePod("test-pod", defaultNamespace.Name).
 				Queue("test-queue").
 				OwnerReference("parent-ray-cluster", rayv1.GroupVersion.WithKind("RayCluster")).
@@ -139,6 +159,7 @@ func TestDefault(t *testing.T) {
 					schema.GroupVersionKind{Group: "kubeflow.org", Version: "v2beta1", Kind: "MPIJob"},
 				).
 				Obj(),
+			enableIntegrations: []string{"kubeflow.org/mpijob"},
 			want: testingpod.MakePod("test-pod", defaultNamespace.Name).
 				Queue("test-queue").
 				OwnerReference(
@@ -158,6 +179,7 @@ func TestDefault(t *testing.T) {
 					schema.GroupVersionKind{Group: "kubeflow.org", Version: "v1", Kind: "PyTorchJob"},
 				).
 				Obj(),
+			enableIntegrations: []string{"kubeflow.org/pytorchjob"},
 			want: testingpod.MakePod("test-pod", defaultNamespace.Name).
 				Queue("test-queue").
 				OwnerReference(
@@ -177,6 +199,7 @@ func TestDefault(t *testing.T) {
 					schema.GroupVersionKind{Group: "kubeflow.org", Version: "v1", Kind: "TFJob"},
 				).
 				Obj(),
+			enableIntegrations: []string{"kubeflow.org/tfjob"},
 			want: testingpod.MakePod("test-pod", defaultNamespace.Name).
 				Queue("test-queue").
 				OwnerReference(
@@ -196,6 +219,7 @@ func TestDefault(t *testing.T) {
 					schema.GroupVersionKind{Group: "kubeflow.org", Version: "v1", Kind: "XGBoostJob"},
 				).
 				Obj(),
+			enableIntegrations: []string{"kubeflow.org/xgboostjob"},
 			want: testingpod.MakePod("test-pod", defaultNamespace.Name).
 				Queue("test-queue").
 				OwnerReference(
@@ -215,6 +239,7 @@ func TestDefault(t *testing.T) {
 					schema.GroupVersionKind{Group: "kubeflow.org", Version: "v1", Kind: "PaddleJob"},
 				).
 				Obj(),
+			enableIntegrations: []string{"kubeflow.org/paddlejob"},
 			want: testingpod.MakePod("test-pod", defaultNamespace.Name).
 				Queue("test-queue").
 				OwnerReference(
@@ -261,6 +286,7 @@ func TestDefault(t *testing.T) {
 
 	for name, tc := range testCases {
 		t.Run(name, func(t *testing.T) {
+			defer jobframework.EnableIntegrationsForTest(tc.enableIntegrations...)()
 			builder := utiltesting.NewClientBuilder()
 			builder = builder.WithObjects(tc.initObjects...)
 			cli := builder.Build()
@@ -374,6 +400,7 @@ func TestGetRoleHash(t *testing.T) {
 }
 
 func TestValidateCreate(t *testing.T) {
+	defer jobframework.EnableIntegrationsForTest("batch/job")()
 	testCases := map[string]struct {
 		pod       *corev1.Pod
 		wantErr   error
@@ -476,6 +503,7 @@ func TestValidateCreate(t *testing.T) {
 }
 
 func TestValidateUpdate(t *testing.T) {
+	defer jobframework.EnableIntegrationsForTest("batch/job")()
 	testCases := map[string]struct {
 		oldPod    *corev1.Pod
 		newPod    *corev1.Pod

test/integration/controller/jobs/job/suite_test.go

@@ -73,6 +73,7 @@ func managerSetup(opts ...jobframework.Option) framework.ManagerSetup {
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
 		err = job.SetupWebhook(mgr, opts...)
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
+		jobframework.EnableIntegration(job.FrameworkName)
 	}
 }
 
@@ -99,6 +100,7 @@ func managerAndControllersSetup(enableScheduler bool, configuration *config.Conf
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
 		err = job.SetupWebhook(mgr, opts...)
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
+		jobframework.EnableIntegration(job.FrameworkName)
 
 		if enableScheduler {
 			sched := scheduler.New(queues, cCache, mgr.GetClient(), mgr.GetEventRecorderFor(constants.AdmissionName))

test/integration/controller/jobs/mpijob/suite_test.go

@@ -73,6 +73,7 @@ func managerSetup(setupJobManager bool, opts ...jobframework.Option) framework.M
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
 		err = mpijob.SetupMPIJobWebhook(mgr, opts...)
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
+		jobframework.EnableIntegration(mpijob.FrameworkName)
 
 		if setupJobManager {
 			jobReconciler := job.NewReconciler(

test/integration/controller/jobs/pod/suite_test.go

@@ -92,6 +92,7 @@ func managerSetup(configuration *config.Configuration, opts ...jobframework.Opti
 			opts...)
 		err = jobReconciler.SetupWithManager(mgr)
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
+		jobframework.EnableIntegration(job.FrameworkName)
 
 		cCache := cache.New(mgr.GetClient())
 		queues := queue.NewManager(mgr.GetClient(), cCache)

test/integration/controller/jobs/raycluster/suite_test.go

@@ -70,6 +70,7 @@ func managerSetup(opts ...jobframework.Option) framework.ManagerSetup {
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
 		err = raycluster.SetupRayClusterWebhook(mgr, opts...)
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
+		jobframework.EnableIntegration(rayjob.FrameworkName)
 	}
 }
 
@@ -123,6 +124,7 @@ func managerWithRayClusterAndRayJobControllersSetup(opts ...jobframework.Option)
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
 		err = rayjob.SetupRayJobWebhook(mgr, opts...)
 		gomega.Expect(err).NotTo(gomega.HaveOccurred())
+		jobframework.EnableIntegration(rayjob.FrameworkName)
 
 		failedWebhook, err := webhooks.Setup(mgr)
 		gomega.Expect(err).ToNot(gomega.HaveOccurred(), "webhook", failedWebhook)