Tao-Run-On-KubernetesSkill tao-run-on-kubernetes

名称	tao-run-on-kubernetes
描述	Kubernetes执行平台 — 将TAO容器作业作为单Pod的k8s Job提交，支持NVIDIA GPU调度。 适用于在EKS / GKE / AKS / 本地集群上运行，且已安装NVIDIA GPU Operator，或将TAO集成到现有的k8s原生ML平台。
开源协议	Apache-2.0 compatibility: 需要带有NVIDIA驱动580或更新版本、CUDA Toolkit 13.0或更新版本、NVIDIA Container Toolkit 1.19.0或更新版本的GPU工作节点，除非所选模型在runtime_requirements.gpu_host中声明了不同的最低版本；还需要带有kubernetes扩展的nvidia-tao-sdk Python包、已认证的集群，以及NVIDIA GPU Operator或设备插件。 metadata:
作者	NVIDIA Corporation
版本	“0.1.0” allowed-tools: Read Bash tags: - kubernetes - k8s - gpu - compute - container

Kubernetes

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

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

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

预检

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

# 0. GPU节点主机运行时。
# 在每个自管理GPU工作节点上运行此命令，或在节点镜像构建中运行。
# 仅当使用托管GPU节点（其驱动/工具包生命周期由云提供商或GPU Operator策略管理）时，设置TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1。
if [ "${TAO_K8S_SKIP_NODE_RUNTIME_CHECK:-0}" != "1" ]; then
  SB="${TAO_SKILL_BANK_PATH:-${TAO_SKILL_BANK_ROOT:-$PWD}}"
  SETUP_SCRIPT="${SB}/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节点运行时未就绪。"
    echo "对于自管理GPU节点，请在用户批准后运行："
    echo "  bash \"$SETUP_SCRIPT\" --backend kubernetes --install --yes"
    echo "对于托管集群，请验证节点镜像/GPU Operator策略满足所选模型的运行时要求，然后设置TAO_K8S_SKIP_NODE_RUNTIME_CHECK=1。"
    exit 1
  }
fi

# 1. SDK + kubernetes扩展已安装。
# nvidia-tao-sdk在公共PyPI上；下面的pin从发布清单中标记。
PIN="nvidia-tao-sdk[kubernetes]==7.1.0rc42"  # versions-key: wheels.tao_sdk_kubernetes
python -c "import tao_sdk" 2>/dev/null || {
  echo "正在安装缺失的Python依赖: $PIN"
  python -m pip install "$PIN"
}
python -c "import kubernetes" 2>/dev/null || {
  echo "正在安装缺失的Python依赖: $PIN"
  python -m pip install "$PIN"
}
python -c "import tao_sdk, kubernetes"

# 2. 集群可访问（kubeconfig或集群内服务账户）
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: ~/.kube/config处无kubeconfig，且未在Pod中运行。"
    echo "配置kubectl（例如，‘aws eks update-kubeconfig --name my-cluster’）或设置\$KUBECONFIG。"
    exit 1
  }

# 3. 存在NVIDIA GPU Operator（软检查——如果可用则警告，不失败）
if command -v kubectl >/dev/null 2>&1; then
  gpu=$(kubectl get nodes -o jsonpath='{range .items[*]}{.status.allocatable.nvidia\.com/gpu}{"
"}{end}' 2>/dev/null | grep -v '^$' | head -1)
  if [ -z "$gpu" ] || [ "$gpu" = "0" ]; then
    echo "WARN: 此集群上没有可分配的nvidia.com/gpu。"
    echo "在提交GPU作业之前安装NVIDIA GPU Operator："
    echo "  https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/getting-started.html"
  fi
fi

对于自管理节点，如果所选模型声明了runtime_requirements.gpu_host，请将相应的--min-*-version标志传递给检查和安装命令。在托管集群上，要求提供商节点镜像或GPU Operator策略满足该模型配置。

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上下文名称。
• NGC_KEY（可选）：用于nvcr.io镜像拉取。如果您已在目标命名空间中预先创建了image-pull secret，请通过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。启动前，验证所选命名空间可以创建Job，数据集/结果路径从Pod可见，并且PVC/挂载文件系统路径已证明挂载到Job容器中；代理主机本地路径不足以作为证明。

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()  # 自动检测认证
job = sdk.create_job(
    image='nvcr.io/nvidia/tao/tao-toolkit:7.1.0-pyt',  # versions-key: images.tao_toolkit.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',                       # 可选覆盖
    image_pull_secret='ngc-pull-secret',         # 可选，预先创建
    node_selector={'gpu-type': 'h100'},          # 可选
)

SDK构造V1Job，包含：

• spec.template.spec.containers[0]: 请求的镜像和command=["/bin/bash", "-c", <command>]。
• resources.limits["nvidia.com/gpu"]: <gpu_count> — 通过NVIDIA设备插件 / 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)  # 连接Job所有Pod的日志

# 对卡在Pending的Job——副本诊断：
for r in sdk.get_job_replicas(job.id):
    issue = r["status"].get("readiness_issue")
    if issue:
        print(issue["reason"], issue["message"])
        # 例如 "ImagePullBackOff" / "Back-off pulling image..."
        # 例如 "Pending"           / "0/3 nodes available: 3 Insufficient nvidia.com/gpu"

# 失败时：
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"

# 对于已确认终止的S3支持的Job，先移除精确UID绑定的Job/pods，然后移除其作业范围的result前缀。
sdk.delete_job_artifacts(job.id)

已完成的Job不会被TTL删除：其UID和终止证据在SDK重启后仍可用于安全恢复。要取消进行中的Job，cancel_job以前台传播删除它及其pods，并且仅在写入者消失后报告成功。对于支持清理的S3作业，delete_job_artifacts首先验证并移除该精确终止Job及其pods，然后验证删除作业范围的S3前缀。绝不删除同名的不同UID的Job。

GPU Operator依赖

SDK拒绝向没有可分配nvidia.com/gpu的集群提交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

多节点训练（分布式）

向create_job()传递num_nodes > 1以跨N个Pod运行分布式训练。SDK提供：

1. 一个以Job命名的headless Service（选择器：job-name=<job-name>，clusterIP: None，publishNotReadyAddresses: true，以便Pod在全部就绪前可以会合）。

2. 一个Indexed Job，parallelism = completions = num_nodes，completionMode: Indexed。每个Pod自动获得k8s注入的JOB_COMPLETION_INDEX（=节点秩）。

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

   环境变量	值	读取者
   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标准会合
   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命令均无需修改即可运行。

job = sdk.create_job(
    image='nvcr.io/nvidia/tao/tao-toolkit:7.1.0-pyt',  # versions-key: images.tao_toolkit.pyt
    command='dino train -e /tmp/spec.yaml',  # TAO入口点读取spec.train.num_nodes；环境变量由容器接线
    gpu_count=8,           # 每个节点的GPU数
    num_nodes=4,           # 4 × 8 = 32个GPU总计
    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',  # 未固定：NGC PyTorch基础镜像示例，非TAO发布节奏
    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。

多节点的集群要求

• k8s 1.28+ 是Indexed Job中稳定Pod主机名（PodIndexLabel特性）所必需的。在较旧的集群上，MASTER_ADDR=<job>-0.<svc> DNS查找会失败。使用kubectl version验证。
• 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 Indexed Job：https://kubernetes.io/docs/concepts/workloads/controllers/job/#completion-mode
• 用于批量ML的Indexed Job：https://kubernetes.io/blog/2022/06/01/indexed-jobs-mpi/
• PyTorch分布式（env-var会合）：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 elastic / 容错训练、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/ — 用于具有内置重启逻辑的elastic PyTorch训练。
• Volcano — https://volcano.sh/ — gang调度、队列、公平共享。在共享多租户集群中有用。
• Kueue — https://kueue.sigs.k8s.io/ — 在上述任何内容之上的配额/队列层。

TAO SDK的Indexed Job路径有意保持简单且无依赖；如果您需要elastic重启或gang调度，请在此基础上分层使用其中一种，并通过operator的CRD提交作业。

常见错误模式

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

ImagePullBackOff / ErrImagePull — 集群无法拉取镜像。对于nvcr.io：在命名空间中预先创建image-pull secret，并通过image_pull_secret参数传递其名称。通过stdin提供密钥——kubectl create secret docker-registry --docker-password=$NGC_KEY会将secret暴露在argv中，在主机进程表和shell历史中可见：<!-- lint-ok: secret-on-argv -->

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
# 不读回secret进行验证：
kubectl get secret ngc-pull-secret -n tao-jobs >/dev/null && echo SECRET_OK

Pod永远卡在Pending — get_job_replicas(job_id)将显示readiness_issue。常见原因：GPU容量不足（Insufficient nvidia.com/gpu）、没有节点匹配node_selector、缺少image-pull secret或PVC挂载失败。

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

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

此技能尚不支持的内容

• Elastic / 容错训练。 Indexed Job的backoff_limit=0——失败会使整个训练运行失败。对于elastic重启（例如，在节点死亡后从检查点恢复），请改用Kubeflow的PyTorchJob operator。
• Gang调度。 Indexed Job Pod独立调度——不提供全有或全无。如果只有部分Pod可以调度，多节点训练将部分启动（rank-0将挂起等待对等节点）。在共享集群上进行全有或全无调度，请使用Volcano或Kueue。
• MPI / Horovod。 使用MPI Operator。此处的Indexed Job路径是PyTorch分布式形状（在MASTER_ADDR:MASTER_PORT上进行env-var会合）。
• 用于共享存储的持久卷。 仅通过script_runner支持S3。PVC支持是后续工作。
• 从$NGC_KEY自动创建image-pull secret。 您在目标命名空间中预先创建secret并传递名称。K8s命名空间约定差异很大，因此我们将secret创建保持显式。