# 通知設計書 27-StreamCallbackRegistry

## 概要

本ドキュメントは、TensorFlowのC++コアランタイム（TFRT）における`StreamCallbackRegistry`の通知設計について記載する。本レジストリは、ストリーミング処理においてコールバックを管理・通知するためのグローバルなレジストリクラスである。

### 本通知の処理概要

StreamCallbackRegistryは、TFRTランタイムにおけるストリーミング推論パイプラインのコールバック管理を担う。`tf.PwStreamResults` Opがストリーミング結果を生成すると、登録されたコールバックが呼び出され、テンソルデータがリアルタイムにクライアントへ配信される。ScopedStreamCallbackによるRAII管理でコールバックのライフサイクルが自動制御される。

**業務上の目的・背景**：大規模言語モデルや生成AIのサービングにおいて、全出力が完成する前に部分的な結果をストリーミング配信する需要がある。StreamCallbackRegistryは、モデル実行中に`tf.PwStreamResults` Opから逐次生成されるテンソル結果を、登録済みのコールバック関数経由で呼び出し元に配信するメカニズムを提供する。callback_id（実行可能モジュール固有）とstep_id（呼び出し固有）のペアでコールバックを一意に識別する。

**通知の送信タイミング**：モデル実行中に`tf.PwStreamResults` Op（もしくはInvokeStreamCallbackOp）が実行されるたびに、StreamCallbackRegistry::Invoke()が呼ばれ、対応するコールバックがスレッドプールで非同期に実行される。

**通知の受信者**：StreamCallbackRegistry::Register()で登録されたコールバック関数が受信者である。コールバックはcallback_idとstep_idのペアで識別される。ScopedStreamCallbackのデストラクタにより自動的にアンレジストされる。

**通知内容の概要**：StreamedResult構造体に含まれるテンソルマップ（`absl::flat_hash_map<std::string, tensorflow::Tensor>`）がコールバックに渡される。各テンソルは名前付きで、モデルの中間出力やストリーミング出力を表す。

**期待されるアクション**：コールバック関数はストリーミングテンソルを受信し、クライアントへの配信（gRPCレスポンスストリーミング等）を行う。コールバックはシングルスレッドで逐次実行されるため、スレッドセーフである必要はない。

## 通知種別

プロセス内コールバック（C++ absl::AnyInvocable呼び出し）/ スレッドプール非同期実行

## 送信仕様

### 基本情報

| 項目 | 内容 |
|-----|------|
| 送信方式 | 非同期（スレッドプール経由でコールバック実行） |
| 優先度 | 高（リアルタイムストリーミング配信） |
| リトライ | 無し（エラー時はStatusを返却） |

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

callback_idとstep_idのペアをキーとしてstream_callbacks_ flat_hash_mapからCallbackStateを検索する。見つからない場合はNotFoundErrorを返却する。見つかった場合はCallbackState::Invoke()でスレッドプールにコールバック実行をスケジュールする。

## 通知テンプレート

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

| 項目 | 内容 |
|-----|------|
| コールバック型 | `absl::AnyInvocable<void(absl::flat_hash_map<std::string, tensorflow::Tensor>)>` |
| 識別キー | (StreamCallbackId, StepId) ペア |
| 実行方式 | tsl::thread::ThreadPoolInterfaceでスケジュール |
| 実行保証 | シリアル実行（シングルスレッド） |

### 本文テンプレート

```cpp
// コールバック関数のシグネチャ
void callback(absl::flat_hash_map<std::string, tensorflow::Tensor> tensors) {
  // tensors: 名前付きテンソルのマップ
  // 例: {"output_token": Tensor(...), "log_prob": Tensor(...)}
  // ストリーミング配信処理を実行
}
```

### 添付ファイル

なし（プロセス内通信）

## テンプレート変数

| 変数名 | 説明 | データ取得元 | 必須 |
|--------|------|-------------|-----|
| callback_id | コールバック識別子（実行可能モジュール固有） | CreateStreamCallbackId()で生成 | Yes |
| step_id | ステップ識別子（呼び出し固有） | リクエストコンテキスト | Yes |
| model_name | モデル名 | サービング設定 | Yes |
| tensors | ストリーミングテンソルマップ | tf.PwStreamResults Op | Yes |
| enqueued_time | キューイング時刻 | StreamedResult構造体 | Yes |

## 送信トリガー・条件

### トリガー一覧

| トリガー種別 | トリガーイベント | 送信条件 | 説明 |
|------------|----------------|---------|------|
| Op実行 | tf.PwStreamResults / InvokeStreamCallbackOp | StreamCallbackRegistry::Invoke()呼び出し | ストリーミング結果生成時 |
| コールバック登録 | StreamCallbackRegistry::Register() | callback_idとstep_idのペアが未登録 | ScopedStreamCallbackを返却 |
| コールバック解除 | ScopedStreamCallback::~ScopedStreamCallback() | デストラクタ呼び出し時 | RAIIによる自動解除 |

### 送信抑止条件

| 条件 | 説明 |
|-----|------|
| コールバック未登録 | Invoke()で(callback_id, step_id)が見つからない場合、NotFoundErrorを返却 |
| コールバック閉鎖済み | CallbackState::closed_がtrueの場合、InternalErrorを返却 |
| 重複登録 | Register()で同一(callback_id, step_id)が既に存在する場合、AlreadyExistsErrorを返却 |

## 処理フロー

### 送信フロー

```mermaid
flowchart TD
    A[tf.PwStreamResults Op実行] --> B[StreamCallbackRegistry::Invoke]
    B --> C[MutexLock取得]
    C --> D[stream_callbacks_.find callback_id, step_id]
    D --> E{CallbackState発見?}
    E -->|No| F[NotFoundError返却]
    E -->|Yes| G[CallbackState::Invoke thread_pool, result]
    G --> H[MutexLock on CallbackState]
    H --> I{closed?}
    I -->|Yes| J[InternalError返却]
    I -->|No| K[num_outstanding++]
    K --> L[thread_pool->Schedule lambda]
    L --> M[InvokeCallback result]
    M --> N[RecordDequeueLatency]
    N --> O[TraceMeプロファイリング]
    O --> P[callback テンソルマップ]
    P --> Q[RecordCallbackLatency]
    Q --> R[num_outstanding--]
```

### コールバックライフサイクル

```mermaid
flowchart TD
    A[CreateStreamCallbackId] --> B[MLIRモジュール書き換え]
    B --> C[Register callback_id, step_id, callback]
    C --> D[ScopedStreamCallback取得]
    D --> E[モデル実行開始]
    E --> F[Invoke 繰り返し]
    F --> G[モデル実行完了]
    G --> H[ScopedStreamCallback破棄]
    H --> I[Unregister callback_id, step_id]
    I --> J[CallbackState::Close]
    J --> K[未完了コールバック待ち]
    K --> L[完了]
```

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

### 参照テーブル一覧

本コンポーネントはデータベースを参照しない。

### 更新テーブル一覧

なし。

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

| データ構造 | 操作 | 概要 |
|-----------|------|------|
| stream_callbacks_ (flat_hash_map) | READ/INSERT/DELETE | (callback_id, step_id)とCallbackStateのマッピング |
| CallbackState.closed_ | READ/WRITE | コールバック閉鎖フラグ |
| CallbackState.num_outstanding_ | READ/WRITE | 未完了コールバック呼び出しカウント |

## エラー処理

### エラーケース一覧

| エラー種別 | 発生条件 | 対処方法 |
|----------|---------|---------|
| NotFoundError | Invoke()でコールバックが見つからない | "does not exist; this usually indicates that a streaming signature was called by a non-streaming request"メッセージを含むエラー返却 |
| AlreadyExistsError | Register()で同一IDが既に存在 | エラーStatusを返却 |
| InternalError | 閉鎖済みコールバックへのInvoke() | "Failed to invoke the callback that is closed."メッセージを含むエラー返却 |
| InternalError | StreamControllerInterfaceが未登録 | GetGlobalStreamCallbackRegistry()初期化時にエラー |

### リトライ仕様

| 項目 | 内容 |
|-----|------|
| リトライ回数 | 0（リトライなし） |
| リトライ間隔 | - |
| リトライ対象エラー | - |

## 配信設定

### レート制限

| 項目 | 内容 |
|-----|------|
| 制限 | なし（tf.PwStreamResults Opの実行頻度に依存） |
| スレッドプール | ThreadPoolInterfaceに委譲 |

### 配信時間帯

制限なし。モデルサービング中に動作する。

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

- ストリーミングテンソルにはモデルの推論結果が含まれるため、適切なアクセス制御が必要
- StreamControllerInterfaceを通じたコントローラとの通信にはアドレス情報が含まれる（controller_address_）
- コールバックIDはtsl::random::New64()で生成されるランダムな64ビット値で、予測困難
- absl::Mutexにより全レジストリ操作がスレッドセーフ
- ScopedStreamCallbackのRAIIパターンにより、リソースリークが防止される

## 備考

- StreamCallbackIdはint64_tベースのSafeId型（TensorFlowはuint64属性をサポートしないため）
- CreateStreamCallbackId()はMLIRモジュールを書き換え、tf.PwStreamResults Opに_controller_address, _model_name, _callback_id属性を付与する（行47-81）
- コールバックはシリアル実行が保証されている（スレッドセーフ不要）
- ScopedStreamCallback::Unregister()ではCallbackState::Close()で未完了コールバックの完了を待つ（行102-111）
- tsl::profiler::TraceMeによるプロファイリングが組み込まれている（行119-125, 207-213）
- GetGlobalStreamCallbackRegistryはグローバルシングルトン
- StreamInterfaceFactoryでコントローラとワーカーのインターフェースを登録可能

---

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

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

### 推奨読解順序

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

ストリーミング通知に関わるデータ型を把握する。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | stream.h | `tensorflow/core/tfrt/runtime/stream.h` | StreamedResult構造体（行46-49）、StreamCallbackId型（行51-53）、ScopedStreamCallbackクラス（行258-280） |
| 1-2 | stream.h | `tensorflow/core/tfrt/runtime/stream.h` | StreamControllerInterface（行57-73）、StreamWorkerInterface（行77-95） |
| 1-3 | step_id.h | `tensorflow/core/tfrt/runtime/step_id.h` | StepId型の定義 |

**読解のコツ**: StreamCallbackIdとStepIdのペアがコールバックの一意識別子である。StreamedResultはテンソルマップとキューイング時刻を保持する。

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

StreamCallbackRegistryの主要APIを把握する。

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | stream.h | `tensorflow/core/tfrt/runtime/stream.h` | StreamCallbackRegistryクラス定義（行155-243）。Register(), Invoke(), stream_interface()のAPI |
| 2-2 | stream.h | `tensorflow/core/tfrt/runtime/stream.h` | CallbackStateネストクラス（行191-232）。Invoke(), Close(), InvokeCallback()のAPI |

**主要処理フロー**:
1. **行175-180**: Register() - コールバック登録、ScopedStreamCallback返却
2. **行182-184**: Invoke() - コールバック呼び出し（スレッドプール経由）
3. **行234-235**: Unregister() - コールバック解除（privateメソッド）

#### Step 3: 実装の詳細を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | stream.cc | `tensorflow/core/tfrt/runtime/stream.cc` | Register()実装（行132-149）。flat_hash_mapへの挿入とCallbackState生成 |
| 3-2 | stream.cc | `tensorflow/core/tfrt/runtime/stream.cc` | Invoke()実装（行151-166）。コールバック検索とCallbackState::Invoke()呼び出し |
| 3-3 | stream.cc | `tensorflow/core/tfrt/runtime/stream.cc` | CallbackState::Invoke()実装（行83-100）。closed_チェック、num_outstanding_管理、スレッドプールスケジューリング |
| 3-4 | stream.cc | `tensorflow/core/tfrt/runtime/stream.cc` | CallbackState::InvokeCallback()実装（行114-130）。デキューレイテンシ記録、TraceMeプロファイリング、コールバック実行 |
| 3-5 | stream.cc | `tensorflow/core/tfrt/runtime/stream.cc` | CallbackState::Close()実装（行102-112）。closed_設定、未完了コールバック待ち |

#### Step 4: コールバックID生成とMLIR書き換えを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 4-1 | stream.cc | `tensorflow/core/tfrt/runtime/stream.cc` | CreateStreamCallbackId()実装（行47-81）。MLIRモジュールのwalk、PwStreamResultsOp属性の設定 |

**主要処理フロー**:
- **行54-55**: module->walk()でPwStreamResultsOpを収集
- **行57-58**: ops.empty()チェック（ストリーミングOpがなければnullopt）
- **行70-71**: tsl::random::New64()でランダムIDを生成（int64_tにキャスト）
- **行74-78**: 各PwStreamResultsOpに_controller_address, _model_name, _callback_id属性を設定

#### Step 5: RAIIライフサイクルを理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 5-1 | stream.cc | `tensorflow/core/tfrt/runtime/stream.cc` | ScopedStreamCallback::Unregister()実装（行202-224）。TraceMeプロファイリング、Unregister呼び出し、Close()呼び出し |
| 5-2 | stream.cc | `tensorflow/core/tfrt/runtime/stream.cc` | ムーブコンストラクタとムーブ代入演算子（行181-200） |

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

```
[モデルロード時]
CreateStreamCallbackId(model_name, module)
    +-- module->walk() [PwStreamResultsOp収集]
    +-- tsl::random::New64() [コールバックID生成]
    +-- op->setAttr() [MLIRモジュール書き換え]

[リクエスト処理時]
StreamCallbackRegistry::Register(model_name, callback_id, step_id, callback)
    +-- stream_callbacks_.insert({callback_id, step_id}, CallbackState)
    +-- ScopedStreamCallback返却

[ストリーミング結果生成時（繰り返し）]
StreamCallbackRegistry::Invoke(thread_pool, callback_id, step_id, result)
    +-- stream_callbacks_.find({callback_id, step_id})
    +-- CallbackState::Invoke(thread_pool, result)
            +-- [closed_チェック]
            +-- num_outstanding_++
            +-- thread_pool->Schedule(lambda)
                    +-- CallbackState::InvokeCallback(result)
                            +-- RecordDequeueLatency()
                            +-- TraceMe("StreamCallbackInvocation")
                            +-- callback_(tensors)
                            +-- RecordCallbackLatency()
                            +-- num_outstanding_--

[リクエスト完了時]
ScopedStreamCallback::~ScopedStreamCallback()
    +-- Unregister()
            +-- registry_->Unregister(callback_id, step_id)
            +-- CallbackState::Close()
                    +-- closed_ = true
                    +-- mu_.Await(num_outstanding_ == 0)
```

### データフロー図

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

MLIRモジュール -----------------> CreateStreamCallbackId() -------> StreamCallbackId
(PwStreamResultsOp)               ランダムID生成、属性設定

ユーザコールバック関数 -----------> Register() ---------------------> ScopedStreamCallback
callback_id + step_id              stream_callbacks_に登録           (RAII管理オブジェクト)

StreamedResult -----------------> Invoke() -----------------------> callback_(tensors)
(テンソルマップ + enqueued_time)   CallbackState::Invoke()           (ユーザ定義処理)
                                  thread_pool->Schedule()
                                  InvokeCallback()
                                  +-- RecordDequeueLatency()
                                  +-- RecordCallbackLatency()
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| stream.h | `tensorflow/core/tfrt/runtime/stream.h` | ソース | StreamCallbackRegistry, ScopedStreamCallback, StreamedResult等の定義 |
| stream.cc | `tensorflow/core/tfrt/runtime/stream.cc` | ソース | StreamCallbackRegistryの全実装 |
| step_id.h | `tensorflow/core/tfrt/runtime/step_id.h` | ソース | StepId型定義 |
| tf_ops.h | `tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h` | ソース | TF::PwStreamResultsOp定義 |
| threadpool_interface.h | `tsl/platform/threadpool_interface.h` | ソース | スレッドプールインターフェース |
| traceme.h | `tsl/profiler/lib/traceme.h` | ソース | プロファイリングトレース |
