Generalize distributed PyTorch training example

[jayeshmahajan]

Summary

This PR adds a new example demonstrating distributed training with PyTorch's Distributed Data Parallel (DDP) on Kubernetes. The example showcases multi-node, multi-GPU training using Kubernetes Jobs with comprehensive support for major cloud providers (GKE, EKS, AKS) and on-premises deployments.

What This Example Demonstrates

• Distributed Data Parallel (DDP) Training: Multi-node, multi-GPU PyTorch training using DDP
• Kubernetes Jobs with Indexed Completion: Coordinated parallel training workers using completionMode: Indexed
• Pod-to-Pod Communication: Headless Services for stable DNS-based worker discovery
• Persistent Storage: PVCs for training data and model checkpoints
• Workload-Aware Scheduling: Integration with Kubernetes v1.35+ workload scheduling (optional)

Key Features

1. Distributed Training Setup

• Uses PyTorch DDP for gradient synchronization across workers
• Automatic rank assignment from Kubernetes Job completion index
• Master worker discovery via headless Service DNS
• DistributedSampler for data sharding across workers

2. Kubernetes Resources

• Job: Indexed completion mode for stable pod naming and rank assignment
• Headless Service: Enables direct pod-to-pod communication
• PersistentVolumeClaims: Separate volumes for training data and outputs
• ConfigMaps: Training script and hyperparameters
• Workload: Workload-aware scheduling support (Kubernetes v1.35+)

3. Multi-Cloud and On-Premises Support

• Base configuration: Generic setup that works across environments
• Kustomize overlays: Provider-specific configurations for:
  • Google Kubernetes Engine (GKE)
  • Amazon Elastic Kubernetes Service (EKS)
  • Azure Kubernetes Service (AKS)
  • On-premises Kubernetes
• Comprehensive comments explaining cloud-specific vs generic configurations
• Storage class guidance for different deployment scenarios

4. Training Script

• Simple CNN model for CIFAR-10 classification
• Automatic dataset download (CIFAR-10)
• Checkpoint saving at each epoch
• TensorBoard logging support
• Proper DDP initialization and cleanup

Files Included

• training-job.yaml - Main Kubernetes Job configuration
• train.py - PyTorch DDP training script
• training-script-configmap.yaml - Training script as ConfigMap
• service.yaml - Headless Service for pod communication
• data-pvc.yaml / output-pvc.yaml - Persistent storage
• train-config.yaml - Training hyperparameters
• workload.yaml - Workload-aware scheduling configuration
• kustomization.yaml - Kustomize base configuration
• README.md - Comprehensive documentation

[k8s-ci-robot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by: jayeshmahajan
Once this PR has been reviewed and has the lgtm label, please assign soltysh for approval. For more information see the Code Review Process.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[soltysh]

This PR has several problems:

1. It uses scheduling.k8s.io/v1alpha1. I prefer our examples use stable APIs that k8s supports.
2. The kustomiztion is pointing to non-existing files. Have you tried if this example works?
3. Why both train.py and training-script-configmap.yaml are needed? I haven't compared them line-by-line, but I'd prefer to have a single copy, and if needed just use kubectl create configmap to populate the necessary resource.

[k8s-triage-robot]

The Kubernetes project currently lacks enough contributors to adequately respond to all PRs.

This bot triages PRs according to the following rules:

• After 90d of inactivity, lifecycle/stale is applied
• After 30d of inactivity since lifecycle/stale was applied, lifecycle/rotten is applied
• After 30d of inactivity since lifecycle/rotten was applied, the PR is closed

You can:

• Mark this PR as fresh with /remove-lifecycle stale
• Close this PR with /close
• Offer to help out with Issue Triage

Please send feedback to sig-contributor-experience at kubernetes/community.

/lifecycle stale