# 機能設計書 102-TPU構成管理

## 概要

本ドキュメントは、TensorFlowにおけるTPU（Tensor Processing Unit）デバイスの初期化・構成・接続管理機能の設計を記述する。分散TPUシステムのセットアップからシャットダウンまでのライフサイクル全体を管理するOp群を提供する。

### 本機能の処理概要

**業務上の目的・背景**：TPUは機械学習ワークロードのための専用アクセラレータであり、大規模モデルの学習・推論に不可欠である。分散TPUシステムを構成するには、複数のホストにまたがるTPUチップの初期化、トポロジ設定、グローバルID割り当てなどの複雑な手順が必要であり、本機能はその一連のプロセスを管理する。

**機能の利用シーン**：TPU Podを使用した大規模分散学習のセットアップ、TPUクラスタの初期化・構成変更・シャットダウン、TPUエンベディング機能の設定、マルチワーカー環境でのTPUシステム構成などで利用される。

**主要な処理内容**：
1. 分散TPUシステムの構成初期化（_ConfigureDistributedTPU）
2. TPUトポロジの待機と検証（_WaitForDistributedTPU）
3. ホストごとのTPU初期化（_InitializeHostForDistributedTPU）
4. グローバルTPU配列の設定（_SetGlobalTPUArray）
5. ホストの分散TPUシステムからの切断（_DisconnectHostFromDistributedTPUSystem）
6. 分散TPUシステムのシャットダウン（_ShutdownDistributedTPU）
7. ユーザー向け高レベルAPI（ConfigureDistributedTPU、ShutdownDistributedTPU）
8. TPUエンベディング構成（ConfigureTPUEmbedding）
9. TPUエンベディング初期化状態確認（IsTPUEmbeddingInitialized）

**関連システム・外部連携**：TPU_SYSTEMデバイス、UberDriverによるTPUチップ間通信、TPUトポロジプロトコルバッファ。

**権限による制御**：TPU_SYSTEMデバイスへのアクセス権限が必要。マスターワーカーのみがConfigureDistributedTPUを実行可能（デフォルト設定）。

## 関連画面

| 画面No | 画面名 | 関連種別 | 関連する操作・処理 |
|--------|--------|----------|------------------|
| - | - | - | 本機能に関連する画面は登録されていない |

## 機能種別

デバイス管理・システム構成

## 入力仕様

### 入力パラメータ

| パラメータ名 | 型 | 必須 | 説明 | バリデーション |
|-------------|-----|-----|------|---------------|
| inputs (Configure) | N * int32 | Yes | 各ホストのTPUチップ数（スカラー） | N >= 1、全入力がスカラー |
| inputs (Wait) | N * int32 | Yes | 各ホストのTPUグローバルID（ベクトル） | 全入力が1次元ベクトル |
| input (InitializeHost) | string | Yes | UberDriverアドレス | スカラー文字列 |
| topology (SetGlobal) | string | Yes | TPUトポロジプロト（シリアライズ） | スカラー文字列 |
| embedding_config | string | No | エンベディング設定（ConfigureDistributedTPU時） | 空文字列許容 |
| tpu_embedding_config | string | No | TPUエンベディング設定 | 空文字列許容 |
| enable_whole_mesh_compilations | bool | No | メッシュ全体コンパイルの有効化（デフォルト: false） | - |
| startup_timeout_sec | int | No | 起動タイムアウト秒数（デフォルト: 20） | - |
| tpu_cancellation_closes_chips | int | No | キャンセル時チップクローズ（0: 未設定、1: 有効、2: 無効） | 0, 1, 2のいずれか |
| config (ConfigureTPUEmbedding) | string | Yes | TPUEmbeddingConfiguration プロト | 有効なプロトバッファ |

### 入力データソース

TPUシステムのハードウェア構成情報、クラスタ設定。

## 出力仕様

### 出力データ

| 項目名 | 型 | 説明 |
|--------|-----|------|
| output (_Configure) | string | TPUHostConfigurationプロトのシリアライズ文字列 |
| topology (_Wait) | string | TPUトポロジプロトのシリアライズ文字列 |
| tpu_ids (_InitializeHost) | int32 vector | 各TPUのグローバルID |
| number_of_tpu_chips (_Disconnect) | int32 | ホスト上のTPUチップ数 |
| topology (Configure) | string | TPUトポロジ文字列 |
| is_tpu_embedding_initialized | bool | TPUエンベディング初期化済みフラグ |

### 出力先

TPU_SYSTEMデバイスの内部状態、TPUトポロジ情報は後続のTPU操作に渡される。

## 処理フロー

### 処理シーケンス

```
1. 各ホストの切断（既存構成がある場合）
   └─ _DisconnectHostFromDistributedTPUSystem を各ホストで実行
2. 分散TPUシステムの構成
   └─ _ConfigureDistributedTPU をワーカー0のTPU_SYSTEMで実行
3. 各ホストの初期化
   └─ _InitializeHostForDistributedTPU を各ホストのTPU_SYSTEMで実行
4. 分散TPUの待機
   └─ _WaitForDistributedTPU でシステム全体の初期化完了を待機
5. グローバルTPU配列の設定
   └─ _SetGlobalTPUArray を各ホストのTPU_SYSTEMで実行
6. （オプション）TPUエンベディングの構成
   └─ ConfigureTPUEmbedding を実行
```

### フローチャート

```mermaid
flowchart TD
    A[開始] --> B[_DisconnectHostFromDistributedTPUSystem<br/>各ホストで実行]
    B --> C[_ConfigureDistributedTPU<br/>ワーカー0で実行]
    C --> D[_InitializeHostForDistributedTPU<br/>各ホストで実行]
    D --> E[_WaitForDistributedTPU<br/>システム初期化完了待機]
    E --> F[_SetGlobalTPUArray<br/>各ホストで実行]
    F --> G{TPUエンベディング使用?}
    G -->|Yes| H[ConfigureTPUEmbedding]
    G -->|No| I[構成完了]
    H --> I
    I --> J[学習/推論処理]
    J --> K[シャットダウン?]
    K -->|Yes| L[_DisconnectHostFromDistributedTPUSystem<br/>各ホストで実行]
    L --> M[_ShutdownDistributedTPU]
    M --> N[終了]
    K -->|No| J
```

## ビジネスルール

### 業務ルール

| ルールNo | ルール名 | 内容 | 適用条件 |
|---------|---------|------|---------|
| BR-102-01 | 構成順序制約 | 5段階の初期化手順は規定された順序で実行する必要がある | 分散TPU構成時 |
| BR-102-02 | マスターワーカー制約 | _ConfigureDistributedTPUと_WaitForDistributedTPUはワーカー0で実行する | デフォルト設定時 |
| BR-102-03 | シャットダウン順序 | 全ホストのDisconnect完了後にShutdownを実行する | システム終了時 |
| BR-102-04 | チップ数一致 | 全ホストのTPUチップ数が一致し、トポロジと整合性がある必要がある | 構成時 |
| BR-102-05 | メッシュコンパイル | enable_whole_mesh_compilations=trueでメッシュ全体でのコンパイル分散が有効 | マルチクライアント環境 |

### 計算ロジック

特段の計算ロジックはない。TPUトポロジの検証はフラグで設定されたトポロジと実際のチップ数の整合性チェックで行われる。

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

本機能はデータベース操作を行わない。

## エラー処理

### エラーケース一覧

| エラーコード | エラー種別 | 発生条件 | 対処方法 |
|------------|----------|---------|---------|
| InvalidArgument | バリデーションエラー | 入力のランクが不正（スカラーでないなど） | 入力形状を確認 |
| Internal | システムエラー | TPUシステム未構成でShutdown実行 | Configure実行後にShutdownを実行 |
| DeadlineExceeded | タイムアウト | startup_timeout_sec内に初期化未完了 | タイムアウト値を増加、ネットワーク確認 |
| InvalidArgument | トポロジ不整合 | ホスト間のチップ数不一致 | ハードウェア構成を確認 |

### リトライ仕様

startup_timeout_secで指定された秒数以内に初期化が完了しない場合はタイムアウトエラーとなる。大規模システムでは数分かかる場合がある。

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

全てのOpsはSetIsStateful()が設定されており、副作用を持つステートフルな操作である。構成操作は冪等ではないため、失敗時は切断からやり直す必要がある。

## パフォーマンス要件

- 大規模TPU Podの初期化には数分を要する可能性がある（startup_timeout_secのデフォルト: 20秒）
- _WaitForDistributedTPUは全ホストの初期化完了をブロッキングで待機する

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

- TPU_SYSTEMデバイスへのアクセス制御が必要
- UberDriverアドレスはネットワーク内部で管理されるべき
- マルチクライアント環境ではクライアント認証が推奨される

## 備考

- ユーザーコードは通常、高レベルのConfigureDistributedTPU Opを使用し、グラフリライトにより内部Opsのサブグラフに展開される
- tpu_cancellation_closes_chips属性はTFRTランタイムでのみ有効
- IsTPUEmbeddingInitialized OpはSetDoNotOptimize()が設定されており、Grappler最適化の対象外

---

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

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

### 推奨読解順序

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

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 1-1 | topology.proto | `tensorflow/core/protobuf/tpu/topology.proto` | TPUトポロジプロトの構造 |
| 1-2 | tpu_embedding_configuration.proto | `tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto` | TPUエンベディング設定プロトの構造 |

**読解のコツ**: TPU構成のデータフローはプロトコルバッファを中心に設計されているため、プロト定義を先に理解することが重要。

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

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 2-1 | tpu_configuration_ops.cc | `tensorflow/core/ops/tpu_configuration_ops.cc` | 全10個のTPU構成Op登録定義 |

**主要処理フロー**:
1. **26-79行目**: 分散TPUシステム構成の5段階手順のコメント（設計意図の解説）
2. **81-110行目**: _ConfigureDistributedTPU Op - 入力はN個のint32スカラー、出力はstring
3. **112-139行目**: _WaitForDistributedTPU Op - startup_timeout_sec属性、出力はトポロジ文字列
4. **141-155行目**: _SetGlobalTPUArray Op - トポロジ文字列を入力として受け取る
5. **157-167行目**: _ShutdownDistributedTPU Op - 出力なし、ステートフル
6. **169-194行目**: _InitializeHostForDistributedTPU Op - UberDriverアドレスを入力、TPU IDベクトルを出力
7. **196-206行目**: _DisconnectHostFromDistributedTPUSystem Op - TPUチップ数を出力
8. **208-219行目**: ConfigureDistributedTPU Op（高レベル） - embedding_config等の属性付き
9. **225-228行目**: ConfigureTPUEmbedding Op - config文字列属性
10. **230-234行目**: IsTPUEmbeddingInitialized Op - bool出力、SetDoNotOptimize

#### Step 3: カーネル実装を理解する

| 順序 | ファイル | パス | 読解ポイント |
|-----|---------|------|-------------|
| 3-1 | tpu_configuration_ops.cc | `tensorflow/core/tpu/kernels/tpu_configuration_ops.cc` | TPU構成カーネルの実装 |

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

```
ConfigureDistributedTPU (高レベルOp)
    │ (グラフリライトで展開)
    ├─ _DisconnectHostFromDistributedTPUSystem (各ホスト)
    │      └─ TPUチップ数を返却
    │
    ├─ _ConfigureDistributedTPU (ワーカー0)
    │      └─ TPUHostConfigurationプロトを生成
    │
    ├─ _InitializeHostForDistributedTPU (各ホスト)
    │      └─ グローバルTPU IDベクトルを返却
    │
    ├─ _WaitForDistributedTPU (ワーカー0)
    │      └─ TPUトポロジプロトを返却
    │
    └─ _SetGlobalTPUArray (各ホスト)
           └─ トポロジ情報を設定

ShutdownDistributedTPU (高レベルOp)
    │
    ├─ _DisconnectHostFromDistributedTPUSystem (各ホスト)
    └─ _ShutdownDistributedTPU (ワーカー0)
```

### データフロー図

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

チップ数(各ホスト) ──▶ _ConfigureDistributedTPU ──▶ HostConfig Proto
                              │
HostConfig Proto ────▶ _InitializeHostForDistributedTPU ──▶ TPU IDs
                              │
TPU IDs (全ホスト) ──▶ _WaitForDistributedTPU ──▶ Topology Proto
                              │
Topology Proto ──────▶ _SetGlobalTPUArray ──▶ (内部状態更新)
```

### 関連ファイル一覧

| ファイル | パス | 種別 | 役割 |
|---------|------|------|------|
| tpu_configuration_ops.cc | `tensorflow/core/ops/tpu_configuration_ops.cc` | ソース | TPU構成Op登録定義（10 Ops） |
| topology.proto | `tensorflow/core/protobuf/tpu/topology.proto` | プロト | TPUトポロジプロト定義 |
| tpu_embedding_configuration.proto | `tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto` | プロト | TPUエンベディング設定プロト |
