diff --git a/internal/controller/application_controller.go b/internal/controller/application_controller.go index eaf2a87c..de0097a4 100644 --- a/internal/controller/application_controller.go +++ b/internal/controller/application_controller.go @@ -649,9 +649,9 @@ func (r *ApplicationReconciler) reconcileJobs(ctx context.Context, app *wandbv2. if jobToReconcile.Labels == nil { jobToReconcile.Labels = make(map[string]string) } - jobToReconcile.Labels["app.kubernetes.io/name"] = app.Name - jobToReconcile.Labels["app.kubernetes.io/instance"] = app.Namespace - jobToReconcile.Labels["app.kubernetes.io/managed-by"] = "application-controller" + jobToReconcile.Labels[common.AppNameLabel] = app.Name + jobToReconcile.Labels[common.AppInstanceLabel] = app.Namespace + jobToReconcile.Labels[common.AppManagedByLabel] = common.ManagedByWandbOperator if err = controllerutil.SetControllerReference(app, jobToReconcile, r.Scheme); err != nil { return err @@ -698,12 +698,12 @@ func (r *ApplicationReconciler) deleteJobs(ctx context.Context, app *wandbv2.App logger.Info("Deleting Jobs", "Application", app.Name) jobList := &batchv1.JobList{} + // Omit managed-by so pre-unification Jobs are still matched. listOpts := []client.ListOption{ client.InNamespace(app.Namespace), client.MatchingLabels{ - "app.kubernetes.io/name": app.Name, - "app.kubernetes.io/instance": app.Namespace, - "app.kubernetes.io/managed-by": "application-controller", + common.AppNameLabel: app.Name, + common.AppInstanceLabel: app.Namespace, }, } @@ -755,9 +755,9 @@ func (r *ApplicationReconciler) reconcileCronJobs(ctx context.Context, app *wand if cronJobToReconcile.Labels == nil { cronJobToReconcile.Labels = make(map[string]string) } - cronJobToReconcile.Labels["app.kubernetes.io/name"] = app.Name - cronJobToReconcile.Labels["app.kubernetes.io/instance"] = app.Namespace - cronJobToReconcile.Labels["app.kubernetes.io/managed-by"] = "application-controller" + cronJobToReconcile.Labels[common.AppNameLabel] = app.Name + cronJobToReconcile.Labels[common.AppInstanceLabel] = app.Namespace + cronJobToReconcile.Labels[common.AppManagedByLabel] = common.ManagedByWandbOperator if err = controllerutil.SetControllerReference(app, cronJobToReconcile, r.Scheme); err != nil { return err @@ -793,12 +793,12 @@ func (r *ApplicationReconciler) deleteCronJobs(ctx context.Context, app *wandbv2 logger.Info("Deleting CronJobs", "Application", app.Name) cronJobList := &batchv1.CronJobList{} + // Omit managed-by so pre-unification CronJobs are still matched. listOpts := []client.ListOption{ client.InNamespace(app.Namespace), client.MatchingLabels{ - "app.kubernetes.io/name": app.Name, - "app.kubernetes.io/instance": app.Namespace, - "app.kubernetes.io/managed-by": "application-controller", + common.AppNameLabel: app.Name, + common.AppInstanceLabel: app.Namespace, }, } diff --git a/internal/controller/common/labels.go b/internal/controller/common/labels.go index 245a830b..03ab448a 100644 --- a/internal/controller/common/labels.go +++ b/internal/controller/common/labels.go @@ -2,16 +2,24 @@ package common import ( apiv2 "github.com/wandb/operator/api/v2" + "github.com/wandb/operator/pkg/utils" ) const ( WandbNameLabel = "weightsandbiases.apps.wandb.com/name" WandbNamespaceLabel = "weightsandbiases.apps.wandb.com/namespace" WandbComponentLabel = "weightsandbiases.apps.wandb.com/component" + + AppNameLabel = "app.kubernetes.io/name" + AppInstanceLabel = "app.kubernetes.io/instance" + AppPartOfLabel = "app.kubernetes.io/part-of" + AppManagedByLabel = "app.kubernetes.io/managed-by" + + PartOfWandb = "wandb" + ManagedByWandbOperator = "wandb-operator" ) -// HasAllLabelKeys reports whether existing contains every key present in desired, -// regardless of value. +// HasAllLabelKeys reports whether existing has every key in desired. func HasAllLabelKeys(existing, desired map[string]string) bool { for k := range desired { if _, ok := existing[k]; !ok { @@ -21,8 +29,7 @@ func HasAllLabelKeys(existing, desired map[string]string) bool { return true } -// BuildWandbLabels returns the standard wandb labels for resources managed -// on behalf of the given WeightsAndBiases CR. +// BuildWandbLabels returns ownership labels; componentName is the service id. func BuildWandbLabels(wandb *apiv2.WeightsAndBiases, componentName string) map[string]string { return map[string]string{ WandbNameLabel: wandb.Name, @@ -30,3 +37,20 @@ func BuildWandbLabels(wandb *apiv2.WeightsAndBiases, componentName string) map[s WandbComponentLabel: componentName, } } + +// BuildIdentityLabels returns name, part-of, and managed-by (not instance). +func BuildIdentityLabels(serviceName string) map[string]string { + return map[string]string{ + AppNameLabel: serviceName, + AppPartOfLabel: PartOfWandb, + AppManagedByLabel: ManagedByWandbOperator, + } +} + +// BuildApplicationLabels merges ownership+identity for MetaTemplate only. +func BuildApplicationLabels(wandb *apiv2.WeightsAndBiases, serviceName string) map[string]string { + return utils.MergeMapsStringString( + BuildWandbLabels(wandb, serviceName), + BuildIdentityLabels(serviceName), + ) +} diff --git a/internal/controller/common/labels_test.go b/internal/controller/common/labels_test.go new file mode 100644 index 00000000..8660458d --- /dev/null +++ b/internal/controller/common/labels_test.go @@ -0,0 +1,51 @@ +package common + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + apiv2 "github.com/wandb/operator/api/v2" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) + +var _ = Describe("BuildIdentityLabels", func() { + It("sets name, part-of, and managed-by without instance", func() { + Expect(BuildIdentityLabels("api")).To(Equal(map[string]string{ + AppNameLabel: "api", + AppPartOfLabel: PartOfWandb, + AppManagedByLabel: ManagedByWandbOperator, + })) + }) +}) + +var _ = Describe("BuildApplicationLabels", func() { + It("merges ownership and identity labels", func() { + wandb := &apiv2.WeightsAndBiases{ + ObjectMeta: metav1.ObjectMeta{Name: "wandb", Namespace: "wandb-ns"}, + } + + Expect(BuildApplicationLabels(wandb, "api")).To(Equal(map[string]string{ + WandbNameLabel: "wandb", + WandbNamespaceLabel: "wandb-ns", + WandbComponentLabel: "api", + AppNameLabel: "api", + AppPartOfLabel: PartOfWandb, + AppManagedByLabel: ManagedByWandbOperator, + })) + }) +}) + +var _ = Describe("HasAllLabelKeys", func() { + It("returns true when every desired key is present", func() { + Expect(HasAllLabelKeys( + map[string]string{"a": "1", "b": "2", "c": "3"}, + map[string]string{"a": "x", "b": "y"}, + )).To(BeTrue()) + }) + + It("returns false when a desired key is missing", func() { + Expect(HasAllLabelKeys( + map[string]string{"a": "1"}, + map[string]string{"a": "1", "b": "2"}, + )).To(BeFalse()) + }) +}) diff --git a/internal/controller/infra/managed/kafka/bufstream/spec.go b/internal/controller/infra/managed/kafka/bufstream/spec.go index fa49171a..5c816fc6 100644 --- a/internal/controller/infra/managed/kafka/bufstream/spec.go +++ b/internal/controller/infra/managed/kafka/bufstream/spec.go @@ -289,6 +289,7 @@ func ToEtcdApplication( ) (*apiv2.Application, error) { infraSpec := wandb.Spec.Kafka.ManagedKafka labels := BuildWandbKafkaLabels(wandb) + metaLabels := utils.MergeMapsStringString(labels, common.BuildIdentityLabels(nsnBuilder.EtcdName())) storageSize := infraSpec.StorageSize if storageSize == "" { @@ -330,7 +331,7 @@ func ToEtcdApplication( Replicas: ptr.To(int32(EtcdReplicas)), ServiceName: nsnBuilder.EtcdName(), MetaTemplate: metav1.ObjectMeta{ - Labels: labels, + Labels: metaLabels, }, PodTemplate: corev1.PodTemplateSpec{ Spec: corev1.PodSpec{ @@ -545,6 +546,7 @@ func ToBufstreamApplication( ) (*apiv2.Application, error) { infraSpec := wandb.Spec.Kafka.ManagedKafka labels := BuildWandbKafkaLabels(wandb) + metaLabels := utils.MergeMapsStringString(labels, common.BuildIdentityLabels(nsnBuilder.BufstreamName())) replicas := effectiveBufstreamReplicas(infraSpec.Replicas) @@ -582,7 +584,7 @@ func ToBufstreamApplication( Kind: "Deployment", Replicas: ptr.To(replicas), MetaTemplate: metav1.ObjectMeta{ - Labels: labels, + Labels: metaLabels, }, PodTemplate: corev1.PodTemplateSpec{ Spec: corev1.PodSpec{ diff --git a/internal/controller/infra/managed/kafka/bufstream/spec_test.go b/internal/controller/infra/managed/kafka/bufstream/spec_test.go index 04551b09..67ad91f2 100644 --- a/internal/controller/infra/managed/kafka/bufstream/spec_test.go +++ b/internal/controller/infra/managed/kafka/bufstream/spec_test.go @@ -5,6 +5,7 @@ import ( "github.com/stretchr/testify/require" apiv2 "github.com/wandb/operator/api/v2" + "github.com/wandb/operator/internal/controller/common" "github.com/wandb/operator/internal/controller/infra/objectstore" "github.com/wandb/operator/pkg/utils" "github.com/wandb/operator/pkg/wandb/manifest" @@ -63,6 +64,11 @@ func TestToEtcdApplication(t *testing.T) { requireKafkaPodSecurityContext(t, app.Spec.PodTemplate.Spec.SecurityContext) requireKafkaContainerSecurityContext(t, app.Spec.PodTemplate.Spec.Containers[0].SecurityContext) require.NotNil(t, app.Spec.ServiceTemplate) + require.Equal(t, BuildWandbKafkaLabels(wandb), app.Labels) + require.Equal(t, "wandb-kafka-etcd", app.Spec.MetaTemplate.Labels[common.AppNameLabel]) + require.Equal(t, common.PartOfWandb, app.Spec.MetaTemplate.Labels[common.AppPartOfLabel]) + require.Equal(t, common.ManagedByWandbOperator, app.Spec.MetaTemplate.Labels[common.AppManagedByLabel]) + require.Equal(t, KafkaModuleName, app.Spec.MetaTemplate.Labels[common.WandbComponentLabel]) } func TestToEtcdApplicationHA(t *testing.T) { @@ -130,6 +136,11 @@ func TestToBufstreamApplication(t *testing.T) { require.Equal(t, "wandb-kafka", app.Name) require.Equal(t, "Deployment", app.Spec.Kind) require.NotNil(t, app.Spec.Replicas) + require.Equal(t, BuildWandbKafkaLabels(wandb), app.Labels) + require.Equal(t, "wandb-kafka", app.Spec.MetaTemplate.Labels[common.AppNameLabel]) + require.Equal(t, common.PartOfWandb, app.Spec.MetaTemplate.Labels[common.AppPartOfLabel]) + require.Equal(t, common.ManagedByWandbOperator, app.Spec.MetaTemplate.Labels[common.AppManagedByLabel]) + require.Equal(t, KafkaModuleName, app.Spec.MetaTemplate.Labels[common.WandbComponentLabel]) require.Len(t, app.Spec.PodTemplate.Spec.InitContainers, 1) ensureBucket := app.Spec.PodTemplate.Spec.InitContainers[0] require.Equal(t, "ensure-bucket", ensureBucket.Name) diff --git a/internal/controller/reconciler/reconcile_v2.go b/internal/controller/reconciler/reconcile_v2.go index c5bdf3a4..d3bb10f9 100644 --- a/internal/controller/reconciler/reconcile_v2.go +++ b/internal/controller/reconciler/reconcile_v2.go @@ -674,6 +674,12 @@ func reconcileApplications( application.Spec.PodTemplate.Spec.Tolerations = *wandb.Spec.Tolerations setCustomCACertsChecksumAnnotation(&application.Spec.PodTemplate, caChecksum) + // MetaTemplate only: ObjectMeta ownership labels would skip pruning. + application.Spec.MetaTemplate.Labels = oputils.MergeMapsStringString( + application.Spec.MetaTemplate.Labels, + common.BuildApplicationLabels(wandb, app.Name), + ) + application.Spec.HpaTemplate = ResolveAutoscaling(app, wandb) // Set shared service account for all W&B applications