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