# 通知設計書 26-CheckpointCallbackManager

## 概要

本ドキュメントは、TensorFlowのC++コアランタイムにおける`CheckpointCallbackManager`の通知設計について記載する。本マネージャは、チェックポイントの保存・復元時にコールバックを発火させ、追加情報をチェックポイントと共に保存・復元するためのクラスである。

### 本通知の処理概要

CheckpointCallbackManagerは、SaveV2/RestoreV2 Opの実行時に登録されたコールバック関数を呼び出し、チェックポイントと共に追加情報をファイルとして保存・復元するメカニズムを提供する。ResourceBaseを継承し、TensorFlowのリソース管理システムに統合されている。

**業務上の目的・背景**：機械学習のトレーニングでは、モデルの重みだけでなく、学習スケジュール、データパイプラインの状態、カスタムメトリクスなどの付随情報もチェックポイントと共に保存・復元したい場合がある。CheckpointCallbackManagerは、チェックポイントのライフサイクルイベントにフックし、任意の追加データを同一ディレクトリにファイルとして保存する拡張ポイントを提供する。

**通知の送信タイミング**：(1) SaveV2 Op実行時（Save()メソッド）に登録済みの全SaveCallbackが呼び出される。(2) RestoreV2 Op実行時（Restore()メソッド）に登録済みの全RestoreCallbackが呼び出される。コールバック登録が保存/復元の後に行われた場合は、最後の保存/復元情報を記録しておき、登録時に遅延実行（lazy trigger）する。

**通知の受信者**：RegisterSaveCallback()およびRegisterRestoreCallback()で登録されたコールバック関数が受信者となる。各コールバックはfile_extensionをキーとして一意に識別される。

**通知内容の概要**：SaveCallbackにはcheckpoint_idが渡され、保存すべき文字列コンテンツを返す。返された文字列は`{checkpoint_id}.{file_extension}`ファイルとして保存される。RestoreCallbackにはcheckpoint_idと既存ファイルの内容が渡される。

**期待されるアクション**：SaveCallbackは保存すべきデータを文字列として返す。RestoreCallbackは渡されたデータを用いて内部状態を復元する。空文字列を返した場合は保存がスキップされる。

## 通知種別

プロセス内コールバック（C++ std::function呼び出し）/ ファイルI/O

## 送信仕様

### 基本情報

| 項目 | 内容 |
|-----|------|
| 送信方式 | 同期（SaveV2/RestoreV2 Op内で直接コールバック呼び出し） |
| 優先度 | 高（チェックポイントの整合性に関わる） |
| リトライ | 無し（エラーはログ警告のみ） |

### 送信先決定ロジック

file_extensionをキーとしてsave_callbacks_またはrestore_callbacks_のflat_hash_mapから対応するコールバックを検索する。Save()時は全登録済みSaveCallbackを呼び出し、Restore()時は全登録済みRestoreCallbackを呼び出す。

## 通知テンプレート

### コールバック呼び出しの場合

| 項目 | 内容 |
|-----|------|
| SaveCallback型 | `std::function<StatusOr<std::string>(absl::string_view checkpoint_id)>` |
| RestoreCallback型 | `std::function<Status(absl::string_view checkpoint_id, absl::string_view content)>` |
| ファイル名形式 | `{checkpoint_id}.{file_extension}` |

### 本文テンプレート

```cpp
// SaveCallback
StatusOr<std::string> save_callback(absl::string_view checkpoint_id) {
  // checkpoint_id: 例 "checkpoint-1"
  // 保存したい文字列コンテンツを返す
  return "saved data content";
}

// RestoreCallback
Status restore_callback(absl::string_view checkpoint_id,
                        absl::string_view content_from_checkpoint) {
  // checkpoint_id: 例 "checkpoint-1"
  // content_from_checkpoint: 保存時に返した文字列
  return OkStatus();
}
```

### 添付ファイル

| ファイル名 | 形式 | 条件 | 説明 |
|----------|------|------|------|
| {checkpoint_id}.{file_extension} | テキスト/バイナリ | SaveCallback登録時 | コールバックが返した文字列をファイルとして保存 |

## テンプレート変数

| 変数名 | 説明 | データ取得元 | 必須 |
|--------|------|-------------|-----|
| checkpoint_id | チェックポイント識別子 | GetCheckpointIdAndPathFromPrefix()で抽出 | Yes |
| checkpoint_dir | チェックポイントディレクトリ | GetCheckpointIdAndPathFromPrefix()で抽出 | Yes |
| file_extension | コールバック識別用拡張子 | RegisterSaveCallback()/RegisterRestoreCallback()の引数 | Yes |
| prefix | SaveV2/RestoreV2に渡されるプレフィックス | Op引数 | Yes |

## 送信トリガー・条件

### トリガー一覧

| トリガー種別 | トリガーイベント | 送信条件 | 説明 |
|------------|----------------|---------|------|
| Op実行 | SaveV2::Compute() | Save()呼び出し時 | チェックポイント保存時に全SaveCallbackを発火 |
| Op実行 | RestoreV2::Compute() | Restore()呼び出し時 | チェックポイント復元時に全RestoreCallbackを発火（同一IDの重複復元はスキップ） |
| 遅延実行 | RegisterSaveCallback() | last_saved_checkpoint_id_and_dir_が空でない場合 | コールバック登録が保存後の場合の遅延実行 |
| 遅延実行 | RegisterRestoreCallback() | last_restored_checkpoint_id_and_dir_が空でない場合 | コールバック登録が復元後の場合の遅延実行 |

### 送信抑止条件

| 条件 | 説明 |
|-----|------|
| checkpoint_id抽出失敗 | GetCheckpointIdAndPathFromPrefix()がNotFoundを返した場合、処理をスキップ |
| ファイル既存（Save時） | 同名ファイルが既に存在する場合、SaveCallbackの呼び出しをスキップ |
| 同一ID復元（Restore時） | 前回復元と同じcheckpoint_idの場合、RestoreCallbackを再呼び出ししない |
| ファイル不在（Restore時） | 対応するファイルが存在しない場合、RestoreCallbackの呼び出しをスキップ |
| 空文字列（Save時） | SaveCallbackが空文字列を返した場合、ファイル保存をスキップ |

## 処理フロー

### 送信フロー（Save）

```mermaid
flowchart TD
    A[SaveV2::Compute] --> B[CheckpointCallbackManager::Save prefix]
    B --> C[GetCheckpointIdAndPathFromPrefix]
    C --> D{ID抽出成功?}
    D -->|No| E[return スキップ]
    D -->|Yes| F[mutex_lock: last_saved更新, callbacks コピー]
    F --> G[各SaveCallbackについて]
    G --> H{ファイル既存?}
    H -->|Yes| I[スキップ]
    H -->|No| J[LOG INFO: Calling save callback]
    J --> K[callback checkpoint_id]
    K --> L{Status OK?}
    L -->|No| M[LOG WARNING]
    L -->|Yes| N{空文字列?}
    N -->|Yes| O[スキップ]
    N -->|No| P[WriteStringToFile]
    P --> Q[LOG INFO: Written to file_path]
    I --> G
    M --> G
    O --> G
    Q --> G
```

### 送信フロー（Restore）

```mermaid
flowchart TD
    A[RestoreV2::Compute] --> B[CheckpointCallbackManager::Restore prefix]
    B --> C[GetCheckpointIdAndPathFromPrefix]
    C --> D{ID抽出成功?}
    D -->|No| E[return スキップ]
    D -->|Yes| F{前回と同一ID?}
    F -->|Yes| G[return 重複スキップ]
    F -->|No| H[mutex_lock: last_restored更新, callbacks コピー]
    H --> I[各RestoreCallbackについて]
    I --> J{ファイル存在?}
    J -->|No| K[スキップ]
    J -->|Yes| L[ReadFileToString]
    L --> M[LOG INFO: Calling restore callback]
    M --> N[callback checkpoint_id, payload]
    N --> O{Status OK?}
    O -->|No| P[LOG WARNING]
    O -->|Yes| Q[成功]
```

## データベース参照・更新仕様

### 参照テーブル一覧

本コンポーネントはデータベースを参照しない。ファイルシステム上のチェックポイントファイルを参照する。

| テーブル名 | 用途 | 備考 |
|-----------|------|------|
| ファイルシステム | チェックポイントファイルの存在確認と読み取り | checkpoint_dir配下 |

### 更新テーブル一覧

| テーブル名 | 操作 | 概要 |
|-----------|------|------|
| ファイルシステム | WRITE | SaveCallbackの返却値を{checkpoint_id}.{file_extension}として保存 |

#### 内部管理データ構造

| データ構造 | 操作 | 概要 |
|-----------|------|------|
| save_callbacks_ (flat_hash_map) | READ/INSERT | file_extensionとSaveCallbackのマッピング |
| restore_callbacks_ (flat_hash_map) | READ/INSERT | file_extensionとRestoreCallbackのマッピング |
| last_saved_checkpoint_id_and_dir_ | READ/WRITE | 最後に保存されたチェックポイント情報（遅延実行用） |
| last_restored_checkpoint_id_and_dir_ | READ/WRITE | 最後に復元されたチェックポイント情報（遅延実行・重複防止用） |

## エラー処理

### エラーケース一覧

| エラー種別 | 発生条件 | 対処方法 |
|----------|---------|---------|
| checkpoint_id抽出失敗 | プレフィックスからIDが抽出できない | NotFoundエラーを返却（Save/Restore内ではreturnでスキップ） |
| コールバック重複登録 | 同一file_extensionで2回RegisterCallback | AlreadyExistsエラーを返却 |
| SaveCallback失敗 | コールバックがエラーStatusを返した | LOG(WARNING)で警告出力し処理を継続 |
| ファイル書き込み失敗 | WriteStringToFileが失敗 | LOG(WARNING)で警告出力し処理を継続 |
| ファイル読み取り失敗 | ReadFileToStringが失敗 | LOG(WARNING)で警告出力し処理を継続 |
| RestoreCallback失敗 | コールバックがエラーStatusを返した | LOG(WARNING)で警告出力し処理を継続 |

### リトライ仕様

| 項目 | 内容 |
|-----|------|
| リトライ回数 | 0（リトライなし、エラーは警告ログのみ） |
| リトライ間隔 | - |
| リトライ対象エラー | - |

## 配信設定

### レート制限

| 項目 | 内容 |
|-----|------|
| 制限 | チェックポイント保存/復元の頻度に依存 |

### 配信時間帯

制限なし。チェックポイント操作時に動作する。

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

- コールバックが返すデータはファイルシステムに平文で保存される
- チェックポイントディレクトリのアクセス権限管理が重要
- コールバックの登録はプロセス内のC++コードからのみ可能
- mutex(mu_)により全操作がスレッドセーフ
- ResourceBaseを継承しているため、TensorFlowのリソース管理メカニズムで管理される

## 備考

- リソース名は`kCheckpointCallbackManagerResourceName = "checkpoint_callback_manager"`
- checkpoint_idの抽出は正規表現ベース（kCheckpointFileRegex, kCheckpointTempDirRegex, kCheckpointDirRegex）
- `_temp`サフィックスの一時ディレクトリにも対応（行128-133）
- コールバックのコピーをmutexの外で操作することで、コールバック実行中のデッドロックを回避（行222-227, 244-253）
- Non-copyable, Non-movable（行52-54）

---

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

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

### 推奨読解順序

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

コールバック型とリソース管理の基盤を把握する。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | checkpoint_callback_manager.h | `tensorflow/core/kernels/checkpoint_callback_manager.h` | SaveCallback型（行38-39）、RestoreCallback型（行43-44）、CheckpointCallbackManagerクラス構造（行47-108） |
| 1-2 | resource_base.h | `tensorflow/core/framework/resource_base.h` | ResourceBase基底クラスの契約 |

**読解のコツ**: SaveCallbackは`StatusOr<string>`を返す関数、RestoreCallbackは`Status`を返す関数であることに注目。SaveCallbackの返却文字列がファイルに保存される。

#### Step 2: エントリーポイントを理解する

Save()とRestore()メソッドがOp実行からどのように呼ばれるかを把握する。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | checkpoint_callback_manager.cc | `tensorflow/core/kernels/checkpoint_callback_manager.cc` | Save()メソッド（行214-234）、Restore()メソッド（行236-260） |

**主要処理フロー**:
1. **行214-219**: Save() - GetCheckpointIdAndPathFromPrefix()でcheckpoint_idとdirを抽出
2. **行221-227**: Save() - mutex_lockでlast_savedを更新しcallbacksをコピー
3. **行229-233**: Save() - 各callbackについてTriggerSaveCallbackIfFileNotExist()を呼び出し
4. **行236-242**: Restore() - IDの抽出と重複チェック
5. **行244-253**: Restore() - mutex_lockでlast_restoredを更新しcallbacksをコピー
6. **行255-259**: Restore() - 各callbackについてTriggerRestoreCallbackIfFileExists()を呼び出し

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

内部ヘルパー関数の動作を把握する。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | checkpoint_callback_manager.cc | `tensorflow/core/kernels/checkpoint_callback_manager.cc` | TriggerSaveCallbackIfFileNotExist()（行46-79）、TriggerRestoreCallbackIfFileExists()（行81-104）、GetCheckpointIdAndPathFromPrefix()（行115-144） |

**主要処理フロー**:
- **行46-79**: ファイル不在確認 -> callback呼び出し -> 返却文字列をファイル書き込み
- **行81-104**: ファイル存在確認 -> ファイル読み取り -> callback呼び出し
- **行115-144**: プレフィックスパスからcheckpoint_idとディレクトリを正規表現で抽出

#### Step 4: コールバック登録と遅延実行を理解する

RegisterSaveCallback/RegisterRestoreCallbackの遅延実行メカニズムを把握する。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 4-1 | checkpoint_callback_manager.cc | `tensorflow/core/kernels/checkpoint_callback_manager.cc` | RegisterSaveCallback()（行146-172）、RegisterRestoreCallback()（行180-206） |

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

```
[チェックポイント保存時]
SaveV2::Compute()
    +-- CheckpointCallbackManager::Save(prefix)
            +-- GetCheckpointIdAndPathFromPrefix(prefix)
            +-- [各SaveCallback]
                    +-- TriggerSaveCallbackIfFileNotExist()
                            +-- Env::Default()->FileExists(file_path)
                            +-- callback(checkpoint_id)
                            +-- WriteStringToFile(file_path, content)

[チェックポイント復元時]
RestoreV2::Compute()
    +-- CheckpointCallbackManager::Restore(prefix)
            +-- GetCheckpointIdAndPathFromPrefix(prefix)
            +-- [各RestoreCallback]
                    +-- TriggerRestoreCallbackIfFileExists()
                            +-- Env::Default()->FileExists(file_path)
                            +-- ReadFileToString(file_path, &payload)
                            +-- callback(checkpoint_id, payload)

[遅延実行パターン]
RegisterSaveCallback(file_extension, callback)
    +-- [last_saved非空] TriggerSaveCallbackIfFileNotExist()
RegisterRestoreCallback(file_extension, callback)
    +-- [last_restored非空] TriggerRestoreCallbackIfFileExists()
```

### データフロー図

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

SaveV2 prefix ----------------> Save() ----------------------> {checkpoint_id}.{ext} ファイル
                                 GetCheckpointIdAndPathFromPrefix()
                                 TriggerSaveCallbackIfFileNotExist()
                                 callback(checkpoint_id)
                                 WriteStringToFile()

RestoreV2 prefix --------------> Restore() ------------------> callback(id, content)
                                  GetCheckpointIdAndPathFromPrefix()
                                  TriggerRestoreCallbackIfFileExists()
                                  ReadFileToString()
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| checkpoint_callback_manager.h | `tensorflow/core/kernels/checkpoint_callback_manager.h` | ソース | CheckpointCallbackManagerクラス定義 |
| checkpoint_callback_manager.cc | `tensorflow/core/kernels/checkpoint_callback_manager.cc` | ソース | CheckpointCallbackManager実装 |
| resource_base.h | `tensorflow/core/framework/resource_base.h` | ソース | ResourceBase基底クラス |
| env.h | `tensorflow/core/platform/env.h` | ソース | ファイルシステムAPI（FileExists, WriteStringToFile, ReadFileToString） |
| path.h | `tensorflow/core/platform/path.h` | ソース | io::JoinPath, io::Dirname, io::Basename |
