feat(api): Add terminationGracePeriodSeconds to PodSpecPatch in TrainJob

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 #3285

Signed-off-by: krishdef7 <gargkrish06@gmail.com>

Author: krishdef7

api/openapi-spec/swagger.json

@@ -3122,6 +3122,14 @@ spec:
                                                         x-kubernetes-validations:
                                                         - message: field is immutable
                                                           rule: self == oldSelf
+                                                      terminationGracePeriodSeconds:
+                                                        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.
+                                                        format: int64
+                                                        minimum: 0
+                                                        type: integer
                                                       tolerations:
                                                         description: tolerations patches
                                                           the Pod's tolerations.

api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_pod_spec_patch.py

@@ -3122,6 +3122,14 @@ spec:
                                                         x-kubernetes-validations:
                                                         - message: field is immutable
                                                           rule: self == oldSelf
+                                                      terminationGracePeriodSeconds:
+                                                        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.
+                                                        format: int64
+                                                        minimum: 0
+                                                        type: integer
                                                       tolerations:
                                                         description: tolerations patches
                                                           the Pod's tolerations.

charts/kubeflow-trainer/crds/trainer.kubeflow.org_trainjobs.yaml

@@ -432,6 +432,13 @@ type PodSpecPatch struct {
 	// +kubebuilder:validation:MaxItems=16
 	// +optional
 	SchedulingGates []corev1.PodSchedulingGate `json:"schedulingGates,omitempty"`
+
+	// 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
+	// +optional
+	TerminationGracePeriodSeconds *int64 `json:"terminationGracePeriodSeconds,omitempty"`
 }
 
 // ContainerPatch represents parameters that can be patched using PodSpecPatch.

manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml

@@ -338,6 +338,15 @@ func (j *JobSetWrapper) SchedulingGates(rJobName string, schedulingGates ...core
 	return j
 }
 
+func (j *JobSetWrapper) TerminationGracePeriodSeconds(rJobName string, seconds int64) *JobSetWrapper {
+	for i, rJob := range j.Spec.ReplicatedJobs {
+		if rJob.Name == rJobName {
+			j.Spec.ReplicatedJobs[i].Template.Spec.Template.Spec.TerminationGracePeriodSeconds = &seconds
+		}
+	}
+	return j
+}
+
 func (j *JobSetWrapper) ImagePullSecrets(rJobName string, imagePullSecrets ...corev1.LocalObjectReference) *JobSetWrapper {
 	for i, rJob := range j.Spec.ReplicatedJobs {
 		if rJob.Name == rJobName {

pkg/apis/trainer/v1alpha1/trainjob_types.go

@@ -301,6 +301,79 @@ var _ = ginkgo.Describe("TrainJob controller", ginkgo.Ordered, func() {
 					g.Expect(k8sClient.Update(ctx, trainJob)).Should(testingutil.BeInvalidError())
 				}, util.Timeout, util.Interval).Should(gomega.Succeed())
 			})
+	ginkgo.It("Should propagate terminationGracePeriodSeconds from RuntimePatches to JobSet pods", func() {
+					ginkgo.By("Creating a TrainingRuntime and TrainJob with terminationGracePeriodSeconds patch")
+					gracePeriodRuntime := testingutil.MakeTrainingRuntimeWrapper(ns.Name, "alpha-grace").
+						RuntimeSpec(
+							testingutil.MakeTrainingRuntimeSpecWrapper(testingutil.MakeTrainingRuntimeWrapper(ns.Name, "alpha-grace").Spec).
+								WithMLPolicy(
+									testingutil.MakeMLPolicyWrapper().
+										WithNumNodes(1).
+										Obj(),
+								).
+								Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+								Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+								Container(constants.Node, constants.Node, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+								Obj()).
+						Obj()
+					gomega.Expect(k8sClient.Create(ctx, gracePeriodRuntime)).Should(gomega.Succeed())
+					gomega.Eventually(func(g gomega.Gomega) {
+						g.Expect(k8sClient.Get(ctx, client.ObjectKeyFromObject(gracePeriodRuntime), gracePeriodRuntime)).Should(gomega.Succeed())
+					}, util.Timeout, util.Interval).Should(gomega.Succeed())
+
+					gracePeriod := int64(300)
+					graceJob := testingutil.MakeTrainJobWrapper(ns.Name, "grace-period-job").
+						Suspend(true).
+						RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "alpha-grace").
+						RuntimePatches([]trainer.RuntimePatch{{
+							Manager: "test.io/manager",
+							TrainingRuntimeSpec: &trainer.TrainingRuntimeSpecPatch{
+								Template: &trainer.JobSetTemplatePatch{
+									Spec: &trainer.JobSetSpecPatch{
+										ReplicatedJobs: []trainer.ReplicatedJobPatch{{
+											Name: constants.Node,
+											Template: &trainer.JobTemplatePatch{
+												Spec: &trainer.JobSpecPatch{
+													Template: &trainer.PodTemplatePatch{
+														Spec: &trainer.PodSpecPatch{
+															TerminationGracePeriodSeconds: &gracePeriod,
+														},
+													},
+												},
+											},
+										}},
+									},
+								},
+							},
+						}}).
+						Trainer(
+							testingutil.MakeTrainJobTrainerWrapper().
+								Container("test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+								Obj()).
+						Obj()
+					graceJobKey := client.ObjectKeyFromObject(graceJob)
+					gomega.Expect(k8sClient.Create(ctx, graceJob)).Should(gomega.Succeed())
+
+					ginkgo.By("Checking that JobSet node pods have terminationGracePeriodSeconds set to 300")
+					gomega.Eventually(func(g gomega.Gomega) {
+						jobSet := &jobsetv1alpha2.JobSet{}
+						g.Expect(k8sClient.Get(ctx, graceJobKey, jobSet)).Should(gomega.Succeed())
+						g.Expect(jobSet).Should(gomega.BeComparableTo(
+							testingutil.MakeJobSetWrapper(ns.Name, graceJobKey.Name).
+								ControllerReference(trainer.SchemeGroupVersion.WithKind(trainer.TrainJobKind), graceJobKey.Name, string(graceJob.UID)).
+								Suspend(true).
+								Replicas(1, constants.Node, constants.DatasetInitializer, constants.ModelInitializer).
+								Parallelism(1, constants.DatasetInitializer, constants.ModelInitializer).
+								Completions(1, constants.DatasetInitializer, constants.ModelInitializer).
+								NumNodes(1).
+								Container(constants.DatasetInitializer, constants.DatasetInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+								Container(constants.ModelInitializer, constants.ModelInitializer, "test:runtime", []string{"runtime"}, []string{"runtime"}, resRequests).
+								Container(constants.Node, constants.Node, "test:trainjob", []string{"trainjob"}, []string{"trainjob"}, resRequests).
+								TerminationGracePeriodSeconds(constants.Node, gracePeriod).
+								Obj(),
+							util.IgnoreObjectMetadata))
+					}, util.Timeout, util.Interval).Should(gomega.Succeed())
+				})
 		})
 
 		ginkgo.Context("Integration tests for the Torch Runtime", func() {