0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

強化学習の自己対戦をマルチプロセス化しようとしたら、次々に謎のクラッシュが出て原因調査に追われた話

0
Posted at

はじめに

通常の機械学習では、データをミニバッチにまとめて処理できるため、学習そのものをマルチプロセス化する場面はあまりありません。
一方、強化学習では、シミュレータが GPU に対応しておらずバッチ処理できないことがあり、その場合は高速化のためにマルチプロセス実行が選択肢に入ってきます。

私も強化学習の自己対戦で学習データ集めを並列化しようとしたのですが、これが一筋縄ではいかず、原因の異なるクラッシュやフリーズに次々と遭遇しました。
しかも py-spygdb といった定番ツールが ptrace の権限制限で使えないサンドボックス環境で、標準的な手段だけで原因を追うことになりました。

この記事は、その調査で使った「プロセスの状態を観測し、原因を切り分けていく方法」を、自分なりにまとめたものです。
題材は私の自己対戦 MCTS ですが、ネイティブライブラリ × マルチプロセスでクラッシュやフリーズに悩んでいる人がそのまま使える手順になっていると思います。

1. 前提:なぜマルチプロセスにするのか、なぜ難しいのか

通常の機械学習はデータをミニバッチにまとめて GPU に流せますが、強化学習では経験を生成するシミュレータが GPU 非対応でバッチ処理できないことがあります。その場合は 1 件ずつ CPU で回すしかなく、素直な高速化の手段がプロセス並列になります。

やっかいなのは、コンペなどで配布されるシミュレータ(共有ライブラリ/ネイティブバイナリ)が、必ずしも並列実行を前提に作られていない点です。プロセスグローバルな状態を持っていたり、スレッド安全でなかったりすると、複数プロセスから同時に叩いたときにデッドロックや SIGSEGV を起こしえます。

「強化学習ライブラリの並列環境機構を使えばいいのでは」と思うかもしれません。Stable-Baselines3 の SubprocVecEnv や Gymnasium の AsyncVectorEnv などはありますが、いずれも「プロセスごとに独立した環境を立てて並列に進める」点は自前実装と同じで、環境側が並列安全でなければ載せ替えても解決しません。

私の場合は、CPU 専用・1 プロセス 1 対局の対戦エンジンで、1 手ごとに数十回探索する MCTS を 1 ラウンド 200 局ほど回す必要があり、逐次では終わらないので multiprocessing で並列化しました。

2. どう並列化するか(最小構成)

対局(エピソード生成)どうしは独立なので、素直には multiprocessing(以下 mp)の Pool に 1 ラウンド分のタスク(ここでは 200 局)を渡し、N 個のワーカープロセスで消化させます。各ワーカーは 1 タスク=1 対局を最後まで進め、学習データを返します。

with mp.Pool(processes=N) as pool:
    results = pool.map(play_one_game, range(200))

play_one_game の中身は「対局を開始 → 1 手ごとに探索して着手 → 終局まで繰り返す → 学習データを返す」だけです。ここまではごく普通の構成ですが、この「対局は独立なはず」という前提が、ネイティブ側の共有状態で崩れることがあります。ここから、実際に出たクラッシュを順に見ていきます。

3. まず「止まっているのか、動いているのか」を読む

観測:ps / wchan / cgroup で読む

並列で「進まない」ときは、まず止まっているのか遅いだけなのかを切り分けます。ps auxSTATS=スリープ / R=実行 / D=割り込み不能スリープ)と %CPUuptime の load average を見ます。なお D はディスクやカーネル内 I/O を待っていることが多く、必ずしも停止やデッドロックではありません。

止まっていそうなら、何を待っているかを見ます。

ps -eo pid,stat,wchan | grep <ワーカーのPID>

wchanfutex_wait_queue(環境によっては futex_wait_queue_me などの別名)が出ていれば、ミューテックスや条件変数などの待機を表します。これ自体でデッドロックとは断定できませんが、全ワーカーが長時間そのままで、CPU 時間も増えていないなら、ロック待ちによる停止を疑えます。停止か低速かは、数分あけて 2 回 ps を撮り、CPU 時間(TIME)が増えているかで判定します。OOM で殺されていないかは cgroup の memory.events で否定できます。

cat /sys/fs/cgroup/memory.events   
# 実行前後で oom_kill が増えていなければ、
# 少なくともこの cgroup での OOM kill は発生していない

最初のフリーズの犯人は fork だった(→ spawn)

multiprocessing は Linux 既定で fork を使いますが(Python 3.14 以降は異なるらしい、未調査)、fork はメモリだけコピーして実行中スレッドを引き継ぎません。torch のように内部スレッドプールやロックを持つライブラリを import 済みだと、fork の瞬間に他スレッドが握っていたロックが、保持されたまま子プロセスへ渡ります。ところが、そのロックを解放するはずのスレッドは子プロセスに存在しないため、ロックは永久に解放されず、子プロセスはそこで固まります(PyTorch も既知問題として明言しています)。対策は spawn に変えるだけです。

ctx = mp.get_context("spawn")
with ctx.Pool(processes=N) as pool:
    ...

私の場合: fork で全ワーカーが futex_wait のまま完全停止していましたが、spawn に変えると全ワーカーが CPU を使って動き出し、60 局が約 140 秒で完走しました(spawn は再 import のぶん起動は遅くなります)。

4. ptrace が使えない環境で、プロセスの中を見る(faulthandler)

spawn で安定したと思ったら、別の状況で今度は完全にフリーズしました。中身を覗きたいところですが、py-spy / gdb / strace はどれも ptrace を必要とし、サンドボックスでは権限で弾かれます(Permission denied)。

そこで、外部からアタッチする代わりに、プロセス自身にスタックを吐かせる faulthandler を使います(標準ライブラリ、ptrace 不要)。各ワーカーの起動時に仕込んでおきます。

import faulthandler
# ファイルはグローバル等で参照を保持し、プロセスの生存中は閉じない
_fault_log = open(f"worker_{os.getpid()}.log", "a", buffering=1)
faulthandler.enable(file=_fault_log, all_threads=True)
faulthandler.dump_traceback_later(30, repeat=True, file=_fault_log)

これで、ハング中でも 30 秒おきに全スレッドのスタックが採れます。読み解き方として、毎回同じ場所で止まっていれば特定のロックや呼び出しを疑います。逆にバラバラの場所に見える場合は、単一の Python レベルのロックだけでは説明しにくく、複数の待機要因やネイティブ側の状態不整合も候補に入れます(単に各ワーカーが別々の処理段階にいるだけ、という可能性も残ります)。なお faulthandler で採れるのは基本的に Python のスタックで、ネイティブライブラリ内部の C/C++ スタックまで読めるわけではない点にも注意します。

私の場合: 停止箇所が NN の forward 中・torch.no_grad の終了処理中・json.load 中とバラバラで、「1 箇所のロック待ち」では説明できないと分かりました。

5. クラッシュを定量化して切り分ける(exitcode + 並列数スイープ)

再実行すると、今度はフリーズではなくワーカーが次々クラッシュしていました。multiprocessing.Pool はワーカー数を保とうとして異常終了したワーカーを再生成することがあり、親プロセスだけを見ていると気づきにくく、処理が進まない・同じクラッシュを繰り返す、といった状態になりがちです。

原因を見るには Pool をやめ、Process で 1 タスク=1 プロセスにして exitcode を拾います。0 は正常終了、正の値は Python 側の異常終了、負の値はシグナルによる終了(-11=SIGSEGV、-6=SIGABRT)、タイムアウトは別枠、と種類を区別できるようにします(ctx は §3 の get_context("spawn") です)。

p = ctx.Process(target=play_one_game, args=(i,))
p.start()
p.join(timeout=300)
if p.is_alive():
    p.terminate()
    p.join()
    status = "timeout"
else:
    status = "ok" if p.exitcode == 0 else f"crash({p.exitcode})"

あとは変数を 1 つずつ動かします。まず単一タスクを逐次実行してロジック自体の健全性を確認し、次に並列数を 2→4→8→… とスイープしてクラッシュ率の変化を見ます。並列数に強く依存して悪化するなら、たまのレースではなく「同時アクセスによるリソース競合」を疑えます。

私の場合: 単一実行では 20 局すべて正常、並列は数を上げるほど悪化し、並列 18 では最大 78% が SIGSEGV/SIGABRT でした。ここで「複数プロセスが同時にネイティブ処理を叩くこと」自体が問題だと絞れました。

6. 「効きそうな仮説」を worktree で実検証する

§4 の結果から、単一の Python レベルのロック待ちだけでは説明しにくく、ネイティブ側の状態不整合やメモリ破壊も候補に入りました。とはいえ確定させる前に、外形から否定できる仮説は先に潰しておきます。まず浮かんだのは「最近ネットワーク構成を変えたせいでは」でした。もっともらしい説明はいくらでも作れますが、こういう仮説は外れることも多いので、決めつける前に確かめます。

git worktree を使うと、変更前のコミットを別ディレクトリに取り出して、同じ診断を同条件でかけられます。

git worktree add ../old-version <古いコミット>

旧コードでも同じように起きれば、その変更は無実だと切り分けられます。

私の場合: §5 と同じ並列数スイープを旧コードにかけたところ、同水準(並列 18 で最大 89%)でクラッシュしたので、直近の変更は無実と確定しました。ただし「では何が原因か」は残ります。

7. ネイティブ並列クラッシュへの打ち手カタログ

ネイティブ処理が並列で壊れるときに取れる、一般的な打ち手をまとめます。

  • ネイティブ呼び出しをロックで直列化する: エンジンを叩く部分だけをプロセス間ロックで囲み、純粋な Python や推論計算はロック外で並列のまま残します。
with engine_lock:
    next_state = engine.search_step(search_id, action)
  • 遅延初期化を先読みする: 「初回だけネイティブを呼ぶ」キャッシュやマスタデータ読み込みは、ロックの外に残りがちです。ワーカー起動時にロック内で先に呼んでおきます。
  • スレッド数を絞る: 各プロセスがさらに内部スレッド(PyTorch/BLAS)を張るため、torch.set_num_threadsOMP_NUM_THREADS でプロセスあたりのスレッド数を制限し、CPU オーバーサブスクリプションを避けます。

私の場合: ネイティブ呼び出しをロックで直列化し、遅延ロードも先読みに直したのですが、それでもクラッシュは収まりませんでした。ロックで塞いだ経路の外にまだ何かが残っている——というより、原因はそもそも別のところにありました。

8. 結局は「自分の使い方」を疑う

ここまでの対策は、すべて「ネイティブライブラリが並列安全でない」という仮説に立っていました。しかし、ロックで直列化しても再現するなら、その仮説自体を疑う番です。ここが山場で、配布物のせいだと決めつける前に、自分の呼び出し側(ライブラリの利用契約を満たしているか)を一から見直します。ネイティブライブラリで特に多いのは次のあたりです。

  • リソースの解放漏れ: 確保系の API(状態・ハンドル・バッファを確保するもの)に対応する解放 API を呼び忘れると、ライブラリ内部の状態が解放されないまま溜まり続けます。単純なメモリリークなら最終的に OOM で落ちますが(これは §3 の cgroup で切り分け済み)、内部実装によっては、メモリ使用量の増加だけでなく、内部スロットの枯渇や不正な状態の再利用につながり、確保箇所とは離れた場所で SIGSEGV として現れることもあります。「回数・並列数が増えるほど悪化する」のが典型的なサインです。確保したら必ず解放します。
  • 呼び出し順序依存: 特定の順序でだけ壊れることもあります。ソースを追っても根本が分からなくても、再現する順序を避ければ回避できます(根本特定と回避は別問題です)。
  • 例外時の後片付け漏れ: 処理の途中で例外が飛ぶと後始末(finish 系)が呼ばれず、状態が壊れます。try/finally で必ず呼びます。

配布物にソースが同梱されていることもあります。「ブラックボックス」と決めつけず、読めるなら読むと裏取りが早いです。

私の場合(Tips): 原因は 2 つでした。1 つは、探索状態を確保する API に対する解放呼び出しの漏れ(=上のリソースリーク)です。ヒープ破壊そのものを直接確認したわけではありませんが、解放を足すと再現しなくなったので、少なくともこのリソース管理がクラッシュに関与していたと判断しました。直したら並列 18 でも 54/54 クリーンになり、並列を復旧できました。もう 1 つは、あるネイティブ呼び出しの直後に大量のファイル読み込みを行うと SIGSEGV になる順序依存で、呼び出し順を入れ替えて回避しました(根本は未特定)。「エンジン固有の不安定さ」と思っていたものが、実は自分のリソース管理と呼び出し順序だった、というオチです。

9. 補足:強化学習でマルチプロセスが詰まりやすい理由

補足として、この手の詰まりを一般化しておきます。「強化学習だからマルチプロセス化できない」わけではなく、正確には、環境・GPU・ネイティブライブラリ・並列処理を組み合わせるぶん問題が起きやすい、というのが実態です。代表的な落とし穴は次のあたりです。

  • 環境が並列を想定していない: 固定パスの一時ファイル・グローバル状態・キャッシュ・ネイティブの静的変数などを内部で共有していると、別プロセスに分けても競合します。
  • fork × CUDA / スレッド: 親で CUDA やスレッドを初期化した後に fork すると、不完全な状態が子に渡ります(いわゆる poison fork)。子で CUDA を使うなら spawn / forkserver を。
  • リソースの多重化: 各プロセスが GPU モデルや内部スレッド(PyTorch/BLAS)を抱えると、GPU メモリや CUDA コンテキストが多重化し、CPU もオーバーサブスクリプションになります。推論を 1 プロセスに集約する・OMP_NUM_THREADS でスレッドを絞る、などで抑えます。
  • spawn の制約: spawn では子へ渡すオブジェクトが pickle 可能である必要があり、ネイティブハンドル等は子プロセス内で作ります。観測の受け渡しに共有メモリや pipe / Queue を使う実装では、それ由来のハングもあります。

ライブラリを使っても「プロセスに分ければ自動で安全」ではありません(Stable-Baselines3 の SubprocVecEnvspawn / forkserver 前提や if __name__ == "__main__" を注意書きし、Gymnasium は並列の AsyncVectorEnv と単一プロセスの SyncVectorEnv を使い分けます)。

なお、私自身の真因(§8 のリソース解放漏れ)は、この一覧の「環境側の罠」とは別種で、自分のコードの問題でした。並列化の一般的な落とし穴と、自分の使い方の問題は別軸として切り分けておくと、原因を見誤りにくくなります。

まとめ

原因の異なるクラッシュに追われましたが、最終的に原因は「自分の使い方」(リソースの解放漏れと呼び出し順序)にあると分かり、並列実行を復旧できました。

使った手順を、道具箱としてまとめておきます。

  • psSTAT / %CPU で「止まりか遅さか」を見る
  • wchanfutex_wait_queue ならデッドロックを疑う
  • 2 回 ps を撮って CPU 時間の増分を見る
  • cgroup の memory.events で OOM を否定する
  • ptrace が使えないなら faulthandler.dump_traceback_later
  • Pool をやめ exitcode でシグナルを区別する
  • 並列数を 1 つずつスイープする
  • 仮説は git worktree で旧コードと実比較する

特別なツールは使わず、標準の psmultiprocessingfaulthandler だけでここまで切り分けられました。

最後に 2 つ、教訓を書いておきます。1 つは、配布物のせいだと決めつける前に、自分の使い方を疑うこと。もう 1 つは、根本原因を完全に特定できなくても、再現条件さえ外せれば前に進めること。原因不明のまま無限に粘るのでも、早々に諦めるのでもなく、分かっている事実を積み上げて判断していくのがデバッグなのだと思います。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?