# 機能設計書 65-非同期チェックポイント

## 概要

本ドキュメントは、TensorFlowにおけるバックグラウンドでの非同期チェックポイント保存を提供する機能の設計を記述する。

### 本機能の処理概要

非同期チェックポイント機能は、AsyncCheckpointHelperクラスを通じて、アクセラレータ（GPU/TPU）上の変数値をホストCPUにコピーし、バックグラウンドスレッドで非同期的にチェックポイントファイルに書き出す機能である。これにより、チェックポイント保存中もアクセラレータでの学習を継続でき、学習のスループットを向上させる。

**業務上の目的・背景**：大規模モデルの学習では、チェックポイント保存に数十秒～数分を要する場合がある。同期的な保存ではこの間アクセラレータがアイドル状態となり、高価な計算リソースが無駄になる。非同期チェックポイントにより、保存処理をバックグラウンドで行いアクセラレータの利用効率を最大化する。

**機能の利用シーン**：GPU/TPUを使用した大規模モデルの学習時、チェックポイント保存の頻度が高い場合、学習スループットの最大化が求められる場合。

**主要な処理内容**：
1. Trackableオブジェクトの走査とCPUコピー対象の特定
2. アクセラレータからホストCPUへの変数値コピー（_copy_to_cpu）
3. バックグラウンドスレッドでのチェックポイント書き出し
4. TPUEmbeddingの特別なハンドリング（_retrieve_variables）
5. スロット変数の非同期コピー対応
6. 非同期スレッドのライフサイクル管理

**関連システム・外部連携**：チェックポイントシステム（tf.train.Checkpoint）との連携。TPUEmbeddingとの連携。

**権限による制御**：チェックポイントディレクトリへの書き込み権限が必要。

## 関連画面

| 画面No | 画面名 | 関連種別 | 関連する操作・処理 |
|--------|--------|----------|------------------|
| - | - | - | 画面機能マッピングに該当なし |

## 機能種別

データ連携（非同期状態シリアライゼーション）

## 入力仕様

### 入力パラメータ

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| checkpointer_impl | class | Yes | 基底のCheckpointクラス | Noneでないこと |
| root | Trackable/WeakRef | No | チェックポイント対象のルートオブジェクト | Trackableまたはその弱参照 |
| **kwargs | Trackable | No | 名前付きTrackableオブジェクト | 各値がTrackable |

### 入力データソース

アクセラレータ（GPU/TPU）上のTrackableオブジェクトのインメモリ状態。

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| checkpoint.index | ファイル | チェックポイントのインデックスファイル |
| checkpoint.data-* | ファイル | チェックポイントのデータファイル |

### 出力先

指定されたsave_pathのディレクトリ配下のファイルシステム。

## 処理フロー

### 処理シーケンス

```
1. AsyncCheckpointHelper の初期化
   └─ Checkpointインスタンスの作成、キュー・スレッドの準備
2. 初回保存時の初期化（_ensure_initialized）
   └─ 2-1. 全Trackableオブジェクトの走査
   └─ 2-2. TPUEmbeddingの特定と分離
   └─ 2-3. スロット変数の特定
   └─ 2-4. 初回のCPUコピーで object_map を構築
   └─ 2-5. 非同期保存スレッドの起動
3. 保存処理（write）
   └─ 3-1. _copy_to_cpu でアクセラレータ→CPU コピー
   └─ 3-2. キューにTrueを追加して非同期スレッドに保存を通知
   └─ 3-3. 非同期スレッドでCheckpoint.writeを実行
4. 終了処理
   └─ キューにFalseを追加してスレッドを終了
   └─ スレッドのjoin（タイムアウト300秒）
```

### フローチャート

```mermaid
flowchart TD
    A[AsyncCheckpointHelper 初期化] --> B{初期化済み?}
    B -->|No| C[_ensure_initialized]
    C --> D[Trackable走査・object_map構築]
    D --> E[非同期スレッド起動]
    B -->|Yes| F[_copy_to_cpu]
    E --> F
    F --> G[キューにTrue追加]
    G --> H[非同期スレッド: Checkpoint.write実行]
    H --> I{次の保存要求?}
    I -->|Yes| F
    I -->|No/終了| J[キューにFalse追加]
    J --> K[スレッドjoin]
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-65-01 | キューサイズ制約 | キューのmaxsizeは1に設定され、前回の保存完了まで次の保存は待機する | 常時 |
| BR-65-02 | TPUEmbedding特別処理 | TPUEmbeddingオブジェクトは_retrieve_variablesで値を取得する | TPUEmbeddingが存在する場合 |
| BR-65-03 | スレッド終了タイムアウト | 非同期スレッドの終了待機は300秒のタイムアウトを持つ | プロセス終了時 |
| BR-65-04 | atexit登録 | プロセス終了時に自動的に非同期スレッドがjoinされる | 常時 |
| BR-65-05 | エラー伝播 | 非同期スレッドで発生したエラーは次の操作時に呼び出し元に伝播される | エラー発生時 |

### 計算ロジック

該当なし。

## データベース操作仕様

### 操作別データベース影響一覧

データベース操作は行わない。

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| AttributeError | checkpointer_impl未指定 | checkpointer_implがNone | 有効なCheckpointクラスを指定 |
| queue.Full | タイムアウト | 非同期スレッドが300秒以内に完了しない | チェックポイントサイズの削減またはI/O性能の改善 |
| NotImplementedError | CPUコピー未実装 | Trackableが_copy_trackable_to_cpuを実装していない | 警告ログが出力され、該当Trackableはスキップされる |

### リトライ仕様

非同期スレッドのエラーは`_check_async_thread_error`で次回操作時に検出・伝播される。自動リトライは行わない。

## トランザクション仕様

非同期スレッドでの書き出しはCheckpointの通常の保存メカニズムに従う。キューサイズ1により同時書き込みは防止される。

## パフォーマンス要件

- アクセラレータからCPUへのコピーはtf.functionでトレースされ最適化される
- キューサイズ1により、前回の保存完了まで次のCPUコピーが待機するため、保存頻度の上限がある

## セキュリティ考慮事項

通常のチェックポイントと同様のセキュリティ要件。

## 備考

- `_ASYNC_CHECKPOINT`メトリクスラベルで使用状況が追跡される
- `_TPU_EMBEDDING_ATTR`でTPUEmbeddingオブジェクトを識別する

---

## コードリーディングガイド

本機能を理解するために参照すべきファイルと、推奨する読み解き順序を以下に示す。

### 推奨読解順序

#### Step 1: データ構造を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | async_checkpoint_helper.py | `tensorflow/python/checkpoint/async_checkpoint_helper.py` | モジュールレベル定数（40-48行目）：共有状態とラベル |

**読解のコツ**: `_END_TIME_OF_LAST_ASYNC_WRITE`はグローバル変数で複数のAsyncCheckpointHelperインスタンス間で共有される。スレッドロックで保護されている。

#### Step 2: ヘルパー関数を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | async_checkpoint_helper.py | `tensorflow/python/checkpoint/async_checkpoint_helper.py` | `_get_all_trackables`関数（68-140行目）：Trackableオブジェクトの分類 |

**主要処理フロー**:
- **84行目**: TrackableViewでdescendantsを取得
- **89-98行目**: exclude_setに含まれるTrackableを除外
- **108-135行目**: `_trackable_needs_to_be_saved`で保存が必要なTrackableを判定（_serialize_to_tensors、_gather_saveables_for_checkpoint、_copy_trackable_to_cpuの有無で判定）

#### Step 3: AsyncCheckpointHelperクラス

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | async_checkpoint_helper.py | `tensorflow/python/checkpoint/async_checkpoint_helper.py` | `AsyncCheckpointHelper.__init__`（146-217行目） |
| 3-2 | async_checkpoint_helper.py | `tensorflow/python/checkpoint/async_checkpoint_helper.py` | `_ensure_initialized`（241-316行目） |
| 3-3 | async_checkpoint_helper.py | `tensorflow/python/checkpoint/async_checkpoint_helper.py` | `_copy_to_cpu`（219-233行目） |
| 3-4 | async_checkpoint_helper.py | `tensorflow/python/checkpoint/async_checkpoint_helper.py` | `_async_save`（348行目以降） |

**主要処理フロー**:
- **161-163行目**: rootオブジェクトの処理（WeakRef対応）
- **171-172行目**: checkpointer_implの保持
- **193-194行目**: デフォルトデバイスの設定
- **207行目**: キュー（maxsize=1）の初期化
- **210行目**: atexitでスレッド終了を登録
- **244行目**: object_mapの初期化
- **252-253行目**: exclude_setにCheckpointとsave_counterを追加
- **294行目**: object_mapをCheckpointの_saverに設定
- **312-313行目**: 非同期スレッドの起動

### プログラム呼び出し階層図

```
AsyncCheckpointHelper(checkpointer_impl, root, **kwargs)
    │
    ├─ _ensure_initialized()
    │      ├─ _get_all_trackables(root, exclude_set)
    │      │      └─ TrackableView.descendants()
    │      ├─ _handle_tpu_embedding(t)  [TPUEmbedding特別処理]
    │      ├─ _copy_trackable_to_cpu()  [初回コピー]
    │      └─ threading.Thread(target=_async_save)  [スレッド起動]
    │
    ├─ write(save_path)
    │      ├─ _copy_to_cpu()  [tf.function]
    │      │      └─ t._copy_trackable_to_cpu(object_map)
    │      └─ self._queue.put(True)  [非同期保存トリガー]
    │
    └─ _join_async_save_thread()  [終了処理]
           ├─ self._queue.put(False, timeout=300)
           └─ self._async_save_thread.join()
```

### データフロー図

```
[入力]                          [処理]                              [出力]

GPU/TPU上の変数 ─────────────▶ _copy_to_cpu() ──▶ CPU上のコピー変数
                                │                    (object_map)
                                │
CPU上のコピー変数 ──────────────▶ _async_save() ──▶ checkpoint files
(バックグラウンドスレッド)            │                    (非同期書き出し)
                                │
TPUEmbedding ──────────────────▶ _retrieve_variables() ──▶ 値のCPUコピー
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| async_checkpoint_helper.py | `tensorflow/python/checkpoint/async_checkpoint_helper.py` | ソース | 非同期チェックポイントのメインロジック |
| checkpoint.py | `tensorflow/python/checkpoint/checkpoint.py` | ソース | 基底のCheckpointクラス |
| trackable_view.py | `tensorflow/python/checkpoint/trackable_view.py` | ソース | Trackableオブジェクトの走査 |
| checkpoint_context.py | `tensorflow/python/checkpoint/checkpoint_context.py` | ソース | チェックポイントコンテキスト管理 |
