tao-run-on-local-docker

ローカル Docker

TAO ジョブをローカルの Docker デーモン上で、名前付きの Docker コンテナとして実行する単一ノード実行プラットフォームです。開発、デバッグ、小規模な実行、およびエージェントホストに必要な GPU、NVIDIA ドライバー、Docker、NVIDIA Container Toolkit がすでに備わっているマシンでの利用に適しています。

データが Docker ホスト上にある場合、またはマウント済みボリュームやクラウド認証情報を介してアクセスできる場合は、ローカル Docker を使用してください。リモートクラスターのスケジューリング、マルチノードトレーニング、SLURM キューイングが必要なジョブには使用しないでください。

事前確認

ワークフローでは、Docker ジョブの開始前にホストの GPU ランタイムを検証する必要があります。確認に失敗した場合は、ユーザーにインストールを承認するよう促し、表示されたインストールコマンドを実行してから、事前確認を再実行してください。

# Host GPU runtime: NVIDIA driver 580, CUDA 13.0, NVIDIA Container Toolkit 1.19.0.
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 docker --check-only || {
  echo "MISSING: TAO GPU host runtime is not ready."
  echo "After user approval, run:"
  echo "  bash \"$SETUP_SCRIPT\" --backend docker --install --yes"
  exit 1
}

# Mode 1 — direct docker (no Python). All you need is docker + the GPU runtime.
docker info >/dev/null 2>&1 || { echo "MISSING: docker daemon not reachable. Start Docker."; exit 1; }
docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi >/dev/null 2>&1 || {
  echo "MISSING: NVIDIA Container Toolkit not installed/configured. See:"
  echo "  bash \"$SETUP_SCRIPT\" --backend docker --install --yes"
  exit 1
}

# Mode 2 — TAO SDK wrapper. Adds Job handles, S3 I/O wrapping, ActionWorkflow.
# Skip this block if Mode 1 is sufficient for the user's request.
# When Mode 2 is in scope, read `tao-skill-bank:tao-run-platform` for the DockerSDK
# kwarg contract, build_entrypoint, and monitoring patterns.
# nvidia-tao-sdk is on public PyPI; pin lives in versions.yaml (wheels.tao_sdk_docker).
PIN=$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_sdk_docker)
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 docker" 2>/dev/null || {
  echo "MISSING: docker Python client not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}

# DockerSDK attaches every job container to ${DOCKER_NETWORK:-tao_default}. If
# the network does not exist, container start fails instantly with
# `network <name> not found` for every create_job.
DOCKER_NETWORK_NAME="${DOCKER_NETWORK:-tao_default}"
docker network ls --format '{{.Name}}' | grep -qx "$DOCKER_NETWORK_NAME" || {
  echo "MISSING: docker network '$DOCKER_NETWORK_NAME' not found. After user approval, run:"
  echo "  docker network create $DOCKER_NETWORK_NAME"
  exit 1
}

確認に失敗した場合、エージェントは続行する前に、Bash 経由でインストールまたは修正を許可するようユーザーに促します。

認証情報

Docker デーモンへのアクセス以外に、プラットフォームの認証情報は必要ありません。

任意の環境変数：

• DOCKER_HOST：任意の Docker デーモン URL です。未設定の場合、SDK は Docker Python クライアントの通常の環境変数／デフォルトソケットの解決方法を使用します。
• DOCKER_NETWORK：ジョブコンテナ用の Docker ネットワークです。デフォルトは tao_default です。
• DOCKER_USERNAME：レジストリのユーザー名です。NGC の場合、デフォルトは $oauthtoken です。
• NGC_KEY：nvcr.io からプライベートイメージをプルする際に使用します。
• HOST_SSH_PATH：リモート SLURM 子ジョブを監視するために SSH キーが必要な場合、AutoML の brain コンテナにマウントされます。
• ACCESS_KEY、SECRET_KEY、S3_ENDPOINT_URL、S3_BUCKET_NAME： ローカルコンテナからクラウドストレージの読み書きを行うジョブ向けの、任意の S3 互換ストレージ設定です。

起動前確認

スクリプトを生成したり、コンテナを起動したりする前に：

1. Docker デーモンに接続できること、および NVIDIA ランタイムから GPU が認識されることを確認します。
2. すべてのローカル／ファイルデータセットのアノテーションとメディアのパスが、Docker ホスト上に存在することを確認します。
3. s3:// データセット／結果の場合は、ACCESS_KEY と SECRET_KEY が設定されていること、および aws s3 ls で正確なパスを読み取れることを確認します。
4. 起動前に、HF_TOKEN などのモデル固有の認証情報を確認します。

マルチ GPU とマルチノード

ローカル Docker ではマルチノードはサポートされていません。 1 つのジョブは、ホスト間の連携なしでローカルの Docker デーモンを実行するホスト上で実行されます。

ローカルホスト上のマルチ GPU は、NVIDIA Container Toolkit の --gpus フラグ（--gpus all または --gpus '"device=0,1,2,3"'）を使用してサポートされます。DockerSDK.create_job(gpu_count=N) は --gpus に値を渡します。単一ホストの分散初期化では localhost を使用し、torchrun --nproc-per-node=N または PyTorch DDP は通常どおり動作します。

バックエンドの詳細

SDK のバックエンド値 local-docker を使用します。ローカルバックエンドのスキーマには追加のバックエンド詳細がないため、ルーティングの大部分は環境変数とジョブパラメーターで制御されます。

{
  "backend_type": "local-docker",
  "num_gpu": 1
}

Lepton/Brev SDK の設計に従い、プラットフォーム／コントロールプレーンの値は SDK の状態と Docker ラベルに保持されます。SDK は BACKEND、HOST_PLATFORM、MONGOSECRET、DOCKER_HOST、DOCKER_NETWORK をトレーニングコンテナに注入しません。

コンテナーの実行

TAO SDK のローカル Docker ハンドラーは、Docker Python クライアントを通じてコンテナーを起動します。

• バックエンドのジョブ名には、SDK ハンドラーで使われる tao-job-<job_id> 形式を使用します。
• コマンドは通常、["/bin/bash", "-c", "<job command>"] です。
• コンテナーはデタッチモードで実行されます。SDK はデフォルトでコンテナーを保持するため、ステータスやログを引き続き確認できます。ただし、DOCKER_AUTO_REMOVE=true の場合は除きます。
• /dev/shm は tmpfs としてマウントされます。
• 設定された Docker ネットワークは、ジョブコンテナーに対して Docker デーモンによって適用されます。プロセスの環境変数として渡されるわけではありません。
• 同じジョブ ID の既存コンテナーがある場合、置き換え前に停止して削除します。

GPU にアクセスするため、ハンドラーはホストの種類を自動検出します。

• Tegra または Jetson ホストでは、runtime="nvidia" に加えて、NVIDIA_VISIBLE_DEVICES と NVIDIA_DRIVER_CAPABILITIES=all を使用します。
• 標準の x86 ホストでは、GPU 機能を指定した Docker の device_requests を使用します。

num_gpus が 0 の場合、GPU は割り当てられません。num_gpus が -1 の場合、認識可能なすべての GPU が要求されます。複数のユーザーが使用する開発マシンでは、GPU 数を明示的に指定することをおすすめします。

ストレージ

コンテナーは同じ Docker ホストで実行されるため、ローカル Docker ではローカルパスと file:// パスを使用できます。spec 内の各パスが次のいずれかであることを確認してください。

• ハンドラーまたは周辺サービスによってコンテナーにマウントされている
• すでにコンテナー内からアクセスできる
• 対応する認証情報を備えたクラウド URI である

リモートファイルシステムや共有ファイルシステムには、そのファイルシステムを管理するプラットフォームを使用してください。たとえば、クラスター上の Lustre パスには、lustre:///... と SLURM を使用します。

モニタリング

• SDK ハンドラーは Docker コンテナーの状態を直接マッピングします。created → Pending、running/restarting → Running、paused → Paused、終了コード 0 → Complete、ゼロ以外の終了コード → Error。
• ログは Docker Python クライアントを通じて、指定されたコンテナーから直接取得します（docker logs tao-job-<job_id>）。

コンテナーが終了した場合、停止した場合、削除中の場合、または見つからない場合、ステータスの調整ではバックエンドのプロセスを終了済みとして扱います。

キャンセル

キャンセルすると、指定されたコンテナーが停止します。GPU の所有権は TAO Core のローカル GPU マネージャーではなく、Docker／NVIDIA ランタイムによって管理されます。

オプション：TAO SDK を使用する

Job ハンドル、SDK の script_runner を介した S3 I/O ラッパー、またはセッションをまたいだ永続性が必要な場合：

from tao_sdk.platforms.docker import DockerSDK

sdk = DockerSDK()  # reads DOCKER_HOST, NGC_KEY, S3 creds from env
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,
    inputs={'/data/train.json': 's3://bucket/coco/train.json'},
    outputs=['/results/'],
)

status = sdk.get_job_status(job.id)
logs = sdk.get_job_logs(job.id, tail=200)

これは同じ docker run の呼び出しを Job ハンドルでラップし、エントリーポイントを script_runner 経由で実行するため、inputs／outputs が S3 から自動的にダウンロード／アップロードされます。これらが不要な場合は、docker run を直接使用してください。SDK のインストールは不要です。

障害の種類

Docker クライアントが初期化されていない：Docker Python パッケージがインストールされていることを確認してください。デフォルトのローカルソケットを使用していない場合は DOCKER_HOST を設定し、プロセスからデーモンに接続できることを確認してください。

GPU の割り当てに失敗した：要求した GPU を使用できない、NVIDIA Container Toolkit が設定されていない、または Docker デーモンが GPU のデバイス要求を作成できない可能性があります。GPU 数を減らす、別のジョブが終わるまで待つ、またはホスト上で docker run --gpus ... が動作することを確認してください。

イメージのプル認証に失敗した：プライベートな nvcr.io イメージには有効な NGC_KEY を設定するか、Docker ホスト上で docker login nvcr.io -u '$oauthtoken' を実行してください。

コンテナーが予期せず終了した：docker logs tao-job-<job_id>、設定された DOCKER_NETWORK、および SDK アクションランナーが生成したコマンドを確認してください。

コンテナー内にパスが存在しない：ホスト上のローカルパスが、ジョブコンテナーにマウントされているとは限りません。アクションランナーがサポートするパスの規則を使用するか、周辺サービスを通じて明示的なボリュームを設定してください。