[slice] NodeSelector in STS webhook

slice/cmd/main.go

@@ -298,6 +298,10 @@ func setupControllers(mgr ctrl.Manager, certsReady chan struct{}, activationTime
 		setupLog.Error(err, "Unable to create webhook", "webhook", "LeaderWorkerSet")
 		os.Exit(1)
 	}
+	if err := webhooks.SetupStatefulSetWebhookWithManager(mgr); err != nil {
+		setupLog.Error(err, "Unable to create webhook", "webhook", "StatefulSet")
+		os.Exit(1)
+	}
 
 	if failedCtrl, err := controller.SetupControllers(mgr, controller.Options{
 		ActivationTimeout: activationTimeout, RetryDelayOnSliceFailure: retryDelay}); err != nil {

slice/config/webhook/manifests.yaml

@@ -61,3 +61,22 @@ webhooks:
     resources:
     - leaderworkersets
   sideEffects: None
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /mutate-apps-v1-statefulset
+  failurePolicy: Fail
+  name: mstatefulset.slice-controller.kb.io
+  rules:
+  - apiGroups:
+    - apps
+    apiVersions:
+    - v1
+    operations:
+    - CREATE
+    resources:
+    - statefulsets
+  sideEffects: None

slice/internal/util/testingjobs/leaderworkerset/wrappers.go

@@ -182,3 +182,72 @@ func (w *Wrapper) LeaderNodeAffinity(key string, values []string) *Wrapper {
 	)
 	return w
 }
+
+func (w *Wrapper) StartupPolicy(policy leaderworkersetv1.StartupPolicyType) *Wrapper {
+	w.Spec.StartupPolicy = policy
+	return w
+}
+
+func (w *Wrapper) WorkerName(name string) *Wrapper {
+	if len(w.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	w.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers[0].Name = name
+	return w
+}
+
+func (w *Wrapper) LeaderName(name string) *Wrapper {
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{}
+	}
+	if len(w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Name = name
+	return w
+}
+
+func (w *Wrapper) LeaderImage(img string) *Wrapper {
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{}
+	}
+	if len(w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Image = img
+	return w
+}
+
+func (w *Wrapper) LeaderArgs(args ...string) *Wrapper {
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{}
+	}
+	if len(w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Args = args
+	return w
+}
+
+func (w *Wrapper) LeaderLimit(resourceName corev1.ResourceName, quantity string) *Wrapper {
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate = &corev1.PodTemplateSpec{}
+	}
+	if len(w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers) == 0 {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers = []corev1.Container{{}}
+	}
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Limits == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Limits = make(corev1.ResourceList)
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Limits[resourceName] = resource.MustParse(quantity)
+	return w
+}
+
+func (w *Wrapper) LeaderRequestAndLimit(resourceName corev1.ResourceName, quantity string) *Wrapper {
+	w.LeaderLimit(resourceName, quantity)
+	if w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Requests == nil {
+		w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Requests = make(corev1.ResourceList)
+	}
+	w.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Resources.Requests[resourceName] = resource.MustParse(quantity)
+	return w
+}

slice/internal/webhooks/statefulset_webhook.go

@@ -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")
+
+	if !core.IsRelevantPodTemplateSpec(sts.Spec.Template) {
+		log.V(5).Info("Skipping non-relevant StatefulSet")
+		return nil
+	}
+
+	log.V(5).Info("Defaulting StatefulSet")
+	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
+}