hanzo-ml/trainer
1
0
Fork 0
mirror of https://github.com/hanzo-ml/trainer.git synced 2026-09-24 00:31:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(api): Add terminationGracePeriodSeconds to PodSpecPatch in TrainJob (#3324)

Browse Source 
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 <[email protected]>
This commit is contained in:
Krish Garg
2026-05-26 17:36:09 +00:00
committed by GitHub
parent cdff07ff48
commit 4f8920c902
13 changed files with 269 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -39,3 +39,5 @@ examples/pytorch/*/data
examples/pytorch/audio-classification/genres_original
examples/pytorch/audio-classification/output
trainer-controller-manager.exe
api/openapi-spec/swagger.json
Generated
+5
View File
@@ -14421,6 +14421,11 @@
"description": "serviceAccountName patches the service account for the Pods in the target job templates.",
"type": "string"
},
"terminationGracePeriodSeconds": {
"description": "terminationGracePeriodSeconds patches the termination grace period for Pods in the target job templates.",
"type": "integer",
"format": "int64"
},
"tolerations": {
"description": "tolerations patches the Pod's tolerations.",
"type": "array",
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_pod_spec_patch.py
Generated
+4 -2
View File
@@ -17,7 +17,7 @@ import pprint
import re # noqa: F401
import json
from pydantic import BaseModel, ConfigDict, Field, StrictStr
from pydantic import BaseModel, ConfigDict, Field, StrictInt, StrictStr
from typing import Any, ClassVar, Dict, List, Optional
from kubeflow_trainer_api.models.io_k8s_api_core_v1_affinity import IoK8sApiCoreV1Affinity
from kubeflow_trainer_api.models.io_k8s_api_core_v1_local_object_reference import IoK8sApiCoreV1LocalObjectReference
@@ -41,9 +41,10 @@ class TrainerV1alpha1PodSpecPatch(BaseModel):
scheduling_gates: Optional[List[IoK8sApiCoreV1PodSchedulingGate]] = Field(default=None, description="schedulingGates patches the scheduling gates for the Pods in the target job templates. More info: https://kubernetes.io/docs/concepts/scheduling-eviction/pod-scheduling-readiness/", alias="schedulingGates")
security_context: Optional[IoK8sApiCoreV1PodSecurityContext] = Field(default=None, description="securityContext patches the Pod's security context. More info: https://kubernetes.io/docs/tasks/configure-pod-container/security-context/", alias="securityContext")
service_account_name: Optional[StrictStr] = Field(default=None, description="serviceAccountName patches the service account for the Pods in the target job templates.", alias="serviceAccountName")
termination_grace_period_seconds: Optional[StrictInt] = Field(default=None, description="terminationGracePeriodSeconds patches the termination grace period for Pods in the target job templates.", alias="terminationGracePeriodSeconds")
tolerations: Optional[List[IoK8sApiCoreV1Toleration]] = Field(default=None, description="tolerations patches the Pod's tolerations.")
volumes: Optional[List[IoK8sApiCoreV1Volume]] = Field(default=None, description="volumes patches the Pod's volumes.")
__properties: ClassVar[List[str]] = ["affinity", "containers", "imagePullSecrets", "initContainers", "nodeSelector", "schedulingGates", "securityContext", "serviceAccountName", "tolerations", "volumes"]
__properties: ClassVar[List[str]] = ["affinity", "containers", "imagePullSecrets", "initContainers", "nodeSelector", "schedulingGates", "securityContext", "serviceAccountName", "terminationGracePeriodSeconds", "tolerations", "volumes"]
model_config = ConfigDict(
populate_by_name=True,
@@ -152,6 +153,7 @@ class TrainerV1alpha1PodSpecPatch(BaseModel):
"schedulingGates": [IoK8sApiCoreV1PodSchedulingGate.from_dict(_item) for _item in obj["schedulingGates"]] if obj.get("schedulingGates") is not None else None,
"securityContext": IoK8sApiCoreV1PodSecurityContext.from_dict(obj["securityContext"]) if obj.get("securityContext") is not None else None,
"serviceAccountName": obj.get("serviceAccountName"),
"terminationGracePeriodSeconds": obj.get("terminationGracePeriodSeconds"),
"tolerations": [IoK8sApiCoreV1Toleration.from_dict(_item) for _item in obj["tolerations"]] if obj.get("tolerations") is not None else None,
"volumes": [IoK8sApiCoreV1Volume.from_dict(_item) for _item in obj["volumes"]] if obj.get("volumes") is not None else None
})
charts/kubeflow-trainer/crds/trainer.kubeflow.org_trainjobs.yaml
+10
View File
@@ -3147,6 +3147,16 @@ 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.
format: int64
minimum: 0
type: integer
x-kubernetes-validations:
- message: field is immutable
rule: self == oldSelf
tolerations:
description: tolerations patches
the Pod's tolerations.
manifests/base/crds/trainer.kubeflow.org_trainjobs.yaml
+10
View File
@@ -3147,6 +3147,16 @@ 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.
format: int64
minimum: 0
type: integer
x-kubernetes-validations:
- message: field is immutable
rule: self == oldSelf
tolerations:
description: tolerations patches
the Pod's tolerations.
pkg/apis/trainer/v1alpha1/trainjob_types.go
+7
View File
@@ -453,6 +453,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.
// +kubebuilder:validation:Minimum=0
// +kubebuilder:validation:XValidation:rule="self == oldSelf",message="field is immutable"
// +optional
TerminationGracePeriodSeconds *int64 `json:"terminationGracePeriodSeconds,omitempty"`
}
// ContainerPatch represents parameters that can be patched using PodSpecPatch.
pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go
Generated
+5
View File
@@ -647,6 +647,11 @@ func (in *PodSpecPatch) DeepCopyInto(out *PodSpecPatch) {
*out = make([]v1.PodSchedulingGate, len(*in))
copy(*out, *in)
}
if in.TerminationGracePeriodSeconds != nil {
in, out := &in.TerminationGracePeriodSeconds, &out.TerminationGracePeriodSeconds
*out = new(int64)
**out = **in
}
return
}
pkg/apis/trainer/v1alpha1/zz_generated.openapi.go
Generated
+7
View File
@@ -1374,6 +1374,13 @@ func schema_pkg_apis_trainer_v1alpha1_PodSpecPatch(ref common.ReferenceCallback)
},
},
},
"terminationGracePeriodSeconds": {
SchemaProps: spec.SchemaProps{
Description: "terminationGracePeriodSeconds patches the termination grace period for Pods in the target job templates.",
Type: []string{"integer"},
Format: "int64",
},
},
},
},
},
pkg/client/applyconfiguration/trainer/v1alpha1/podspecpatch.go
Generated
+11
View File
@@ -48,6 +48,9 @@ type PodSpecPatchApplyConfiguration struct {
// schedulingGates patches the scheduling gates for the Pods in the target job templates.
// More info: https://kubernetes.io/docs/concepts/scheduling-eviction/pod-scheduling-readiness/
SchedulingGates []corev1.PodSchedulingGate `json:"schedulingGates,omitempty"`
// terminationGracePeriodSeconds patches the termination grace period for Pods
// in the target job templates.
TerminationGracePeriodSeconds *int64 `json:"terminationGracePeriodSeconds,omitempty"`
}
// PodSpecPatchApplyConfiguration constructs a declarative configuration of the PodSpecPatch type for use with
@@ -165,3 +168,11 @@ func (b *PodSpecPatchApplyConfiguration) WithSchedulingGates(values ...corev1.Po
}
return b
}
// WithTerminationGracePeriodSeconds sets the TerminationGracePeriodSeconds field in the declarative configuration to the given value
// and returns the receiver, so that objects can be built by chaining "With" function invocations.
// If called multiple times, the TerminationGracePeriodSeconds field is set to the value of the last call.
func (b *PodSpecPatchApplyConfiguration) WithTerminationGracePeriodSeconds(value int64) *PodSpecPatchApplyConfiguration {
b.TerminationGracePeriodSeconds = &value
return b
}
pkg/runtime/core/trainingruntime_test.go
+2
View File
@@ -430,6 +430,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
Effect: corev1.TaintEffectNoSchedule,
},
},
TerminationGracePeriodSeconds: ptr.To(int64(300)),
Volumes: []corev1.Volume{
{
Name: "node_secret",
@@ -575,6 +576,7 @@ func TestTrainingRuntimeNewObjects(t *testing.T) {
Operator: corev1.TolerationOpExists,
Effect: corev1.TaintEffectNoSchedule,
}).
TerminationGracePeriodSeconds(constants.Node, 300).
Volumes(constants.DatasetInitializer,
corev1.Volume{
Name: "initializer_claim",
pkg/util/testing/wrapper.go
+9
View File
@@ -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 {
test/integration/controller/trainjob_controller_test.go
+73
View File
@@ -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() {
test/integration/webhooks/trainjob_test.go
+124
View File
@@ -508,6 +508,96 @@ var _ = ginkgo.Describe("TrainJob marker validations and defaulting", ginkgo.Ord
Obj()
},
testingutil.BeForbiddenError()),
ginkgo.Entry("Should succeed to create TrainJob with valid terminationGracePeriodSeconds in RuntimePatches",
func() *trainer.TrainJob {
return testingutil.MakeTrainJobWrapper(ns.Name, "valid-termination-grace-period").
RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "acme.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: ptr.To(int64(300)),
},
},
},
},
}},
},
},
},
},
}).
Obj()
},
gomega.Succeed()),
ginkgo.Entry("Should fail to create TrainJob with negative terminationGracePeriodSeconds",
func() *trainer.TrainJob {
return testingutil.MakeTrainJobWrapper(ns.Name, "invalid-termination-grace-period").
RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "acme.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: ptr.To(int64(-1)),
},
},
},
},
}},
},
},
},
},
}).
Obj()
},
testingutil.BeInvalidError()),
ginkgo.Entry("Should succeed to create TrainJob with zero terminationGracePeriodSeconds",
func() *trainer.TrainJob {
return testingutil.MakeTrainJobWrapper(ns.Name, "zero-termination-grace-period").
RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "acme.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: ptr.To(int64(0)),
},
},
},
},
}},
},
},
},
},
}).
Obj()
},
gomega.Succeed()),
)
ginkgo.DescribeTable("Defaulting TrainJob on creation", func(trainJob func() *trainer.TrainJob, wantTrainJob func() *trainer.TrainJob) {
created := trainJob()
@@ -688,6 +778,40 @@ var _ = ginkgo.Describe("TrainJob marker validations and defaulting", ginkgo.Ord
return job
},
gomega.Succeed()),
ginkgo.Entry("Should fail to update TrainJob terminationGracePeriodSeconds after creation",
func() *trainer.TrainJob {
return testingutil.MakeTrainJobWrapper(ns.Name, "immutable-termination-grace-period").
RuntimeRef(trainer.GroupVersion.WithKind(trainer.TrainingRuntimeKind), "testing").
RuntimePatches([]trainer.RuntimePatch{
{
Manager: "acme.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: ptr.To(int64(300)),
},
},
},
},
}},
},
},
},
},
}).
Obj()
},
func(job *trainer.TrainJob) *trainer.TrainJob {
job.Spec.RuntimePatches[0].TrainingRuntimeSpec.Template.Spec.ReplicatedJobs[0].Template.Spec.Template.Spec.TerminationGracePeriodSeconds = ptr.To(int64(600))
return job
},
testingutil.BeInvalidError()),
)
})
})