fix(jobset): address CodeRabbit review: compare fully-built spec and improve test assertions

- Move RHOAIENG-48867 delete guard to after the full JobSet build so that
  all builder mutations (Initializer, Trainer, PodLabels, PodAnnotations)
  are reflected in the comparison; this replaces the trainerImage workaround
- Expand replicatedJobsSpecChanged() to detect removed containers and
  initContainer image changes (bi-directional comparison for both lists)
- Assert DeletePropagationBackground in unit test using client interceptor
- Add Eventually wait in integration test AfterEach to ensure resources are
  fully removed before the next test runs

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Ambient Code Bot

pkg/runtime/framework/plugins/jobset/jobset.go

@@ -280,6 +280,32 @@ func (j *JobSet) Build(ctx context.Context, info *runtime.Info, trainJob *traine
 		}
 	}
 
+	// Init the JobSet apply configuration from the runtime template spec.
+	// The full build must happen before the spec comparison below (RHOAIENG-48867)
+	// so that all builder mutations (Initializer, Trainer, PodLabels, PodAnnotations)
+	// are reflected in the desired state used for the comparison.
+	jobSetBuilder := NewBuilder(jobsetv1alpha2ac.JobSet(trainJob.Name, trainJob.Namespace).
+		WithLabels(maps.Clone(info.Labels)).
+		WithAnnotations(maps.Clone(info.Annotations)).
+		WithSpec(jobSetSpec))
+
+	// TODO (andreyvelich): Refactor the builder with wrappers for PodSpec.
+	// TODO: Once we remove deprecated runtime.Info.Trainer, we should remove JobSet Builder with DeprecatedTrainer().
+	jobSet := jobSetBuilder.
+		Initializer(trainJob).
+		Trainer(info, trainJob).
+		PodLabels(info.Scheduler.PodLabels).
+		PodAnnotations(info.Scheduler.PodAnnotations).
+		Suspend(trainJob.Spec.Suspend).
+		Build().
+		WithOwnerReferences(metav1ac.OwnerReference().
+			WithAPIVersion(trainer.GroupVersion.String()).
+			WithKind(trainer.TrainJobKind).
+			WithName(trainJob.Name).
+			WithUID(trainJob.UID).
+			WithController(true).
+			WithBlockOwnerDeletion(true))
+
 	// RHOAIENG-48867: If the existing JobSet is suspended and the TrainJob has been
 	// resumed (Kueue admitted it, Suspend=false) but spec.replicatedJobs has changed
 	// (e.g., container images updated during a ClusterTrainingRuntime upgrade), delete
@@ -292,18 +318,13 @@ func (j *JobSet) Build(ctx context.Context, info *runtime.Info, trainJob *traine
 	// while suspended). Delete-recreate is only needed when the TrainJob transitions
 	// from suspended to running and the runtime spec has changed.
 	//
-	// trainerImage corrects the node container comparison: jobSetSpec holds the raw
-	// runtime template image for the node container at this point (the builder applies
-	// TrainJob.Spec.Trainer.Image override later). Without this, every normal admission
-	// would incorrectly detect a node image change and trigger an unnecessary deletion.
-	var trainerImage *string
-	if trainJob.Spec.Trainer != nil {
-		trainerImage = trainJob.Spec.Trainer.Image
-	}
+	// The comparison uses the fully-built desired JobSet spec so that all builder
+	// mutations (Initializer, Trainer, PodLabels, PodAnnotations) are accounted for,
+	// covering container and initContainer images, added and removed containers.
 	if oldJobSet != nil &&
 		ptr.Deref(oldJobSet.Spec.Suspend, false) &&
 		!ptr.Deref(trainJob.Spec.Suspend, false) &&
-		replicatedJobsSpecChanged(oldJobSet.Spec.ReplicatedJobs, jobSetSpec.ReplicatedJobs, trainerImage) {
+		replicatedJobsSpecChanged(oldJobSet.Spec.ReplicatedJobs, jobSet.Spec.ReplicatedJobs) {
 		j.logger.Info("Deleting stale suspended JobSet: spec.replicatedJobs changed post-upgrade, will recreate",
 			"jobSet", client.ObjectKeyFromObject(trainJob))
 		// Use background propagation: the JobSet is suspended so there are no
@@ -318,29 +339,6 @@ func (j *JobSet) Build(ctx context.Context, info *runtime.Info, trainJob *traine
 		return nil, nil
 	}
 
-	// Init the JobSet apply configuration from the runtime template spec
-	jobSetBuilder := NewBuilder(jobsetv1alpha2ac.JobSet(trainJob.Name, trainJob.Namespace).
-		WithLabels(maps.Clone(info.Labels)).
-		WithAnnotations(maps.Clone(info.Annotations)).
-		WithSpec(jobSetSpec))
-
-	// TODO (andreyvelich): Refactor the builder with wrappers for PodSpec.
-	// TODO: Once we remove deprecated runtime.Info.Trainer, we should remove JobSet Builder with DeprecatedTrainer().
-	jobSet := jobSetBuilder.
-		Initializer(trainJob).
-		Trainer(info, trainJob).
-		PodLabels(info.Scheduler.PodLabels).
-		PodAnnotations(info.Scheduler.PodAnnotations).
-		Suspend(trainJob.Spec.Suspend).
-		Build().
-		WithOwnerReferences(metav1ac.OwnerReference().
-			WithAPIVersion(trainer.GroupVersion.String()).
-			WithKind(trainer.TrainJobKind).
-			WithName(trainJob.Name).
-			WithUID(trainJob.UID).
-			WithController(true).
-			WithBlockOwnerDeletion(true))
-
 	return []apiruntime.ApplyConfiguration{jobSet}, nil
 }
 
@@ -350,15 +348,12 @@ func (j *JobSet) Build(ctx context.Context, info *runtime.Info, trainJob *traine
 // scenario where a ClusterTrainingRuntime change updates the runtime container image.
 // RHOAIENG-48867
 //
-// trainerImage is the TrainJob.Spec.Trainer.Image override (may be nil). The
-// builder applies this override to the node container after this function is
-// called, so the raw jobSetSpec still holds the runtime template image for that
-// container. Passing trainerImage here ensures we compare the effective image
-// that will end up in the JobSet, avoiding false-positive deletions.
+// desired must be the fully-built desired spec — all builder mutations (Initializer,
+// Trainer, PodLabels, PodAnnotations) must be applied before calling this function
+// so that the comparison reflects the effective state that will be applied to the cluster.
 func replicatedJobsSpecChanged(
 	existing []jobsetv1alpha2.ReplicatedJob,
 	desired []jobsetv1alpha2ac.ReplicatedJobApplyConfiguration,
-	trainerImage *string,
 ) bool {
 	if len(existing) != len(desired) {
 		return true
@@ -380,23 +375,48 @@ func replicatedJobsSpecChanged(
 			d.Template.Spec.Template == nil || d.Template.Spec.Template.Spec == nil {
 			continue
 		}
-		existingImages := make(map[string]string, len(e.Template.Spec.Template.Spec.Containers))
-		for _, c := range e.Template.Spec.Template.Spec.Containers {
+		dSpec := d.Template.Spec.Template.Spec
+		eSpec := &e.Template.Spec.Template.Spec
+
+		// Check containers: detect changed/added images (desired → existing direction).
+		existingImages := make(map[string]string, len(eSpec.Containers))
+		for _, c := range eSpec.Containers {
 			existingImages[c.Name] = c.Image
 		}
-		for _, c := range d.Template.Spec.Template.Spec.Containers {
+		desiredContainerNames := sets.New[string]()
+		for _, c := range dSpec.Containers {
+			if c.Name == nil || c.Image == nil {
+				continue
+			}
+			desiredContainerNames.Insert(*c.Name)
+			if img, found := existingImages[*c.Name]; !found || img != *c.Image {
+				return true
+			}
+		}
+		// Detect containers removed from desired (existing → desired direction).
+		for _, c := range eSpec.Containers {
+			if !desiredContainerNames.Has(c.Name) {
+				return true
+			}
+		}
+
+		// Repeat the same checks for initContainers.
+		existingInitImages := make(map[string]string, len(eSpec.InitContainers))
+		for _, c := range eSpec.InitContainers {
+			existingInitImages[c.Name] = c.Image
+		}
+		desiredInitNames := sets.New[string]()
+		for _, c := range dSpec.InitContainers {
 			if c.Name == nil || c.Image == nil {
 				continue
 			}
-			// For the node (trainer) container, use the TrainJob's trainer image
-			// override as the effective desired image. The builder applies this
-			// override after replicatedJobsSpecChanged is called, so jobSetSpec
-			// still has the raw runtime template image at this point.
-			effectiveImage := c.Image
-			if *c.Name == constants.Node && trainerImage != nil {
-				effectiveImage = trainerImage
+			desiredInitNames.Insert(*c.Name)
+			if img, found := existingInitImages[*c.Name]; !found || img != *c.Image {
+				return true
 			}
-			if img, found := existingImages[*c.Name]; !found || img != *effectiveImage {
+		}
+		for _, c := range eSpec.InitContainers {
+			if !desiredInitNames.Has(c.Name) {
 				return true
 			}
 		}

pkg/runtime/framework/plugins/jobset/jobset_test.go

@@ -122,21 +122,23 @@ func TestBuild_ImmutableJobSetUpgrade(t *testing.T) {
 	}
 
 	cases := map[string]struct {
-		trainJob          *trainer.TrainJob
-		existingJobSet    *jobsetv1alpha2.JobSet
-		info              *runtime.Info
-		wantNilResult     bool // true if Build() should return no apply configs
-		wantJobSetDeleted bool // true if the existing JobSet should be deleted
-		wantError         error
+		trainJob               *trainer.TrainJob
+		existingJobSet         *jobsetv1alpha2.JobSet
+		info                   *runtime.Info
+		wantNilResult          bool // true if Build() should return no apply configs
+		wantJobSetDeleted      bool // true if the existing JobSet should be deleted
+		wantBackgroundDeletion bool // true if the delete should use Background propagation
+		wantError              error
 	}{
 		"suspended JobSet with changed image is deleted on upgrade (RHOAIENG-48867)": {
 			trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
 				Suspend(false).
 				Obj(),
-			existingJobSet:    makeSuspendedJobSet(oldImage),
-			info:              makeInfo(newImage),
-			wantNilResult:     true,
-			wantJobSetDeleted: true,
+			existingJobSet:         makeSuspendedJobSet(oldImage),
+			info:                   makeInfo(newImage),
+			wantNilResult:          true,
+			wantJobSetDeleted:      true,
+			wantBackgroundDeletion: true,
 		},
 		"suspended JobSet with unchanged image is not deleted": {
 			trainJob: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test").
@@ -191,10 +193,24 @@ func TestBuild_ImmutableJobSetUpgrade(t *testing.T) {
 			ctx, cancel = context.WithCancel(ctx)
 			t.Cleanup(cancel)
 
+			var capturedPropagation *metav1.DeletionPropagation
 			clientBuilder := utiltesting.NewClientBuilder()
 			if tc.existingJobSet != nil {
 				clientBuilder = clientBuilder.WithObjects(tc.existingJobSet)
 			}
+			// Intercept Delete calls to capture the deletion propagation policy.
+			clientBuilder = clientBuilder.WithInterceptorFuncs(interceptor.Funcs{
+				Delete: func(ctx context.Context, cli client.WithWatch, obj client.Object, opts ...client.DeleteOption) error {
+					if _, ok := obj.(*jobsetv1alpha2.JobSet); ok {
+						deleteOpts := &client.DeleteOptions{}
+						for _, o := range opts {
+							o.ApplyToDelete(deleteOpts)
+						}
+						capturedPropagation = deleteOpts.PropagationPolicy
+					}
+					return cli.Delete(ctx, obj, opts...)
+				},
+			})
 			cli := clientBuilder.Build()
 
 			p, err := New(ctx, cli, nil)
@@ -224,6 +240,14 @@ func TestBuild_ImmutableJobSetUpgrade(t *testing.T) {
 					t.Errorf("Expected existing JobSet to remain, but it was deleted")
 				}
 			}
+			// Verify deletion propagation policy for cases that expect background deletion.
+			if tc.wantBackgroundDeletion {
+				if capturedPropagation == nil {
+					t.Errorf("Expected Delete to be called with Background propagation, but Delete was not called")
+				} else if *capturedPropagation != metav1.DeletePropagationBackground {
+					t.Errorf("Expected DeletePropagationBackground, got %v", *capturedPropagation)
+				}
+			}
 		})
 	}
 }

test/integration/controller/trainjob_controller_test.go

@@ -1455,6 +1455,17 @@ alpha-node-0-1.alpha slots=8
 		ginkgo.AfterEach(func() {
 			gomega.Expect(k8sClient.DeleteAllOf(ctx, &trainer.TrainJob{}, client.InNamespace(ns.Name))).Should(gomega.Succeed())
 			gomega.Expect(k8sClient.DeleteAllOf(ctx, &trainer.ClusterTrainingRuntime{})).Should(gomega.Succeed())
+			// Wait for resources to be fully removed to avoid name collisions in subsequent tests.
+			gomega.Eventually(func(g gomega.Gomega) {
+				trainJobList := &trainer.TrainJobList{}
+				g.Expect(k8sClient.List(ctx, trainJobList, client.InNamespace(ns.Name))).Should(gomega.Succeed())
+				g.Expect(trainJobList.Items).Should(gomega.BeEmpty())
+			}, util.Timeout, util.Interval).Should(gomega.Succeed())
+			gomega.Eventually(func(g gomega.Gomega) {
+				runtimeList := &trainer.ClusterTrainingRuntimeList{}
+				g.Expect(k8sClient.List(ctx, runtimeList)).Should(gomega.Succeed())
+				g.Expect(runtimeList.Items).Should(gomega.BeEmpty())
+			}, util.Timeout, util.Interval).Should(gomega.Succeed())
 		})
 
 		ginkgo.BeforeEach(func() {