tao-run-on-kubernetes

Kubernetes

Standalone install? If this session was not initialized by the TAO skill bank plugin, run the tao-setup skill first (host preflight, credentials, cross-skill discovery).

Submits TAO container jobs as Kubernetes Jobs. Works on any cluster reachable via kubeconfig (EKS / GKE / AKS / on-prem) or in-cluster service account (when the SDK runs inside a pod).

Single-pod by default; opt into multi-node distributed training via num_nodes > 1 (uses Indexed Job + headless Service, see Multi-node training below).

Preflight

Four checks: GPU host runtime ready, SDK installed, cluster reachable, GPU Operator/device plugin present.

For self-managed nodes, if the selected model declares runtime_requirements.gpu_host, pass the corresponding --min-*-version flags to check and install commands. On managed clusters, require the provider node image or GPU Operator policy to satisfy that model profile.

The GPU node runtime check is mandatory for self-managed nodes. For managed clusters where the client is not running on a GPU worker, verify the provider node image or GPU Operator policy and set TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1 instead of running the installer on the client. The final GPU capacity check is a warning rather than a hard fail — kubectl isn't always installed. The SDK does a hard guard inside KubernetesSDK.create_job() that uses the kubernetes Python client to verify GPU capacity before submitting.

Credentials & configuration

• Kubeconfig (one of):
  • ~/.kube/config — default discovery path
  • $KUBECONFIG — alternate path
  • In-cluster service account — used when running inside a pod (no kubeconfig needed)
• TAO_K8S_NAMESPACE (optional): default namespace for Job submission. Defaults to default.
• TAO_K8S_CONTEXT (optional): kubeconfig context name to switch clusters.
• NGC_KEY (optional): for nvcr.io image pulls. If you've pre-created an image-pull secret in the target namespace, pass its name to create_job via the image_pull_secret argument.
• ACCESS_KEY / SECRET_KEY / S3_BUCKET_NAME / S3_ENDPOINT_URL (optional): for S3 dataset I/O via the SDK's inputs/outputs script_runner wrapping.

Do not ask for Brev or SLURM credentials for Kubernetes runs. Ask for S3 credentials only when the selected workflow uses s3:// inputs or outputs, and ask for model-specific credentials such as HF_TOKEN only when the selected model requires them. Before launch, verify the selected namespace can create Jobs, dataset/result paths are visible from the pod, and PVC/mounted filesystem paths are proven to be mounted into the job container; an agent-host local path is not sufficient proof.

SDK API

K8s is SDK-only — there is no kubectl-only launch path. Read tao-skill-bank:tao-run-platform before drafting create_job calls; it covers build_entrypoint, the shared kwarg contract, monitoring, and ActionWorkflow.

The SDK constructs a V1Job with:

• spec.template.spec.containers[0]: the requested image and command=["/bin/bash", "-c", <command>].
• resources.limits["nvidia.com/gpu"]: <gpu_count> — schedules onto GPU nodes via the NVIDIA Device Plugin / GPU Operator.
• env_vars flowed through, plus auto-injected S3/NGC/HF credentials for script_runner.
• restart_policy=Never and backoff_limit=0 — failures surface to the user instead of silently retrying.
• ttl_seconds_after_finished=3600 — Job auto-cleans 1 hour after terminal state.

Status & monitoring

Cancel & cleanup

Completed Jobs are not TTL-deleted: their UID and terminal evidence remain available for safe recovery after an SDK restart. To cancel an in-flight Job, cancel_job deletes it and its pods with foreground propagation and reports success only after the writers are absent. For cleanup-supported S3 jobs, delete_job_artifacts first verifies and removes that exact terminal Job and its pods, then verifies deletion of the job-scoped S3 prefix. A same-name Job with a different UID is never deleted.

GPU Operator dependency

The SDK refuses to submit GPU jobs to a cluster with no nvidia.com/gpu allocatable. For self-managed clusters, first run the tao-setup-nvidia-gpu-host install action on every GPU worker node or bake the same package set into the node image:

Then install the NVIDIA GPU Operator or device plugin:

Full guide: https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/getting-started.html

Multi-node training (distributed)

Pass num_nodes > 1 to create_job() to run distributed training across N pods. The SDK provisions:

1. A headless Service named after the Job (selector: job-name=<job-name>, clusterIP: None, publishNotReadyAddresses: true so pods can rendezvous before they're all Ready).

2. An Indexed Job with parallelism = completions = num_nodes, completionMode: Indexed. Each pod gets JOB_COMPLETION_INDEX injected by k8s automatically (= the node rank).

3. A command wrapper that exports the rendezvous env vars before invoking the user command. Two naming conventions are exported simultaneously:

   Env var	Value	Read by
   WORLD_SIZE	num_nodes	TAO PyTorch container's nvidia_tao_pytorch/core/entrypoint.py (uses this to mean node count, even though PyTorch's own convention is total processes)
   NUM_GPU_PER_NODE	gpu_count	TAO PyTorch container's entrypoint
   NNODES	num_nodes	torchrun and PyTorch-standard rendezvous
   NPROC_PER_NODE	gpu_count	torchrun
   NODE_RANK	$JOB_COMPLETION_INDEX	both
   MASTER_ADDR	<job-name>-0.<job-name> (pod-0's DNS)	both
   MASTER_PORT	29500	both (TAO's default)

   Both naming conventions are set so TAO entrypoints (dino train, etc.) and raw torchrun commands work without modification.

For raw torchrun-based commands (non-TAO containers):

The capacity check sums across nodes: gpu_count × num_nodes ≤ cluster's allocatable nvidia.com/gpu.

Cluster requirements for multi-node

• k8s 1.28+ is required for stable pod hostnames in Indexed Jobs (the PodIndexLabel feature). On older clusters the MASTER_ADDR=<job>-0.<svc> DNS lookup fails. Verify with kubectl version.
• Pod-to-pod networking must be open on port 29500 (PyTorch default; configurable via MASTER_PORT env var). Most CNIs (Calico, Cilium, AWS VPC CNI) allow this by default; restrictive NetworkPolicies must be relaxed.
• NCCL in the container talks GPU-to-GPU; if the cluster has multi-NIC nodes or RDMA, set NCCL_SOCKET_IFNAME / NCCL_IB_HCA via env_vars.

Reference reading

• Kubernetes Indexed Job: https://kubernetes.io/docs/concepts/workloads/controllers/job/#completion-mode
• Indexed Job for batch ML: https://kubernetes.io/blog/2022/06/01/indexed-jobs-mpi/
• PyTorch distributed (env-var rendezvous): https://pytorch.org/docs/stable/elastic/run.html
• NCCL networking tuning (NCCL_SOCKET_IFNAME, NCCL_IB_HCA): https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html

Kubernetes operator alternatives

For more sophisticated topologies (gang scheduling, PyTorch elastic / fault-tolerant training, MPI / Horovod, RDMA setup), reach for an operator instead of plain Indexed Job:

• MPI Operator — https://github.com/kubeflow/mpi-operator — for MPI / Horovod workloads.
• Kubeflow Training Operator (PyTorchJob, TFJob) — https://www.kubeflow.org/docs/components/training/ — for elastic PyTorch training with built-in restart logic.
• Volcano — https://volcano.sh/ — gang scheduling, queues, fair-share. Useful in shared multi-tenant clusters.
• Kueue — https://kueue.sigs.k8s.io/ — quota / queue layer on top of any of the above.

The TAO SDK's Indexed Job path is intentionally simple and dependency-free; if you need elastic restart or gang scheduling, layer one of these on top and submit jobs through the operator's CRD instead.

Common error patterns

No nvidia.com/gpu resources allocatable on the cluster — the GPU Operator (or NVIDIA Device Plugin) isn't installed. Install per the link above; verify with kubectl get nodes -o jsonpath='{.items[*].status.allocatable}'.

ImagePullBackOff / ErrImagePull — the cluster can't pull the image. For nvcr.io: pre-create an image-pull secret in the namespace and pass its name via the image_pull_secret argument. Feed the key over stdin — kubectl create secret docker-registry --docker-password=$NGC_KEY would put the secret in argv, visible in the host's process table and shell history:

Pod stays Pending forever — get_job_replicas(job_id) will show the readiness_issue. Common causes: insufficient GPU capacity (Insufficient nvidia.com/gpu), no node matches node_selector, missing image-pull secret, or PVC mount failure.

OOMKilled (exit 137) — container exceeded memory. Reduce batch size, lower max_length, or add a memory request/limit and target a larger node.

CredentialError: Could not authenticate to a Kubernetes cluster — neither kubeconfig nor in-cluster auth worked. Run kubectl get nodes to verify your config, or set $KUBECONFIG to the right path.

What this skill does NOT support (yet)

• Elastic / fault-tolerant training. Indexed Job has backoff_limit=0 — failures fail the whole training run. For elastic restart (e.g., resume from checkpoint after a node death), use Kubeflow's PyTorchJob operator instead.
• Gang scheduling. Indexed Job pods are scheduled independently — no all-or-nothing. Multi-node training will partially start if only some pods can be scheduled (rank-0 will hang waiting for peers). For all-or-nothing scheduling on shared clusters, use Volcano or Kueue.
• MPI / Horovod. Use the MPI Operator. The Indexed Job path here is PyTorch-distributed-shaped (env-var rendezvous on MASTER_ADDR:MASTER_PORT).
• Persistent volumes for shared storage. S3 only via the script_runner. PVC support is a follow-up.
• Auto-creating image-pull secrets from $NGC_KEY. You pre-create the secret in the target namespace and pass the name. K8s namespace conventions vary widely, so we keep secret creation explicit.*