Azure OpenAI をネットワーク機能で負荷分散する

Azure OpenAI をそれなりの規模の環境で利用する場合、リソースの上限があるため、負荷分散を行いたいといった要望があります。Azure の機能で実現する場合、以下の方法が考えられます。実際に利用可能であるかを検証して確認してみましたので、ご利用の際はご一読いただければと思います。

• API Management (※ポリシーを利用して実現。サンプルコード。 本記事では詳細は記載してません)
• Application Gateway
• Front Door
• Traffic Manager

なお、Azure Load Balancer で PaaS や Private Endpoint を負荷分散することはサポートされてませんので、このシナリオでは利用できません。

負荷分散の機能のおさらい

こちらに詳細はまとまっておりますが、本件に関連しそうなところとして、それぞれのネットワーク機能は以下のような違いがありますので、要件に応じて選択する必要があります。

サービス	範囲	ルーティング	分散方式	対応プロトコル	WAFの有無	URL,ヘッダー書き換え	正常性プローブの応答	閉域利用
Application Gateway	リージョン内	ラウンドロビンのみ	L7 ロードバランサー	HTTP, HTTPS, HTTP/2, WebSocket	〇	〇	Status コードの修正可能	〇
Front Door	グローバル	優先順位/レイテンシ/重み	L7 ロードバランサー	HTTP, HTTPS, HTTP/2	〇	〇	Status 200 のみ	×
Traffic Manager	グローバル	優先順位/レイテンシ/重み/地理/サブネット	DNS ベースのロードバランサー	TCP, UDP	×	×	Status コードの修正可能	×

Azure OpenAI の認証方式

API キーまたは AAD 認証の 2 種類がサポートされています。

• API キー認証: この種類の認証の場合、すべての API 要求で、api-key HTTP ヘッダーに API キーを含める必要があります。 クイック スタートに、この種類の認証を使用して呼び出しを行う方法に関するガイダンスが用意されています。

• Azure Active Directory 認証: Azure Active Directory トークンを使用して API 呼び出しを認証できます。 認証トークンは、Authorization ヘッダーとして要求に含まれます。 指定するトークンの前に Bearer を付ける必要があります (例: Bearer YOUR_AUTH_TOKEN)。 Azure Active Directory での認証に関する攻略ガイドをお読みください。

構成検討

Azure OpenAI にアクセスするためには認証が必須という点と、リソースごとに異なる API キーを持っているという点を考慮して構成を検討する必要があります。
また、各サービスには可用性を担保する機能も備わっているため、最低限、以下の点を考慮して、構成したほうが良いと思います。

• 各負荷分散サービスからの正常性プローブへの応答をどうするか。(※認証がないため、Status 404 等を許容するか) ※2023/11/27 追記。 こちらに正常性プローブの参考情報がありました。 モデルによって動作が異なる可能性があるため事前検証をお勧めします。
• Azure OpenAI への認証方式は API Key, AAD どちらを利用するか。
• アクセスする URL を統一するため、各 Azure OpenAI のデプロイ名は合わせる。

各機能の検証

以下のような構成で各機能で Azure OpenAI の負荷分散を実施します。
クライアントは検証しやすかったため、curl コマンドで実施しています。

Application Gateway の場合

正常性プローブの構成

Application Gateway の正常性プローブでヘッダーを付与する機能はないため、Status 404 を許容するといった対応が必要となります。Status 200 ではないため、API へのリクエストが成功するような状態であるか？までは監視できませんが、ネットワークレベル (http リクエストが届き、status 404 が返る) での監視は可能です。※2023/11/27 追記。 こちらに正常性プローブの参考情報がありました。"/status-0123456789abcdef" のパスを指定することで status 200 にできます。

この例ではバックエンドプールのホスト名に対してパス / の HTTP リクエスト https://xxx.openai.azure.com/ を投げる構成としています。

正常性プローブが健全となっていることを確認できます。

認証方式1 (API Key 認証)

Application Gateway では振り分け先のバックエンドに応じてヘッダーを付与するような機能はありません。HTTP リクエストのサーバー変数 (IP, Port 等) やヘッダーを元に振り分け先を指定し、ヘッダーを付与することはできます。バックエンドの冗長化はできませんが、負荷分散のみの要件であれば以下のような構成は可能です。
ただし、テンプレート等で平文で API Key が参照可能となってしまいますので、一時的な動作確認程度で本番環境での利用はお勧めしません。

どのようにヘッダー書き換えを設定すればよいかは以下が参考になると思います。

接続元ポートが 0-4 であれば URL の先頭に /path1 を付与し、api-key (West Europe) も付与します。また、パスマップの再評価も行います。(変更後の URL パスで再度ルーティングを行う)

一致させるパターンは ^\d*[01234]$ の正規表現を利用。

接続元ポートが 5-9 であれば URL の先頭に /path2 を付与し、api-key (France Central) も付与します。また、パスマップの再評価も行います。

以下のようなコマンドで Application Gateway に割り当てた FQDN に対して POST リクエストを投げます。ヒアドキュメントで JSON を読み込んで、jq で見やすくフォーマットしています。

curl -v https://xxxx.hiyamanet2.org/openai/deployments/gpt35model/chat/completions?api-version=2023-03-15-preview \
-H "Content-Type: application/json" \
-d @- <<EOM | jq
{
    "messages": [
        {
            "role": "user",
            "content": "Hi OpenAI"
        }
    ]
}
EOM

以下の様に Azure OpenAI から応答があることを確認できました。

診断ログのアクセスログを見てみると clientPort の番号により、想定通り West Europe (weu), France Central (crf) に負荷分散されています。

ただし、前述のとおりこの構成は 1 バックエンドプール、1 エンドポイントとなっているため、バックエンドの冗長化はできません。どちらかの Azure OpenAI が異常となった時に、そのエンドポイントに振り分けられたリクエストのレスポンスは 502 となります。

認証方式2 (AAD 認証)

書き換え規則を削除して、2 つの Azure OpenAI リソースを同じバックエンドプールに登録し、Basic ルールと関連付けを行います。

以下のドキュメントを参考に実施します。

curl の実行端末で Azure CLI を使えるようにしておきます。また、jq コマンドも利用できるようにしておきます。
検証で利用してる端末は CentOS 8.5 です。

Azure AD のユーザーで Token を取得し、リクエストを投げるため、各 Azure OpenAI のリソースに Cognitive Service ユーザーロールを割り当てておきます。(ユーザーはなくサービスプリンシパルの場合も似たような手順になるはずです)

az login コマンドを実行し、RBAC を割り当てたユーザーでログインし、token を取得します。

az login

export accessToken=$(az account get-access-token --resource https://cognitiveservices.azure.com -o json | jq -r .accessToken)

-H "Authorization: Bearer $accessToken" を付け足してコマンドを実行します。

curl -v https://xxxx.hiyamanet2.org/openai/deployments/gpt35model/chat/completions?api-version=2023-03-15-preview -H "Content-Type: application/json" -H "Authorization: Bearer $accessToken" -d @- <<EOM | jq
{
    "messages": [
        {
            "role": "user",
            "content": "Hi OpenAI"
        }
    ]
}
EOM

何回か実行し、レスポンスに含まれるリージョンが想定通りに切り替わっていることを確認できました。こちらのレスポンスヘッダーはリソースのリージョンと同じになるようです。

また、試しに疑似的な障害として France Central のプライベートエンドポイントを削除してみましたが、想定通りに正常性プローブが失敗し、West Europe のみにリクエストが振り分けられました。

OpenAI サービスのファイアウォール

Application Gateway からプライベートエンドにアクセス可能ですので、ファイアウォールを無効化してしまい、プライベートエンドのみが有効化されていれば問題ありません。

Front Door の場合

Azure Front Door (Standard または Premium) を利用して負荷分散を行います。

正常性プローブの構成

Front Door の正常性プローブは Status 200 のみが成功とみなされます。そのため、認証が必須となる Azure OpenAI のようなリソースの場合、正常性プローブを使ってバックエンドを切り替えることができません。
Azure OpenAI の負荷分散のために、Azure Front Door を使う場合はバックエンドの冗長化ができないことになります。 ※2023/11/27 追記。 こちらの "/status-0123456789abcdef" のパスを指定する方法で正常性を確認できるようです。正常性プローブの参考情報

200 OK 状態コードは、配信元が正常であることを示します。 その他の状態コードは失敗と見なされます。

各配信グループの設定で正常性プローブを無効化しておきます。

認証方式1 (API Key 認証)

Application Gateway と同じような方法で、クライアントポートが 1-5 で終わる場合、api-key (West Europe) を付与して、West Europe の配信グループに振り分けし、クライアントポートが 1-5 で終わらない場合、api-key (France Central) を付与して、France Central の配信グループに振り分けするルールセットを作成します。

作成したルールセットをルールに関連付けします。

Application Gateway と同じように curl コマンドを実行したところ、想定通りに振り分けられ応答が得られることを確認できました。

ただし、こちらもテンプレート等で平文で API Key が参照可能となってしまいますので、一時的な動作確認程度で本番環境での利用はお勧めしません。

認証方式2 (AAD 認証)

ルールセットから API Key を削除しておきます。

想定通りに送信元ポートに応じて負荷分散されました。

※2023/11/27 追記。こちらの "/status-0123456789abcdef" のパスを指定する方法でも問題なく負荷分散可能でした。 正常性プローブの参考情報

OpenAI サービスのファイアウォール

Azure OpenAI へのアクセスを "AzureFrontDoor.Backend" タグに限定します。サービスタグの機能はないのでアドレスを記載する必要があります。
様々なやり方があると思いますが、サービスタグの一覧ファイルから AzureFrontDoor.Backend タグの IPv4 アドレスを取得し、テキストエディタ等で { "value": "xxx" } のような JSON 形式に変更後、az rest コマンドでヒアドキュメントで JSON を読み込んでいくのが簡単でした。/32 がある場合は /32 部分の削除が必要だったりしますので、JSON の中身は適宜修正してください。

サービスタグの一覧は更新がありますので定期的なメンテが必要となります。

az rest --method PATCH --url "/subscriptions/ad982283-xxxx-xxxx-xxxx-8043c3586b04/resourceGroups/rg-testxxxx/providers/Microsoft.CognitiveServices/accounts/cog-xxxxxx?api-version=2022-10-01" --body @- << EOM
{
    "properties": {
        "networkAcls": {
            "defaultAction": "Deny",
                "ipRules": [
                    { "value": "4.232.98.120/29" },
                                  ※省略
                    { "value": "158.23.108.56/29" },
                    { "value": "191.233.9.120/29" },
                    { "value": "191.235.225.128/29" }
                ],
                    "virtualNetworkRules": []
        }
    }
}
EOM

Traffic Manager の場合

正常性プローブの構成

Status 404 を許可します。

認証方式1 (API Key 認証)

Traffic Manager にはエンドポイントごとに API Key 付与するような機能がないため実施できません。

認証方式2 (AAD 認証)

AAD 認証の場合は実施可能だが、Azure OpenAI にはカスタムドメインを割り当てる機能はないため、Traffic Manager の既定のホストでアクセスしてみます。
証明書の警告は出てしまうため、-k オプションで証明書の検証を無視してアクセスしてみます。

curl -kv https://testoaixxxx.trafficmanager.net/openai/deployments/gpt35model/chat/completions?api-version=2023-03-15-preview -H "Content-Type: application/json" -H "Authorization: Bearer $accessToken" -d @- <<EOM | jq
{
    "messages": [
        {
            "role": "user",
            "content": "Hi OpenAI"
        }
    ]
}
EOM

以下のようなエラーが出て、Status 404 でアクセスすることはできませんでした。このことから Azure OpenAI サービスは異なるホスト名でのアクセスは受け付けないようです。 Traffic Manager での負荷分散は難しいことがわかりました。

まとめ

検証結果をふまえると以下の様になります。

サービス名	負荷分散	バックエンド冗長化	ファイアウォール	備考
Application Gateway	〇	〇	〇
FrontDoor	〇	〇	△
Traffic Manager	×	×	-	動作せず

※2023/11/27 追記。こちらの "/status-0123456789abcdef" のパスを指定する方法にて Azure Front Door でも負荷分散、バックエンドの冗長化が可能でした。正常性プローブの参考情報