tao-run-on-lepton

Lepton

DGX Cloud 上のマネージド GPU コンピューティングプラットフォームです。ジョブは、専用または共有の GPU ノードグループで実行されるコンテナワークロードとして送信します。Lepton はスケジューリング、イメージのプル、ログの収集、ジョブのライフサイクルを管理します。

Kubernetes や SLURM のインフラストラクチャを直接管理せずに、クラウドベースの GPU コンピューティングが必要な場合に Lepton を使用します。

事前確認

Lepton は API ファーストのため、docker-run による代替手段はありません。このスキルには Lepton エクストラ付きの TAO SDK が必要です。nvidia-tao-sdk はパブリック PyPI にあり、固定バージョンは versions.yaml（wheels.tao_sdk_lepton）に記載されています。scripts/resolve_versions_key.py で解決します。

PIN=$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_sdk_lepton)
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 leptonai" 2>/dev/null || {
  echo "MISSING: lepton extra not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}

見つからない場合、エージェントは Bash 経由でインストールを許可するようユーザーに求め、続行前に事前確認を再実行します。

認証情報

• LEPTON_WORKSPACE_ID（必須）：ジョブの実行先となるクラスターと請求アカウントを決定します。
• LEPTON_AUTH_TOKEN（必須）：Lepton コントロールプレーンでの認証に使う API トークンです。
• NGC_KEY（任意）：nvcr.io から TAO コンテナイメージをプルするためのイメージプルシークレットの作成に使います。
• ACCESS_KEY／SECRET_KEY（任意）：データセットおよびチェックポイントの URI 用の S3 互換ストレージキーです。
• S3_ENDPOINT_URL（任意）：カスタム S3 エンドポイントです（MinIO や AWS 以外の S3 など）。
• S3_BUCKET_NAME（任意）：ジョブの出力アーティファクト用バケットです。
• CLOUD_REGION（任意）：ストレージのリージョンです（us-east-1 など）。

起動前の事前確認

スクリプトの生成やジョブの送信を行う前に、以下を確認します。

1. LEPTON_WORKSPACE_ID と LEPTON_AUTH_TOKEN が設定されていることを確認します。
2. 同梱のヘルパーを使って、ワークスペース API にアクセスできることを確認します。 scripts/check_tao_launch_preflight.py --platform lepton ...。
3. s3:// のデータセット／結果の場合、ACCESS_KEY と SECRET_KEY が設定されていること、および正確なパスを aws s3 ls で読み取れることを確認します。
4. NFS／Lustre のマウントパスの場合、そのパスがジョブにマウントされることを、Lepton のボリューム／ストレージ権限の証拠で確認する必要があります。エージェントホスト上でローカルファイルシステムに対して test -e を実行しても、Lepton のジョブにおける証明とは見なさないでください。
5. 起動前に、HF_TOKEN などモデル固有の認証情報を確認します。

バックエンドの詳細

LeptonSDK.create_job は、プラットフォームに依存しない引数（image、command、gpu_count、env_vars、inputs、outputs、hooks）に加えて、以下の Lepton 固有のキーワード引数を受け取ります。

• resource_shape：明示的な GPU リソースシェイプ ID（例："gpu.8xh100-sxm"）です。設定すると、gpu_count による自動解決をスキップします。形式は不透明です（Lepton API がインスタンスメタデータの id として返す値がそのまま使われます）。有効な ID は sdk.list_resource_shapes() で確認します。
• dedicated_node_group：GPU の割り当てが保証されるノードグループ ID です（プリエンプションなし）。共有リソースの場合は省略します。
• num_nodes：分散トレーニングに使用するノード数です。デフォルトは 1 です。1 より大きい値にすると、ジョブ内通信と PyTorch の分散初期化が有効になります（マルチノードトレーニングを参照）。
• mounts：NFS／Lustre 用の、事前に作成された Mount オブジェクトです。設定されていない場合は、ノードグループから自動検出されます。

ワークスペースのシェイプ／ボリュームの確認

shapes = sdk.list_resource_shapes()
# {<platform_id>: {"cluster": ..., "gpu_type": "gpu.8xh100-sxm",
#                   "gpu_count": 8, "instance_type": ..., ...}, ...}

volumes = sdk.get_volumes(node_group_id="my-h100-pool")
# [{"name": "lustre", "from_path": "/lustre", "type": "Lustre"}, ...]

prefixes = sdk.get_storage_permissions("lustre", "my-h100-pool")
# ["/lustre/fsw/portfolios/edgeai/...", ...]

マルチノードトレーニング（分散）

マルチノード分散トレーニングを行うには、create_job に num_nodes > 1 を渡します。Lepton ハンドラー（tao_sdk/platforms/lepton/handler.py）は、intra_job_communication=True（Pod 間ネットワークを有効化）、parallelism=num_nodes および completions=num_nodes（Lepton が N 個のレプリカをスケジュール）を設定し、コンテナーの環境変数として WORLD_SIZE=num_nodes をエクスポートすることで、基盤となる LeptonJob を構成します。

Lepton のレプリカごとのネイティブ環境変数には、Lepton 固有の名前（LEPTON_JOB_WORKER_INDEX、LEPTON_JOB_TOTAL_WORKERS、LEPTON_JOB_WORKER_PREFIX、LEPTON_SUBDOMAIN）が使われます。そのため、ハンドラーは起動処理の先頭に、Lepton 公式の変換スクリプトを読み込む処理を追加します。

wget -O init.sh https://raw.githubusercontent.com/leptonai/scripts/main/lepton_env_to_pytorch.sh
chmod +x init.sh
source init.sh
# user command runs here

読み込み後、次の環境変数が設定されます。

環境変数	設定元	値
MASTER_ADDR	スクリプト	${LEPTON_JOB_WORKER_PREFIX}-0.${LEPTON_SUBDOMAIN}
MASTER_PORT	スクリプト	29400
NNODES	スクリプト	${LEPTON_JOB_TOTAL_WORKERS}
NODE_RANK	スクリプト	${LEPTON_JOB_WORKER_INDEX}
WORKER_ADDRS	スクリプト	マスター以外のワーカーホスト名をカンマ区切りで並べたリスト
WORLD_SIZE	TAO SDK ハンドラー	num_nodes（TAO コンテナーの規約。NNODES と同じ値）
NUM_GPU_PER_NODE	TAO SDK ハンドラー	gpu_count（TAO コンテナーのエントリーポイントが読み取る値）

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 WORLD_SIZE + NUM_GPU_PER_NODE
    gpu_count=8,                          # GPUs per node
    num_nodes=4,                          # 4 × 8 = 32 GPUs total
    dedicated_node_group='my-h100-pool',
    inputs={'/data/train.json': 's3://bucket/coco/train.json'},
    outputs=['/results/'],
)

生の torchrun ベースのコマンド（TAO 以外のコンテナー）の場合：

command='torchrun --nnodes=$NNODES --nproc-per-node=8 --node-rank=$NODE_RANK '
        '--master-addr=$MASTER_ADDR --master-port=$MASTER_PORT train.py'

Lepton で分散ジョブを実行する 2 つの方法

方法	使用する状況
TAO SDK create_job(num_nodes=N)（このスキル）	エージェントコードからプログラムで送信する場合。SDK による S3 のラッピング、監視、失敗分析、JobStore を利用したい場合。
Lepton の「Torchrun」ジョブタイプ（Lepton UI／lep CLI）	Lepton コンソールから手動で送信する場合。Lepton UI には、ランデブーを自動で設定する専用の「Torchrun」モードがあり、起動スクリプトは不要です。公式の例を参照してください。

参考資料

• NVIDIA の Lepton マルチノード PyTorch の例（UI／Torchrun モード）：https://docs.nvidia.com/dgx-cloud/lepton/examples/batch-job/distributed-training-with-pytorch/
• SDK が読み込む変換スクリプト：https://github.com/leptonai/scripts/blob/main/lepton_env_to_pytorch.sh
• PyTorch 分散処理（環境変数によるランデブー）：https://pytorch.org/docs/stable/elastic/run.html
• NCCL ネットワークのチューニング：https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html

注意事項

• レプリカを同じ低レイテンシーのインターコネクト（NVLink／InfiniBand）上に配置するため、マルチノードでは dedicated_node_group を優先してください。
• 共有ノードグループでレプリカがプリエンプトされると、ジョブ全体が失敗します。v1 の Lepton では、弾力的な再起動は行われません。長時間実行する場合は、専用ノードグループを使用してください。
• Lustre を利用するデータセットでは、同じマウントがすべてのレプリカに公開されるため、レプリカごとの I/O ラッピングは不要です。

クラウドストレージ

プラットフォームは Lepton ですが、ストレージレイヤーは S3 互換です。データセットと results_dir の両方で、cloud_metadata キーには必ず aws を、URI プロトコルには s3:// を使用してください。

• 正しい例：s3://bucket-name/path
• 誤った例：lepton://bucket-name/path

コンテナーの get_cloud_storage_class_object() は URI プロトコルを解析し、CLOUD_METADATA[protocol][bucket] から認証情報を取得します。

共有ストレージ（NFS/Lustre）

ノードグループには、NFS または Lustre のボリュームを接続できます。SDK はこれらを自動検出し、ジョブをまたいでデータを永続的に共有できるよう、コンテナーにマウントします。

SDK 関数

• sdk.get_volumes(node_group_id=None) — ノードグループの仕様から、利用可能なボリューム（名前、from_path、タイプ）を返します
• sdk.get_storage_permissions(volume_name, node_group_id) — ボリュームで許可されているパスプレフィックスを返します

LeptonSDK.create_job() はこれらを自動的に呼び出してマウントを検出し、ジョブ仕様に適した Mount オブジェクトを構築します。

スクリプトランナーによるマウントの使用方法

Lustre マウントが利用可能な場合：

• 入力：S3 パスは Lustre にマッピングされます（s3://bucket/path → /mnt/lustre/bucket/path）。ファイルが Lustre 上に存在する場合は、そのまま使用されます（ダウンロード不要）。存在しない場合は、S3 から Lustre にダウンロードされ、以降のジョブでも保持されます。
• 出力：結果はまず Lustre に書き込まれ（高速で永続的）、その後 S3 にアップロードされます（耐久性を確保）。後続のジョブ（ギャップ分析など）は、S3 を経由せずに Lustre から直接結果を読み取れます。

ボリュームの優先順位

lustre > filestore > 最初に利用可能なもの

Lustre キャッシュの無効化

Lustre はジョブ間でファイルを永続的にキャッシュします。組み込みの無効化機能はありません。上流のデータが変更されても S3 パスが同じ場合、Lustre は古いキャッシュ済みバージョンを返します。キャッシュを確実にミスさせるには：

• S3 上のファイル名を変更します（例：prompt.txt を上書きせずに prompt_v2.txt を使用する）
• イテレーション間で新しい storage_root を使用し、イテレーションをまたいだ古いデータの参照を防ぎます
• 再生成するアーティファクトには新しいパスを使用します

モニタリング

ジョブのステータス

概要レベルのステータス（Pending、Running、Complete、Error）を確認するには、sdk.get_job_status(job_id) を使用します。

レプリカのステータス

起動時に詳細なレプリカレベルの情報を確認するには、sdk.get_job_replicas(job_id) を使用します。各レプリカは辞書です：

replicas = sdk.get_job_replicas(job_id)
for r in replicas:
    node = r["status"]["node"]["name"]           # e.g., "node-ip-10-50-111-24"
    node_group = r["status"]["node"]["node_group_id"]
    cpu = r["status"]["cpu"]                      # e.g., 2
    memory_mb = r["status"]["memory_in_mb"]       # e.g., 8192
    readiness = r["status"].get("readiness_issue")
    if readiness:
        reason = readiness["reason"]   # "InProgress", "Failed", "ConfigError"
        message = readiness["message"] # "Pulling image", "Mount point not found", etc.

主な readiness_issue のパターン：

• reason="InProgress"、message="Pulling image" — イメージのプルが進行中です（大きなイメージでは通常の動作です）
• reason="Failed" — イメージのプルに失敗しました（NGC_KEY を確認してください）
• reason="ConfigError" — ノードの問題です（マウントの失敗、GPU エラー）
• readiness_issue がない — レプリカは実行中です

ジョブが Pending のまま停止している場合、レプリカのステータスが特に役立ちます。問題がイメージのプル、リソースのスケジューリング、ノードの状態のいずれにあるかが分かります。

ジョブログ

直近の N 行のログを取得するには、sdk.get_job_logs(job_id, tail=N) を使用します。ログは Lepton のログ収集サービスから取得されます。

並列ジョブ

並列で実行されるワークフローのステージ（例：動画生成 x8）では：

1. 起動： 各分割に対して execute_step(plan, step_id, extra_args={"split_id": i}) を呼び出します。各呼び出しは job_id を返してすぐに終了します。
2. モニタリング： すべてのジョブをポーリングします。各ジョブで sdk.get_job_status(job_id) を使用します。起動時の診断には get_job_replicas(job_id) を使用します。
3. 完了： すべてのステータスが Complete または Error になれば、すべてのジョブが完了です。
4. 部分的な失敗： 失敗した分割のみを再試行します。成功した分割を再実行する必要はありません。execute_step には同じ split_id を渡します。

失敗の分析

ジョブが失敗した場合は、自動的に根本原因を検出する sdk.get_failure_analysis(job_id) を使用します：

analysis = sdk.get_failure_analysis(job_id)
if analysis:
    print(analysis["err_class"])    # e.g., "ERR_PROGRAM"
    print(analysis["suggestion"])   # Human-readable fix
    for event in analysis.get("job_failure_by_node_event", []):
        print(event["node_event_name"], event["message"])
        # e.g., "OOM", "OOM encountered, victim process: cosmos-rl-evalu, pid: 3368483"

返される内容：

• err_class：エラーの分類（ERR_PROGRAM、ERR_INFRA など）
• suggestion：考えられる原因と修正方法
• job_failure_by_node_event：ノードレベルのイベント（OOM による強制終了、GPU エラー、マウントの失敗）
• log_streams：エラーのコンテキストを含む関連ログの抜粋

再試行する前に、失敗したジョブに対して必ずこれを呼び出してください。ユーザー側のエラー（不正な設定、OOM）とインフラストラクチャの問題（ノードの障害、エビクション）を区別できます。

失敗モード

OOM による強制終了：コンテナが GPU またはシステムメモリを超過しました。検出方法：get_failure_analysis() が node_event_name: "OOM" を返します。よくある原因：evaluation.batch_size が大きすぎる、または max_length が利用可能な KV キャッシュに対して大きすぎる。復旧方法：batch_size を減らす、テンソル並列処理で GPU を追加する、または max_length を減らします。

イメージのプル失敗：TAO コンテナイメージを nvcr.io からプルできません。通常、イメージプルシークレットがないか、有効期限が切れていることが原因です。SDK は NGC_KEY からシークレットを自動的にプロビジョニングしますが、NGC_KEY が無効な場合、ジョブは失敗します。検出方法：get_job_replicas() を確認します。readiness_issue.reason が、長時間にわたって message = "Pulling image" のまま InProgress を示すか、プルに失敗すると Failed を示します。復旧方法：NGC_KEY が有効であることを確認します。

リソースを利用できない：要求した GPU の構成を利用できません。ジョブは Queueing 状態のままになります。検出方法：Pending 状態が 15 分を超えても、レプリカにノードの割り当てが表示されません。復旧方法：別の resource_shape または dedicated_node_group を試すか、リソースが利用可能になるまで待ちます。

認証失敗：LEPTON_AUTH_TOKEN が無効か、有効期限が切れています。すべての API 呼び出しが 401/403 で失敗します。検出方法：ジョブの作成時にすぐ例外が発生します。復旧方法：トークンを更新し、SDK を再初期化します。

ノードの異常：割り当てられたノードに、マウントの失敗、GPU エラー、ネットワークの問題などのインフラストラクチャ上の問題があります。検出方法：get_job_replicas() を確認します。readiness_issue.reason = "ConfigError" とともに、"Mount point not found" のようなメッセージが表示されます。問題のあるノードでは、ジョブが Pending 状態のままになります。復旧方法：ジョブをキャンセルして再送信します。Lepton により別のノードにスケジュールされます。問題が再発する場合は、別の dedicated_node_group または resource_shape を試します。

ジョブのエビクション：共有ノードグループでは、リソースの逼迫時に Lepton がジョブをエビクトする場合があります。検出方法：ジョブが予期せず Running から Error に遷移します。復旧方法：再試行するか、dedicated_node_group を使用します。