feat: add Flux integration and e2e tests (#3561)

* test: add Flux framework integration and e2e tests Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * fix: fix integration and e2e Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * fix: EOF linting issue Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * chore: add logs for E2E ClusterTrainingRuntimes in Github Actions Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * feat: add Flux runtime to Helm chart Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * fix: fix EOF linting issue Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * fix: install Flux runtime only for E2E tests Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * chore: regenerate Helm chart docs and generated assets Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * fix: retry failed deepspeed CI dependency Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * revert back to the previous commit Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * fix(e2e): install Flux runtime only during E2E setup Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * fix: fix linting issue Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * fix: fix runtimeRef for CI Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> * fix: fix flux examples runtimeref Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> --------- Signed-off-by: Amir380-A <62997533+Amir380-A@users.noreply.github.com> Signed-off-by: Amir Ibrahim <62997533+Amir380-A@users.noreply.github.com>
2026-07-27 06:55:07 +00:00 · 2026-07-02 15:21:31 +00:00
parent 6d920fd6cb
commit a9f4c62d7f
11 changed files with 463 additions and 11 deletions
@@ -10,7 +10,7 @@ metadata:
 spec:
  # Reference the pre-defined runtime by name
  runtimeRef:
-    name: flux-runtime
+    name: flux-distributed
  trainer:
    numNodes: 4
    image: ghcr.io/converged-computing/metric-lammps:latest
@@ -1,7 +1,7 @@
 apiVersion: trainer.kubeflow.org/v1alpha1
 kind: ClusterTrainingRuntime
 metadata:
-  name: flux-runtime
+  name: flux-distributed
  labels:
    trainer.kubeflow.org/framework: flux
 spec:
@@ -20,7 +20,7 @@ metadata:
 spec:
  # Reference the pre-defined runtime by name
  runtimeRef:
-    name: flux-runtime
+    name: flux-distributed
  trainer:
    numNodes: 2
    numProcPerNode: 1
@@ -15,7 +15,7 @@ metadata:
 spec:
  # Reference the pre-defined runtime by name
  runtimeRef:
-    name: flux-runtime
+    name: flux-distributed
  trainer:
    numNodes: 2
    numProcPerNode: 4
@@ -235,6 +235,10 @@ elif [ "${INSTALL_METHOD}" = "helm" ]; then
    --wait
 fi

+echo "Installing Flux runtime for E2E"
+kubectl apply --server-side \
+  -f examples/flux/flux-runtime.yaml
+
 if [ "${CLUSTER_TYPE}" = "gpu" ]; then
  # hotfix: patch CRDs to run on GPU nodes (Check #3067)
  echo "Patch CRDs to run on GPU nodes"
@@ -110,9 +110,13 @@ func (f *Flux) Validate(_ context.Context, runtimeInfo *runtime.Info, _, newJobO
 	fluxPolicy := runtimeInfo.RuntimePolicy.MLPolicySource.Flux

 	// We require at least 1 proc per node. Zero or fewer does not make sense.
-	if fluxPolicy.NumProcPerNode != nil && *fluxPolicy.NumProcPerNode < 1 {
+	numProcPerNode := fluxPolicy.NumProcPerNode
+	if newJobObj.Spec.Trainer != nil && newJobObj.Spec.Trainer.NumProcPerNode != nil {
+		numProcPerNode = newJobObj.Spec.Trainer.NumProcPerNode
+	}
+	if numProcPerNode != nil && *numProcPerNode < 1 {
 		numProcPerNodePath := field.NewPath("spec").Child("trainer").Child("numProcPerNode")
-		allErrs = append(allErrs, field.Invalid(numProcPerNodePath, *fluxPolicy.NumProcPerNode, "must be greater than or equal to 1 for Flux TrainJob"))
+		allErrs = append(allErrs, field.Invalid(numProcPerNodePath, *numProcPerNode, "must be greater than or equal to 1 for Flux TrainJob"))
 	}

 	// Iterate through Trainer's internal PodSet abstraction
@@ -28,6 +28,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	apiruntime "k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/validation/field"
 	batchv1ac "k8s.io/client-go/applyconfigurations/batch/v1"
 	corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
 	"k8s.io/klog/v2/ktesting"
@@ -227,8 +228,9 @@ func TestFlux(t *testing.T) {

 func TestValidate(t *testing.T) {
 	cases := map[string]struct {
-		info   *runtime.Info
-		newObj *trainer.TrainJob
+		info      *runtime.Info
+		newObj    *trainer.TrainJob
+		wantError field.ErrorList
 	}{
 		"valid when flux policy is nil": {
 			info:   &runtime.Info{},
@@ -244,6 +246,76 @@ func TestValidate(t *testing.T) {
 			},
 			newObj: &trainer.TrainJob{},
 		},
+		"invalid when runtime policy numProcPerNode is less than one": {
+			info: &runtime.Info{
+				RuntimePolicy: runtime.RuntimePolicy{
+					MLPolicySource: &trainer.MLPolicySource{
+						Flux: &trainer.FluxMLPolicySource{
+							NumProcPerNode: ptr.To[int32](0),
+						},
+					},
+				},
+			},
+			newObj: &trainer.TrainJob{},
+			wantError: field.ErrorList{
+				field.Invalid(
+					field.NewPath("spec").Child("trainer").Child("numProcPerNode"),
+					int32(0),
+					"must be greater than or equal to 1 for Flux TrainJob",
+				),
+			},
+		},
+		"invalid when trainJob numProcPerNode is less than one": {
+			info: &runtime.Info{
+				RuntimePolicy: runtime.RuntimePolicy{
+					MLPolicySource: &trainer.MLPolicySource{
+						Flux: &trainer.FluxMLPolicySource{
+							NumProcPerNode: ptr.To[int32](1),
+						},
+					},
+				},
+			},
+			newObj: utiltesting.MakeTrainJobWrapper(metav1.NamespaceDefault, "test-job").
+				Trainer(utiltesting.MakeTrainJobTrainerWrapper().
+					NumProcPerNode(0).
+					Obj(),
+				).
+				Obj(),
+			wantError: field.ErrorList{
+				field.Invalid(
+					field.NewPath("spec").Child("trainer").Child("numProcPerNode"),
+					int32(0),
+					"must be greater than or equal to 1 for Flux TrainJob",
+				),
+			},
+		},
+		"invalid when node podSet includes reserved flux installer init container": {
+			info: &runtime.Info{
+				RuntimePolicy: runtime.RuntimePolicy{
+					MLPolicySource: &trainer.MLPolicySource{
+						Flux: &trainer.FluxMLPolicySource{},
+					},
+				},
+				TemplateSpec: runtime.TemplateSpec{
+					PodSets: []runtime.PodSet{
+						{
+							Name: constants.Node,
+							InitContainers: []runtime.Container{
+								{Name: constants.FluxInstallerContainerName},
+							},
+						},
+					},
+				},
+			},
+			newObj: &trainer.TrainJob{},
+			wantError: field.ErrorList{
+				field.Invalid(
+					field.NewPath("spec", "trainer", "initContainers"),
+					constants.FluxInstallerContainerName,
+					"InitContainer 'flux-installer' is reserved",
+				),
+			},
+		},
 	}

 	for name, tc := range cases {
@@ -251,9 +323,9 @@ func TestValidate(t *testing.T) {
 			_, ctx := ktesting.NewTestContext(t)
 			p, _ := New(ctx, utiltesting.NewClientBuilder().Build(), nil, nil)

-			_, errs := p.(framework.CustomValidationPlugin).Validate(ctx, tc.info, nil, tc.newObj)
-			if len(errs) > 0 {
-				t.Errorf("Unexpected validation error: %v", errs)
+			_, gotError := p.(framework.CustomValidationPlugin).Validate(ctx, tc.info, nil, tc.newObj)
+			if diff := gocmp.Diff(tc.wantError, gotError); diff != "" {
+				t.Errorf("Unexpected validation errors (-want, +got): %s", diff)
 			}
 		})
 	}
@@ -1293,6 +1293,14 @@ func (m *MLPolicySourceWrapper) MPIPolicy(numProcPerNode *int32, MPImplementatio
 	return m
 }

+func (m *MLPolicySourceWrapper) FluxPolicy(numProcPerNode *int32) *MLPolicySourceWrapper {
+	if m.Flux == nil {
+		m.Flux = &trainer.FluxMLPolicySource{}
+	}
+	m.Flux.NumProcPerNode = numProcPerNode
+	return m
+}
+
 func (m *MLPolicySourceWrapper) Obj() *trainer.MLPolicySource {
 	return &m.MLPolicySource
 }
@@ -41,6 +41,7 @@ const (
 	deepSpeedRuntime = "deepspeed-distributed"
 	jaxRuntime       = "jax-distributed"
 	xgboostRuntime   = "xgboost-distributed"
+	fluxRuntime      = "flux-distributed"
 )

 //go:embed testdata/status_update.py
@@ -320,6 +321,74 @@ var _ = ginkgo.Describe("TrainJob e2e", func() {
 		})
 	})

+	ginkgo.When("Creating TrainJob to perform Flux workload", func() {
+		ginkgo.It("should create TrainJob with Flux runtime reference", func() {
+			trainJob := testingutil.MakeTrainJobWrapper(ns.Name, "e2e-test-flux").
+				RuntimeRef(
+					trainer.SchemeGroupVersion.WithKind(trainer.ClusterTrainingRuntimeKind),
+					fluxRuntime,
+				).
+				Trainer(
+					testingutil.MakeTrainJobTrainerWrapper().
+						NumNodes(2).
+						Container(
+							"ubuntu:22.04",
+							[]string{"flux", "getattr", "rank"},
+							nil,
+							corev1.ResourceList{},
+						).
+						Obj(),
+				).
+				Obj()
+
+			ginkgo.By("Create a TrainJob with flux-distributed runtime reference", func() {
+				gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
+			})
+
+			ginkgo.By("Wait for TrainJob jobs to become active", func() {
+				gomega.Eventually(func(g gomega.Gomega) {
+					gotTrainJob := &trainer.TrainJob{}
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainJob), gotTrainJob)).Should(gomega.Succeed())
+					g.Expect(gotTrainJob.Status.JobsStatus).Should(gomega.BeComparableTo([]trainer.JobStatus{
+						{
+							Name:      constants.Node,
+							Ready:     ptr.To(int32(0)),
+							Succeeded: ptr.To(int32(0)),
+							Failed:    ptr.To(int32(0)),
+							Active:    ptr.To(int32(1)),
+							Suspended: ptr.To(int32(0)),
+						},
+					}, util.SortJobsStatus))
+				}, util.TimeoutE2E, util.Interval).Should(gomega.Succeed())
+			})
+
+			ginkgo.By("Wait for TrainJob to be in Succeeded status with all jobs succeeded", func() {
+				gomega.Eventually(func(g gomega.Gomega) {
+					gotTrainJob := &trainer.TrainJob{}
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainJob), gotTrainJob)).Should(gomega.Succeed())
+					g.Expect(gotTrainJob.Status.Conditions).Should(gomega.BeComparableTo([]metav1.Condition{
+						{
+							Type:    trainer.TrainJobComplete,
+							Status:  metav1.ConditionTrue,
+							Reason:  jobsetconsts.AllJobsCompletedReason,
+							Message: jobsetconsts.AllJobsCompletedMessage,
+						},
+					}, util.IgnoreConditions))
+					g.Expect(gotTrainJob.Status.JobsStatus).Should(gomega.BeComparableTo([]trainer.JobStatus{
+						{
+							Name:      constants.Node,
+							Ready:     ptr.To(int32(0)),
+							Succeeded: ptr.To(int32(1)),
+							Failed:    ptr.To(int32(0)),
+							Active:    ptr.To(int32(0)),
+							Suspended: ptr.To(int32(0)),
+						},
+					}, util.SortJobsStatus))
+				}, util.TimeoutE2E, util.Interval).Should(gomega.Succeed())
+			})
+		})
+	})
+
 	ginkgo.When("Creating a TrainJob with RuntimePatches", func() {
 		ginkgo.It("should preserve user-provided manager fields", func() {
 			userTime := metav1.NewTime(time.Now().Add(-time.Hour).Truncate(time.Second))
@@ -2283,5 +2283,256 @@ alpha-node-0-1.alpha slots=8
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})
 		})
+		ginkgo.Context("Integration Tests for the Flux Runtime", func() {
+			var (
+				configMapKey client.ObjectKey
+				secKey       client.ObjectKey
+			)
+
+			makeFluxObjects := func(suspend bool) {
+				trainJobWrapper := testingutil.MakeTrainJobWrapper(ns.Name, "alpha").
+					RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "alpha").
+					Trainer(
+						testingutil.MakeTrainJobTrainerWrapper().
+							NumNodes(2).
+							Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+							Obj(),
+					)
+				if suspend {
+					trainJobWrapper.Suspend(true)
+				}
+				trainJob = trainJobWrapper.Obj()
+				trainJobKey = client.ObjectKeyFromObject(trainJob)
+
+				configMapKey = client.ObjectKey{
+					Name:      fmt.Sprintf("%s-flux-entrypoint", trainJobKey.Name),
+					Namespace: trainJobKey.Namespace,
+				}
+				secKey = client.ObjectKey{
+					Name:      fmt.Sprintf("%s-flux-curve", trainJobKey.Name),
+					Namespace: trainJobKey.Namespace,
+				}
+
+				trainingRuntime = testingutil.MakeTrainingRuntimeWrapper(ns.Name, "alpha").
+					RuntimeSpec(
+						testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(ns.Name, "alpha").Spec).
+							WithMLPolicy(
+								testingutil.MakeMLPolicyWrapper().
+									WithNumNodes(2).
+									WithMLPolicySource(*testingutil.MakeMLPolicySourceWrapper().
+										FluxPolicy(ptr.To[int32](1)).
+										Obj(),
+									).
+									Obj(),
+							).
+							Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+							Obj(),
+					).
+					Obj()
+			}
+
+			ginkgo.It("Should succeed to create TrainJob with Flux TrainingRuntime", func() {
+				ginkgo.By("Creating Flux TrainingRuntime and TrainJob")
+				makeFluxObjects(false)
+				gomega.Expect(k8sClient.Create(ctx, trainingRuntime)).Should(gomega.Succeed())
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainingRuntime), trainingRuntime)).Should(gomega.Succeed())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+				gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
+
+				ginkgo.By("Checking if the appropriate Flux JobSet is created")
+				gomega.Eventually(func(g gomega.Gomega) {
+					jobSet := &jobsetv1alpha2.JobSet{}
+					g.Expect(k8sClient.Get(ctx, trainJobKey, jobSet)).Should(gomega.Succeed())
+					g.Expect(jobSet.Spec.ReplicatedJobs).Should(gomega.HaveLen(3))
+
+					var nodeJob *jobsetv1alpha2.ReplicatedJob
+					for i := range jobSet.Spec.ReplicatedJobs {
+						if jobSet.Spec.ReplicatedJobs[i].Name == constants.Node {
+							nodeJob = &jobSet.Spec.ReplicatedJobs[i]
+						}
+					}
+					g.Expect(nodeJob).ShouldNot(gomega.BeNil())
+					g.Expect(nodeJob.Replicas).Should(gomega.Equal(int32(1)))
+					g.Expect(nodeJob.Template.Spec.Parallelism).Should(gomega.Equal(ptr.To[int32](2)))
+					g.Expect(nodeJob.Template.Spec.Completions).Should(gomega.Equal(ptr.To[int32](2)))
+
+					podSpec := nodeJob.Template.Spec.Template.Spec
+					var fluxInstaller *corev1.Container
+					for i := range podSpec.InitContainers {
+						if podSpec.InitContainers[i].Name == constants.FluxInstallerContainerName {
+							fluxInstaller = &podSpec.InitContainers[i]
+						}
+					}
+					g.Expect(fluxInstaller).ShouldNot(gomega.BeNil())
+					g.Expect(fluxInstaller.Image).Should(gomega.Equal(constants.FluxInstallerImage))
+					g.Expect(fluxInstaller.Command).Should(gomega.Equal([]string{"/bin/bash", "/etc/flux-config/init.sh"}))
+					g.Expect(fluxInstaller.VolumeMounts).Should(gomega.ConsistOf(
+						corev1.VolumeMount{Name: constants.FluxInstallVolumeName, MountPath: constants.FluxVolumePath},
+						corev1.VolumeMount{Name: configMapKey.Name, MountPath: constants.FluxConfigVolumeName, ReadOnly: true},
+					))
+					g.Expect(podSpec.Volumes).Should(gomega.ContainElements(
+						corev1.Volume{
+							Name: constants.FluxSpackViewVolumeName,
+							VolumeSource: corev1.VolumeSource{
+								EmptyDir: &corev1.EmptyDirVolumeSource{},
+							},
+						},
+						corev1.Volume{
+							Name: constants.FluxInstallVolumeName,
+							VolumeSource: corev1.VolumeSource{
+								EmptyDir: &corev1.EmptyDirVolumeSource{},
+							},
+						},
+						corev1.Volume{
+							Name: constants.FluxMemoryVolumeName,
+							VolumeSource: corev1.VolumeSource{
+								EmptyDir: &corev1.EmptyDirVolumeSource{
+									Medium: corev1.StorageMediumMemory,
+								},
+							},
+						},
+					))
+					var configVolume, curveVolume *corev1.Volume
+					for i := range podSpec.Volumes {
+						switch podSpec.Volumes[i].Name {
+						case configMapKey.Name:
+							configVolume = &podSpec.Volumes[i]
+						case constants.FluxCurveVolumeName:
+							curveVolume = &podSpec.Volumes[i]
+						}
+					}
+					g.Expect(configVolume).ShouldNot(gomega.BeNil())
+					g.Expect(configVolume.ConfigMap).ShouldNot(gomega.BeNil())
+					g.Expect(configVolume.ConfigMap.Name).Should(gomega.Equal(configMapKey.Name))
+					g.Expect(curveVolume).ShouldNot(gomega.BeNil())
+					g.Expect(curveVolume.Secret).ShouldNot(gomega.BeNil())
+					g.Expect(curveVolume.Secret.SecretName).Should(gomega.Equal(secKey.Name))
+
+					var nodeContainer *corev1.Container
+					for i := range podSpec.Containers {
+						if podSpec.Containers[i].Name == constants.Node {
+							nodeContainer = &podSpec.Containers[i]
+						}
+					}
+					g.Expect(nodeContainer).ShouldNot(gomega.BeNil())
+					g.Expect(nodeContainer.Image).Should(gomega.Equal("test:trainjob"))
+					g.Expect(nodeContainer.Command).Should(gomega.Equal([]string{"/bin/bash", "/etc/flux-config/entrypoint.sh", "trainjob trainjob"}))
+					g.Expect(nodeContainer.VolumeMounts).Should(gomega.ContainElements(
+						corev1.VolumeMount{Name: constants.FluxInstallVolumeName, MountPath: constants.FluxVolumePath},
+						corev1.VolumeMount{Name: constants.FluxSpackViewVolumeName, MountPath: constants.FluxSpackViewVolumePath},
+						corev1.VolumeMount{Name: configMapKey.Name, MountPath: constants.FluxConfigVolumeName, ReadOnly: true},
+						corev1.VolumeMount{Name: constants.FluxCurveVolumeName, MountPath: constants.FluxCurveVolumePath, ReadOnly: true},
+						corev1.VolumeMount{Name: constants.FluxMemoryVolumeName, MountPath: constants.FluxMemoryVolumePath, ReadOnly: true},
+					))
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+
+				ginkgo.By("Checking if the Flux ConfigMap and Secret are created")
+				gomega.Eventually(func(g gomega.Gomega) {
+					cm := &corev1.ConfigMap{}
+					g.Expect(k8sClient.Get(ctx, configMapKey, cm)).Should(gomega.Succeed())
+					g.Expect(cm.Data).Should(gomega.HaveKey("entrypoint.sh"))
+					g.Expect(cm.Data).Should(gomega.HaveKey("init.sh"))
+
+					sec := &corev1.Secret{}
+					g.Expect(k8sClient.Get(ctx, secKey, sec)).Should(gomega.Succeed())
+					g.Expect(sec.Data).Should(gomega.HaveKey("curve.cert"))
+					g.Expect(string(sec.Data["curve.cert"])).Should(gomega.ContainSubstring("public-key"))
+					g.Expect(string(sec.Data["curve.cert"])).Should(gomega.ContainSubstring("secret-key"))
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+			})
+
+			ginkgo.It("Should succeed to reconcile TrainJob conditions with Complete condition", func() {
+				ginkgo.By("Creating Flux TrainingRuntime and suspended TrainJob")
+				makeFluxObjects(true)
+				gomega.Expect(k8sClient.Create(ctx, trainingRuntime)).Should(gomega.Succeed())
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainingRuntime), trainingRuntime)).Should(gomega.Succeed())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+				gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
+
+				ginkgo.By("Checking if the JobSet was created")
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, trainJobKey, &jobsetv1alpha2.JobSet{})).Should(gomega.Succeed())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+
+				ginkgo.By("Unsuspending TrainJob")
+				gomega.Eventually(func(g gomega.Gomega) {
+					gotTrainJob := &trainer.TrainJob{}
+					g.Expect(k8sClient.Get(ctx, trainJobKey, gotTrainJob)).Should(gomega.Succeed())
+					gotTrainJob.Spec.Suspend = ptr.To(false)
+					g.Expect(k8sClient.Update(ctx, gotTrainJob)).Should(gomega.Succeed())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+
+				ginkgo.By("Updating JobSet with completed status")
+				gomega.Eventually(func(g gomega.Gomega) {
+					jobSet := &jobsetv1alpha2.JobSet{}
+					g.Expect(k8sClient.Get(ctx, trainJobKey, jobSet)).Should(gomega.Succeed())
+					meta.SetStatusCondition(&jobSet.Status.Conditions, metav1.Condition{
+						Type:    string(jobsetv1alpha2.JobSetCompleted),
+						Reason:  jobsetconsts.AllJobsCompletedReason,
+						Message: jobsetconsts.AllJobsCompletedMessage,
+						Status:  metav1.ConditionTrue,
+					})
+					jobSet.Status.ReplicatedJobsStatus = []jobsetv1alpha2.ReplicatedJobStatus{
+						{Name: constants.Node, Ready: 0, Succeeded: 1, Failed: 0, Active: 0, Suspended: 0},
+					}
+					g.Expect(k8sClient.Status().Update(ctx, jobSet)).Should(gomega.Succeed())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+
+				ginkgo.By("Checking Complete=True condition")
+				gomega.Eventually(func(g gomega.Gomega) {
+					gotTrainJob := &trainer.TrainJob{}
+					g.Expect(k8sClient.Get(ctx, trainJobKey, gotTrainJob)).Should(gomega.Succeed())
+					g.Expect(meta.IsStatusConditionTrue(gotTrainJob.Status.Conditions, trainer.TrainJobComplete)).Should(gomega.BeTrue())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+			})
+
+			ginkgo.It("Should succeed to reconcile TrainJob conditions with Failed condition", func() {
+				ginkgo.By("Creating Flux TrainingRuntime and suspended TrainJob")
+				makeFluxObjects(true)
+				gomega.Expect(k8sClient.Create(ctx, trainingRuntime)).Should(gomega.Succeed())
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(trainingRuntime), trainingRuntime)).Should(gomega.Succeed())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+				gomega.Expect(k8sClient.Create(ctx, trainJob)).Should(gomega.Succeed())
+
+				ginkgo.By("Checking if the JobSet was created")
+				gomega.Eventually(func(g gomega.Gomega) {
+					g.Expect(k8sClient.Get(ctx, trainJobKey, &jobsetv1alpha2.JobSet{})).Should(gomega.Succeed())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+
+				ginkgo.By("Unsuspending TrainJob")
+				gomega.Eventually(func(g gomega.Gomega) {
+					gotTrainJob := &trainer.TrainJob{}
+					g.Expect(k8sClient.Get(ctx, trainJobKey, gotTrainJob)).Should(gomega.Succeed())
+					gotTrainJob.Spec.Suspend = ptr.To(false)
+					g.Expect(k8sClient.Update(ctx, gotTrainJob)).Should(gomega.Succeed())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+
+				ginkgo.By("Updating JobSet with failed condition")
+				gomega.Eventually(func(g gomega.Gomega) {
+					jobSet := &jobsetv1alpha2.JobSet{}
+					g.Expect(k8sClient.Get(ctx, trainJobKey, jobSet)).Should(gomega.Succeed())
+					meta.SetStatusCondition(&jobSet.Status.Conditions, metav1.Condition{
+						Type:    string(jobsetv1alpha2.JobSetFailed),
+						Reason:  jobsetconsts.FailedJobsReason,
+						Message: jobsetconsts.FailedJobsMessage,
+						Status:  metav1.ConditionTrue,
+					})
+					jobSet.Status.ReplicatedJobsStatus = []jobsetv1alpha2.ReplicatedJobStatus{
+						{Name: constants.Node, Ready: 0, Succeeded: 0, Failed: 1, Active: 0, Suspended: 0},
+					}
+					g.Expect(k8sClient.Status().Update(ctx, jobSet)).Should(gomega.Succeed())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+
+				ginkgo.By("Checking Failed=True condition")
+				gomega.Eventually(func(g gomega.Gomega) {
+					gotTrainJob := &trainer.TrainJob{}
+					g.Expect(k8sClient.Get(ctx, trainJobKey, gotTrainJob)).Should(gomega.Succeed())
+					g.Expect(meta.IsStatusConditionTrue(gotTrainJob.Status.Conditions, trainer.TrainJobFailed)).Should(gomega.BeTrue())
+				}, util.Timeout, util.Interval).Should(gomega.Succeed())
+			})
+		})
 	})
 })
@@ -196,6 +196,50 @@ var _ = ginkgo.Describe("TrainJob Webhook", ginkgo.Ordered, func() {
 						Obj()
 				},
 				testingutil.BeForbiddenError()),
+
+			ginkgo.Entry("Should fail in creating TrainJob with Flux numProcPerNode < 1",
+				func() *trainer.TrainJob {
+					trainingRuntime.Spec.MLPolicy = &trainer.MLPolicy{
+						MLPolicySource: trainer.MLPolicySource{
+							Flux: &trainer.FluxMLPolicySource{},
+						},
+					}
+					gomega.Expect(k8sClient.Update(ctx, trainingRuntime)).To(gomega.Succeed())
+
+					return testingutil.MakeTrainJobWrapper(ns.Name, jobName).
+						RuntimeRef(
+							trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind),
+							runtimeName,
+						).
+						Trainer(&trainer.Trainer{
+							NumProcPerNode: ptr.To[int32](0),
+						}).
+						Obj()
+				},
+				testingutil.BeForbiddenError(),
+			),
+			ginkgo.Entry("Should fail in creating TrainJob with reserved flux-installer init container",
+				func() *trainer.TrainJob {
+					trainingRuntime.Spec.MLPolicy = &trainer.MLPolicy{
+						MLPolicySource: trainer.MLPolicySource{
+							Flux: &trainer.FluxMLPolicySource{},
+						},
+					}
+					trainingRuntime.Spec = testingutil.MakeTrainingRuntimeSpecWrapper(trainingRuntime.Spec).
+						InitContainer(constants.Node, constants.FluxInstallerContainerName, "ubuntu:22.04").
+						Obj()
+
+					gomega.Expect(k8sClient.Update(ctx, trainingRuntime)).To(gomega.Succeed())
+
+					return testingutil.MakeTrainJobWrapper(ns.Name, jobName).
+						RuntimeRef(
+							trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind),
+							runtimeName,
+						).
+						Obj()
+				},
+				testingutil.BeForbiddenError(),
+			),
 			ginkgo.Entry("Should fail with invalid dataset storageUri",
 				func() *trainer.TrainJob {
 					return testingutil.MakeTrainJobWrapper(ns.Name, jobName).