JobSet

JobSet is a Kubernetes-native API designed to manage a group of Kubernetes Jobs as a single unit. It is primarily used for deploying HPC (e.g., MPI) and AI/ML training workloads (PyTorch, Jax, TensorFlow, etc.).

Key capabilities include:

• Multi-template jobs: Specify different pod templates for different groups of pods (e.g., separate templates for a leader, workers, and parameter servers) within a single JobSet.
• Automatic networking: Automatically configures and manages headless services using IndexedJobs to provide stable pod hostnames for distributed training.
• Configurable failure policies: Define how the JobSet should react to job failures. You can specify a maximum number of restarts or use the extended failure policy API (v0.6.0+) to handle different error types differently.
• Configurable success policies: Define when a JobSet is considered complete by targeting specific ReplicatedJobs using Any or All operators.
• Exclusive Placement: Use annotations to request a 1:1 mapping between a child job and a topology domain (like a rack or zone), ensuring all pods in that job are colocated and exclusive to that domain.
• Startup Sequencing: (v0.6.0+) Define a startup order for ReplicatedJobs to support patterns like the "leader-worker" paradigm, ensuring the leader is running before workers attempt to connect.
• Kueue Integration: (v0.2.3+ for JobSet, v0.6.0+ for Kueue) Supports oversubscription and multi-tenancy via Kueue.

JobSet provides a way to track the total number of execution iterations a JobSet has undergone, including both failure-triggered restarts and suspension-driven resumptions (e.g., when evicted by a queueing system like Kueue).

This is useful for external logging and telemetry systems (like FluentBit) to isolate log streams for specific execution cycles. While the existing jobset.sigs.k8s.io/restart-attempt annotation only tracks failures, the new jobset.sigs.k8s.io/execution-attempt annotation tracks every time the JobSet is resumed from a suspended state.

### Example Lifecycle
1. **Initial Run**: `restart-attempt` = "0", `execution-attempt` = "0"
2. **Preemption (Suspension)**: JobSet is suspended (`suspend: true`). Pods are deleted.
3. **Resumed Execution**: JobSet is unsuspended (`suspend: false`).
   * `restart-attempt` remains "0"
   * `execution-attempt` increments to "1"

JobSet's volumeClaimPolicies API allows you to automatically create and manage shared PersistentVolumeClaims (PVCs) across multiple ReplicatedJobs within a single JobSet. This is used for stateful workloads requiring persistent storage for datasets, models, or checkpoints.

To implement this:

1. Define volumeClaimPolicies in the JobSet spec.
2. Each policy contains one or more PVC templates.
3. JobSet automatically creates PVCs using the naming convention: <claim-name>-<jobset-name>.

Example of a JobSet using different retention policies for an initializer and checkpoints:

# Note: This is a conceptual representation based on the documentation example
volumeClaimPolicies:
  - templates:
      - metadata:
          name: initializer-volume
        spec:
          accessModes: ["ReadWriteOnce"]
          resources:
            requests:
              storage: 10Gi
    retentionPolicy:
      whenDeleted: Delete
  - templates:
      - metadata:
          name: checkpoints
        spec:
          accessModes: ["ReadWriteOnce"]
          resources:
            requests:
              storage: 50Gi
    retentionPolicy:
      whenDeleted: Retain

Exclusive Topology is a feature used to ensure that all Pods belonging to a single TPU slice are scheduled onto the exact same physical Node Pool. This is critical for TPU Multislice workloads to prevent breaking the high-bandwidth Inter-Chip Interconnect (ICI) links, which rely on physical wiring within a specific hardware boundary (represented in GKE as a Node Pool).

JobSet implements this using a Leader/Follower scheduling mechanism:

1. Leader Scheduling: The JobSet controller ensures that only 'leader' Pods (those with batch.kubernetes.io/job-completion-index: 0) are scheduled first.
2. Follower Binding: 'Follower' Pods (those with a non-zero batch.kubernetes.io/job-completion-index) are only allowed to schedule once their leader is placed. The JobSet webhook intercepts these follower Pods and dynamically injects nodeSelectors to force them into the exact same Node Pool as their leader.

To enable this feature, add the following annotation to your JobSet: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool

JobSet supports configurable failure policies that allow you to distinguish between failures that should trigger a JobSet restart and failures that should cause the JobSet to fail immediately. This is particularly useful for handling maintenance events (like SIGTERM from host maintenance) differently from application errors.

Common Usage Patterns

• Fail on non-maintenance errors: Configure the policy to fail the JobSet immediately if a child job fails for any reason other than a SIGTERM (which typically indicates graceful node shutdown/maintenance). Otherwise, allow the JobSet to restart up to a specified limit (e.g., 10 times).
• Restart on all failures: Configure the policy to restart the JobSet upon any failure type for all replicated jobs, up to a specified limit (e.g., 10 times).

This mechanism relies on the interaction between the JobSet's failure policy and the podFailurePolicy defined on individual pods. By defining which exit codes or signals are considered 'retriable' via podFailurePolicy, you can control whether the JobSet controller treats a pod's termination as a reason to restart the entire JobSet or as a terminal failure.

JobSet supports "elasticity" by allowing the parallelism and completions fields of the underlying Job template to be mutated during execution. This enables Pod-level scaling (adjusting the number of worker pods within a Job) to support workloads like PyTorch Elastic (Torchrun) that need to handle varying worker counts due to hardware failures or resource availability.

Key Constraints:

• Supported Fields: Only parallelism and completions within the template.spec of a replicatedJob are mutable.
• Unsupported (Deferred): Mutating replicas (Job-level scaling), the JobTemplate itself, or container images is not supported in this implementation.
• Terminal State Protection: You cannot mutate parallelism or completions once the JobSet has reached a terminal state (Completed or Failed). Attempting to do so will result in a validation error: Forbidden: Cannot mutate a JobSet in a terminal state.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: pytorch-elastic
spec:
  replicatedJobs:
  - name: workers
    replicas: 1
    template:
      spec:
        completionMode: Indexed
        parallelism: 2      # This field can be scaled (e.g., 2 -> 8)
        completions: 2      # This field can be scaled (e.g., 2 -> 8)
        template:
          spec:
            containers:
            - name: trainer
              image: pytorch/torchserve
              command: ["torchrun", "--nnodes=2:8", "--nproc_per_node=1", "train.py"]
              resources:
                limits:
                  nvidia.com/gpu: 1