tao-run-on-kubernetes

Kubernetes

将 TAO 容器作业提交为 Kubernetes Jobs。适用于任何可通过 kubeconfig 访问的集群（EKS / GKE / AKS / 本地部署），或使用集群内服务账号（当 SDK 在 Pod 内运行时）。

默认使用单个 Pod；通过 num_nodes > 1 可启用多节点分布式训练（使用 Indexed Job + 无头 Service，详见下文的多节点训练（分布式））。

预检

四项检查：GPU 主机运行时已就绪、SDK 已安装、集群可访问、GPU Operator / 设备插件已就位。

# 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}/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 "Installing missing Python requirement: $PIN"
  python -m pip install "$PIN"
}
python -c "import kubernetes" 2>/dev/null || {
  echo "Installing missing Python requirement: $PIN"
  python -m pip install "$PIN"
}
python -c "import tao_sdk, kubernetes"

# 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

对于自行管理的节点，GPU 节点运行时检查是必需的。对于客户端不在 GPU 工作节点上运行的托管集群，请验证提供商的节点映像或 GPU Operator 策略，并设置 TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1，而不是在客户端运行安装程序。最终的 GPU 容量检查是警告而非硬性失败——kubectl 不一定已安装。SDK 会在 KubernetesSDK.create_job() 内执行硬性保护：使用 kubernetes Python 客户端，在提交前验证 GPU 容量。

凭据与配置

• Kubeconfig（以下任选其一）：
  • ~/.kube/config — 默认发现路径
  • $KUBECONFIG — 替代路径
  • 集群内服务账号 — 在 Pod 内运行时使用（无需 kubeconfig）
• TAO_K8S_NAMESPACE（可选）：提交 Job 时使用的默认命名空间。默认为 default。
• TAO_K8S_CONTEXT（可选）：要切换到的集群对应的 kubeconfig context 名称。
• NGC_KEY（可选）：用于拉取 nvcr.io 镜像。如果已在目标命名空间中预先创建镜像拉取密钥，请通过 image_pull_secret 参数将其名称传递给 create_job。
• ACCESS_KEY / SECRET_KEY / S3_BUCKET_NAME / S3_ENDPOINT_URL（可选）：通过 SDK 的 inputs/outputs script_runner 封装进行 S3 数据集 I/O。

对于 Kubernetes 运行，不要索要 Brev 或 SLURM 凭据。仅当所选工作流使用 s3:// 输入或输出时，才索要 S3 凭据；仅当所选模型需要时，才索要模型专用凭据，例如 HF_TOKEN。启动前，请确认所选命名空间有权创建 Jobs、Pod 可以访问数据集 / 结果路径，并已验证 PVC / 挂载的文件系统路径确实挂载到了作业容器中；Agent 主机上的本地路径不足以作为证明。

SDK API

K8s 仅支持 SDK — 没有仅使用 kubectl 的启动路径。编写 create_job 调用前，请阅读 tao-skill-bank:tao-run-platform；其中介绍了 build_entrypoint、共享的 kwarg 契约、监控和 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
)

SDK 会构建一个 V1Job，其中包含：

• spec.template.spec.containers[0]：请求的镜像以及 command=["/bin/bash", "-c", <command>]。
• resources.limits["nvidia.com/gpu"]: <gpu_count> — 通过 NVIDIA Device Plugin / GPU Operator 将作业调度到 GPU 节点。
• 传入的 env_vars，以及为 script_runner 自动注入的 S3 / NGC / HF 凭据。
• restart_policy=Never 和 backoff_limit=0 — 失败会直接呈现给用户，而不是静默重试。
• ttl_seconds_after_finished=3600 — Job 进入终止状态 1 小时后自动清理。

状态与监控

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", ...}]}

取消与清理

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

ttl_seconds_after_finished=3600 表示已完成的 Jobs 会在 1 小时后自动删除。若要取消正在运行的 Job，cancel_job 会立即删除该 Job 及其 Pod。

GPU Operator 依赖

如果集群没有可分配的 nvidia.com/gpu，SDK 会拒绝向其提交 GPU 作业。对于自行管理的集群，请先在每个 GPU 工作节点上运行 tao-setup-nvidia-gpu-host 安装操作，或将相同的软件包集成到节点映像中：

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

然后安装 NVIDIA GPU Operator 或设备插件：

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

完整指南：https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/getting-started.html

多节点训练（分布式）

将 num_nodes > 1 传入 create_job()，即可在 N 个 Pod 上运行分布式训练。SDK 会配置：

1. 一个以 Job 命名的无头 Service（选择器：job-name=<job-name>，clusterIP: None，publishNotReadyAddresses: true，这样即使 Pod 尚未全部 Ready，也能进行 rendezvous）。

2. 一个索引式 Job，其中 parallelism = completions = num_nodes，completionMode: Indexed。k8s 会自动为每个 Pod 注入 JOB_COMPLETION_INDEX（即节点 rank）。

3. 一个命令包装器，在调用用户命令前导出 rendezvous 环境变量。它会同时导出两种命名约定：

   环境变量	值	读取方
   WORLD_SIZE	num_nodes	TAO PyTorch 容器的 nvidia_tao_pytorch/core/entrypoint.py（该容器将其作为节点数使用，尽管 PyTorch 自身的约定是进程总数）
   NUM_GPU_PER_NODE	gpu_count	TAO PyTorch 容器的 entrypoint
   NNODES	num_nodes	torchrun 和 PyTorch 标准 rendezvous
   NPROC_PER_NODE	gpu_count	torchrun
   NODE_RANK	$JOB_COMPLETION_INDEX	两者
   MASTER_ADDR	<job-name>-0.<job-name>（pod-0 的 DNS）	两者
   MASTER_PORT	29500	两者（TAO 的默认值）

   同时设置两种命名约定，这样 TAO entrypoint（dino train 等）和原始 torchrun 命令都无需修改即可运行。

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/'],
)

对于基于原始 torchrun 的命令（非 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,
)

容量检查会汇总所有节点：gpu_count × num_nodes ≤ 集群可分配的 nvidia.com/gpu。

多节点集群要求

• 稳定的索引式 Job Pod 主机名需要 k8s 1.28+（PodIndexLabel 功能）。在较旧的集群中，MASTER_ADDR=<job>-0.<svc> DNS 查询会失败。请使用 kubectl version 验证。
• 必须开放 Pod 到 Pod 的网络通信，端口为 29500（PyTorch 默认值；可通过 MASTER_PORT 环境变量配置）。大多数 CNI（Calico、Cilium、AWS VPC CNI）默认允许此通信；必须放宽限制严格的 NetworkPolicy。
• 容器中的 NCCL 会进行 GPU 到 GPU 通信；如果集群有多网卡节点或 RDMA，请通过 env_vars 设置 NCCL_SOCKET_IFNAME / NCCL_IB_HCA。

参考资料

• Kubernetes 索引式 Job：https://kubernetes.io/docs/concepts/workloads/controllers/job/#completion-mode
• 用于批处理 ML 的索引式 Job：https://kubernetes.io/blog/2022/06/01/indexed-jobs-mpi/
• PyTorch 分布式（环境变量 rendezvous）：https://pytorch.org/docs/stable/elastic/run.html
• NCCL 网络调优（NCCL_SOCKET_IFNAME、NCCL_IB_HCA）：https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html

Kubernetes Operator 替代方案

对于更复杂的拓扑（gang 调度、PyTorch 弹性 / 容错训练、MPI / Horovod、RDMA 设置），请使用 Operator，而不是普通的索引式 Job：

• MPI Operator — https://github.com/kubeflow/mpi-operator — 适用于 MPI / Horovod 工作负载。
• Kubeflow Training Operator（PyTorchJob、TFJob）— https://www.kubeflow.org/docs/components/training/ — 适用于内置重启逻辑的弹性 PyTorch 训练。
• Volcano — https://volcano.sh/ — gang 调度、队列、公平共享。适用于共享的多租户集群。
• Kueue — https://kueue.sigs.k8s.io/ — 在上述任一方案之上提供配额 / 队列层。

TAO SDK 的索引式 Job 路径有意保持简单且无依赖；如果需要弹性重启或 gang 调度，可在其上叠加其中一种方案，并通过 Operator 的 CRD 提交 Job。

常见错误模式

No nvidia.com/gpu resources allocatable on the cluster — 未安装 GPU Operator（或 NVIDIA Device Plugin）。请按照上方链接安装；使用 kubectl get nodes -o jsonpath='{.items[*].status.allocatable}' 验证。

ImagePullBackOff / ErrImagePull — 集群无法拉取镜像。对于 nvcr.io：请在命名空间中预先创建镜像拉取密钥，并通过 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 永远处于 Pending 状态 — get_job_replicas(job_id) 会显示 readiness_issue。常见原因：GPU 容量不足（Insufficient nvidia.com/gpu）、没有节点匹配 node_selector、缺少镜像拉取密钥，或 PVC 挂载失败。

OOMKilled（退出代码 137） — 容器超出内存限制。请减小批次大小、降低 max_length，或添加内存请求/限制并选择更大的节点。

CredentialError: Could not authenticate to a Kubernetes cluster — kubeconfig 和集群内身份验证均未成功。运行 kubectl get nodes 验证配置，或将 $KUBECONFIG 设置为正确的路径。

此技能目前不支持的功能

• 弹性／容错训练。 Indexed Job 的 backoff_limit=0——失败会导致整个训练运行失败。若要弹性重启（例如节点故障后从检查点恢复），请改用 Kubeflow 的 PyTorchJob Operator。
• Gang 调度。 Indexed Job 的 Pod 会独立调度——不支持全有或全无。如果只能调度部分 Pod，多节点训练会部分启动（rank-0 会一直等待其他节点）。若要在共享集群中实现全有或全无调度，请使用 Volcano 或 Kueue。
• MPI / Horovod。 请使用 MPI Operator。此处的 Indexed Job 路径采用 PyTorch 分布式模式（通过 MASTER_ADDR:MASTER_PORT 上的环境变量进行会合）。
• 用于共享存储的持久卷。 仅通过 script_runner 支持 S3。PVC 支持将在后续添加。
• 根据 $NGC_KEY 自动创建镜像拉取密钥。 你需要在目标命名空间中预先创建密钥，并传入其名称。K8s 命名空间约定差异很大，因此我们将密钥创建设为显式操作。