Fixed admission and scheduler PodDisruptionBudget reconciliation

Signed-off-by: dttung2905 <ttdao.2015@accountancy.smu.edu.sg>

Author: dttung2905

.changes/unreleased/Fixed-20260718-212943.yaml

@@ -0,0 +1,6 @@
+kind: Fixed
+body: Fixed admission and scheduler PodDisruptionBudget reconciliation, including drift recovery and cleanup when no longer desired.
+time: 2026-07-18T21:29:43.130860557+01:00
+custom:
+    Author: dttung2905
+    Issue: "1477"

pkg/operator/controller/integration_tests/config_controller_test.go

@@ -8,14 +8,18 @@ import (
 	"os"
 
 	kaiv1 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/kai/v1"
+	kaiv1admission "github.com/kai-scheduler/KAI-scheduler/pkg/apis/kai/v1/admission"
 	kaiv1binder "github.com/kai-scheduler/KAI-scheduler/pkg/apis/kai/v1/binder"
+	kaiv1scheduler "github.com/kai-scheduler/KAI-scheduler/pkg/apis/kai/v1/scheduler"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 
 	nvidiav1 "github.com/kai-scheduler/KAI-scheduler/third_party/nvidia/gpu-operator/api/nvidia/v1"
 
 	appsv1 "k8s.io/api/apps/v1"
+	policyv1 "k8s.io/api/policy/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/utils/ptr"
@@ -116,4 +120,169 @@ var _ = Describe("KAIConfigController", Ordered, func() {
 		})
 	})
 
+	Context("Reconciling the admission PodDisruptionBudget", Ordered, func() {
+		const pdbName = "admission"
+
+		updateAdmissionPDBConfig := func(ctx context.Context, replicas int32, enabled bool) {
+			Eventually(func() error {
+				currentConfig := &kaiv1.Config{}
+				if err := k8sClient.Get(
+					ctx,
+					types.NamespacedName{Name: constants.DefaultKAIConfigSingeltonInstanceName},
+					currentConfig,
+				); err != nil {
+					return err
+				}
+				currentConfig.Spec.Admission = &kaiv1admission.Admission{
+					Replicas: ptr.To(replicas),
+					Service: &v1common.Service{
+						Enabled: ptr.To(true),
+						PodDisruptionBudget: &v1common.PodDisruptionBudget{
+							Enabled:        ptr.To(enabled),
+							MaxUnavailable: ptr.To(int32(1)),
+						},
+					},
+				}
+				return k8sClient.Update(ctx, currentConfig)
+			}, "10s", "200ms").Should(Succeed())
+		}
+
+		getPDB := func(ctx context.Context) (*policyv1.PodDisruptionBudget, error) {
+			pdb := &policyv1.PodDisruptionBudget{}
+			err := k8sClient.Get(ctx, types.NamespacedName{
+				Name:      pdbName,
+				Namespace: kaiConfig.Spec.Namespace,
+			}, pdb)
+			return pdb, err
+		}
+
+		It("creates and watches the PDB", func(ctx context.Context) {
+			updateAdmissionPDBConfig(ctx, 2, true)
+
+			Eventually(func(g Gomega) {
+				pdb, err := getPDB(ctx)
+				g.Expect(err).NotTo(HaveOccurred())
+				g.Expect(metav1.GetControllerOf(pdb)).NotTo(BeNil())
+				g.Expect(metav1.GetControllerOf(pdb).Kind).To(Equal("Config"))
+			}, "10s", "200ms").Should(Succeed())
+
+			pdb, err := getPDB(ctx)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(k8sClient.Delete(ctx, pdb)).To(Succeed())
+
+			Eventually(func() error {
+				_, err := getPDB(ctx)
+				return err
+			}, "10s", "200ms").Should(Succeed())
+		})
+
+		It("removes the PDB after scaling admission down to one replica", func(ctx context.Context) {
+			updateAdmissionPDBConfig(ctx, 1, true)
+
+			Eventually(func() bool {
+				_, err := getPDB(ctx)
+				return apierrors.IsNotFound(err)
+			}, "10s", "200ms").Should(BeTrue())
+		})
+
+		It("removes the PDB when it is disabled", func(ctx context.Context) {
+			updateAdmissionPDBConfig(ctx, 2, true)
+			Eventually(func() error {
+				_, err := getPDB(ctx)
+				return err
+			}, "10s", "200ms").Should(Succeed())
+
+			updateAdmissionPDBConfig(ctx, 2, false)
+			Eventually(func() bool {
+				_, err := getPDB(ctx)
+				return apierrors.IsNotFound(err)
+			}, "10s", "200ms").Should(BeTrue())
+		})
+	})
+
+	Context("Reconciling the scheduler PodDisruptionBudget", Ordered, func() {
+		const (
+			shardName = "pdb-test"
+			pdbName   = "kai-scheduler-" + shardName
+		)
+
+		updateSchedulerPDBConfig := func(ctx context.Context, replicas int32, enabled bool) {
+			Eventually(func() error {
+				currentConfig := &kaiv1.Config{}
+				if err := k8sClient.Get(
+					ctx,
+					types.NamespacedName{Name: constants.DefaultKAIConfigSingeltonInstanceName},
+					currentConfig,
+				); err != nil {
+					return err
+				}
+				currentConfig.Spec.Scheduler = &kaiv1scheduler.Scheduler{
+					Replicas: ptr.To(replicas),
+					Service: &v1common.Service{
+						Enabled: ptr.To(true),
+						PodDisruptionBudget: &v1common.PodDisruptionBudget{
+							Enabled:        ptr.To(enabled),
+							MaxUnavailable: ptr.To(int32(1)),
+						},
+					},
+				}
+				return k8sClient.Update(ctx, currentConfig)
+			}, "10s", "200ms").Should(Succeed())
+		}
+
+		getPDB := func(ctx context.Context) (*policyv1.PodDisruptionBudget, error) {
+			pdb := &policyv1.PodDisruptionBudget{}
+			err := k8sClient.Get(ctx, types.NamespacedName{
+				Name:      pdbName,
+				Namespace: kaiConfig.Spec.Namespace,
+			}, pdb)
+			return pdb, err
+		}
+
+		BeforeAll(func(ctx context.Context) {
+			updateSchedulerPDBConfig(ctx, 2, true)
+			Expect(k8sClient.Create(ctx, &kaiv1.SchedulingShard{
+				ObjectMeta: metav1.ObjectMeta{Name: shardName},
+			})).To(Succeed())
+		})
+
+		AfterAll(func(ctx context.Context) {
+			shard := &kaiv1.SchedulingShard{}
+			err := k8sClient.Get(ctx, types.NamespacedName{Name: shardName}, shard)
+			if err == nil {
+				Expect(k8sClient.Delete(ctx, shard)).To(Succeed())
+			} else {
+				Expect(apierrors.IsNotFound(err)).To(BeTrue())
+			}
+		})
+
+		It("creates and watches one PDB for the shard", func(ctx context.Context) {
+			Eventually(func(g Gomega) {
+				pdb, err := getPDB(ctx)
+				g.Expect(err).NotTo(HaveOccurred())
+				g.Expect(metav1.GetControllerOf(pdb)).NotTo(BeNil())
+				g.Expect(metav1.GetControllerOf(pdb).Kind).To(Equal("SchedulingShard"))
+				g.Expect(metav1.GetControllerOf(pdb).Name).To(Equal(shardName))
+			}, "10s", "200ms").Should(Succeed())
+
+			pdb, err := getPDB(ctx)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(k8sClient.Delete(ctx, pdb)).To(Succeed())
+
+			Eventually(func() error {
+				_, err := getPDB(ctx)
+				return err
+			}, "10s", "200ms").Should(Succeed())
+		})
+
+		It("removes the shard PDB after scheduler scales down", func(ctx context.Context) {
+			updateSchedulerPDBConfig(ctx, 1, true)
+
+			Eventually(func() bool {
+				_, err := getPDB(ctx)
+				return apierrors.IsNotFound(err)
+			}, "10s", "200ms").Should(BeTrue())
+		})
+	})
+
 })

pkg/operator/controller/integration_tests/suite_test.go

@@ -33,7 +33,10 @@ import (
 	monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1"
 	admissionv1 "k8s.io/api/admissionregistration/v1"
 	appsv1 "k8s.io/api/apps/v1"
+	coordinationv1 "k8s.io/api/coordination/v1"
 	v1 "k8s.io/api/core/v1"
+	discoveryv1 "k8s.io/api/discovery/v1"
+	policyv1 "k8s.io/api/policy/v1"
 	apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -73,6 +76,9 @@ var _ = BeforeSuite(func() {
 	Expect(appsv1.AddToScheme(scheme)).To(Succeed())
 	Expect(apiserver.AddToScheme(scheme)).To(Succeed())
 	Expect(admissionv1.AddToScheme(scheme)).To(Succeed())
+	Expect(coordinationv1.AddToScheme(scheme)).To(Succeed())
+	Expect(discoveryv1.AddToScheme(scheme)).To(Succeed())
+	Expect(policyv1.AddToScheme(scheme)).To(Succeed())
 	Expect(apiextensionsv1.AddToScheme(scheme)).To(Succeed())
 
 	By("bootstrapping test environment")
@@ -125,6 +131,14 @@ var _ = BeforeSuite(func() {
 	err = configReconciler.SetupWithManager(k8sManager)
 	Expect(err).NotTo(HaveOccurred())
 
+	shardReconciler := controller.NewSchedulingShardReconciler(
+		k8sManager.GetClient(),
+		k8sManager.GetScheme(),
+	)
+	shardReconciler.SetOperands(controller.OperandsForShard)
+	err = shardReconciler.SetupWithManager(k8sManager)
+	Expect(err).NotTo(HaveOccurred())
+
 	go func() {
 		defer GinkgoRecover()
 		err = k8sManager.Start(testContext)

pkg/operator/operands/deployable/deployable_test.go

@@ -19,7 +19,9 @@ import (
 	monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1"
 	appsv1 "k8s.io/api/apps/v1"
 	v1 "k8s.io/api/core/v1"
+	policyv1 "k8s.io/api/policy/v1"
 	apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	vpav1 "k8s.io/autoscaler/vertical-pod-autoscaler/pkg/apis/autoscaling.k8s.io/v1"
 	"k8s.io/client-go/kubernetes/scheme"
@@ -54,6 +56,7 @@ var _ = Describe("Deployable", func() {
 		Expect(kaiv1.AddToScheme(testScheme)).To(Succeed())
 		Expect(apiextensionsv1.AddToScheme(testScheme)).To(Succeed())
 		Expect(monitoringv1.AddToScheme(testScheme)).To(Succeed())
+		Expect(policyv1.AddToScheme(testScheme)).To(Succeed())
 		Expect(vpav1.AddToScheme(testScheme)).To(Succeed())
 
 		fakeClientBuilder = fake.NewClientBuilder().
@@ -215,6 +218,77 @@ var _ = Describe("Deployable", func() {
 				Expect(updateCalls).To(Equal(1))
 			})
 		})
+
+		Context("PodDisruptionBudget lifecycle", func() {
+			var existingPDB *policyv1.PodDisruptionBudget
+
+			BeforeEach(func() {
+				existingPDB = &policyv1.PodDisruptionBudget{
+					TypeMeta: metav1.TypeMeta{
+						Kind:       "PodDisruptionBudget",
+						APIVersion: policyv1.SchemeGroupVersion.String(),
+					},
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "admission",
+						Namespace: "kai-scheduler",
+						OwnerReferences: []metav1.OwnerReference{{
+							APIVersion: kaiv1.GroupVersion.String(),
+							Kind:       "Config",
+							Name:       kaiConfig.Name,
+							UID:        kaiConfig.UID,
+							Controller: ptr.To(true),
+						}},
+					},
+				}
+			})
+
+			It("collects an existing PDB instead of attempting to create it again", func() {
+				pdbCreateCalls := 0
+				builder := fakeClientBuilder.
+					WithObjects(existingPDB).
+					WithInterceptorFuncs(interceptor.Funcs{
+						Create: func(
+							ctx context.Context,
+							runtimeClient client.WithWatch,
+							obj client.Object,
+							opts ...client.CreateOption,
+						) error {
+							if _, ok := obj.(*policyv1.PodDisruptionBudget); ok {
+								pdbCreateCalls++
+							}
+							return runtimeClient.Create(ctx, obj, opts...)
+						},
+					})
+				fakeClient := getFakeClient(builder, known_types.KAIConfigRegisteredCollectible)
+				deployable := New(
+					[]operands.Operand{&fakePDBOperand{enabled: true}},
+					known_types.KAIConfigRegisteredCollectible,
+				)
+
+				Expect(deployable.Deploy(context.Background(), fakeClient, kaiConfig, kaiConfig)).To(Succeed())
+				Expect(pdbCreateCalls).To(BeZero())
+			})
+
+			It("deletes an owned PDB when it is no longer desired", func() {
+				fakeClient := getFakeClient(
+					fakeClientBuilder.WithObjects(existingPDB),
+					known_types.KAIConfigRegisteredCollectible,
+				)
+				deployable := New(
+					[]operands.Operand{&fakePDBOperand{enabled: false}},
+					known_types.KAIConfigRegisteredCollectible,
+				)
+
+				Expect(deployable.Deploy(context.Background(), fakeClient, kaiConfig, kaiConfig)).To(Succeed())
+
+				err := fakeClient.Get(
+					context.Background(),
+					client.ObjectKeyFromObject(existingPDB),
+					&policyv1.PodDisruptionBudget{},
+				)
+				Expect(apierrors.IsNotFound(err)).To(BeTrue())
+			})
+		})
 	})
 
 	Describe("IsDeployed", func() {
@@ -442,6 +516,53 @@ type fakeOperand struct {
 	name                    string
 }
 
+type fakePDBOperand struct {
+	enabled bool
+}
+
+func (f *fakePDBOperand) DesiredState(
+	ctx context.Context,
+	runtimeClient client.Reader,
+	_ *kaiv1.Config,
+) ([]client.Object, error) {
+	if !f.enabled {
+		return nil, nil
+	}
+
+	pdb := &policyv1.PodDisruptionBudget{}
+	err := runtimeClient.Get(ctx, client.ObjectKey{Name: "admission", Namespace: "kai-scheduler"}, pdb)
+	if err != nil && !apierrors.IsNotFound(err) {
+		return nil, err
+	}
+	pdb.TypeMeta = metav1.TypeMeta{
+		Kind:       "PodDisruptionBudget",
+		APIVersion: policyv1.SchemeGroupVersion.String(),
+	}
+	pdb.Name = "admission"
+	pdb.Namespace = "kai-scheduler"
+	return []client.Object{pdb}, nil
+}
+
+func (f *fakePDBOperand) IsDeployed(context.Context, client.Reader) (bool, error) {
+	return true, nil
+}
+
+func (f *fakePDBOperand) IsAvailable(context.Context, client.Reader) (bool, error) {
+	return true, nil
+}
+
+func (f *fakePDBOperand) Monitor(context.Context, client.Reader, *kaiv1.Config) error {
+	return nil
+}
+
+func (f *fakePDBOperand) HasMissingDependencies(context.Context, client.Reader, *kaiv1.Config) (string, error) {
+	return "", nil
+}
+
+func (f *fakePDBOperand) Name() string {
+	return "fakePDBOperand"
+}
+
 func (f *fakeOperand) DesiredState(_ context.Context, _ client.Reader, _ *kaiv1.Config) ([]client.Object, error) {
 	return []client.Object{
 		&v1.Pod{

pkg/operator/operands/known_types/known_types.go

@@ -49,6 +49,7 @@ func init() {
 	registerCustomResourceDefinitions()
 	registerPrometheus()
 	registerVerticalPodAutoscalers()
+	registerPodDisruptionBudgets()
 }
 
 func SetupKAIConfigOwned(fn *Collectable) {