fix(webhook): set the workload's owner Reference to the pod's root owner (#158)

- Enhance pod_webhook_test.go to verify workload owner reference propagation in tests

Author: 0x5457

internal/utils/owner_ref_utils.go

@@ -0,0 +1,37 @@
+package utils
+
+import (
+	context "context"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+)
+
+// FindRootOwnerReference recursively finds the root owner reference for a given object (e.g. Pod).
+func FindRootOwnerReference(ctx context.Context, c client.Client, namespace string, obj metav1.Object) (*metav1.OwnerReference, error) {
+	current := obj
+	for {
+		owners := current.GetOwnerReferences()
+		if len(owners) == 0 {
+			return nil, nil // no owner, this is root
+		}
+		ownerRef := owners[0]
+		// Try to get the owner object as unstructured
+		unObj := &unstructured.Unstructured{}
+		unObj.SetAPIVersion(ownerRef.APIVersion)
+		unObj.SetKind(ownerRef.Kind)
+		key := client.ObjectKey{Name: ownerRef.Name, Namespace: namespace}
+		err := c.Get(ctx, key, unObj)
+		if err != nil {
+			// If not found, treat this ownerRef as root
+			return &ownerRef, nil
+		}
+		// Cast back to metav1.Object if possible
+		if metaObj, ok := any(unObj).(metav1.Object); ok {
+			current = metaObj
+		} else {
+			return &ownerRef, nil
+		}
+	}
+}

internal/webhook/v1/pod_webhook.go

@@ -160,6 +160,12 @@ func (m *TensorFusionPodMutator) createOrUpdateWorkload(ctx context.Context, pod
 		if !errors.IsNotFound(err) {
 			return fmt.Errorf("failed to get workload: %w", err)
 		}
+		// find root owner references of pod
+		rootOwnerRef, err := utils.FindRootOwnerReference(ctx, m.Client, pod.Namespace, pod)
+		if err != nil {
+			return fmt.Errorf("failed to find root owner reference: %w", err)
+		}
+
 		// Create a new workload
 		replicas := tfInfo.Replicas
 		workload = &tfv1.TensorFusionWorkload{
@@ -176,6 +182,10 @@ func (m *TensorFusionPodMutator) createOrUpdateWorkload(ctx context.Context, pod
 			},
 		}
 
+		if rootOwnerRef != nil {
+			workload.OwnerReferences = []metav1.OwnerReference{*rootOwnerRef}
+		}
+
 		if err := m.Client.Create(ctx, workload); err != nil {
 			return fmt.Errorf("failed to create workload: %w", err)
 		}

internal/webhook/v1/pod_webhook_test.go

@@ -33,6 +33,7 @@ import (
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 )
@@ -133,6 +134,13 @@ var _ = Describe("TensorFusionPodMutator", func() {
 					Labels: map[string]string{
 						constants.TensorFusionEnabledLabelKey: "true",
 					},
+					OwnerReferences: []metav1.OwnerReference{{
+						APIVersion: "apps/v1",
+						Kind:       "Deployment",
+						Name:       "owner",
+						UID:        "owner-uid",
+						Controller: ptr.To(true),
+					}},
 					Annotations: map[string]string{
 						constants.GpuPoolKey:                "mock",
 						constants.WorkloadProfileAnnotation: "test-profile-handle",
@@ -179,6 +187,15 @@ var _ = Describe("TensorFusionPodMutator", func() {
 			resp := mutator.Handle(ctx, req)
 			Expect(resp.Allowed).To(BeTrue())
 			Expect(resp.Patches).NotTo(BeEmpty())
+
+			// Check workload created
+			workload := &tfv1.TensorFusionWorkload{}
+			err = k8sClient.Get(ctx, client.ObjectKey{Name: "test-workload", Namespace: "default"}, workload)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(*workload.Spec.Replicas).To(Equal(int32(1)))
+			// check workload owner reference
+			Expect(workload.OwnerReferences).To(HaveLen(1))
+			Expect(workload.OwnerReferences[0].Name).To(Equal("owner"))
 		})
 
 		It("should handle pods without TF requirements", func() {