tao-run-on-kubernetes

Kubernetes

TAO コンテナジョブを Kubernetes Job として送信します。kubeconfig からアクセスできる任意のクラスター（EKS／GKE／AKS／オンプレミス）またはクラスター内サービスアカウント（SDK が Pod 内で実行される場合）で動作します。

デフォルトではシングル Pod です。num_nodes > 1 を指定すると、マルチノード分散トレーニングを利用できます（Indexed Job とヘッドレス Service を使用します。以下のマルチノードトレーニング（分散）を参照してください）。

事前確認

確認項目は 4 つです。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 は、送信前に Kubernetes Python クライアントを使用して GPU 容量を確認するハードガードを KubernetesSDK.create_job() 内で実行します。

認証情報と設定

• Kubeconfig（次のいずれか）：
  • ~/.kube/config — デフォルトの検出パス
  • $KUBECONFIG — 代替パス
  • クラスター内サービスアカウント — Pod 内で実行する場合に使用します（kubeconfig は不要です）
• TAO_K8S_NAMESPACE（任意）：Job の送信先となるデフォルトの namespace です。デフォルトは default です。
• TAO_K8S_CONTEXT（任意）：クラスターを切り替えるための kubeconfig context 名です。
• NGC_KEY（任意）：nvcr.io からイメージをプルする場合に使用します。対象 namespace にイメージプルシークレットを事前作成してある場合は、image_pull_secret 引数を使って create_job にその名前を渡してください。
• ACCESS_KEY / SECRET_KEY / S3_BUCKET_NAME / S3_ENDPOINT_URL（任意）：SDK の inputs／outputs による script_runner ラッピングを介した S3 データセットの入出力に使用します。

Kubernetes の実行では、Brev または SLURM の認証情報を求めないでください。選択したワークフローで s3:// の入力または出力を使用する場合にのみ S3 の認証情報を求め、HF_TOKEN などのモデル固有の認証情報は、選択したモデルで必要な場合にのみ求めてください。起動前に、選択した namespace で Job を作成できること、データセット／結果のパスが Pod から見えること、PVC／マウントされたファイルシステムのパスがジョブコンテナにマウントされていることを確認してください。エージェントホスト上のローカルパスだけでは、マウントされている証明にはなりません。

SDK API

K8s は SDK 専用です。kubectl のみで起動する方法はありません。create_job の呼び出しを作成する前に tao-skill-bank:tao-run-platform を読み、build_entrypoint、共通の kwargs 契約、モニタリング、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 の場合、完了した Job は 1 時間後に自動削除されます。実行中の Job をキャンセルするには、cancel_job を実行します。Job とその Pod が直ちに削除されます。

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 の名前を付けた ヘッドレス Service（セレクター：job-name=<job-name>、clusterIP: None、すべての Pod が Ready になる前に相互検出できるよう、publishNotReadyAddresses: true）。

2. parallelism = completions = num_nodes、completionMode: Indexed の インデックス付き Job。各 Pod には、k8s によって JOB_COMPLETION_INDEX が自動的に設定されます（＝ノードランク）。

3. ユーザーコマンドを呼び出す前にランデブー用の環境変数をエクスポートする コマンドラッパー。2 つの命名規則を同時にエクスポートします。

   環境変数	値	読み取り元
   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 を直接使うコマンドの両方を変更なしで実行できるよう、2 つの命名規則を設定します。

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

TAO 以外のコンテナーで torchrun を直接使うコマンドの場合：

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。

マルチノードのクラスター要件

• Indexed Job で安定した Pod ホスト名を使うには、k8s 1.28 以降が必要です（PodIndexLabel 機能）。古いクラスターでは、MASTER_ADDR=<job>-0.<svc> の DNS ルックアップに失敗します。kubectl version で確認してください。
• Pod 間のネットワークでポート 29500（PyTorch のデフォルト。MASTER_PORT 環境変数で変更可能）を開放する必要があります。ほとんどの CNI（Calico、Cilium、AWS VPC CNI）ではデフォルトで許可されますが、制限の厳しい NetworkPolicy は緩和する必要があります。
• コンテナー内の NCCL は GPU 間の通信を行います。クラスターにマルチ NIC ノードまたは 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 分散処理（環境変数によるランデブー）：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 オペレーターの代替手段

より高度なトポロジー（ギャングスケジューリング、PyTorch elastic／フォールトトレラントトレーニング、MPI／Horovod、RDMA のセットアップ）には、通常の 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/ — ギャングスケジューリング、キュー、公平なリソース配分を提供します。複数のテナントが共有するクラスターで役立ちます。
• Kueue — https://kueue.sigs.k8s.io/ — 上記のいずれにも適用できるクォータ／キューレイヤーです。

TAO SDK の Indexed Job パスは、意図的にシンプルで依存関係のない設計になっています。弾力的な再起動やギャングスケジューリングが必要な場合は、これらのいずれかを上に重ね、オペレーターの 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 オペレーターを使用してください。
• ギャングスケジューリング。 Indexed Job の Pod は個別にスケジュールされるため、全体をまとめて実行する仕組みはありません。一部の Pod だけがスケジュール可能な場合、マルチノードトレーニングは一部だけ開始し（rank-0 はピアを待って停止します）。共有クラスターで全体をまとめてスケジュールするには、Volcano または Kueue を使用してください。
• MPI／Horovod。 MPI Operator を使用してください。ここでの Indexed Job の方式は、PyTorch 分散処理の形式（MASTER_ADDR:MASTER_PORT で環境変数を使ってランデブー）です。
• 共有ストレージ用の永続ボリューム。 S3 は script_runner 経由のみです。PVC のサポートは今後追加予定です。
• $NGC_KEY からのイメージプルシークレットの自動作成。 対象の名前空間にシークレットを事前に作成し、その名前を渡してください。K8s の名前空間の慣習は大きく異なるため、シークレットの作成は明示的に行う方式にしています。