Kubernetes for ML
Intermediate 16 min read Kubernetes for ML

K8s for Training

Training as a first-class cluster workload

Training on Kubernetes means Job-based semantics: a training run is a declarative workload with backoffs, checkpoints and log history — single-pod by default, distributed via Kubeflow-style operators (TFJob, PyTorchJob) that handle leader-workers and fault recovery.

trainingjobkubeflow operatorpytorchjobbackoffdistributed training
Before we start

What you will learn

  • Run training as Jobs with backoff and restart policy
  • Distinguish single-pod vs distributed training objects
  • Handle checkpoints and cleanup for training pods
Introduction

What is this?

K8s for training models a train as a Job or operator workload: single-pod jobs for small runs; Kubeflow operators (PyTorchJob, TFJob) coordinate leader/workers, gang scheduling and fault restart for distributed runs. Checkpoints, logs and artifacts become cluster-managed resources.

Motivation

Why do you need it?

  • Jobs give retries/backoff to flaky training that plain pods never had.
  • Distributed topologies (worker/ps) are encoded in the workload, not scripts.
  • Cluster-managed storage and cleanup keep training reproducible and tidy.
  • Operators handle the hard parts — failure, eviction, gang scheduling.
Think about it

Mental model

Imagine

A construction crew with a site foreman

Think of

Distributed training on K8s

Individual workers (pods) do shifts; the foreman (operator) assigns roles (leader/worker), coordinates concrete pours (gradient sync) and replaces any worker who doesn't show. Burned-out crew get rebuilt overnight, and the schedule (checkpoint) continues.

Mechanics

How it works

  • A Job YAML declares the image, replicas, resources and restart/backoff policy.
  • For distributed runs, an operator (e.g. Kubeflow) creates leader + worker pods with a spec.
  • The scheduler co-schedules partition-workers for gang-style parallelism.
  • Checkpoints stream to persistent storage; failures resk from the checkpoint.
  • Completed runs clean up pods while preserving logs and artifacts.

A distributed training topology on K8s.

TrainingJob
Operator
Gang schedule
Train + checkpoint
Complete
Real world

In production

  • A nightly fine-tune runs as a Job with backoffLimit=3 and auto-captured logs.
  • A PyTorchJob spreads 8 GPUs across workers with gang scheduling via a queue.
  • Checkpoints on shared storage let an evicted worker resk instead of restart.
Guidelines

Best practices

  • Declare backoffLimit and activeDeadline so runaway jobs can't burn forever.
  • Choose operator-per-framework (PyTorchJob, TFJob) over hand-rolled distributed scripts.
  • Stream checkpoints to persistent storage; use the pod only for compute.
  • Set cleanup TTLs / finalizers so finished training doesn't litter the namespace.
Watch out

Common mistakes

  • Running distributed training as N independent pods with no topology — 'the loss didn't sync'.
  • No restart policy — an evicted worker silently halves the world.
  • Checkpoints on local disk — an eviction or reschedule loses the training run.
  • No deadline: a hang trains forever and bills the whole namespace.
When things break

Debugging tips

  • Workers lag or desync — check the shared-rank/backend config, not the model code.
  • Job stuck in backoff loop — read the failed pod's exit reason and logs; often a config or image issue.
  • Gang job won't start — too few schedulable GPUs; verify quota and node availability.
Interviews

Interview checkpoint

Practice

Hands-on lab

Declare a Training Job
35 min · intermediate
In progress
Objectives
  • Author a training Job with backoff and deadline
  • Add checkpoint-on-storage semantics
  • Model distributed worker topology

Declare a Training Job

35 min · intermediate

0/3 steps
0%
sandbox · Linux sandbox
$
Try these commands
Test yourself

Quick quiz

0/2
1

A training Job's backoffLimit governs…

2

Distributed training on K8s belongs in framework operators, not hand-rolled scripts.

Go deeper

Official documentation