@yag_ays

pueueによるGPU実験のジョブ管理

KaggleやatmaCupなどのデータコンペにおいて、GPUを使った学習/予測などの実験を効率的に行うには、とにかくGPUを休ませずに働き続けてもらう必要があります。手動で一つずつ実験を投入するのでは、どうしても待ち時間が発生してしまいます。複数の実験を事前にキューに入れておき、前の実験が終わったら自動で次の実験が始まるようにすることで、GPUの稼働率を最大化できます。

ここでは、pueueというジョブ管理ツールによる、ローカル環境で完結するGPUの割当管理の方法について説明します。なお、ここでは単一サーバでのGPU学習を対象にしたものであり、クラウド環境やColab等での複数インスタンスが柔軟に利用できる状況は対象としていません。

pueueについて

pueueはRustで書かれたジョブ管理ツールです。あくまで単一ユーザ向けに単独で動作するよう設計されており、シンプルな機能で軽量に動作するのが特徴です。pueuedというデーモンを動作させておき、ユーザはpueueコマンドでジョブを投入すると、優先度や並列度に応じて自動で実行されます。

https://github.com/Nukesor/pueue

pueueの主な特徴は以下の通りです。

  • ターミナル非依存: デーモンとして動作するため、ターミナルを閉じてもジョブは継続して実行されます。別のターミナルからタスクの追加や状態確認が可能です。
  • 並列実行数の制御: 同時に実行するタスク数を指定できます。GPUを使う実験では1並列に制限することで、複数タスクによるGPUの競合やVRAMのOOMを防げます。
  • グループ機能: タスクをグループに分類し、グループごとに並列数を制御できます。複数GPUがある場合には、GPUごとにキューを分けて管理できます。
  • 優先度とタスク順序の操作: タスクの優先度を設定したり、キュー内の順序を入れ替えたりできます。

実際に使っているところ

実際に私が実験管理に使っている様子です。ラベルには実験のIDと概要を記載することで、実験の進捗具合をわかりやすくしています。

pueue status

基本コマンド

タスクの追加 (add)

タスクをキューに追加するには add コマンドを使います。

pueue add "uv run python experiments/exp001/run.py"

複数のタスクを一度に追加することで、自動的に順次実行されます。

pueue add "uv run python experiments/exp001/run.py"
pueue add "uv run python experiments/exp002/run.py"
pueue add "uv run python experiments/exp003/run.py"

状態確認 (status)

現在のタスクの状態を確認するには status コマンドを使います。

pueue status

各タスクのID、状態（Queued、Running、Doneなど）、コマンドが表示されます。

タスクの削除 (remove)

キューからタスクを削除するには remove コマンドを使います。

pueue remove 3  # タスクID 3を削除

実行結果の確認 (log)

完了したタスクの出力を確認するには log コマンドを使います。

pueue log 1     # タスクID 1のログを表示
pueue log       # 最新のログを表示

リアルタイム監視 (follow)

実行中のタスクの出力をリアルタイムで監視するには follow コマンドを使います。tail -f のような挙動で、tqdm等のprogress barもきちんと表示されます。

pueue follow 1  # タスクID 1の出力をリアルタイム表示

完了タスクの削除 (clean)

完了したタスクをリストから削除するには clean コマンドを使います。

pueue clean     # 完了したタスクをすべて削除

GPUの並列制御

各GPUに対応するグループを作成してそれぞれの並列数を1に制限することで、複数GPUへのジョブ投入をコントロールできます。タスク追加時に CUDA_VISIBLE_DEVICES と組み合わせることで、GPUごとにジョブを振り分けることもできます。なお、私は複数GPUを持っていないので、このあたりは未検証です。

priorityによる実行順序の制御

実験を進める中で「この実験を先に試したい」という状況がよく発生します。pueueでは優先度の設定やタスク順序の入れ替えが可能です。

優先度を指定してタスクを追加

--priority オプションで優先度を指定できます。値が大きいほど優先的に実行されます。

pueue add --priority 10 "uv run python experiments/exp004/run.py"
pueue add --priority 1 "uv run python experiments/exp005/run.py"

タスク順序の入れ替え (switch)

すでにキューに入っているタスクの順序を入れ替えるには switch コマンドを使います。

pueue switch 3 5  # タスクID 3と5の位置を入れ替え

スタッシュと復帰

一時的に未実行のタスクを保留にしたい場合は stash を使い、復帰させるには enqueue を使います。

pueue stash 4     # タスクID 4を保留
pueue enqueue 4   # 保留を解除してキューに戻す

callbackによる通知

pueueではcallbackを設定することで、ジョブが終わったタイミングで好きなコマンドを実行できます。ジョブに関する情報も同時に取得することが出来るため、実験IDや評価値を含めてslackに通知することが可能です。ただし、pueue全体で1つのコマンドしか対応していないので、複数のコンペに対応するには、一度振り分けのスクリプトを挟むやり方が良いと思います。

私の場合はSlack Webhookで実験完了時に以下のような通知が飛ぶようにしていました。

pueueのslack通知

まとめ

pueueはローカル環境下でのジョブ管理にちょうどよいシンプルさで扱いやすく気に入っています。CLIなのでClaude Code等のAI Agentからも容易に操作できるため、親和性が高いのも嬉しいところです。