Amazon Bedrockをマルチリージョンで使おうとした話

## 前提

* 以下の内容は2023/11/03時点で試した内容なので、最新では異なる情報となっている場合があります。
* Bedrockのモデルアクセスは設定済みの状態で話を始めています。
* Bedrockで利用可能な生成AIモデルの精度については本記事では扱っていません。
* 本記事に記載の内容・設定を実施して生じた問題については責任を負いかねます。

## 何をしようとしたのか

* Amazon BedrockではClaudeやTitanなどの大規模言語モデルが利用できますが、エンドポイントはリージョン単位で用意されています。
    * 例えば、us-east-1であれば、bedrock-runtime.us-east-1.amazonaws.comのようなエンドポイントが用意されています。
* us-east-1のエンドポイントでRate Limit（Quota）上限に引っかかった場合は、us-west-2などの他のリージョンのエンドポイントに切り替える必要があります。
* アプリケーションの中でエラーが出たら別のリージョンのエンドポイントを使うようにするのも1つの手ですが、**インフラレベルで複数のリージョンのエンドポイントの自動切り替え（正確にはリクエストの分散＝ロードバランシング）ができないか**を試してみました。
    * 「Azureでいうところの、各リージョンにデプロイしたOpenAI ServiceのモデルをAPIManagementを使って負荷分散する仕組み」と同じことができないかというのを試してみました。
    * Azureの参考：https://logico-jp.io/2023/06/08/request-load-balancing-for-azure-openai-service/

## 結論

**リクエストの分散はできなくはないが、アプリケーションレベルでの制御の方が楽かもしれない**です。

## 最終的なアーキテクチャ

![aws_bedrock.drawio.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1062325/3cc8a209-e985-e442-5665-0ce935d04952.png)
※AZやInternet Gatewayなど細かい部分は割愛しています。

* Bedrockが利用可能なリージョンとしてus-east-1とus-west-2を選び、各リージョン内でBedrock用のVPCエンドポイントを作成しました。
    * Bedrockのエンドポイント自体はインターネット経由でも利用できますが、閉域網での利用も想定したかったので今回はVPCエンドポイントを用意してAWS内のNWから利用できるようにしました。
* ap-northeast-1にロードバランサー用のEC2インスタンス（アーキテクチャ図の左側の方）を立てて、us-east-1とus-west-2のリージョンと通信できるようにVPCピアリングを設定しています。
    * ロードバランシングは自前で作ったFlaskのアプリケーションサーバで行うようにしています。
    * ELBやAPI Gatewayなどを使って分散（ロードバランシング）できないか調査・検証してみましたが、（自身の知識不足もあるのか、、）できなかったため、EC2インスタンスで自前でロードバランシングするようにしました。
    * なお、API Gatewayはリクエストタイムアウトが最大30秒となっているので、処理に時間がかかりやすい大規模言語モデルを使うには不向きです。

このようなアーキテクチャになった経緯は後述します。

<!--
原因はエンドポイントへのリクエスト時に生成されるシグネチャ（署名）の不整合でした。
もう少し具体的かつざっくりいうと、クライアントであるPython SDK（boto3）からエンドポイントへのリクエストの過程でリクエストヘッダ（Authorizationヘッダ）をいじる必要があったのですが、それをいじるとSDKで生成されたシグネチャとの整合性が取れなくなりエラーになる、という感じでした。
-->


## 試したことその1

ロードバランシングにApacheが使えないか試してみました。
※Nginxにもロードバランシングの機能はありますが、試している中でうまく動かなかった（リクエストがタイムアウトしてしまう）ので、今回はApacheを利用しました。

結論を先にいうと、**Apacheではうまくいきませんでした**。
リクエストヘッダのAuthorizationを書き換える必要があったのですが、[AWS APIのシグネチャ（署名）の検証の仕様](https://docs.aws.amazon.com/ja_jp/IAM/latest/UserGuide/reference_aws-signing.html)により、書き換えを行うことができず断念しました。

<!--
### アーキテクチャ

![aws_bedrock.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1062325/0915938e-bc8b-387c-f02f-98332e8c1f6a.png)
※AZやInternet Gatewayなど細かい部分は割愛しています。
-->

### Apacheの設定

* Apacheをインストール後、必要なApacheのモジュールを有効にします。
    ```
    $ sudo a2enmod headers lbmethod_byrequests proxy proxy_balancer proxy_http slotmem_shm socache_shmcb ssl status
    ```
* ロードバランシング用の設定ファイルを`/etc/apache2/conf-available/bedrock-balancer.conf`に作成します。
    ```apache
    ProxyRequests Off
    ProxyPreserveHost On
    SSLProxyEngine On
    SSLProxyCheckPeerCN Off
    SSLProxyCheckPeerName Off
    
    <Proxy balancer://bedrock-runtime-cluster>
      BalancerMember https://vpce-XXXXX.bedrock-runtime.us-east-1.vpce.amazonaws.com route=us-east-1 loadfactor=1
      BalancerMember https://vpce-XXXXX.bedrock-runtime.us-west-2.vpce.amazonaws.com route=us-west-2 loadfactor=1
    </Proxy>
    
    <Location />
      ProxyPass balancer://bedrock-runtime-cluster/ lbmethod=byrequests timeout=1
    </Location>
    ```
    * 今回はSSL Proxyを使うことになるので、`SSLProxyEngine On`を入れています。
    * また、リクエスト先のホスト名が変わるとSSL証明書の検証に失敗するので、`SSLProxyCheckPeerCN`と`SSLProxyCheckPeerName`を`Off`にしています。
* コンフィグチェックをして問題ないことを確認し、Apacheサービスの再起動をします。
    ```
    $ apache2ctl configtest
    Syntax OK
    $ sudo systemctl restart apache2
    ```

### 動かしてみる

* 以下のようなテストスクリプトをEC2上に用意して、Apacheサーバ宛にリクエストを飛ばしてみます。
    ```python
    import json
    import boto3
    
    ENDPOINT_URL = "http://localhost"
    bedrock_runtime = boto3.client(
        "bedrock-runtime",
        endpoint_url=ENDPOINT_URL,
        region_name="us-east-1",
        aws_access_key_id="XXXXXXXXXXXXXXXXXXX",
        aws_secret_access_key="XXXXXXXXXXXXXXXXXXXXXX")
    
    res = bedrock_runtime.invoke_model(
        modelId="ai21.j2-ultra-v1",
        body=json.dumps({"prompt": "大谷翔平について教えて"}),
        contentType="application/json",
        accept="application/json",
    )
    answer = res["body"].read().decode()
    print(json.loads(answer)["completions"])
    ```
* （内容はともかく）動いた…！
    ```
    $ python3 test.py
    [{'data': {'text': '下さい\n大谷翔平は、日本の男性', ...}}]
    ```
* もう一度実行すると違うエラーが出ました…
    ```
    Traceback (most recent call last):
      File "/home/ubuntu/test.py", line 23, in <module>
        resp = bedrock_runtime.invoke_model(
      File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/client.py", line 535, in _api_call
        return self._make_api_call(operation_name, kwargs)
      File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/client.py", line 980, in _make_api_call
        raise error_class(parsed_response, operation_name)
    botocore.exceptions.ClientError: An error occurred (InvalidSignatureException) when calling the InvokeModel operation: Credential should be scoped to a valid region.
    ```

### Apacheのログを確認してみる
※ApacheのLogFormatに`%{BALANCER_WORKER_NAME}e`を追加して、どのVPCエンドポイントが選ばれたか出力するようにしています。

* 成功したときのログ
    ```
    "POST /model/ai21.j2-ultra-v1/invoke HTTP/1.1" 200 4484 "-" "https://vpce-XXXXX.bedrock-runtime.us-east-1.vpce.amazonaws.com" "Boto3/1.28.76 md/Botocore#1.31.76 ua/2.0 os/linux#6.2.0-1012-aws md/arch#x86_64 lang/python#3.10.12 md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.31.76"
    ```
    ちゃんとus-east-1のVPCエンドポイントが選択されています。
* エラーが出たときのログ
    ```
    "POST /model/ai21.j2-ultra-v1/invoke HTTP/1.1" 403 368 "-" "https://vpce-XXXXX.bedrock-runtime.us-west-2.vpce.amazonaws.com" "Boto3/1.28.76 md/Botocore#1.31.76 ua/2.0 os/linux#6.2.0-1012-aws md/arch#x86_64 lang/python#3.10.12 md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.31.76"
    ```
    us-west-2のVPCエンドポイントが選択されているので、ロードバランシング自体はちゃんとできていそうです。

### なぜエラーが起きたか

「テストスクリプトを書いている時点で気づけ」というところもありますが、スクリプトでは
```python
bedrock_runtime = boto3.client(
    "bedrock-runtime",
    endpoint_url=ENDPOINT_URL,
    region_name="us-east-1",
    aws_access_key_id="XXXXXXXXXXXXXXXXXXX",
    aws_secret_access_key="XXXXXXXXXXXXXXXXXXXXXX")
```
このようにリージョンとして`us-east-1`を指定しています。
一方で実際にリクエストしにいったエンドポイントは`us-west-2`と別のリージョンなのでクレデンシャル情報に不整合が起きてエラーになったということになります。

### Apacheの設定を変えてみる

* リクエストの中にリージョン関係の情報が含まれていると考え、Apacheのログにリクエストヘッダも出力させてみたところ、Authorizationヘッダ中に以下のようにリージョン情報が含まれていました。
    ```
    AWS4-HMAC-SHA256 Credential=<AWSアクセスキー>/20231103/us-east-1/bedrock/aws4_request, SignedHeaders=accept;content-type;host;x-amz-date, Signature=<シグネチャ>
    ```
* このリージョン情報をApacheの中で書き換えて上げれば良いのではと考え、ひとまず実験的にus-west-2に固定で書き換えるようにApacheの設定を修正しました。
    ```apache
    <Location />
      RequestHeader edit* Authorization "us-east-1" "us-west-2"
      ProxyPass balancer://bedrock-runtime-cluster/ lbmethod=byrequests timeout=1
    </Location>
    ```
    ※[RequestHeaderディレクティブ](https://httpd.apache.org/docs/2.4/en/mod/mod_headers.html#requestheader)でリクエストヘッダをいじることができます。
* Apacheのサービスを再起動して再度テストスクリプトを実行したところ…
    us-east-1のエンドポイントの方：リクエストヘッダをus-west-2に書き換えているので予想通りのエラー
    ```
    Traceback (most recent call last):
      File "/home/ubuntu/test.py", line 23, in <module>
        resp = bedrock_runtime.invoke_model(
      File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/client.py", line 535, in _api_call
        return self._make_api_call(operation_name, kwargs)
      File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/client.py", line 980, in _make_api_call
        raise error_class(parsed_response, operation_name)
    botocore.exceptions.ClientError: An error occurred (InvalidSignatureException) when calling the InvokeModel operation: Credential should be scoped to a valid region.
    ```
    
    us-west-2のエンドポイントの方：さっきと違うエラーが…
    ```
    Traceback (most recent call last):
      File "/home/ubuntu/test.py", line 23, in <module>
        resp = bedrock_runtime.invoke_model(
      File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/client.py", line 535, in _api_call
        return self._make_api_call(operation_name, kwargs)
      File "/home/ubuntu/.local/lib/python3.10/site-packages/botocore/client.py", line 980, in _make_api_call
        raise error_class(parsed_response, operation_name)
    botocore.exceptions.ClientError: An error occurred (InvalidSignatureException) when calling the InvokeModel operation: The request signature we calculated does not match the signature you provided. Check your AWS Secret Access Key and signing method. Consult the service documentation for details.
    ```

### 新たなエラーの原因

* エラーメッセージには「The request signature we calculated does not match the signature you provided.」とあるので、Authorizationヘッダの中身とセットで考えると、Authorizationヘッダにあるシグネチャとエンドポイント側で計算したシグネチャがマッチしないということになります。
* ということでPython SDK（boto3, botocore）の中身を見てシグネチャの生成方法を確認したところ、[こちら](https://github.com/boto/botocore/blob/975b30af43979372b069ff6c42a061992870ee5b/botocore/auth.py#L406)にありました。
    * サービス名（今回はbedrock）やリクエスト年月日、**リージョン名**、シークレットキーなどの情報からハッシュ値として生成しているようでした。
    * [AWSの公式ドキュメント](https://docs.aws.amazon.com/ja_jp/IAM/latest/UserGuide/reference_aws-signing.html)にもAPIリクエストの署名に関する記述がありました。
* AuthorizationヘッダにあるシグネチャはApacheで上書きする前のリージョン名から作られている＆エンドポイント側ではAuthorizationヘッダにあるCredentialの情報からシグネチャを計算していると思われるので、Authorizationヘッダにあるリージョン名を変えてしまうとマッチしなくなるのは確かにそう…
 
Apacheでシグネチャの再計算をして上書きすれば良いのですが、シークレットキーの情報はリクエストにない（あるのはアクセスキーのみ）ので、これ以上どうしようもなくなってしまいました…

## 試したことその2

「新たにBedrockのエンドポイントにリクエストを作る機能を用意すればシグネチャも新たに作られるのでいけるのでは？」と考え、以下の仕様を満たすFlaskのアプリケーションサーバを作りました。
* クライアントからのリクエストボディに含まれるプロンプトを抜き出す
* Python SDK（boto3）で新たにBedrockエンドポイントへのリクエストを生成し、前項のプロンプトを与える
* 新たに作ったリクエストに対するレスポンスをそのままクライアントに返す

### Flaskアプリを作る

boto3のパーサーに合う形でレスポンスを返す必要があったので、そこはbotocoreのソースコードを参考にしました。
※エラーレスポンスについては発生した例外のクラスに応じて変えるようにしていますが、ステータスコードの設定は合ってない可能性があります。

```python
import json
import logging
import random
import traceback
from typing import Any

import boto3
from botocore.exceptions import BotoCoreError, ClientError
from flask import Flask, Response, request


# 分散先のBedrock runtimeエンドポイント
ENDPOINT_URLS = [
    {
        "region": "us-east-1",
        "url": "https://vpce-XXXXX.bedrock-runtime.us-east-1.vpce.amazonaws.com"
    },
    {
        "region": "us-west-2",
        "url": "https://vpce-XXXXX.bedrock-runtime.us-west-2.vpce.amazonaws.com"
    }
]


def get_endpoint() -> dict[str, str]:
    return random.choice(ENDPOINT_URLS)


def invoke_bedrock(endpoint_url: str, region: str, model_id: str, body: str) -> dict[str, Any]:
    # EC2インスタンスにBedrockの操作が可能なロールをあてているので、ここではAPIの認証は不要
    client = boto3.client("bedrock-runtime", endpoint_url=endpoint_url, region_name=region)
    res = client.invoke_model(
        modelId=model_id,
        body=body,
        contentType="application/json",
        accept="application/json",
    )

    return res


app = Flask(__name__)
logging.basicConfig(level=logging.INFO)


# AWS APIのエラーレスポンスのパースコードは以下
# https://github.com/boto/botocore/blob/1.31.77/botocore/parsers.py#L686
@app.route("/model/<model_id>/invoke", methods=["POST"])
def invoke(model_id: str) -> Response:
    endpoint = get_endpoint()
    app.logger.info(f'selected region: {endpoint["region"]}')
    try:
        res = invoke_bedrock(endpoint["url"], endpoint["region"], model_id, request.data.decode())
        obj = Response(
            res["body"].read().decode(),
            status=res["ResponseMetadata"]["HTTPStatusCode"],
            headers=res["ResponseMetadata"]["HTTPHeaders"]
        )
    except ClientError as e:
        traceback.print_exc()
        res = e.response
        obj = Response(
            json.dumps({"message": res["Error"]["Message"]}, ensure_ascii=False),
            status=res["ResponseMetadata"]["HTTPStatusCode"],
            headers=res["ResponseMetadata"]["HTTPHeaders"]
        )
    except BotoCoreError as e:
        traceback.print_exc()
        obj = Response(
            json.dumps({"message": "Some error happened.", "__type": e.__class__.__name__}),
            status=400
        )
    except Exception as e:
        traceback.print_exc()
        obj = Response(
            json.dumps({"message": "Some error happened.", "__type": e.__class__.__name__}),
            status=500
        )

    return obj
```
* `invoke`関数が実際にクライアントからリクエストを受け付ける関数になっており、リクエストボディだけ取り出しています。
* リクエスト先のBedrockのエンドポイントは`get_endpoint`関数内でランダムに選択しています。
* `invoke_bedrock`関数内で選択したエンドポイントと抜き出したリクエストボディ（プロンプト情報）で新たにBedrockエンドポイントにリクエストをしています。

なお、今回はこのFlaskアプリをDockerで動かすようにしました。
docker-compose.ymlとDockerfile、requirements.txtは以下になります。
<details><summary>docker-compose.yml</summary>

```yaml
version: '3'

services:
  app:
    build: .
    image: bedrock-flask-proxy:latest
    container_name: bedrock-flask-proxy
    restart: always
    ports:
      - "80:8000"
```
</details>
<details><summary>Dockerfile</summary>

```dockerfile
FROM python:3.10.13-slim-bullseye

WORKDIR /app

RUN set -ex \
    && useradd -d /app -s /bin/bash flask \
    && chown -R flask:flask /app

COPY --chown=flask:flask . .

USER flask

RUN set -ex \
    && pip install --no-cache-dir -r requirements.txt

ENV PATH=/app/.local/bin:$PATH

ENTRYPOINT ["gunicorn", "-w", "2", "-b", "0.0.0.0:8000", "app:app"]
```
</details>
<details><summary>requirements.txt</summary>

```
Flask==3.0.0
boto3==1.28.77
gunicorn==21.2.0
```
</details>

### Flaskアプリを動かしてみる

`docker compose up -d`でFlaskアプリ（のコンテナ）を起動し、Apacheのときに作ったテストスクリプトを実行してみます。
```
$ python test_request.py
[{'data': {'text': 'くだ\nさいたしました', 'tokens': ... }}]
```
ちゃんと結果は返ってきていますね。
コンテナのログを確認すると…
```
bedrock-flask-proxy  | INFO:app:selected region: us-east-1
bedrock-flask-proxy  | INFO:botocore.credentials:Found credentials from IAM Role: ExecuteBedrockFromEC2
```
選択されたリージョンのエンドポイントがちゃんと実行できていそうです。
何回かテストスクリプトを実行しましたが、以下のようにus-west-2の方も問題なく選択・実行できていたので大丈夫そうでした。
```
bedrock-flask-proxy  | INFO:app:selected region: us-west-2
bedrock-flask-proxy  | INFO:botocore.credentials:Found credentials from IAM Role: ExecuteBedrockFromEC2
```

ということで新たにリクエストを作るアプリケーションサーバを用意すればロードバランシングできることが確認できました。


## おわりに

* Bedrockのエンドポイントをマルチリージョンで使うときはロードバランシング用のアプリケーションサーバを用意すれば可能なことがわかりました。
    * が、力技感が否めないので個人的には微妙な感じがしています。（もう少しスマートにやりたかった…）
* ロギングやレスポンス性能、可用性、セキュリティ（認証・認可）など非機能面は考慮してないので、実際に利用する際はそのあたりの検証・設計も必要になります。
* ApacheやNginxのロードバランシングに関する情報が思っていたより少なくて大変でした…
* こういうアーキテクチャ・設定ならできるのでは？といったコメントがあればぜひお願いします。