Add pickle-based checkpoint support for exact suspend/resume

docs/sphinx/manual/en/source/checkpoint.rst

@@ -0,0 +1,132 @@
+Checkpoint and restart
+==================================
+
+PHYSBO provides two ways to persist a search:
+
+- ``policy.save()`` / ``policy.load()`` save the search *results*
+  (history, training data, and predictor) into portable files.
+  They are suitable for analyzing results later or for warm-starting a
+  new search, and the files can be loaded by runs with a different
+  number of MPI processes. They do **not** save the random number
+  generator state, so a restarted run does not reproduce the
+  uninterrupted one.
+- The checkpoint API described here saves the *complete execution
+  state*, including the random number generator, so that a suspended
+  run can be resumed **bit-exactly**: the continued search selects
+  exactly the same candidates as if it had never stopped.
+
+Random number generator modes
+----------------------------------
+
+The checkpoint mechanism interacts with the ``rng`` argument of the
+policies:
+
+.. code-block:: python
+
+    # legacy mode (default): uses the global numpy.random state.
+    # set_seed() seeds the global state, as in previous versions.
+    policy = physbo.search.discrete.Policy(test_X=X)
+
+    # Generator mode: the policy owns a numpy.random.Generator.
+    # The state is stored on the policy itself.
+    policy = physbo.search.discrete.Policy(test_X=X, rng=12345)
+
+In the Generator mode the RNG state is part of the policy object, so
+pickling the policy captures everything needed for exact resumption.
+In the legacy mode the RNG state lives in the global ``numpy.random``
+module; the checkpoint API captures and restores it explicitly (note
+that loading a legacy-mode checkpoint therefore *sets the global
+numpy.random state* as a side effect).
+
+Saving and restoring a checkpoint
+----------------------------------
+
+.. code-block:: python
+
+    import physbo
+
+    policy = physbo.search.discrete.Policy(test_X=X, rng=12345)
+    policy.random_search(max_num_probes=10, simulator=simulator)
+    policy.bayes_search(max_num_probes=20, simulator=simulator, score="TS",
+                        num_rand_basis=500)
+
+    # save the complete execution state into a single file
+    policy.save_checkpoint("search.ckpt")
+
+    # ... later, possibly in a new process ...
+
+    policy = physbo.search.discrete.Policy.load_checkpoint("search.ckpt")
+    # continues exactly as if the run had never stopped
+    policy.bayes_search(max_num_probes=20, simulator=simulator, score="TS",
+                        num_rand_basis=500)
+
+``load_checkpoint`` is a class method of the policy class that saved
+the checkpoint; loading with a different policy class raises an error.
+The checkpoint file records the PHYSBO version (a mismatch emits a
+warning: bit-exact resumption is guaranteed only within the same
+version) and the checkpoint format version.
+
+Usage under MPI
+----------------------------------
+
+``save_checkpoint`` and ``load_checkpoint`` are *collective*
+operations: every rank must call them. The rank-local state (the
+remaining candidates of each rank and its RNG state) is gathered, and
+rank 0 writes a single file; on load, rank 0 reads the file and
+broadcasts it.
+
+.. code-block:: python
+
+    policy = physbo.search.discrete.Policy(test_X=X, comm=comm, rng=12345)
+    policy.random_search(max_num_probes=10, simulator=simulator)
+    policy.save_checkpoint("search.ckpt")     # all ranks call this
+
+    # ... restart with the SAME number of MPI processes ...
+
+    policy = physbo.search.discrete.Policy.load_checkpoint(
+        "search.ckpt", comm=comm)             # all ranks call this
+
+The resuming run must use the same number of MPI processes as the run
+that saved the checkpoint; ``load_checkpoint`` raises an error
+otherwise.
+
+.. note::
+
+   With the BLM predictor (``num_rand_basis > 0``), Thompson sampling
+   draws its posterior sample on rank 0 and broadcasts it, so the
+   search result is independent of the number of ranks. With the GP
+   predictor (``num_rand_basis == 0``), Thompson sampling under MPI is
+   a rank-local approximation and the result depends on the number of
+   ranks; the BLM predictor is recommended for TS under MPI.
+
+Embedding PHYSBO in another application
+----------------------------------------
+
+Applications that embed PHYSBO (such as ODAT-SE) and have their own
+checkpoint mechanism can simply pickle the policy as part of their own
+state. The MPI communicator is excluded from the pickled state
+automatically and must be re-attached with ``set_comm()`` after
+restoring:
+
+.. code-block:: python
+
+    # saving (inside the host application's own checkpoint routine)
+    state = {
+        "step": step,
+        "policy": self.policy,     # picklable; the communicator is excluded
+        # ... other host state ...
+    }
+    with open(filename, "wb") as f:
+        pickle.dump(state, f)
+
+    # restoring
+    with open(filename, "rb") as f:
+        state = pickle.load(f)
+    self.policy = state["policy"]
+    self.policy.set_comm(self.mpicomm)   # re-attach the communicator
+
+``set_comm()`` validates that the size and rank of the communicator
+match the stored state. If the policy uses the Generator mode
+(``rng=`` given), its RNG state is included in the pickle
+automatically; in the legacy mode the host must additionally save and
+restore ``numpy.random.get_state()`` itself (as ODAT-SE already does).

docs/sphinx/manual/en/source/index.rst

@@ -23,6 +23,7 @@ Welcome to PHYSBO's documentation!
    :maxdepth: 2
 
    algorithm
+   checkpoint
 
 .. toctree::
    :maxdepth: 1

docs/sphinx/manual/ja/source/checkpoint.rst

@@ -0,0 +1,126 @@
+チェックポイントと再開
+==================================
+
+PHYSBO には探索を永続化する方法が 2 つあります。
+
+- ``policy.save()`` / ``policy.load()`` は探索の *結果*
+  (history、訓練データ、predictor)を可搬なファイルに保存します。
+  後から結果を解析する場合や、新しい探索のウォームスタートに適して
+  おり、保存時と異なる MPI プロセス数の実行からも読み込めます。
+  ただし乱数生成器の状態は保存されないため、再開した実行は中断しな
+  かった場合の実行を再現しません。
+- 本節で説明するチェックポイント API は、乱数生成器を含む *実行状態
+  の全体* を保存します。中断した探索を **bit-exact に再開**でき、
+  再開後の探索は中断しなかった場合と完全に同じ候補を選択します。
+
+乱数生成器のモード
+----------------------------------
+
+チェックポイント機構は policy の ``rng`` 引数と連動します。
+
+.. code-block:: python
+
+    # legacy モード(デフォルト): グローバルな numpy.random 状態を使用。
+    # set_seed() は従来どおりグローバル状態に seed を設定します。
+    policy = physbo.search.discrete.Policy(test_X=X)
+
+    # Generator モード: policy が numpy.random.Generator を保持。
+    # 乱数状態は policy 自身に格納されます。
+    policy = physbo.search.discrete.Policy(test_X=X, rng=12345)
+
+Generator モードでは乱数状態が policy オブジェクトの一部となるため、
+policy を pickle するだけで厳密な再開に必要な情報がすべて保存されま
+す。legacy モードでは乱数状態はグローバルな ``numpy.random`` モジュ
+ールにあるため、チェックポイント API が明示的に捕捉・復元します
+(このため legacy モードのチェックポイントを読み込むと、副作用として
+*グローバルな numpy.random 状態が書き換わる* ことに注意してください)。
+
+チェックポイントの保存と復元
+----------------------------------
+
+.. code-block:: python
+
+    import physbo
+
+    policy = physbo.search.discrete.Policy(test_X=X, rng=12345)
+    policy.random_search(max_num_probes=10, simulator=simulator)
+    policy.bayes_search(max_num_probes=20, simulator=simulator, score="TS",
+                        num_rand_basis=500)
+
+    # 実行状態の全体を単一ファイルに保存
+    policy.save_checkpoint("search.ckpt")
+
+    # ... 後で(別プロセスでも可)...
+
+    policy = physbo.search.discrete.Policy.load_checkpoint("search.ckpt")
+    # 中断しなかった場合と完全に同じように継続します
+    policy.bayes_search(max_num_probes=20, simulator=simulator, score="TS",
+                        num_rand_basis=500)
+
+``load_checkpoint`` は保存に使った policy クラスのクラスメソッドです。
+異なる policy クラスで読み込もうとするとエラーになります。チェック
+ポイントファイルには PHYSBO のバージョンが記録され、不一致の場合は
+警告が出ます(bit-exact な再開が保証されるのは同一バージョン内のみ
+です)。チェックポイント形式のバージョンも検証されます。
+
+MPI での利用
+----------------------------------
+
+``save_checkpoint`` と ``load_checkpoint`` は *集団操作* であり、
+全 rank が呼び出す必要があります。rank ローカルな状態(各 rank の
+残り候補と乱数状態)が集約され、rank 0 が単一ファイルを書き込みます。
+読み込み時は rank 0 がファイルを読んで全 rank に配布します。
+
+.. code-block:: python
+
+    policy = physbo.search.discrete.Policy(test_X=X, comm=comm, rng=12345)
+    policy.random_search(max_num_probes=10, simulator=simulator)
+    policy.save_checkpoint("search.ckpt")     # 全 rank が呼ぶ
+
+    # ... 保存時と同じ MPI プロセス数で再開 ...
+
+    policy = physbo.search.discrete.Policy.load_checkpoint(
+        "search.ckpt", comm=comm)             # 全 rank が呼ぶ
+
+再開時の MPI プロセス数は保存時と同じでなければなりません。異なる
+場合は ``load_checkpoint`` がエラーを送出します。
+
+.. note::
+
+   BLM predictor(``num_rand_basis > 0``)の場合、Thompson sampling は
+   事後分布からのサンプルを rank 0 で生成して全 rank に配布するため、
+   探索結果は rank 数に依存しません。GP predictor
+   (``num_rand_basis == 0``)の場合、MPI での Thompson sampling は
+   rank ローカルな近似であり結果が rank 数に依存します。MPI で TS を
+   使う場合は BLM predictor を推奨します。
+
+他のアプリケーションへの組み込み
+----------------------------------------
+
+PHYSBO を組み込むアプリケーション(ODAT-SE など)が独自のチェック
+ポイント機構を持つ場合は、自身の状態の一部として policy をそのまま
+pickle できます。MPI コミュニケータは pickle 時に自動的に除外される
+ため、復元後に ``set_comm()`` で再アタッチします。
+
+.. code-block:: python
+
+    # 保存(ホストアプリケーション自身のチェックポイント処理内)
+    state = {
+        "step": step,
+        "policy": self.policy,     # pickle 可能(コミュニケータは除外される)
+        # ... その他のホスト側の状態 ...
+    }
+    with open(filename, "wb") as f:
+        pickle.dump(state, f)
+
+    # 復元
+    with open(filename, "rb") as f:
+        state = pickle.load(f)
+    self.policy = state["policy"]
+    self.policy.set_comm(self.mpicomm)   # コミュニケータを再アタッチ
+
+``set_comm()`` はコミュニケータの size と rank が保存された状態と一致
+することを検証します。policy が Generator モード(``rng=`` 指定)で
+あれば乱数状態は pickle に自動的に含まれます。legacy モードの場合は、
+ホスト側で ``numpy.random.get_state()`` の保存・復元も併せて行って
+ください(ODAT-SE は既にこれを行っています)。

docs/sphinx/manual/ja/source/index.rst

@@ -23,6 +23,7 @@ Welcome to PHYSBO's documentation!
    :maxdepth: 2
 
    algorithm
+   checkpoint
 
 .. toctree::
    :maxdepth: 1