このエントリは2023/11/06現在の情報に基づいています。将来の機能追加や変更に伴い、記載内容からの乖離が発生する可能性があります。

問い合わせ

いつもの人とは違う人から、次のような問い合わせがあった。

現在Azure OpenAI Service (以後、AOAI) を社内で利用してもらうようにAPI管理しているのだが、特定部署が大量に使わないように、利用トークン数に上限を設けてアクセスを制御したい。具体的には、1日あたり○○個のトークンまで利用可能にしておき、上限を超えた場合には429（もしくは403）を返すようにしたいのだが、どうすればよいか？

エンタープライズでの利用であれば結構出てきそうなユースケースである。とはいえ、AOAIならびにAPI Management (以下APIM) の標準機能にはこのような仕組みは存在しないが、どうすればよいのだろうか。

解決策

2024/05/27時点で、新たなAPI ManagementのポリシーであるToken Limitポリシーが出たので、これを使えば簡単に制御できる。レートリミットを超えると429を呼び出し元に返す。

例えば1分あたり5kトークンとする制限を、呼び出し元IPアドレスをキーとして指定すると、以下のように構成できる。各ポリシーの実行後、その期間中にその呼び出し元 IP アドレスに対して許可されている残りのトークン数を remainingTokens 変数に格納している。これを呼び出し元へのレスポンスのHTTP Headerあたりに設定しておく、という使い方もできる。

<policies>
 <inbound>
  <base />
  <azure-openai-token-limit
   counter-key="@(context.Request.IpAddress)"
   tokens-per-minute="5000" 
   estimate-prompt-tokens="false" 
   remaining-tokens-variable-name="remainingTokens" />
 </inbound>
 <outbound>
  <base />
 </outbound>
</policies>

このポリシーは以前の手組みの方式とは異なり、SSEを使うStreamingモード時にも問題なく計算してくれる。詳細は以下のドキュメントを参照。

Azure OpenAI API トークンの使用を制限する / Limit Azure OpenAI API token usage
https://learn.microsoft.com/azure/api-management/azure-openai-token-limit-policy

【2024/07/02】
以下は過去の経緯として残しておく。もはや手作業で構成する必要はない。

実のところ、このような要件を満たすための仕組みがGitHubにすでに上がっている。Streamが有効な場合、無効な場合で方法が異なる。

Open AI Cost Gateway Pattern
https://github.com/ThePreston/Custom-Rate-Limiter-API

[1] Streamが無効な場合

APIMのOutboundセクションでAOAIからのレスポンスに含まれるUsageの要素をsend-one-way-requestポリシーを使ってAzure Cache for Redisに書き込み、Cacheの値を加算していく。APIMのInboundセクションでは、send-requestポリシーでCacheに保持されている利用トークン数をチェックし、上限を超えていればreturn-responseポリシーで429（もしくは403）を返す。

1 方向の要求を送信する / Send one way request
https://learn.microsoft.com/azure/api-management/send-one-way-request-policy
要求の送信 / Send request
https://learn.microsoft.com/azure/api-management/send-request-policy
応答を返す / Return response
https://learn.microsoft.com/azure/api-management/return-response-policy

イメージは以下のような感じ。

[2] Streamが有効な場合

以下のエントリにも記載したように、APIMのOutboundセクションでチェックできるAOAIからのレスポンスのUsageは常にNullであり、APIレスポンスの値から直接トークン数を把握できない。

そのため、APIMのバックエンドサービスとしてFunction Appを立てて、APIMがSSE (Server Sent Events) を受けないようにするか、もしくはAPI Client（APIMでホストするAPIを呼び出すクライアント）で利用トークン数を計算する必要がある。利用トークン数計算後はStream無効時と同じで、Azure Cache for Redisに書き込んでおく。InboundセクションでのチェックはStream無効時と同じ。

注意点

この説明をしたときにいろいろ更問があったので、メモとして残しておく。

In-memory cacheでなければだめなのか？

結論から言うと非機能要件次第。なので、Cosmos DBでもかまわないし、それこそTable storageでもSQL Databaseでもかまわない。ただ、より高速なアクセスを求めるならAzure Cache for Redisが推奨。

利用トークン数書き込み時にRace conditionが発生しないのか？

実装次第ではあるが、シンプルにFunction AppからAzure Cache for Redisを更新するのだとすればRace conditionが発生しうる。そのため、Redis StreamsやEvent Hubs (要件次第ではCosmos DBのChange Feedも候補になり得る) を使ってFIFOを実現しつつ、サブスクライバーであるFunction AppがSingletonになるよう構成する必要がある。

以下はコンポーネント配置案のひとつ。

【Stream無効時】

• Outbound section
  • Event Hubsを使う場合
    • APIMのOutboundセクションでsend-one-way-requestポリシーではなくlog-to-eventhubポリシーでEvent Hubsに投げ込む
    • Event Hubsをサブスクライブする1個だけのFunction Appがメッセージを読み取り、順次Azure Cache for Redisに対して書き込む
  • Redis Streamsを使う場合 (Redis CacheではRedis Streamsを構成済みとする)
    • APIMのOutboundセクションでsend-one-way-requestポリシーでCacheに投げ込む
    • Redis Streamsをサブスクライブして、Event Hubsをサブスクライブする1個だけのFunction Appがメッセージを読み取り、順次Cache (もしくは別の永続領域) に書き込む
• Inbound section
  • ReadではRace conditionは発生しないので、Function appで直接Cache (もしくは別の永続領域) を確認してもよい。

イベント ハブにログを記録する / Log to event hub
https://learn.microsoft.com/azure/api-management/log-to-eventhub-policy
Streams in Azure Cache for Redis
https://azure.github.io/redis-on-azure-workshop/labs/04-streams-in-azure-cache-for-redis.html
Implement Pub/Sub and Streams in Azure Cache for Redis
https://learn.microsoft.com/training/modules/azure-redis-publish-subscribe-streams/

【Stream有効時】

• トークン数計算後にEvent HubsもしくはRedis Streamsに渡すように構成する。

Service Busではだめなのか？

Event HubsやRedis Streams以外に、Service Busでももちろん同様に構成できるが、大量イベントを扱う必要がある場合にはEvent Hubsが適している。Service Busを使う場合は、APIMのポリシーは存在しないのでFunction Appで実装する必要が出てくる（APIMのポリシーで扱えるのがEvent Hubsだけ、という理由でEvent Hubsを選択しているわけではない）。Messaging Serviceの使い分けや比較は以下のドキュメントを参照。

Azure メッセージング サービスの中から選択する – Azure Event Grid、Event Hubs、および Service Bus / Choose between Azure messaging services – Event Grid, Event Hubs, and Service Bus
https://learn.microsoft.com/azure/service-bus-messaging/compare-messaging-services

どの部署からのリクエストであるかを判断する方法はどういったものがあるか？

一つはJWTの要素をチェックする、もう一つはサブスクリプションキーを利用する方法が考えられる。JWTがリクエストに付随する場合は、APIMでvalidate-jwtポリシーを使って識別に使える値を取り出し、その値でAzure Cache for Redisを検索し、トークン数を減算することになるだろう。

JWT を検証する / Validate JWT
https://learn.microsoft.com/azure/api-management/validate-jwt-policy

もう一つのサブスクリプションキーは、いわゆる「APIキー」と言われるもので、呼び出し元であるAPI Clientを識別するためのものでJWT validationよりも簡便である。サブスクリプションキーを使う場合には、キーと利用トークン数をEvent Hubsに渡す。書き込みアプリケーションはEvent Hubs経由で渡されたサブスクリプションキーでAzure Cache for Redisを検索し、トークン数を減算することになるだろう。

Azure API Management のサブスクリプション / Subscriptions in Azure API Management
https://learn.microsoft.com/azure/api-management/api-management-subscriptions

サブスクリプションキーはアクセス制御には利用出来るが、認証機構ではない。そのため、「サブスクリプションキーがあれば認証は不要」という判断は厳に慎む必要がある。

• AI
• API Management
• Azure
• Cosmos DB
• Functions
• Redis
• Technology