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 running 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

Three checks: GPU host runtime ready, cluster reachable via kubectl, GPU Operator/device plugin present.

# 0. GPU node host runtime.
# Run this on each self-managed GPU worker node or in the node image build.
# Set TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1 only when using managed GPU nodes whose
# driver/toolkit lifecycle is owned by the cloud provider or GPU Operator policy.
if [ "${TAO_K8S_SKIP_NODE_RUNTIME_CHECK:-0}" != "1" ]; then
  TAO_SKILL_BANK_ROOT="${TAO_SKILL_BANK_ROOT:-$PWD}"
  SETUP_SCRIPT="${TAO_SKILL_BANK_ROOT}/skills/platform/tao-setup-nvidia-gpu-host/scripts/setup-nvidia-gpu-host.sh"

  bash "$SETUP_SCRIPT" --backend kubernetes --check-only || {
    echo "MISSING: TAO Kubernetes GPU node runtime is not ready."
    echo "For self-managed GPU nodes, run after user approval:"
    echo "  bash \"$SETUP_SCRIPT\" --backend kubernetes --install --yes"
    echo "For managed clusters, verify the node image/GPU Operator policy installs driver 580 and toolkit 1.19.0, then set TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1."
    exit 1
  }
fi

# 1. Cluster reachable (kubeconfig OR in-cluster service account)
command -v kubectl >/dev/null 2>&1 || {
  echo "MISSING: kubectl not found on PATH. Install kubectl to submit Jobs."
  exit 1
}
kubectl cluster-info >/dev/null 2>&1 || {
  echo "MISSING: no reachable cluster (kubeconfig at ~/.kube/config, \$KUBECONFIG, or in-pod service account)."
  echo "Configure kubectl for your cluster, or set \$KUBECONFIG:"
  echo "  EKS: aws eks update-kubeconfig --name <cluster> --region <region>"
  echo "  GKE: gcloud container clusters get-credentials <cluster> --region <region>"
  echo "  AKS: az aks get-credentials --resource-group <rg> --name <cluster>"
  echo "  local: minikube start   (see 'Local cluster' below)"
  exit 1
}

# 2. NVIDIA GPU Operator present (soft check — warn, don't fail)
gpu=$(kubectl get nodes -o jsonpath='{range .items[*]}{.status.allocatable.nvidia\.com/gpu}{"\n"}{end}' 2>/dev/null | grep -v '^$' | head -1)
if [ -z "$gpu" ] || [ "$gpu" = "0" ]; then
  echo "WARN: no nvidia.com/gpu allocatable on this cluster."
  echo "Install the NVIDIA GPU Operator before submitting GPU jobs:"
  echo "  https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/getting-started.html"
fi

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 GPU-capacity warning here is a soft check; the submit verb re-checks allocatable nvidia.com/gpu and hard-fails before applying the manifest (there is no gang scheduling, so a too-big Job would sit Pending forever).

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, reference its name in the rendered manifest's imagePullSecrets.
• AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY / S3_BUCKET_NAME / S3_ENDPOINT_URL (optional): for S3 dataset I/O (storage tier C), injected into the pod via the per-job Secret (envFrom.secretRef), never inline. Legacy ACCESS_KEY/SECRET_KEY are mapped by tao-data-io.

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.

Execution — the four verbs

tao-run-on-kubernetes is a platform consumer: it runs a spec-bundle via kubectl, mutating only the job-record. No nvidia-tao-sdk, no tao_sdk import — jobs are submitted with plain kubectl apply. $BANK = ${TAO_SKILL_BANK_PATH}.

submit

1. GPU-capacity gate — hard-fail first (no gang scheduling → a too-big Job sits Pending forever):

   ALLOC=$(kubectl get nodes -o jsonpath='{range .items[*]}{.status.allocatable.nvidia\.com/gpu}{"\n"}{end}' | awk '{s+=$1} END{print s+0}')
   [ "${ALLOC:-0}" -ge "$NUM_GPUS" ] || { echo "insufficient GPUs: need $NUM_GPUS, allocatable $ALLOC"; exit 1; }

2. Storage tier (via tao-data-io): A = mount a bound PVC/NFS holding the data (author the mount paths, no fetch — the air-gap answer, and what the packaged template does); C = ephemeral: an initContainer fetches from S3 into a shared emptyDir and a final step uploads results to S3 before TTL.

   Tier C holds the GPU while it downloads. A pod reserves nvidia.com/gpu for its whole lifetime, initContainers included, so a large tier-C fetch — or a first-time multi-GB image pull — is billed and reaper-eligible idle GPU time, exactly like pulling inside a SLURM allocation. Prefer tier A when the data is already on a PVC; choose tier C knowingly, for small inputs.

   A producer action request may declare several mounts, including duplicate- source aliases for logical and embedded absolute paths. Read references/action-request.md and use its staging map plus packaged renderer. For mode=config, materialize the producer's nested spec first, stage that exact generated file, and pass it back to the renderer so the pod receives a verified read-only config mount. The legacy single-root template cannot represent that contract.

3. Open the record — mints the id, binds results_dir, before launch:

   JOB_ID=$("$BANK/scripts/tao_job_record.py" open --platform kubernetes --image "$IMAGE" \
     --network-arch "$ARCH" --action "$ACTION" --storage-tier "$TIER" --results-dir "$RESULTS_DIR")

   results_dir must be a mounted (surviving) volume path or an S3 prefix — ttlSecondsAfterFinished deletes the Job and its logs after it ends, so nothing is recoverable from the Job object later.

4. Render, gate, apply, and record RUNNING. For a producer action request, follow references/action-request.md; it owns backend-name normalization, conditional Secret references, native argv rendering, server dry-run, and binding the applied object name to the job-record. For a simple one-root spec-bundle, render templates/k8s/single-pod-job.yaml.tmpl, run redact_secrets.py lint plus kubectl apply --dry-run=server, apply it, and mark the record with backend-ref=<namespace>/<actual-object-name>.

A submit that skipped the gate or the open has no id — so it cannot launch.

status

Keep K8S_JOB_NAME from submit. On reattach, read the job-record's backend_ref=<namespace>/<name> and recover both values from that field; do not assume the Kubernetes name equals the record id.

kubectl get job "$K8S_JOB_NAME" -n "$NAMESPACE" \
  -o jsonpath='{.status.conditions[0].type} {.status.active} {.status.succeeded} {.status.failed}'

logs

kubectl logs -n "$NAMESPACE" -l "job-name=$K8S_JOB_NAME" --tail "${N:-200}"

cancel

kubectl delete job "$K8S_JOB_NAME" -n "$NAMESPACE" --cascade=foreground
if [ -n "${CRED_SECRET:-}" ]; then
  kubectl delete secret "$CRED_SECRET" -n "$NAMESPACE" --ignore-not-found
fi
"$BANK/scripts/tao_job_record.py" mark "$JOB_ID" --state CANCELED --source agent

Multi-node (nodes > 1)

Same four verbs, plus:

1. Version gate: require k8s ≥ 1.28 (kubectl version -o json) — the pod hostname <job>-<index> (PodIndexLabel) that MASTER_ADDR=<job>-0.<svc> resolves to needs it; on older clusters rank-0 hangs at rendezvous.
2. Capacity gate ×nodes: hard-fail unless allocatable GPUs ≥ gpus_per_node × nodes (no gang scheduling → a partial start leaves rank-0 waiting forever).
3. Render templates/k8s/indexed-job.yaml.tmpl — the headless Service + Indexed Job + rendezvous env (WORLD_SIZE = node count, NODE_RANK from JOB_COMPLETION_INDEX, MASTER_ADDR=<job>-0.<svc>, /dev/shm 16Gi so NCCL doesn't silently hang). kubectl apply -f creates the Service and Job together; cancel deletes the Job (Foreground) and the Service.
4. NCCL probe first (as SLURM) — a 2-node all-reduce with a timeout; on hang, set the cluster NCCL env and re-probe; cache per cluster.

Local cluster (development, CI, and evals)

A throwaway minikube/kind cluster exercises admission, the four verbs, job-record wiring, and log plumbing without cluster quota — and is what an agent-driven eval should provision for itself. kubectl and minikube are single static binaries needing no root, so a non-root CI container can install them itself.

Two prerequisites keep a rendered Job Pending, and the first masks the second: the PVC the template mounts must exist (persistentvolumeclaim "<name>" not found fires before any GPU complaint), then a Job requesting nvidia.com/gpu on a GPU-less cluster reports Insufficient nvidia.com/gpu and waits forever. Render NUM_GPUS=0 for a lifecycle-only run and say GPU scheduling was not verified; on a Linux GPU host, minikube start --driver=docker --gpus all passes real GPUs through, so one GPU box suffices for a GPU-real smoke.

Install commands, driver choice, the container/host-networking caveat, and the fake-device-plugin middle option: references/local-cluster.md.

Container shell

The simple single-pod template invokes its command via /bin/sh -c (POSIX sh, present in busybox/distroless as well as TAO images). For producer action requests, an args-mode command and its arguments are native container argv; a producer that needs a shell declares the shell and its script explicitly. A simple config-mode command also becomes native argv after {config_path} is substituted. A producer-owned config command that is itself a multi-line shell script is preserved verbatim under /bin/sh -c; the renderer never constructs shell text from config values.