tao-run-on-kubernetes

Kubernetes

独立安装？ 如果此会话不是由 TAO skill bank 插件初始化的，请先运行 tao-setup 技能（主机预检、凭证、跨技能发现）。

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

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

预检

检查三项：GPU 主机运行时已就绪、可通过 kubectl 访问集群、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}/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

对于自管节点，GPU 节点运行时检查是强制要求。对于客户端未运行在 GPU 工作节点上的托管集群，请验证提供商的节点映像或 GPU Operator 策略，并设置 TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1，而不是在客户端运行安装程序。此处的 GPU 容量警告属于软性检查；submit 动词会重新检查可分配的 nvidia.com/gpu，并在应用清单前硬性失败（不支持 gang 调度，因此资源需求过大的 Job 会永远处于 Pending 状态）。

凭证与配置

• Kubeconfig（以下之一）：
  • ~/.kube/config — 默认发现路径
  • $KUBECONFIG — 备用路径
  • 集群内服务账号 — 在 Pod 内运行时使用（无需 kubeconfig）
• TAO_K8S_NAMESPACE（可选）：提交 Job 的默认命名空间。默认为 default。
• TAO_K8S_CONTEXT（可选）：要切换集群的 kubeconfig context 名称。
• NGC_KEY（可选）：用于拉取 nvcr.io 映像。如果已在目标命名空间中预先创建映像拉取密钥，请在渲染后的清单中的 imagePullSecrets 处引用其名称。
• AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY / S3_BUCKET_NAME / S3_ENDPOINT_URL（可选）：用于 S3 数据集 I/O（存储层级 C），通过每个 Job 的 Secret（envFrom.secretRef）注入 Pod，绝不以内联方式传入。tao-data-io 会映射旧版 ACCESS_KEY / SECRET_KEY。

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

执行 — 四个动词

tao-run-on-kubernetes 是一个平台消费者：它通过 kubectl 运行 spec-bundle，仅修改 job-record。不使用 nvidia-tao-sdk，也不导入 tao_sdk — 任务通过普通的 kubectl apply 提交。 $BANK = ${TAO_SKILL_BANK_PATH}。

submit

1. GPU 容量门控——首先硬失败（没有 Gang 调度 → 超出容量的 Job 会永远处于 Pending 状态）：

   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. 存储层级（通过 tao-data-io）：A = 挂载一个已绑定、存有 数据的 PVC/NFS（编写挂载路径，不抓取数据——这是隔离网络环境下的方案， 也是打包模板的做法）；C = 临时存储：initContainer 从 S3 抓取数据到 共享的 emptyDir，最后一步在 TTL 到期前将结果上传到 S3。

   层级 C 会在下载期间占用 GPU。 Pod 在整个生命周期（包括 initContainer）中都会预留 nvidia.com/gpu，因此大型层级 C 数据抓取—— 或首次拉取多 GB 镜像——都会产生计费且符合回收条件的 GPU 空闲时间， 与在 SLURM 分配资源期间拉取的情况完全一样。如果数据已在 PVC 上，优先 选择层级 A；层级 C 仅适用于小型输入，且应审慎选择。

   生产者 action 请求可以声明多个挂载，包括用于逻辑路径和嵌入式绝对路径 的重复源别名。阅读 references/action-request.md，并使用其中的暂存映射 和打包渲染器。对于 mode=config，先实例化生产者的嵌套 spec，再暂存 该精确生成文件，并将其传回渲染器，以便 Pod 接收经过验证的只读配置 挂载。旧版单根模板无法表示此契约。

3. 打开记录——在启动前生成 id，并绑定 results_dir：

   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 必须是已挂载（可保留）的卷路径或 S3 前缀—— ttlSecondsAfterFinished 会在 Job 结束后删除该 Job 及其日志，因此之后 无法从 Job 对象中恢复任何内容。

4. 渲染、门控、应用并记录为 RUNNING。 对于生产者 action 请求， 请遵循 references/action-request.md；该文档负责后端名称规范化、 条件式 Secret 引用、本机 argv 渲染、服务器端 dry-run，以及将已应用对象 名称绑定到 job-record。对于简单的单根 spec-bundle，请渲染 templates/k8s/single-pod-job.yaml.tmpl，运行 redact_secrets.py lint 和 kubectl apply --dry-run=server，应用该对象， 并将记录标记为 backend-ref=<namespace>/<actual-object-name>。

跳过门控或打开记录的提交没有 id，因此无法启动。

status

提交时保留 K8S_JOB_NAME。重新连接时，读取 job-record 的 backend_ref=<namespace>/<name>，并从该字段恢复这两个值；不要假设 Kubernetes 名称与记录 id 相同。

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

kubectl 信号	术语
未调度任何 Pod	PENDING（kubectl get pods -n "$NAMESPACE" -l job-name="$K8S_JOB_NAME" → message 中的 ImagePullBackOff / Insufficient nvidia.com/gpu）
active ≥ 1	RUNNING
condition Complete	COMPLETE
condition Failed	ERROR（根据 Pod 的终止原因分类——OOMKilled → ERR_INFRA）
未找到 Job/Pod	UNKNOWN（可能已被 TTL 删除——job-record 才是事实来源）

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

多节点（nodes > 1）

相同的四种动词，另加：

1. **版本门控：**要求 k8s ≥ 1.28（kubectl version -o json）——需要此版本 才能解析 Pod 主机名 <job>-<index>（PodIndexLabel）到 MASTER_ADDR=<job>-0.<svc>；在较旧集群上，rank-0 会在 rendezvous 时挂起。
2. **容量门控 ×nodes：**除非可分配 GPU 数量 ≥ gpus_per_node × nodes，否则硬失败（没有 Gang 调度 → 部分启动会让 rank-0 永远等待）。
3. 渲染 templates/k8s/indexed-job.yaml.tmpl——headless Service + Indexed Job + rendezvous 环境变量（WORLD_SIZE = 节点数， NODE_RANK 来自 JOB_COMPLETION_INDEX，MASTER_ADDR=<job>-0.<svc>， /dev/shm 为 16Gi，以免 NCCL 悄无声息地挂起）。 kubectl apply -f 会同时创建 Service 和 Job； cancel 会删除 Job（Foreground）和 Service。
4. 先探测 NCCL（与 SLURM 相同）——执行带超时的 2 节点 all-reduce； 如果挂起，设置集群 NCCL 环境变量并重新探测；按集群缓存结果。

本地集群（开发、CI 和评估）

临时的 minikube/kind 集群可以在不占用集群配额的情况下，测试准入、四种动词、作业记录接线和日志管道；这也是由 Agent 驱动的评估应自行配置的环境。kubectl 和 minikube 是无需 root 权限的单个静态二进制文件，因此非 root 的 CI 容器可以自行安装。

有两个先决条件会让渲染后的 Job 保持 Pending，而第一个会掩盖第二个：必须存在模板挂载的 PVC（persistentvolumeclaim "<name>" not found 会在任何 GPU 报错之前触发），然后，在没有 GPU 的集群上，请求 nvidia.com/gpu 的 Job 会报告 Insufficient nvidia.com/gpu 并一直等待。仅测试生命周期时，使用 NUM_GPUS=0 进行渲染，并说明未验证 GPU 调度；在 Linux GPU 主机上，minikube start --driver=docker --gpus all 可以直通真实 GPU，因此一台 GPU 主机就足以进行 GPU 实测冒烟测试。

安装命令、驱动选择、容器与主机网络的注意事项，以及使用模拟设备插件的折中方案：references/local-cluster.md。

容器 Shell

简单的单 Pod 模板通过 /bin/sh -c 调用命令（POSIX sh，busybox/distroless 和 TAO 镜像中均有）。对于生产方的 action 请求，args 模式的命令及其参数会作为原生容器 argv；需要 shell 的生产方会显式声明 shell 及其脚本。简单的配置模式命令在替换 {config_path} 后也会成为原生 argv。生产方自有配置命令如果本身是多行 shell 脚本，则会原样保留在 /bin/sh -c 下；渲染器绝不会根据配置值构造 shell 文本。

GPU Operator 依赖项

如果集群没有可分配的 nvidia.com/gpu，submit 动词会拒绝启动 GPU 作业。对于自行管理的集群，先在每个 GPU 工作节点上运行 tao-setup-nvidia-gpu-host 安装 action，或者将相同的软件包集打包到节点镜像中：

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（参见上文多节点（nodes > 1）动词步骤），即可在 N 个 Pod 上运行分布式训练。渲染 templates/k8s/indexed-job.yaml.tmpl 会配置：

1. 一个以 Job 命名的无头 Service（选择器：job-name=<job-name>，clusterIP: None，publishNotReadyAddresses: true，让 Pod 在全部就绪前即可相互 rendezvous）。

2. 一个Indexed 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 容器的入口点
   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 入口点（dino train 等）和原始 torchrun 命令都无需修改即可运行。

对于 TAO 入口点，容器会读取 spec.train.num_nodes 和接线配置的环境变量，例如：gpu_count=8、num_nodes=4 时，执行 dino train -e /tmp/spec.yaml（共 4 × 8 = 32 个 GPU）。

对于基于原始 torchrun 的命令（非 TAO 容器），包装器会调用：

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 × num_nodes ≤ 集群可分配的 nvidia.com/gpu。

多节点的集群要求

• 需要 k8s 1.28+，以确保 Indexed Jobs 中 Pod 主机名稳定（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 里设置 NCCL_SOCKET_IFNAME / NCCL_IB_HCA。

参考资料

• Kubernetes Indexed Job：https://kubernetes.io/docs/concepts/workloads/controllers/job/#completion-mode
• 用于批处理机器学习的 Indexed Job：https://kubernetes.io/blog/2022/06/01/indexed-jobs-mpi/
• PyTorch 分布式（环境变量会合）：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，而不是普通的 Indexed 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/ — 在上述任一方案之上提供配额／队列层。

本技能的 Indexed Job 路径刻意设计得简单且无依赖；如果需要弹性重启或 gang 调度，请在其上叠加其中一个方案，并通过 Operator 的 CRD 提交作业。

常见错误模式

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：在命名空间中预先创建镜像拉取密钥，并在渲染后的清单中将其设为 Pod 的 imagePullSecrets： 通过标准输入传递密钥——--docker-password=$NGC_KEY 会将密钥放入 argv， 使其在主机的进程表和 Shell 历史记录中可见：

set -a; source /path/to/.env; set +a   # omit if already exported
kubectl create secret generic ngc-pull-secret -n tao-jobs \
  --type=kubernetes.io/dockerconfigjson \
  --from-file=.dockerconfigjson=/dev/stdin <<EOF
{"auths": {"nvcr.io": {"username": "\$oauthtoken", "password": "${NGC_KEY}"}}}
EOF
# Verify without reading the secret back:
kubectl get secret ngc-pull-secret -n tao-jobs >/dev/null && echo SECRET_OK

Pod 一直处于 Pending 状态 — kubectl describe pod -l job-name=$JOB_ID 会在 Events 中显示调度原因。常见原因：GPU 容量不足（Insufficient nvidia.com/gpu）、没有节点匹配 Pod 的 nodeSelector、缺少镜像拉取密钥，或 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 上进行 rendezvous）。
• 根据 $NGC_KEY 自动创建镜像拉取密钥。 你需要在目标命名空间中预先创建密钥，并传入其名称。K8s 命名空间的约定差异很大，因此我们明确要求创建密钥。