Skip to content

Add terminationGracePeriodSeconds to PodSpecPatch in TrainJob #3285

Description

@efazal

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:

  1. Kubelet sends SIGTERM to the pod
  2. TorchElastic agent receives SIGTERM and propagates it to worker processes
  3. Workers perform JIT checkpointing (saving model state to prevent training progress loss)
  4. 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:

  1. Set terminationGracePeriodSeconds on the pod to allow Kubernetes to wait longer before SIGKILL
  2. 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

Metadata

Metadata

Assignees

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions