feat(api): Add terminationGracePeriodSeconds to PodSpecPatch in TrainJob by krishdef7 · Pull Request #3324 · kubeflow/trainer

krishdef7 · 2026-03-13T10:55:45Z

What this PR does / why we need it

Adds terminationGracePeriodSeconds to PodSpecPatch so users can configure pod termination grace period per TrainJob via RuntimePatches, without requiring cluster-admin access to modify the TrainingRuntime.

This is needed for distributed training with PyTorch Elastic (torchrun). When a TrainJob is paused or a node is drained, Kubelet sends SIGTERM to the pod. The TorchElastic agent propagates this to worker processes, which perform JIT checkpointing before shutdown. For large models (70B+ parameters), the default 30-second grace period is insufficient to complete checkpoint saves to disk or remote storage (S3/PVC).

Users need to set terminationGracePeriodSeconds alongside TORCH_ELASTIC_SHUTDOWN_TIMEOUT (configurable since pytorch/pytorch#172596) to give workers enough time to save state before SIGKILL.

Previously, this field could only be set in the TrainingRuntime template, affecting all jobs using that runtime. This PR exposes it as a per-TrainJob override consistent with other PodSpecPatch fields like nodeSelector and tolerations.

No changes to merge logic in trainingruntime.go are required — the existing StrategicMergePatch applied at batchv1.JobTemplateSpec level already handles this field automatically.

Which issue(s) this PR fixes

Fixes #3285

Changes

Add TerminationGracePeriodSeconds *int64 to PodSpecPatch with // +kubebuilder:validation:Minimum=0 marker
Regenerate zz_generated.deepcopy.go, zz_generated.openapi.go, podspecpatch.go applyconfiguration, Python client model, and CRD
Add TerminationGracePeriodSeconds helper to JobSetWrapper in test utilities
Add integration test verifying propagation to JobSet pods
Add webhook tests for valid (300s) and invalid (-1) values

Checklist

Docs included if any changes are user facing

github-actions · 2026-03-13T10:55:55Z

🎉 Welcome to the Kubeflow Trainer! 🎉

Thanks for opening your first PR! We're happy to have you as part of our community 🚀

Here's what happens next:

If you haven't already, please check out our Contributing Guide for repo-specific guidelines and the Kubeflow Contributor Guide for general community standards.
Our team will review your PR soon! cc @kubeflow/kubeflow-trainer-team

Join the community:

Slack: Join our #kubeflow-trainer Slack channel.
Meetings: Attend the Kubeflow AutoML and Training Working Group bi-weekly meetings.

Feel free to ask questions in the comments if you need any help or clarification!
Thanks again for contributing to Kubeflow! 🙏

Copilot

Pull request overview

This PR extends the TrainJob RuntimePatches API to allow per-TrainJob overrides of terminationGracePeriodSeconds via PodSpecPatch, enabling longer graceful shutdowns for workloads like PyTorch Elastic that need time to checkpoint on SIGTERM.

Changes:

Add TerminationGracePeriodSeconds *int64 (Minimum=0) to PodSpecPatch and propagate it through generated Go/OpenAPI/CRD artifacts.
Update the Python client model to include terminationGracePeriodSeconds.
Add integration + webhook coverage to validate schema enforcement and verify the field propagates into JobSet pod specs.

Reviewed changes

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

Show a summary per file

File	Description
`pkg/apis/trainer/v1alpha1/trainjob_types.go`	Adds the new `PodSpecPatch` field with kubebuilder minimum validation.
`manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml`	Exposes the field in the TrainJob CRD schema (min 0).
`pkg/client/applyconfiguration/trainer/v1alpha1/podspecpatch.go`	Extends server-side apply configuration with a `WithTerminationGracePeriodSeconds` builder.
`pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go`	Regenerates deepcopy logic to include the new pointer field.
`pkg/apis/trainer/v1alpha1/zz_generated.openapi.go`	Regenerates OpenAPI schema to include the new field.
`api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_pod_spec_patch.py`	Updates the Python client model to support the new field.
`pkg/util/testing/wrapper.go`	Adds a JobSetWrapper helper to set `terminationGracePeriodSeconds` in expected JobSet pod specs.
`test/integration/webhooks/trainjob_test.go`	Adds webhook validation tests for valid and invalid (negative) values.
`test/integration/controller/trainjob_controller_test.go`	Adds integration coverage ensuring the patch propagates into the created JobSet pod specs.

andreyvelich

Thanks @krishdef7, the API changes look good.
/assign @tenzen-y @astefanutti @efazal @kaisoz

andreyvelich · 2026-05-12T02:17:53Z

+	// terminationGracePeriodSeconds patches the termination grace period for Pods
+	// in the target job templates. This allows users to configure sufficient time
+	// for checkpoint saving on TrainJob completion or node drain.
+	// +kubebuilder:validation:Minimum=0


Should it be 1 as min?

Keeping minimum: 0, Kubernetes itself treats 0 as a valid value meaning immediate termination via SIGKILL with no grace period. Setting minimum: 1 would prevent users from disabling the grace period entirely, which is a valid use case (e.g. forcing immediate pod termination on drain).

google-oss-prow · 2026-05-12T02:25:27Z

@andreyvelich: GitHub didn't allow me to assign the following users: efazal.

Note that only kubeflow members with read permissions, repo collaborators and people who have commented on this issue/PR can be assigned. Additionally, issues/PRs can only have 10 assignees at the same time.
For more information please see the contributor guide

Details

In response to this:

Thanks @krishdef7, the API changes looks good.
/assign @tenzen-y @astefanutti @efazal @kaisoz

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

andreyvelich · 2026-05-12T02:31:26Z

+	// in the target job templates. This allows users to configure sufficient time
+	// for checkpoint saving on TrainJob completion or node drain.
+	// +kubebuilder:validation:Minimum=0
+	// +optional


Should this field be immutable too?

Added immutability via +kubebuilder:validation:XValidation:rule="self == oldSelf" consistent with other PodSpecPatch fields like serviceAccountName. Also added a webhook integration test ("Should fail to update TrainJob terminationGracePeriodSeconds after creation") to verify the update is rejected.

andreyvelich

Thanks for this contribution! The implementation looks solid — correct type, appropriate validation, and good test coverage. Review performed with AI tools.

andreyvelich · 2026-05-12T02:33:08Z

+	// for checkpoint saving on TrainJob completion or node drain.
+	// +kubebuilder:validation:Minimum=0
+	// +optional
+	TerminationGracePeriodSeconds *int64 `json:"terminationGracePeriodSeconds,omitempty"`


Suggestion: Consider adding terminationGracePeriodSeconds to the existing unit test in pkg/runtime/core/trainingruntime_test.go (the "succeeded to build JobSet with TrainJob's RuntimePatches" test case around line 191) alongside the other PodSpecPatch fields tested there (ServiceAccountName, Tolerations, etc.). This would provide faster regression detection at the unit level without needing the full integration environment.

Done, added TerminationGracePeriodSeconds: ptr.To(int64(300)) to the input patch and TerminationGracePeriodSeconds(constants.Node, 300) to the expected output in the "succeeded to build JobSet with TrainJob's RuntimePatches" test case.

andreyvelich · 2026-05-12T02:33:11Z

+						Obj()
+				},
+				testingutil.BeInvalidError()),
 		)


Suggestion: Consider adding a boundary test for ptr.To(int64(0)) with gomega.Succeed(). Zero is a valid value per the CRD (minimum: 0) but has special Kubernetes semantics (immediate termination with no grace period). Testing this boundary documents the intent and guards against accidental Minimum=1 refactors.

Done, added "Should succeed to create TrainJob with zero terminationGracePeriodSeconds" entry with ptr.To(int64(0)) and gomega.Succeed().

kaisoz · 2026-05-19T05:00:17Z

+                                                        description: |-
+                                                          terminationGracePeriodSeconds patches the termination grace period for Pods
+                                                          in the target job templates. This allows users to configure sufficient time
+                                                          for checkpoint saving on TrainJob completion or node drain.


IMO this description should be more generic and follow the convention set by the other PodSpecPatch fields, so there's no need to have a use case example here

@krishdef7 Are you still working on this PR? Or we can find a new contributors to finish it?

Updated, removed the use-case example from the description, now reads: "terminationGracePeriodSeconds patches the termination grace period for Pods in the target job templates."

kaisoz · 2026-05-19T05:36:20Z

would it make sense to delegate PodSpec fields validation to the api server? If I'm not mistaken, this is the JobSet behaviour. Here's) the API server validation for the terminationGracePeriodSeconds field

WDYT @andreyvelich ?

andreyvelich · 2026-05-19T12:59:11Z

would it make sense to delegate PodSpec fields validation to the api server? If I'm not mistaken, this is the JobSet behaviour. Here's) the API server validation for the terminationGracePeriodSeconds field

WDYT @andreyvelich ?

I think, we have discussion about it here: kubernetes-sigs/jobset#1211 (comment)
As @kannon92 pointed out, we might have scalability concern if we validate PodSpec via API calls.
We can explore kubectl-validation SDK for that, but this is not ideal too.

Adds terminationGracePeriodSeconds field to PodSpecPatch so users can configure pod termination grace period per TrainJob via RuntimePatches. This is needed for distributed training with PyTorch Elastic (torchrun) where large models (70B+ parameters) require more than the default 30s to complete JIT checkpointing before SIGKILL on node drain or TrainJob pause. No changes to merge logic in trainingruntime.go are required since the existing StrategicMergePatch applied at batchv1.JobTemplateSpec level already handles this field automatically. Closes kubeflow#3285 Signed-off-by: krishdef7 <[email protected]>

krishdef7 · 2026-05-20T06:23:49Z

@kaisoz @andreyvelich @astefanutti @tenzen-y, all comments are addressed, the CRD description is now generic, the immutability XValidation rule is in place, the boundary test for ptr.To(int64(0)) is added, and the unit test in trainingruntime_test.go covers the RuntimePatches path. All CI checks are passing. Could someone issue a /lgtm and /approve when you get a chance?

kaisoz · 2026-05-20T09:12:11Z

/lgtm

All the comments were addressed and I agree on keeping the minimum set to 0 (reason). I'll leave it to @andreyvelich for the final approval

andreyvelich

Thanks for this @krishdef7!
/lgtm
/approve
/hold

/assign @astefanutti @akshaychitneni @robert-bell @tenzen-y in case you have any followup questions.

google-oss-prow · 2026-05-20T14:12:50Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: andreyvelich

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

The pull request process is described here

Details

Needs approval from an approver in each of these files:

~~OWNERS~~ [andreyvelich]

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

google-oss-prow · 2026-05-20T14:59:53Z

@robert-bell: changing LGTM is restricted to collaborators

Details

In response to this:

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

andreyvelich · 2026-05-26T17:33:23Z

/hold cancel

Copilot AI review requested due to automatic review settings March 13, 2026 10:55

google-oss-prow Bot requested review from akshaychitneni and jinchihe March 13, 2026 10:55

google-oss-prow Bot added the size/L label Mar 13, 2026

Copilot started reviewing on behalf of krishdef7 March 13, 2026 10:56 View session

krishdef7 force-pushed the feat/termination-grace-period-patch branch from 5d3b20c to 73143f5 Compare March 13, 2026 10:56

Copilot AI reviewed Mar 13, 2026

View reviewed changes

krishdef7 force-pushed the feat/termination-grace-period-patch branch 5 times, most recently from 8a3aa76 to 6095a15 Compare March 14, 2026 05:11

andreyvelich reviewed May 12, 2026

View reviewed changes

google-oss-prow Bot assigned astefanutti and kaisoz May 12, 2026

google-oss-prow Bot assigned tenzen-y May 12, 2026

andreyvelich reviewed May 12, 2026

View reviewed changes

kaisoz reviewed May 19, 2026

View reviewed changes

krishdef7 force-pushed the feat/termination-grace-period-patch branch 5 times, most recently from 6dc5d79 to 0516d95 Compare May 19, 2026 14:43

krishdef7 force-pushed the feat/termination-grace-period-patch branch from 0516d95 to 5ab2b1e Compare May 19, 2026 14:53

google-oss-prow Bot added the lgtm label May 20, 2026

andreyvelich reviewed May 20, 2026

View reviewed changes

google-oss-prow Bot added the do-not-merge/hold label May 20, 2026

google-oss-prow Bot assigned andreyvelich May 20, 2026

google-oss-prow Bot added the approved label May 20, 2026

robert-bell approved these changes May 20, 2026

View reviewed changes

google-oss-prow Bot removed the do-not-merge/hold label May 26, 2026

google-oss-prow Bot merged commit 7c161d0 into kubeflow:master May 26, 2026
30 checks passed

google-oss-prow Bot added this to the v2.3 milestone May 26, 2026

krishdef7 deleted the feat/termination-grace-period-patch branch June 4, 2026 12:42

Uh oh!

Conversation

krishdef7 commented Mar 13, 2026

What this PR does / why we need it

Which issue(s) this PR fixes

Changes

Checklist

Uh oh!

github-actions Bot commented Mar 13, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

andreyvelich left a comment • 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!

krishdef7 May 19, 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!

google-oss-prow Bot commented May 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andreyvelich left a comment

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kaisoz commented May 19, 2026

Uh oh!

andreyvelich commented May 19, 2026

Uh oh!

krishdef7 commented May 20, 2026

Uh oh!

kaisoz commented May 20, 2026

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

google-oss-prow Bot commented May 20, 2026

Uh oh!

google-oss-prow Bot commented May 20, 2026

Uh oh!

andreyvelich commented May 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

andreyvelich left a comment •

edited

Loading

krishdef7 May 19, 2026 •

edited

Loading