tao-run-on-slurm

SLURM

SLURM で管理されるクラスター向けのリモート GPU コンピューティングプラットフォームです。ジョブは TAO サービスまたは SDK のホストから SSH 経由でログインノードに送信され、共有ファイルシステムに配置された後、sbatch で登録され、srun のコンテナーサポートを使って実行されます。

ユーザーが管理対象の GPU クラスター、共有 Lustre ストレージ、スケジューラーが所有する GPU 割り当てを利用できる場合は、SLURM を使用してください。エージェントマシンにしか存在しないローカルファイルには SLURM を使用しないでください。データと出力はクラスターからアクセスできる必要があります。

事前確認

# 1. SSH to the login node works without a password prompt
SLURM_HOST="${SLURM_HOSTNAME%%,*}"
[ -n "$SLURM_USER" ] && [ -n "$SLURM_HOST" ] || {
  echo "MISSING: set SLURM_USER and SLURM_HOSTNAME (comma-separated for failover) in your env (~/.config/tao/.env)."
  exit 1
}
ssh -o BatchMode=yes -o ConnectTimeout=10 "${SLURM_USER}@${SLURM_HOST}" "true" 2>/dev/null || {
  echo "MISSING: passwordless SSH to ${SLURM_USER}@${SLURM_HOST} not working. See references/ssh-setup.md."
  exit 1
}

# 2. Optional: TAO SDK wrapper for Job handles + S3 wrapping.
# nvidia-tao-sdk is on public PyPI; pin lives in versions.yaml (wheels.tao_sdk_slurm).
PIN=$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_sdk_slurm)
python -c "import tao_sdk" 2>/dev/null || {
  echo "MISSING: nvidia-tao-sdk not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}

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

3 つ目の事前確認は、プライベートな nvcr.io イメージにのみ適用されます。計算ノード上の Pyxis では、クラスターの ~/.config/enroot/.credentials に永続的な enroot 認証情報が必要です（ジョブ環境の NGC_KEY は読み込まれません）。認証が必要なイメージのプルでは、この設定がないとジョブの起動時に "Could not process JSON input" というエラーが発生します。この設定は（クラスター、ユーザー）の組み合わせごとに 1 回実行します。完全な確認手順と、NGC_KEY を履歴、ファイル、チャット出力に残さない printf | ssh のインストールパターンについては、references/ssh-setup.md を参照してください。パブリックイメージの場合は省略してください。

前提条件

ジョブを送信する前に、TAO サービスまたは SDK を実行しているホストから、SLURM_HOSTNAME に指定されたホストの少なくとも 1 つへ SSH でログインできる必要があります。対話的なパスワードプロンプトは表示されないようにしてください。ハンドラーは sbatch、squeue、sacct、scancel、ログの末尾取得を非対話的に実行するため、パスワードや 2FA のプロンプトが表示されると、ジョブの送信時または状態確認時に失敗します。

（ホスト、ログインノード、ユーザー）の組み合わせごとに 1 回設定します。SSH キーペアを作成し、各ログインホストに公開鍵をインストールして、ホスト鍵を信頼し、秘密鍵の権限を chmod 600 に制限したうえで、ssh -o BatchMode=yes ... で確認してください。~/.ssh/config のエイリアス、コンテナーへの鍵のマウントに関する注意事項、2FA / SSH_AUTH_SOCK の代替方法を含む詳しい手順は、references/ssh-setup.md を参照してください。同じファイルには、パスワードなしの SSH が失敗した場合にユーザーへ提示する SSH 失敗時の修正プロンプト も記載されています。

認証情報

• SLURM_USER（必須）：ログインノードの SSH ユーザー名です。マイクロサービスのワークスペースメタデータでは cloud_specific_details.slurm_user です。
• SLURM_HOSTNAME（必須）：フェイルオーバー用の、カンマ区切りのログインホスト名です。マイクロサービスのスキーマでは、リストフィールド cloud_specific_details.slurm_hostname に格納されます。
• SLURM_PARTITION（必須）：GPU ジョブの送信用パーティションリストです。必須の SLURM 情報収集リストで確認してください。パッケージに含まれるデフォルト値は polar,polar3,polar4,grizzly で、4 時間のキューとして扱われます。
• SSH_KEY_PATH（推奨。起動前に設定してください）：ログインノードへの非対話的な公開鍵認証に使う秘密鍵のパスです。パスワードなしの SSH が失敗した場合は、ユーザーに SSH_KEY_PATH=/path/to/private_key を尋ね、references/ssh-setup.md の設定手順を提示してください。複数の代替案の中に埋もれさせないでください。
• SSH_AUTH_SOCK（高度な代替手段）：承認済みの鍵がすでに読み込まれている SSH エージェントのソケットです。ユーザー向けの修正プロンプトでは SSH_KEY_PATH を優先してください。
• SLURM_BASE_RESULTS_DIR（任意）：共有ファイルシステムのベースパスです。tao-core のデフォルト規則は /lustre/fsw/portfolios/edgeai/<your-dir> です。<your-dir> はクラスター上のユーザーごとのディレクトリです。
• SLURM_ACCOUNT（サイトのポリシーにより通常は必須）：#SBATCH --account に指定する課金対象のアカウントです。

サイトでアカウントが必要だとユーザーが伝えた場合、カスタムの結果保存先ルートを希望した場合、またはデフォルトを上書きしないとワークフローを進められない場合を除き、最初の情報収集では SLURM_ACCOUNT や SLURM_BASE_RESULTS_DIR を尋ねないでください。

バックエンドの詳細

このプラットフォームにジョブをルーティングする場合は、backend_details.backend_type = "slurm" を使用してください。マイクロサービスのスキーマでサポートされるバックエンド詳細は次のとおりです。

{
  "backend_type": "slurm",
  "partition": "polar,polar3,polar4,grizzly",
  "cluster_name": "optional-name"
}

ランタイムメタデータは backend_details.slurm_metadata 以下に保存されます。特に slurm_job_id と job_dir が対象です。これらの値を作り出さないでください。sbatch がスケジューラーのジョブ ID を返した後に書き込まれます。

ストレージ

SLURM ジョブはクラスター上で実行されるため、API ホスト上のローカルパスは有効なデータセットパスではありません。共有ファイルシステムの URI を使用してください。

• Lustre 上のユーザー提供データセットには、lustre:///absolute/path を使用します。
• slurm:// パスはマイクロサービスのメタデータに含まれる場合があり、コンテナーの起動前に実際の Lustre パスへ変換されます。
• SLURM のデータセット URI には、単独の /local/path や file:// を使用しないでください。tao-core の検証では、リモートバックエンド用のローカルパスとファイルパスが拒否されます。

データセットのルートパスと、仕様キーを直接指定するパスのどちらも使用できます。

• ルートモード：/lustre/.../<model>/train。モデルスキルによって、<root>/annotations.json などの必須ファイルや、メディアパスとしての<root>に対応付けられます。
• 仕様直接指定モード：custom.train_dataset.annotation_path=/lustre/.../train.json や custom.train_dataset.media_path=/lustre/.../videos.tar.gz など、正確なフィールドを指定します。

パスワードなしの SSH 接続が成功した後、スクリプトを生成する前に、ログインホストから各必須データセットファイル／パスを検証します。

ssh -o BatchMode=yes <SLURM_USER>@<working-login-host> \
  'test -e /lustre/.../annotations.json && test -e /lustre/.../media_or_archive'

リモートの test -e が失敗した場合は処理を中止し、パスの修正を依頼するか、共有クラスター ストレージへのデータの配置を依頼してください。最初のトレーニングジョブ内で失敗するランナースクリプトを作成しないでください。

結果のデフォルトの保存先は次のとおりです。

/lustre/fsw/portfolios/edgeai/<your-dir>/results/<job_id>

<your-dir> は、クラスター上のユーザーごとのディレクトリです。

コンテナーのコードはステータスとアーティファクトの書き込み時にジョブ ID を追加するため、ランナーは TAO_API_RESULTS_DIR に結果の親ディレクトリを設定します。

SLURM ジョブの入力には S3 ではなく Lustre を使用してください。 SLURM のスケジューラーは GPU のアイドルタイムアウトを適用します。スクリプトの冒頭での s3:// からのダウンロードに時間がかかると、トレーニング開始前に割り当て時間を使い切り、スケジューラーによってジョブが終了される可能性があります。トレーニングデータはあらかじめ Lustre に配置してください。S3／HF／NGC からの事前取得は、チェックポイントや設定ファイルなど、小さな補助入力に限り使用できます。詳しくは、完全な説明が記載された references/sdk-usage.md を参照してください。

コンテナーの実行

tao-core は SLURM ハンドラーを使用し、Pyxis／Enroot を通じて TAO コンテナーを実行します。

1. 仕様、環境、クラウドメタデータ用のコンパクトな JSON ファイルを、<job_dir>/specs、<job_dir>/env、<job_dir>/meta の下に配置します。
2. 必要に応じて、srun -n1 -p <conversion_partition> enroot import を使用し、Docker イメージをキャッシュされた SQSH イメージに変換します。
3. <job_dir>/sbatch/job_<job_id>.sbatch の下に sbatch スクリプトを書き込みます。
4. sbatch --export=ALL <script> を送信します。
5. srun --container-image=<image> --container-mounts=/lustre を使用してコンテナーを実行します。

ハンドラーが受け付けるイメージ形式は次のとおりです。

• /path/to/image.sqsh
• registry#image:tag
• docker://registry#image:tag
• 通常の registry/image:tag。必要に応じて Pyxis 形式に変換されます。

SQSH への変換結果は、イメージ名ごとにキャッシュされます。:latest イメージでは、force_reconvert_latest が有効でない限り、キャッシュされた SQSH が使用されます。

リソースのマッピング

tao-core のデフォルト値：

• num_nodes：1
• num_gpus：4
• max_num_gpus_per_node：8
• cpus_per_task：16
• time_hours：4
• timeout_hours：3.8
• max_time_hours：4
• container_mounts：/lustre
• use_requeue：true
• use_sqsh：true

SLURM 用のランチャーやラッパースクリプトを生成する際は、パッケージに含まれるプラットフォームリソースのデフォルト値を使って、ウォールタイムのデフォルト値を明示的に設定します。

export SLURM_TIME_HOURS="${SLURM_TIME_HOURS:-4}"
export SLURM_TIMEOUT_HOURS="${SLURM_TIMEOUT_HOURS:-3.8}"

SLURM で 12 時間をデフォルトにしないでください。ユーザーがより長い SLURM_TIME_HOURS を指定した場合は、送信前に選択したパーティションがその時間に対応していることを確認します。パッケージに含まれるデフォルトのパーティションリスト polar,polar3,polar4,grizzly では、4 時間を超える要求を拒否し、ユーザーが実際により長いウォールタイムを希望している場合に限り、別のパーティションを指定するよう求めます。

num_gpus が max_num_gpus_per_node 以上の場合、ハンドラーはノードを占有する要求として扱い、必要に応じて GPU の総数から追加ノード数を算出します。

マルチノードジョブ（num_nodes > 1）の場合、sbatch スクリプトは WORLD_SIZE、MASTER_ADDR、MASTER_PORT、NODE_RANK、NUM_GPU_PER_NODE をエクスポートし、Cosmos-RL はコントローラー、ポリシー、ロールアウトのワーカーに対して特別なマルチノードロール処理を行います。sbatch ディレクティブ、ランデブー環境変数の表と契約、クラスター要件の詳細については、references/multi-node.md を参照してください。

監視

• スケジューラーのステータスは、保存された SLURM ジョブ ID を使い、squeue または sacct から取得します。
• TAO の終端ステータスは、共有結果フォルダー内の status.json から取得します。
• ユーザーがチャット監視を有効にしている場合、ジョブが PENDING、RUNNING、またはその他の非終端状態である間は、指定された間隔でポーリングを続けます。30 分など、一定時間が経過したことを理由に停止しないでください。共有 GPU パーティションでは、キューの待ち時間が長くなるのは通常のことです。
• チャット監視が有効な場合、非終端状態の SLURM ジョブに対して最終応答を送信しないでください。最終応答はデタッチの操作です。ユーザーがデタッチまたは停止を求めた場合、あるいはジョブが終端状態に達した場合にのみ使用します。
• ログは SSH 経由で以下から読み込みます。

<job_dir>/slurm-logs/<slurm_job_name>-<slurm_job_id>/main.out
<job_dir>/slurm-logs/<slurm_job_name>-<slurm_job_id>/main.err

ステータスの対応関係：

• PENDING → Pending
• RUNNING または COMPLETING → Running
• COMPLETED → status.json を確認します
• FAILED、BOOT_FAIL、DEADLINE、OUT_OF_MEMORY、NODE_FAIL → ログが再試行可能なインフラストラクチャのパターンに一致する場合は再試行し、それ以外は Error
• CANCELLED、PREEMPTED、REVOKED → Canceled
• TIMEOUT → Error
• SUSPENDED、STOPPED → Paused

キャンセル

backend_details.slurm_metadata.slurm_job_id を参照し、SSH 経由で scancel <slurm_job_id> を実行してキャンセルします。存在しない SLURM ジョブ、またはすでに終了した SLURM ジョブも、キャンセル成功として扱います。

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

大規模なマルチノード実行には SLURM が最適なプラットフォームです。num_nodes > 1 を渡すと、SDK が sbatch ディレクティブと PyTorch 分散環境変数を自動的に処理します。実行可能な create_job の例、生成される sbatch ディレクティブ、ランデブー環境変数の表（WORLD_SIZE、NUM_GPU_PER_NODE、NODE_RANK、MASTER_ADDR、MASTER_PORT）、Cosmos-RL のロールに関する注意事項、クラスター要件（Pyxis／Enroot、InfiniBand／NVLink、Lustre）、上流の参考リンクについては、references/multi-node.md を参照してください。

TAO SDK を使用した実行

SDK のインストールについては、事前準備の pip install 'nvidia-tao-sdk[slurm]' をご覧ください。ジョブハンドル、sbatch／squeue／sacct の処理を任せたい場合、ActionWorkflow による実行フォルダーの永続化を利用したい場合、またはクラウドストレージの I/O（s3://、hf_model://、ngc://）を簡単に利用したい場合に使用します。SDK を使わない場合は、sbatch と srun を自分で実行してください。

自動リトライは完全に自動です。バックグラウンドモニターが squeue／sacct をポーリングし、インフラ障害と思われるエラーが発生した場合は、ステージング済みのスクリプトを sbatch で再送信します。これは MAX_JOB_RETRIES = 10 まで行われます。一方、通常のトレーニング失敗はすぐに通知されます。また、デフォルトで #SBATCH --requeue が設定されます（SLURM_USE_REQUEUE のデフォルト値は true です）。SlurmSDK／build_entrypoint のコード例、Lustre を使用して S3 は使用しないというルール、リトライ可能な失敗の分類、自動リトライと再キューイングの動作全体については、references/sdk-usage.md をご覧ください。

失敗モード

よくある失敗には、SSH 認証の失敗、ローカルデータセットパスの拒否、SQSH 変換のタイムアウト、Pyxis／Enroot を利用できない場合、不良ノードや一時的な GPU 障害（ハンドラーが設定された上限までリトライします）があります。それぞれの診断方法と対処方法については、references/troubleshooting.md をご覧ください。