tao-run-on-kubernetes

Par nvidia · skills

Plateforme d'exécution Kubernetes — soumet les jobs de conteneurs TAO sous forme de Jobs k8s mono-pod avec ordonnancement GPU NVIDIA.

npx skills add https://github.com/nvidia/skills --skill tao-run-on-kubernetes

Kubernetes

Soumet les jobs de conteneurs TAO en tant que Jobs Kubernetes. Fonctionne sur n'importe quel cluster accessible via kubeconfig (EKS / GKE / AKS / on-prem) ou compte de service in-cluster (quand le SDK s'exécute dans un pod).

Single-pod par défaut ; optez pour l'entraînement distribué multi-nœud via num_nodes > 1 (utilise Indexed Job + Service headless, voir Multi-node training ci-dessous).

Preflight

Quatre vérifications : GPU host runtime prêt, SDK installé, cluster accessible, GPU Operator/device plugin présent.

# 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/tao-setup-nvidia-gpu-host/scripts/setup-nvidia-gpu-host.sh"
  [ -x "$SETUP_SCRIPT" ] || SETUP_SCRIPT="${TAO_SKILL_BANK_ROOT}/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. SDK + kubernetes extra installed.
# nvidia-tao-sdk is on public PyPI; pin lives in versions.yaml (wheels.tao_sdk_kubernetes).
PIN=$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_sdk_kubernetes)
python -c "import tao_sdk" 2>/dev/null || {
  echo "MISSING: nvidia-tao-sdk not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}
python -c "import kubernetes" 2>/dev/null || {
  echo "MISSING: kubernetes extra not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}

# 2. Cluster reachable (kubeconfig OR in-cluster service account)
python -c "from kubernetes import config; config.load_kube_config()" 2>/dev/null || \
  python -c "from kubernetes import config; config.load_incluster_config()" 2>/dev/null || {
    echo "MISSING: no kubeconfig at ~/.kube/config and not running in a pod."
    echo "Configure kubectl (e.g., 'aws eks update-kubeconfig --name my-cluster') or set \$KUBECONFIG."
    exit 1
  }

# 3. NVIDIA GPU Operator present (soft check — warn if kubectl available, don't fail)
if command -v kubectl >/dev/null 2>&1; then
  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
fi

La vérification du GPU node runtime est obligatoire pour les nœuds auto-gérés. Pour les clusters gérés où le client ne s'exécute pas sur un nœud GPU worker, vérifiez l'image de nœud du fournisseur ou la politique du GPU Operator et définissez TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1 au lieu d'exécuter le programme d'installation sur le client. La vérification finale de la capacité GPU est un avertissement plutôt qu'une défaillance matérielle — kubectl n'est pas toujours installé. Le SDK effectue une vérification stricte dans KubernetesSDK.create_job() qui utilise le client Python kubernetes pour vérifier la capacité GPU avant la soumission.

Credentials & configuration

  • Kubeconfig (l'un des suivants) :
    • ~/.kube/config — chemin de découverte par défaut
    • $KUBECONFIG — chemin alternatif
    • Compte de service in-cluster — utilisé lors de l'exécution dans un pod (pas besoin de kubeconfig)
  • TAO_K8S_NAMESPACE (optionnel) : espace de noms par défaut pour la soumission du Job. Par défaut default.
  • TAO_K8S_CONTEXT (optionnel) : nom du contexte kubeconfig pour basculer les clusters.
  • NGC_KEY (optionnel) : pour les pulls d'image nvcr.io. Si vous avez pré-créé un secret image-pull dans l'espace de noms cible, passez son nom à create_job via l'argument image_pull_secret.
  • ACCESS_KEY / SECRET_KEY / S3_BUCKET_NAME / S3_ENDPOINT_URL (optionnel) : pour l'I/O de dataset S3 via le wrapping inputs/outputs du script_runner du SDK.

Ne demandez pas d'identifiants Lepton, Brev ou SLURM pour les exécutions Kubernetes. Demandez les identifiants S3 uniquement quand le workflow sélectionné utilise des entrées ou sorties s3://, et demandez les identifiants propres au modèle tels que HF_TOKEN uniquement quand le modèle sélectionné l'exige. Avant le lancement, vérifiez que l'espace de noms sélectionné peut créer des Jobs, que les chemins de dataset/résultat sont visibles depuis le pod, et que les chemins PVC/filesystem montés sont prouvés être montés dans le conteneur du job ; un chemin local de l'agent-host n'est pas une preuve suffisante.

SDK API

K8s est SDK-only — il n'y a pas de chemin de lancement kubectl-only. Lisez tao-skill-bank:tao-run-platform avant de rédiger les appels create_job ; elle couvre build_entrypoint, le contrat kwarg partagé, le monitoring et ActionWorkflow.

from tao_sdk.platforms.kubernetes import KubernetesSDK

sdk = KubernetesSDK()  # auto-detects auth
job = sdk.create_job(
    image='nvcr.io/nvidia/tao/tao-toolkit:6.26.3-pyt',
    command='dino train -e /tmp/spec.yaml',
    gpu_count=1,
    env_vars={'NGC_KEY': os.environ['NGC_KEY']},
    inputs={'/data/train.json': 's3://bucket/coco/train.json'},
    outputs=['/results/'],
    namespace='tao-jobs',                       # optional override
    image_pull_secret='ngc-pull-secret',         # optional, pre-created
    node_selector={'gpu-type': 'h100'},          # optional
)

Le SDK construit un V1Job avec :

  • spec.template.spec.containers[0] : l'image demandée et command=["/bin/bash", "-c", <command>].
  • resources.limits["nvidia.com/gpu"]: <gpu_count> — ordonnance sur les nœuds GPU via le Device Plugin NVIDIA / GPU Operator.
  • env_vars passés, plus l'injection automatique des identifiants S3/NGC/HF pour script_runner.
  • restart_policy=Never et backoff_limit=0 — les défaillances remontent à l'utilisateur au lieu de réessayer silencieusement.
  • ttl_seconds_after_finished=3600 — le Job se nettoie automatiquement 1 heure après l'état terminal.

Status & monitoring

status = sdk.get_job_status(job.id)
# status.status ∈ {"Pending", "Running", "Complete", "Error", "Canceled", "Unknown"}

logs = sdk.get_job_logs(job.id, tail=200)  # concatenates logs from all pods of the Job

# For stuck-Pending jobs — replica diagnostics:
for r in sdk.get_job_replicas(job.id):
    issue = r["status"].get("readiness_issue")
    if issue:
        print(issue["reason"], issue["message"])
        # e.g. "ImagePullBackOff" / "Back-off pulling image..."
        # e.g. "Pending"           / "0/3 nodes available: 3 Insufficient nvidia.com/gpu"

# On failure:
analysis = sdk.get_failure_analysis(job.id)
# {"err_class": "ERR_PROGRAM" | "ERR_INFRA",
#  "suggestion": "Container OOM-killed. Reduce batch size...",
#  "job_failure_by_node_event": [{"node_event_name": "OOMKilled", ...}]}

Cancel & cleanup

sdk.cancel_job(job.id)  # delete_namespaced_job with propagation_policy="Foreground"

ttl_seconds_after_finished=3600 signifie que les Jobs complétés se suppriment automatiquement après 1 h. Pour annuler un Job en cours, cancel_job le supprime ainsi que ses pods immédiatement.

GPU Operator dependency

Le SDK refuse de soumettre les jobs GPU à un cluster sans nvidia.com/gpu allouable. Pour les clusters auto-gérés, exécutez d'abord l'action d'installation tao-setup-nvidia-gpu-host sur chaque nœud GPU worker ou intégrez le même ensemble de paquets dans l'image de nœud :

bash skills/platform/tao-setup-nvidia-gpu-host/scripts/setup-nvidia-gpu-host.sh --backend kubernetes --install --yes

Ensuite, installez le NVIDIA GPU Operator ou le device plugin :

helm repo add nvidia https://helm.ngc.nvidia.com/nvidia
helm repo update
helm install --wait gpu-operator -n gpu-operator --create-namespace nvidia/gpu-operator

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

Multi-node training (distributed)

Passez num_nodes > 1 à create_job() pour exécuter l'entraînement distribué sur N pods. Le SDK provisionne :

  1. Un Service headless nommé d'après le Job (sélecteur : job-name=<job-name>, clusterIP: None, publishNotReadyAddresses: true afin que les pods puissent se rencontrer avant d'être tous Ready).

  2. Un Indexed Job avec parallelism = completions = num_nodes, completionMode: Indexed. Chaque pod reçoit JOB_COMPLETION_INDEX injecté automatiquement par k8s (= le rang du nœud).

  3. Un command wrapper qui exporte les variables d'environnement de rendezvous avant d'invoquer la commande utilisateur. Deux conventions de nommage sont exportées simultanément :

    Variable d'environnement Valeur Lue par
    WORLD_SIZE num_nodes entrypoint du conteneur TAO PyTorch nvidia_tao_pytorch/core/entrypoint.py (utilise ceci pour désigner le node count, bien que la convention propre de PyTorch soit total processes)
    NUM_GPU_PER_NODE gpu_count entrypoint du conteneur TAO PyTorch
    NNODES num_nodes torchrun et rendezvous standard PyTorch
    NPROC_PER_NODE gpu_count torchrun
    NODE_RANK $JOB_COMPLETION_INDEX les deux
    MASTER_ADDR <job-name>-0.<job-name> (DNS du pod-0) les deux
    MASTER_PORT 29500 les deux (défaut TAO)

    Les deux conventions de nommage sont définie afin que les entrypoints TAO (dino train, etc.) et les commandes raw torchrun fonctionnent sans modification.

job = sdk.create_job(
    image='nvcr.io/nvidia/tao/tao-toolkit:6.26.3-pyt',
    command='dino train -e /tmp/spec.yaml',  # TAO entrypoint reads spec.train.num_nodes; env vars are wired by the container
    gpu_count=8,           # GPUs per node
    num_nodes=4,           # 4 × 8 = 32 GPUs total
    inputs={'/data/train.json': 's3://bucket/coco/train.json'},
    outputs=['/results/'],
)

Pour les commandes raw torchrun (conteneurs non-TAO) :

job = sdk.create_job(
    image='nvcr.io/nvidia/pytorch:25.08-py3',
    command='torchrun --nnodes=$NNODES --nproc-per-node=$NPROC_PER_NODE --node-rank=$NODE_RANK '
            '--master-addr=$MASTER_ADDR --master-port=$MASTER_PORT train.py',
    gpu_count=8,
    num_nodes=4,
)

La vérification de capacité somme sur les nœuds : gpu_count × num_nodes ≤ nvidia.com/gpu allouable du cluster.

Cluster requirements for multi-node

  • k8s 1.28+ est requis pour des noms de pods stables dans les Indexed Jobs (la fonctionnalité PodIndexLabel). Sur les anciens clusters, la recherche DNS MASTER_ADDR=<job>-0.<svc> échoue. Vérifiez avec kubectl version.
  • Pod-to-pod networking doit être ouvert sur le port 29500 (défaut PyTorch ; configurable via la variable d'environnement MASTER_PORT). La plupart des CNI (Calico, Cilium, AWS VPC CNI) autorisent ceci par défaut ; les NetworkPolicies restrictives doivent être assouplies.
  • NCCL dans le conteneur communique GPU-to-GPU ; si le cluster a des nœuds multi-NIC ou RDMA, définissez NCCL_SOCKET_IFNAME / NCCL_IB_HCA via env_vars.

Reference reading

When to use a Kubernetes operator instead

Pour les topologies plus sophistiquées (gang scheduling, PyTorch elastic / fault-tolerant training, MPI / Horovod, setup RDMA), préférez un operator au lieu de plain Indexed Job :

Le chemin Indexed Job du TAO SDK est intentionnellement simple et sans dépendances ; si vous avez besoin d'elastic restart ou de gang scheduling, couchez l'un de ceux-ci par-dessus et soumettez les jobs via le CRD de l'operator à la place.

Common error patterns

No nvidia.com/gpu resources allocatable on the cluster — le GPU Operator (ou NVIDIA Device Plugin) n'est pas installé. Installez via le lien ci-dessus ; vérifiez avec kubectl get nodes -o jsonpath='{.items[*].status.allocatable}'.

ImagePullBackOff / ErrImagePull — le cluster ne peut pas tirer l'image. Pour nvcr.io : pré-créez un secret image-pull dans l'espace de noms et passez son nom via l'argument image_pull_secret :

kubectl create secret docker-registry ngc-pull-secret \
  --docker-server=nvcr.io \
  --docker-username='$oauthtoken' \
  --docker-password=$NGC_KEY -n tao-jobs

Pod reste Pending indéfiniment — get_job_replicas(job_id) affichera le readiness_issue. Causes courantes : capacité GPU insuffisante (Insufficient nvidia.com/gpu), aucun nœud ne correspond au node_selector, secret image-pull manquant, ou défaillance du montage PVC.

OOMKilled (exit 137) — le conteneur a dépassé la mémoire. Réduisez la batch size, réduisez max_length, ou ajoutez une memory request/limit et ciblez un nœud plus gros.

CredentialError: Could not authenticate to a Kubernetes cluster — ni kubeconfig ni l'auth in-cluster n'ont fonctionné. Exécutez kubectl get nodes pour vérifier votre config, ou définissez $KUBECONFIG vers le bon chemin.

What this skill does NOT support (yet)

  • Elastic / fault-tolerant training. Indexed Job a backoff_limit=0 — les défaillances font échouer toute la run d'entraînement. Pour elastic restart (p. ex., reprendre à partir d'un checkpoint après une mort de nœud), utilisez plutôt l'operator PyTorchJob de Kubeflow.
  • Gang scheduling. Les pods Indexed Job sont ordonnancés indépendamment — pas d'all-or-nothing. L'entraînement multi-nœud démarrera partiellement si seulement certains pods peuvent être ordonnancés (rank-0 restera suspendu en attente de pairs). Pour l'ordonnancement all-or-nothing sur les clusters partagés, utilisez Volcano ou Kueue.
  • MPI / Horovod. Utilisez l'MPI Operator. Le chemin Indexed Job ici est en forme PyTorch-distributed (env-var rendezvous sur MASTER_ADDR:MASTER_PORT).
  • Persistent volumes pour le stockage partagé. S3 uniquement via le script_runner. Le support PVC est un suivi.
  • Auto-création de secrets image-pull à partir de $NGC_KEY. Vous pré-créez le secret dans l'espace de noms cible et passez le nom. Lepton le fait automatiquement ; nous ne le faisons pas ici car les conventions d'espace de noms k8s varient largement.

Skills similaires

tao-run-on-local-docker
nvidia / skills

Exécuter des jobs TAO comme conteneurs Docker locaux sur GPU NVIDIA.

1 116
multi-node-slurm
nvidia / skills

Convertir des commandes distribuées PyTorch en scripts Slurm multi-nœuds avec conteneurs Enroot.

1 116
tao-run-platform
nvidia / skills

Gérer des jobs TAO avec handles, I/O S3 et support multi-plateformes GPU.

1 116
tao-run-on-brev
nvidia / skills

Gérer des instances GPU NVIDIA on-demand via le CLI Brev avec authentification automatisée.

1 116
launch-nemo-rl
nvidia / skills

Lancer et gérer des recettes NeMo-RL sur Kubernetes via la CLI nrl-k8s.

1 116