🗿
ローカルからのSLURMジョブの投入・監視
はじめに
モデルを計算機サーバーで学習させるためにSlurmがよく使われると思います。
そんな中でよく直面するのが
- ローカルで書いたスクリプトをサーバにアップロードするのが面倒（Cursorとかあんまり直接繋がないでね...）
- ジョブの実行状況をターミナルで
squeueを何度も叩いて確認 - ジョブが失敗した時のログ探しに時間を取られる
- 環境変数の設定を間違える
こうした課題を解決するために、ssh-slurmというPythonライブラリを開発しました。
本記事では、ローカル環境からちゃっとSSH経由でSLURMジョブを投入・監視する方法を解説します。
ssh-slurmの特徴
1. ローカルファイルからのsbatch実行
ローカルの .shや.slurm ファイルを指定するだけで、自動的にリモートサーバにアップロードして sbatch で実行できます。
# ローカルのtrain.slurmがサーバに自動アップロードされて実行される
# /tmp/ssh-slurmに転送
ssb train.slurm --host gpu-server
2. リアルタイムモニタリングとエラーログの取得
ジョブの実行状況をリアルタイムでモニタリングできます。
また、SLURM_LOG_DIRを設定しておき（詳細は後ほど）、エラーファイルのを同名のディレクトリに%x_%j.logの形式で吐くようにしておくと下記のようにエラーを確認できます。
3. 環境変数の自動検出・転送
ML開発でよく使う環境変数（WANDB_API_KEY、HF_TOKENなど）を自動で検出し、リモートサーバに転送します。
4. プロファイル管理
複数の計算クラスタの接続情報などの設定をプロファイルとして保存し、簡単に切り替えられます。
インストール
# uv
uv add ssh-slurm
# pip
pip install ssh-slurm
要件: Python 3.11+
クイックスタート
1. SSH設定の準備
まず、~/.ssh/config でリモートサーバの接続設定を行います。
Host gpu-cluster
HostName gpu-server.company.com
User your-username
Port 22
IdentityFile ~/.ssh/id_ed25519
事前にデフォルトのコンフィグを登録。
# クラスタを登録
ssb profile add gpu-cluster --ssh-host gpu-cluster
# デフォルトのクラスタを指定
ssb profile set gpu-cluster
2. 基本的な使い方
- ローカルからジョブを投入
# ローカルのtrain.slurmを実行（接続先は上で指定したgpu-cluster）
ssb train.slurm
# デフォルト以外を使用する際
ssb train.slurm --host gpu-cluster2
# 直接指定する方法
ssb train.slurm --hostname gpu-server.company.com --key-file ~/.ssh/id_rsa
実行すると、以下のような出力が表示されます：
3. ファイルの扱い方
-
ローカルファイル (
./train.slurm,train.slurm): 自動アップロード -
リモートファイル (
/shared/scripts/train.slurm): サーバ上のファイルを直接実行
他の機能
環境変数の設定
ML系でよく使われる環境変数はローカルのものを読んで、自動で検出・転送されます。
# ローカルのWANDB_API_KEYとHF_TOKENが自動で転送される
export WANDB_API_KEY="your-wandb-key"
export HF_TOKEN="your-hf-token"
ssb train.slurm --host gpu-cluster
対応している環境変数:
-
WANDB_API_KEY,WANDB_PROJECT,WANDB_ENTITY -
HF_TOKEN,HUGGING_FACE_HUB_TOKEN -
OPENAI_API_KEY,ANTHROPIC_API_KEY CUDA_VISIBLE_DEVICES
ローカルの環境変数をオーバーライド
ssb profile env dev set WANDB_PROJECT "dev-experiments"
ssb profile env dev set SLURM_LOG_DIR "/tmp/slurm_logs"
ssb profile env dev set DEBUG "true"
さいごに
そんなわけで、ssh-slurmの使い方について紹介してきました。
自分のHPC研究開発体験は少しはこれで良くなった気がします。
合わせて、Slurmジョブをワークフローで定義投入できるようにしたsrunxなんてライブラリもあるので良ければ是非。
良いモデルの開発を！
Discussion