Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Merged
merged 6 commits into from
Jul 12, 2024
Merged

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

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e15a9fe
[jobframework] Only check if the owner is managed for enabled integra…
trasc Jun 27, 2024
237531e
[jobframework] Check Enabled frameworks in unit tests
trasc Jul 9, 2024
c8ba8b5
Review Remarks
trasc Jul 10, 2024
b610515
Review Remarks
trasc Jul 11, 2024
f5866f6
Review Remarks
trasc Jul 12, 2024
a31adae
Review Remarks
trasc Jul 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 32 additions & 2 deletions pkg/controller/jobframework/integrationmanager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,13 @@ import (
"errors"
"fmt"
"sort"
"testing"

metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"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"
Expand Down Expand Up @@ -72,6 +74,7 @@ type IntegrationCallbacks struct {
type integrationManager struct {
names []string
integrations map[string]IntegrationCallbacks
enabledIntegrations set.Set[string]
externalIntegrations map[string]runtime.Object
}

Expand Down Expand Up @@ -144,6 +147,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)
Expand All @@ -152,8 +163,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
}
}
Expand Down Expand Up @@ -186,6 +198,24 @@ 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(tb testing.TB, names ...string) func() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe for a follow up, but I would prefer that, instead, we get rid of the manager as a global variable and pass it to the functions or store it in the structs that need it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not a fan of global states, however in this case when we are doing init bases framework registration this is our best option.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It can still be done differently. There could be one DefaultIntegrationManager that the implementations can register to.

But, you still can pass a different manager to the objects that need it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure it could have been done in a different way, but it was done like this a long time ago and reviewed at that time.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, but, at that time, we were not modifying it in tests.

tb.Helper()
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) {
Expand Down
32 changes: 30 additions & 2 deletions pkg/controller/jobframework/integrationmanager_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,26 @@ import (
"k8s.io/client-go/tools/record"
ctrl "sigs.k8s.io/controller-runtime"
"sigs.k8s.io/controller-runtime/pkg/client"
ctrlmgr "sigs.k8s.io/controller-runtime/pkg/manager"
"sigs.k8s.io/controller-runtime/pkg/reconcile"
)

func testNewReconciler(client.Client, record.EventRecorder, ...Option) JobReconcilerInterface {
type testReconciler struct{}

func (t *testReconciler) Reconcile(context.Context, reconcile.Request) (reconcile.Result, error) {
return reconcile.Result{}, nil
}

func (t *testReconciler) SetupWithManager(mgr ctrlmgr.Manager) error {
return nil
}

var _ JobReconcilerInterface = (*testReconciler)(nil)

func testNewReconciler(client.Client, record.EventRecorder, ...Option) JobReconcilerInterface {
return &testReconciler{}
}

func testSetupWebhook(ctrl.Manager, ...Option) error {
return nil
}
Expand Down Expand Up @@ -363,18 +377,28 @@ func TestGetJobTypeForOwner(t *testing.T) {
externalK3 := func() runtime.Object {
return &metav1.PartialObjectMetadata{TypeMeta: metav1.TypeMeta{Kind: "K3"}}
}()
manageK4 := func() IntegrationCallbacks {
trasc marked this conversation as resolved.
Show resolved Hide resolved
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
Expand All @@ -396,6 +420,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 {
Expand Down
1 change: 1 addition & 0 deletions pkg/controller/jobframework/reconciler_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ func TestIsParentJobManaged(t *testing.T) {
}
for name, tc := range cases {
t.Run(name, func(t *testing.T) {
defer EnableIntegrationsForTest(t, "kubeflow.org/mpijob")()
builder := utiltesting.NewClientBuilder(kubeflow.AddToScheme)
if tc.parentJob != nil {
builder = builder.WithObjects(tc.parentJob)
Expand Down
7 changes: 6 additions & 1 deletion pkg/controller/jobframework/setup.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,14 +45,18 @@ var (
// until the webhooks are operating, and the webhook won't work until the
// certs are all in place.
func SetupControllers(mgr ctrl.Manager, log logr.Logger, opts ...Option) error {
return manager.setupControllers(mgr, log, opts...)
}

func (m *integrationManager) setupControllers(mgr ctrl.Manager, log logr.Logger, opts ...Option) error {
options := ProcessOptions(opts...)

for fwkName := range options.EnabledExternalFrameworks {
if err := RegisterExternalJobType(fwkName); err != nil {
return err
}
}
return ForEachIntegration(func(name string, cb IntegrationCallbacks) error {
return m.forEach(func(name string, cb IntegrationCallbacks) error {
logger := log.WithValues("jobFrameworkName", name)
fwkNamePrefix := fmt.Sprintf("jobFrameworkName %q", name)

Expand Down Expand Up @@ -83,6 +87,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)
}
m.enableIntegration(name)
logger.Info("Set up controller and webhook for job framework")
return nil
}
Expand Down
50 changes: 46 additions & 4 deletions pkg/controller/jobframework/setup_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,10 +43,38 @@ import (
)

func TestSetupControllers(t *testing.T) {
availableIntegrations := map[string]IntegrationCallbacks{
"batch/job": {
NewReconciler: testNewReconciler,
SetupWebhook: testSetupWebhook,
JobType: &batchv1.Job{},
SetupIndexes: testSetupIndexes,
AddToScheme: testAddToScheme,
CanSupportIntegration: testCanSupportIntegration,
},
"kubeflow.org/mpijob": {
NewReconciler: testNewReconciler,
SetupWebhook: testSetupWebhook,
JobType: &kubeflow.MPIJob{},
SetupIndexes: testSetupIndexes,
AddToScheme: testAddToScheme,
CanSupportIntegration: testCanSupportIntegration,
},
"pod": {
NewReconciler: testNewReconciler,
SetupWebhook: testSetupWebhook,
JobType: &corev1.Pod{},
SetupIndexes: testSetupIndexes,
AddToScheme: testAddToScheme,
CanSupportIntegration: testCanSupportIntegration,
},
}

cases := map[string]struct {
opts []Option
mapperGVKs []schema.GroupVersionKind
wantError error
opts []Option
mapperGVKs []schema.GroupVersionKind
wantError error
wantEnabledIntegrations []string
}{
"setup controllers succeed": {
opts: []Option{
Expand All @@ -60,6 +88,7 @@ func TestSetupControllers(t *testing.T) {
batchv1.SchemeGroupVersion.WithKind("Job"),
kubeflow.SchemeGroupVersionKind,
},
wantEnabledIntegrations: []string{"batch/job", "kubeflow.org/mpijob"},
},
"mapper doesn't have kubeflow.org/mpijob, but no error occur": {
opts: []Option{
Expand All @@ -68,10 +97,19 @@ func TestSetupControllers(t *testing.T) {
mapperGVKs: []schema.GroupVersionKind{
batchv1.SchemeGroupVersion.WithKind("Job"),
},
wantEnabledIntegrations: []string{"batch/job"},
},
}
for name, tc := range cases {
t.Run(name, func(t *testing.T) {
manager := integrationManager{}
for name, cbs := range availableIntegrations {
err := manager.register(name, cbs)
if err != nil {
t.Fatalf("Unexpected error while registering %q: %s", name, err)
}
}

_, logger := utiltesting.ContextWithLog(t)
k8sClient := utiltesting.NewClientBuilder(jobset.AddToScheme, kubeflow.AddToScheme, kftraining.AddToScheme, rayv1.AddToScheme).Build()

Expand All @@ -97,10 +135,14 @@ func TestSetupControllers(t *testing.T) {
t.Fatalf("Failed to setup manager: %v", err)
}

gotError := SetupControllers(mgr, logger, tc.opts...)
gotError := manager.setupControllers(mgr, logger, tc.opts...)
if diff := cmp.Diff(tc.wantError, gotError, cmpopts.EquateErrors()); len(diff) != 0 {
t.Errorf("Unexpected error from SetupControllers (-want,+got):\n%s", diff)
}

if diff := cmp.Diff(tc.wantEnabledIntegrations, manager.enabledIntegrations.SortedList()); len(diff) != 0 {
t.Errorf("Unexpected enabled integrations (-want,+got):\n%s", diff)
}
})
}
}
Expand Down
1 change: 1 addition & 0 deletions pkg/controller/jobs/job/job_controller_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,6 +396,7 @@ var (
)

func TestReconciler(t *testing.T) {
defer jobframework.EnableIntegrationsForTest(t, FrameworkName)()
baseJobWrapper := utiltestingjob.MakeJob("job", "ns").
Suspend(true).
Queue("foo").
Expand Down
1 change: 1 addition & 0 deletions pkg/controller/jobs/pod/pod_controller_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4879,6 +4879,7 @@ func TestReconciler_ErrorFinalizingPod(t *testing.T) {
}

func TestIsPodOwnerManagedByQueue(t *testing.T) {
defer jobframework.EnableIntegrationsForTest(t, "batch/job", "ray.io/raycluster")()
trasc marked this conversation as resolved.
Show resolved Hide resolved
testCases := map[string]struct {
ownerReference metav1.OwnerReference
wantRes bool
Expand Down
28 changes: 28 additions & 0 deletions pkg/controller/jobs/pod/pod_webhook_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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": {
Expand Down Expand Up @@ -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{},
Expand All @@ -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")).
Expand All @@ -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")).
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand All @@ -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(
Expand Down Expand Up @@ -261,6 +286,7 @@ func TestDefault(t *testing.T) {

for name, tc := range testCases {
t.Run(name, func(t *testing.T) {
defer jobframework.EnableIntegrationsForTest(t, tc.enableIntegrations...)()
builder := utiltesting.NewClientBuilder()
builder = builder.WithObjects(tc.initObjects...)
cli := builder.Build()
Expand Down Expand Up @@ -374,6 +400,7 @@ func TestGetRoleHash(t *testing.T) {
}

func TestValidateCreate(t *testing.T) {
defer jobframework.EnableIntegrationsForTest(t, "batch/job")()
testCases := map[string]struct {
pod *corev1.Pod
wantErr error
Expand Down Expand Up @@ -476,6 +503,7 @@ func TestValidateCreate(t *testing.T) {
}

func TestValidateUpdate(t *testing.T) {
defer jobframework.EnableIntegrationsForTest(t, "batch/job")()
testCases := map[string]struct {
oldPod *corev1.Pod
newPod *corev1.Pod
Expand Down
Loading