webhook for STS

Author: pajakd

slice/cmd/main.go

@@ -298,8 +298,8 @@ func setupControllers(mgr ctrl.Manager, certsReady chan struct{}, activationTime
 		setupLog.Error(err, "Unable to create webhook", "webhook", "LeaderWorkerSet")
 		os.Exit(1)
 	}
-	if err := webhooks.SetupPodWebhookWithManager(mgr); err != nil {
-		setupLog.Error(err, "Unable to create webhook", "webhook", "Pod")
+	if err := webhooks.SetupStatefulSetWebhookWithManager(mgr); err != nil {
+		setupLog.Error(err, "Unable to create webhook", "webhook", "StatefulSet")
 		os.Exit(1)
 	}
 

slice/config/dev/manager_config_patch.yaml

@@ -4,7 +4,7 @@
   value: --zap-devel
 - op: add
   path: /spec/template/spec/containers/0/args/0
-  value: --zap-log-level=3
+  value: --zap-log-level=5
 - op: add
   path: /spec/template/spec/containers/0/args/0
   value: --retry-delay-on-slice-failure=0s

slice/config/webhook/kustomization.yaml

@@ -2,8 +2,5 @@ resources:
 - manifests.yaml
 - service.yaml
 
-patches:
-- path: pod_webhook_patch.yaml
-
 configurations:
 - kustomizeconfig.yaml

slice/config/webhook/manifests.yaml

@@ -67,16 +67,16 @@ webhooks:
     service:
       name: webhook-service
       namespace: system
-      path: /mutate--v1-pod
+      path: /mutate-apps-v1-statefulset
   failurePolicy: Fail
-  name: mpod.slice-controller.kb.io
+  name: mstatefulset.slice-controller.kb.io
   rules:
   - apiGroups:
-    - ""
+    - apps
     apiVersions:
     - v1
     operations:
     - CREATE
     resources:
-    - pods
+    - statefulsets
   sideEffects: None

slice/config/webhook/pod_webhook_patch.yaml

@@ -0,0 +1,59 @@
+/*
+Copyright The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package webhooks
+
+import (
+	"context"
+
+	appsv1 "k8s.io/api/apps/v1"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
+
+	"tpu-slice-controller/internal/core"
+)
+
+type StatefulSetWebhook struct{}
+
+func SetupStatefulSetWebhookWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewWebhookManagedBy(mgr, &appsv1.StatefulSet{}).
+		WithDefaulter(&StatefulSetWebhook{}).
+		Complete()
+}
+
+// +kubebuilder:webhook:path=/mutate-apps-v1-statefulset,mutating=true,failurePolicy=fail,sideEffects=None,groups=apps,resources=statefulsets,verbs=create,versions=v1,name=mstatefulset.slice-controller.kb.io,admissionReviewVersions=v1
+
+var _ admission.Defaulter[*appsv1.StatefulSet] = &StatefulSetWebhook{}
+
+func (r *StatefulSetWebhook) Default(ctx context.Context, sts *appsv1.StatefulSet) error {
+	log := ctrl.LoggerFrom(ctx).WithName("statefulset-webhook")
+	log.V(5).Info("Defaulting StatefulSet")
+
+	if !core.IsRelevantPodTemplateSpec(sts.Spec.Template) {
+		log.V(5).Info("Skipping non-TPUv7 StatefulSet")
+		return nil
+	}
+
+	tpuTopology := core.GetTPUTopology(sts.Spec.Template)
+	if sts.Spec.Template.Spec.NodeSelector == nil {
+		sts.Spec.Template.Spec.NodeSelector = make(map[string]string)
+	}
+	if _, ok := sts.Spec.Template.Spec.NodeSelector[core.TPUTopologyAnnotation]; !ok {
+		sts.Spec.Template.Spec.NodeSelector[core.TPUTopologyAnnotation] = tpuTopology
+	}
+
+	return nil
+}