feat(api): Fix immutability of the TrainJob APIs by andreyvelich · Pull Request #3157 · kubeflow/trainer

andreyvelich · 2026-02-03T01:53:33Z

As discussed during our Trainer calls, we want to properly enforce immutability for the TrainJob API fields. It’s best to address this now, before the API graduates to beta, when we can't introduce breaking changes.

Currently, when user tries to update .spec.trainer.image, the JobSet controller fails due to immutability of this field.

I would prefer we gradually introduce mutability when we get more use-cases from users.

/hold for review
/assign @tenzen-y @astefanutti @akshaychitneni @kaisoz
cc @mimowo @kannon92

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

google-oss-prow · 2026-02-03T01:53:39Z

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please ask for approval from andreyvelich. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

Copilot

Pull request overview

This PR aims to enforce immutability rules for TrainJob spec fields (primarily via CRD-level XValidation) ahead of the API’s beta graduation.

Changes:

Add/extend CRD XValidation immutability rules for TrainJobSpec fields (including conditional mutability for podTemplateOverrides).
Update integration tests to cover immutability of initializer, trainer, and conditional updates to podTemplateOverrides.
Adjust OpenAPI/CRD schema details for PodTemplateOverride.targetJobs list semantics.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
`pkg/apis/trainer/v1alpha1/trainjob_types.go`	Adds `XValidation` rules to enforce immutability and conditional mutability for `podTemplateOverrides`.
`test/integration/webhooks/trainjob_test.go`	Adds webhook-focused immutability tests for `initializer`, `trainer`, and `podTemplateOverrides`.
`test/integration/controller/trainjob_controller_test.go`	Updates controller integration test to validate `podTemplateOverrides` behavior across suspend/unsuspend.
`manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml`	Propagates the new schema validations into the base CRD manifest.
`charts/kubeflow-trainer/crds/trainer.kubeflow.org_trainjobs.yaml`	Propagates the new schema validations into the Helm CRD manifest.
`pkg/apis/trainer/v1alpha1/zz_generated.openapi.go`	Updates OpenAPI schema (notably list semantics for `targetJobs`) and descriptions.
`api/openapi-spec/swagger.json`	Updates published OpenAPI spec to match schema changes.
`api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_train_job_spec.py`	Updates generated Python client model descriptions to match OpenAPI changes.
`pkg/client/applyconfiguration/trainer/v1alpha1/trainjobspec.go`	Updates applyconfiguration comments/metadata to align with schema changes.

Copilot · 2026-02-03T01:57:27Z

pkg/apis/trainer/v1alpha1/trainjob_types.go

 }

 // TrainJobSpec represents specification of the desired TrainJob.
+// +kubebuilder:validation:XValidation:rule="(has(oldSelf.suspend) && oldSelf.suspend) || self.podTemplateOverrides == oldSelf.podTemplateOverrides", message="podTemplateOverrides are immutable when suspend is false"


The TrainJobSpec XValidation for podTemplateOverrides gates mutability on oldSelf.suspend, which prevents updating podTemplateOverrides in the same request that sets suspend=true (oldSelf.suspend is false) and diverges from the webhook/plugin logic that checks the new spec’s suspend value.

Suggested change

// +kubebuilder:validation:XValidation:rule="(has(oldSelf.suspend) && oldSelf.suspend) || self.podTemplateOverrides == oldSelf.podTemplateOverrides", message="podTemplateOverrides are immutable when suspend is false"

// +kubebuilder:validation:XValidation:rule="(has(self.suspend) && self.suspend) || self.podTemplateOverrides == oldSelf.podTemplateOverrides", message="podTemplateOverrides are immutable when suspend is false"

Copilot · 2026-02-03T01:57:28Z

test/integration/webhooks/trainjob_test.go

+					return job
+				},
+				testingutil.BeInvalidError()),
+			ginkgo.Entry("Should success to update podTemplateOverride when suspend is true",


The test case name uses ungrammatical wording (“Should success to update…”), which makes intent harder to read when scanning failures; rename to “Should succeed to update …”.

Suggested change

ginkgo.Entry("Should success to update podTemplateOverride when suspend is true",

ginkgo.Entry("Should succeed to update podTemplateOverride when suspend is true",

Copilot · 2026-02-03T01:57:28Z

test/integration/controller/trainjob_controller_test.go

 					g.Expect(k8sClient.Get(ctx, trainJobKey, trainJob)).Should(gomega.Succeed())
-					trainJob.Spec.Trainer.Image = &originImageName
-					g.Expect(k8sClient.Update(ctx, trainJob)).Should(gomega.Succeed())
-				}, util.Timeout, util.Interval).Should(gomega.Succeed())
-
-				ginkgo.By("Checking if JobSet keep having updated Trainer image")
-				gomega.Consistently(func(g gomega.Gomega) {
-					jobSet := &jobsetv1alpha2.JobSet{}
-					g.Expect(k8sClient.Get(ctx, trainJobKey, jobSet)).Should(gomega.Succeed())
-					for _, rJob := range jobSet.Spec.ReplicatedJobs {
-						if rJob.Name == constants.Node {
-							g.Expect(rJob.Template.Spec.Template.Spec.Containers[0].Image).Should(gomega.Equal(updatedImageName))
-						}
-					}
-				}, util.ConsistentDuration, util.Interval).Should(gomega.Succeed())
-
-				ginkgo.By("Trying to re-suspend TrainJob and restore Trainer image")
-				gomega.Eventually(func(g gomega.Gomega) {
-					g.Expect(k8sClient.Get(ctx, trainJobKey, trainJob))
-					trainJob.Spec.Suspend = ptr.To(true)
-					trainJob.Spec.Trainer.Image = &originImageName
-					g.Expect(k8sClient.Update(ctx, trainJob)).Should(gomega.Succeed())
-				}, util.Timeout, util.Interval).Should(gomega.Succeed())
-
-				ginkgo.By("Checking if JobSet image is restored")
-				gomega.Eventually(func(g gomega.Gomega) {
-					jobSet := &jobsetv1alpha2.JobSet{}
-					g.Expect(k8sClient.Get(ctx, trainJobKey, jobSet)).Should(gomega.Succeed())
-					g.Expect(jobSet.Spec.Suspend).ShouldNot(gomega.BeNil())
-					g.Expect(*jobSet.Spec.Suspend).Should(gomega.BeTrue())
-					for _, rJob := range jobSet.Spec.ReplicatedJobs {
-						if rJob.Name == constants.Node {
-							g.Expect(rJob.Template.Spec.Template.Spec.Containers[0].Image).Should(gomega.Equal(originImageName))
-						}
-					}
+					trainJob.Spec.PodTemplateOverrides = emptyPodTemplateOverrides
+					g.Expect(k8sClient.Update(ctx, trainJob)).ShouldNot(gomega.Succeed())
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())


This assertion only checks that the update fails, so the test could pass due to unrelated transient errors (e.g., conflicts) rather than the intended immutability validation; assert the error type/message (e.g., apierrors.IsInvalid / forbidden with expected field path) to make the test deterministic.

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

coveralls · 2026-02-03T02:59:37Z

Pull Request Test Coverage Report for Build 21719581477

Details

0 of 0 changed or added relevant lines in 0 files are covered.
No unchanged relevant lines lost coverage.
Overall coverage remained the same at 51.217%

Totals
Change from base Build 21715897523:	0.0%
Covered Lines:	1241
Relevant Lines:	2423

💛 - Coveralls

astefanutti · 2026-02-04T08:38:10Z

pkg/apis/trainer/v1alpha1/trainjob_types.go

+	// +listType=map
+	// +listMapKey=name


Are we specifying names for PodTemplateOverride?

Don't we want the key to me the manager as introduced after #3020?

Good point, that will be introduced for the top level struct: #3020 (comment) after @kaisoz changes.

// +listType=map // +listMapKey=manager TemplateOverrides []TemplateOverride `json:"templateOverrides,omitempty"`

I am curious if we want to keep // +listMapKey=name for PodTemplateOverrideTargetJob or we can remove it?

IIRC, we introduced PodTemplateOverrideTargetJob struct instead of just []string since @tenzen-y wanted to support label selectors in the future. Context: #2296 (comment)

I guess, in that case we should keep this, right?

// +listType=atomic TargetJobs []PodTemplateOverrideTargetJob `json:"targetJobs,omitempty"`

Right, I think we can keep +listType=atomic for targetJobs.

astefanutti · 2026-02-04T08:39:04Z

pkg/apis/trainer/v1alpha1/trainjob_types.go

 // PodTemplateOverride represents a custom PodTemplateSpec override that will be applied to the TrainJob's training runtime.
 type PodTemplateOverride struct {
 	// targetJobs is the list of replicated jobs in the training runtime template to apply the overrides.
-	// +listType=atomic


I remember we had a discussion on this w.r.t. Kueue manager field.
@tenzen-y @kaisoz @mimowo if you have views on this.

No strong view. Ideally (long term) we use +listType=map, but for that we need to have a unique identifier. This could be the manager field, but for that we would either have a webhook or MAP.

akshaychitneni · 2026-02-04T17:31:26Z

pkg/apis/trainer/v1alpha1/trainjob_types.go

 	// targetJobs is the list of replicated jobs in the training runtime template to apply the overrides.
-	// +listType=atomic
+	// +listType=map
+	// +listMapKey=name


Do we also validate if job name is valid? i.e it exists in the runtime jobs

Yeah, this is good point.
Let me create followup issue for this, since we need to use Validation Webhook for it.

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

astefanutti · 2026-02-05T08:08:21Z

pkg/apis/trainer/v1alpha1/trainjob_types.go

 	Initializer *Initializer `json:"initializer,omitempty"`

 	// trainer defines the configuration of the trainer.
+	// +kubebuilder:validation:XValidation:rule="self == oldSelf", message="field is immutable"


Would that make sense to be consistent with what we did in #2683 for PodTemplateOverrides that can be mutated when the TrainJob is suspended or you think we it's better to be more conservative?

Currently, JobSet/Job doesn't allow us to modify those fields even if TrainJob is suspended (image, command, env, numNodes, etc.) Which leads to a controller error if users try to update these fields.

andreyvelich · 2026-02-05T13:30:41Z

pkg/apis/trainer/v1alpha1/trainjob_types.go

 }

 // TrainJobSpec represents specification of the desired TrainJob.
+// +kubebuilder:validation:XValidation:rule="self.suspend == true || (has(self.podTemplateOverrides) == has(oldSelf.podTemplateOverrides) && (!has(self.podTemplateOverrides) || self.podTemplateOverrides == oldSelf.podTemplateOverrides))", message="podTemplateOverrides are immutable when suspend is false"


One of the problem I see with this rule is that we also block updates to PodTemplateMetadata until TrainJob is suspended. Not sure if that is right approach.

trainer/pkg/apis/trainer/v1alpha1/trainjob_types.go

Lines 273 to 276 in e0b3ef4

// metadata overrides the Pod template metadata.

// These values will be merged with the TrainingRuntime's Pod template metadata.

// +optional

Metadata *metav1.ObjectMeta `json:"metadata,omitempty"`

I would prefer to remove this CEL rule for now, and introduce validation for PodTemplateOverrides/TemplateOverrides in the followup PR after refactoring the API

WDYT @kaisoz @astefanutti @mimowo @tenzen-y @akshaychitneni ?

Yeah, so Kueue will need to be able to modify the podTemplateOverrides when:

old.suspend=true -> new.suspend=false - resuming a Job - Kueue will want to inject nodeSelectors

old.suspend=false -> new.suspend=true -> request to suspend a Job -- Kueue will want to revert the nodeSelectors

Also, any requests which preserve suspend=true -> suspend=true are fine.

I think the rules should only reject requests when old.suspend=false -> new.suspend=false.

Let me remove this CEL for now, and we can iterate with validation in the next PRs.

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

andreyvelich added 2 commits February 3, 2026 00:39

feat(api): Fix immutability for TrainJob APIs

e0fe0e7

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

Only validate immutability of podTemplateOverride API

6d3bf81

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

Copilot AI review requested due to automatic review settings February 3, 2026 01:53

google-oss-prow bot assigned akshaychitneni Feb 3, 2026

google-oss-prow bot added the do-not-merge/hold label Feb 3, 2026

google-oss-prow bot assigned astefanutti, kaisoz and tenzen-y Feb 3, 2026

google-oss-prow bot requested review from akshaychitneni and kuizhiqing February 3, 2026 01:53

google-oss-prow bot added the size/L label Feb 3, 2026

Copilot started reviewing on behalf of andreyvelich February 3, 2026 01:53 View session

Copilot AI reviewed Feb 3, 2026

View reviewed changes

Fix CEL rule

6ac623c

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

andreyvelich force-pushed the fix-immutable-apis branch from 32b1dc7 to 6ac623c Compare February 3, 2026 02:55

astefanutti reviewed Feb 4, 2026

View reviewed changes

akshaychitneni reviewed Feb 4, 2026

View reviewed changes

Use listType atomic for PodTemplateOverrideTargetJob

e0b3ef4

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

astefanutti reviewed Feb 5, 2026

View reviewed changes

andreyvelich commented Feb 5, 2026

View reviewed changes

andreyvelich force-pushed the fix-immutable-apis branch from c5a2345 to db2bb3a Compare February 5, 2026 16:26

Remove PodTemplateOverrides CEL validation

727d99d

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

andreyvelich force-pushed the fix-immutable-apis branch from db2bb3a to 727d99d Compare February 5, 2026 16:27

	// +kubebuilder:validation:XValidation:rule="(has(oldSelf.suspend) && oldSelf.suspend) \|\| self.podTemplateOverrides == oldSelf.podTemplateOverrides", message="podTemplateOverrides are immutable when suspend is false"
	// +kubebuilder:validation:XValidation:rule="(has(self.suspend) && self.suspend) \|\| self.podTemplateOverrides == oldSelf.podTemplateOverrides", message="podTemplateOverrides are immutable when suspend is false"

	ginkgo.Entry("Should success to update podTemplateOverride when suspend is true",
	ginkgo.Entry("Should succeed to update podTemplateOverride when suspend is true",

	// metadata overrides the Pod template metadata.
	// These values will be merged with the TrainingRuntime's Pod template metadata.
	// +optional
	Metadata *metav1.ObjectMeta `json:"metadata,omitempty"`

Conversation

andreyvelich commented Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

google-oss-prow bot commented Feb 3, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

coveralls commented Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 21719581477

Details

💛 - Coveralls

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andreyvelich Feb 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mimowo Feb 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mimowo Feb 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

andreyvelich commented Feb 3, 2026 •

edited

Loading

coveralls commented Feb 3, 2026 •

edited

Loading

andreyvelich Feb 4, 2026 •

edited

Loading

mimowo Feb 5, 2026 •

edited

Loading

mimowo Feb 5, 2026 •

edited

Loading