What you would like to be added?
Add terminationGracePeriodSeconds field to PodSpecPatch so users can configure the pod termination grace period per TrainJob via RuntimePatches.
Depends on: #3199 (RuntimePatches API) must be merged first.
Currently, PodSpecPatch (introduced in #3199) supports these fields:
serviceAccountName, nodeSelector, affinity, tolerations, volumes, initContainers, containers, schedulingGates, imagePullSecrets, securityContext
But it does not expose terminationGracePeriodSeconds, which is a standard corev1.PodSpec field.
The field already exists in the TrainingRuntime CRD as part of the full embedded corev1.PodSpec (e.g., line 7507 in trainer.kubeflow.org_trainingruntimes.yaml), so runtime authors can set it. However, users cannot override it per-TrainJob.
Proposed API change in pkg/apis/trainer/v1alpha1/trainjob_types.go:
type PodSpecPatch struct {
// ... existing fields ...
// terminationGracePeriodSeconds patches the termination grace period for Pods
// in the target job templates.
// +optional
TerminationGracePeriodSeconds *int64 `json:"terminationGracePeriodSeconds,omitempty"`
}
No changes needed to the merge logic in pkg/runtime/core/trainingruntime.go since the existing strategic merge patch already handles arbitrary PodTemplateSpec fields from PodSpecPatch.
Why is this needed?
In distributed training with PyTorch Elastic (torchrun), when a TrainJob is paused or a node is drained:
- Kubelet sends SIGTERM to the pod
- TorchElastic agent receives SIGTERM and propagates it to worker processes
- Workers perform JIT checkpointing (saving model state to prevent training progress loss)
- After the grace period expires, Kubelet sends SIGKILL
The problem: TorchElastic currently has a hardcoded 30-second shutdown timeout (pytorch/pytorch#172596 makes this configurable via TORCH_ELASTIC_SHUTDOWN_TIMEOUT env var). For large models (70B+ parameters), 30 seconds is insufficient to save a checkpoint to disk or upload to S3/PVC.
Users need to:
- Set
terminationGracePeriodSeconds on the pod to allow Kubernetes to wait longer before SIGKILL
- Set
TORCH_ELASTIC_SHUTDOWN_TIMEOUT env var to match, so TorchElastic doesn't kill workers prematurely
Today, users cannot configure terminationGracePeriodSeconds via the TrainJob API, they can only set it in the TrainingRuntime template, which requires cluster-admin access and affects all jobs using that runtime. This field should be user-configurable per TrainJob, just like nodeSelector, tolerations, and other PodSpec fields already exposed in PodSpecPatch.
/kind feature
/area api
What you would like to be added?
Add
terminationGracePeriodSecondsfield toPodSpecPatchso users can configure the pod termination grace period per TrainJob viaRuntimePatches.Currently,
PodSpecPatch(introduced in #3199) supports these fields:serviceAccountName,nodeSelector,affinity,tolerations,volumes,initContainers,containers,schedulingGates,imagePullSecrets,securityContextBut it does not expose
terminationGracePeriodSeconds, which is a standardcorev1.PodSpecfield.The field already exists in the TrainingRuntime CRD as part of the full embedded
corev1.PodSpec(e.g., line 7507 intrainer.kubeflow.org_trainingruntimes.yaml), so runtime authors can set it. However, users cannot override it per-TrainJob.Proposed API change in
pkg/apis/trainer/v1alpha1/trainjob_types.go:No changes needed to the merge logic in
pkg/runtime/core/trainingruntime.gosince the existing strategic merge patch already handles arbitraryPodTemplateSpecfields fromPodSpecPatch.Why is this needed?
In distributed training with PyTorch Elastic (torchrun), when a TrainJob is paused or a node is drained:
The problem: TorchElastic currently has a hardcoded 30-second shutdown timeout (pytorch/pytorch#172596 makes this configurable via
TORCH_ELASTIC_SHUTDOWN_TIMEOUTenv var). For large models (70B+ parameters), 30 seconds is insufficient to save a checkpoint to disk or upload to S3/PVC.Users need to:
terminationGracePeriodSecondson the pod to allow Kubernetes to wait longer before SIGKILLTORCH_ELASTIC_SHUTDOWN_TIMEOUTenv var to match, so TorchElastic doesn't kill workers prematurelyToday, users cannot configure
terminationGracePeriodSecondsvia the TrainJob API, they can only set it in the TrainingRuntime template, which requires cluster-admin access and affects all jobs using that runtime. This field should be user-configurable per TrainJob, just likenodeSelector,tolerations, and other PodSpec fields already exposed inPodSpecPatch./kind feature
/area api