レートリミットとトークン会計のゲートウェイを作る

目標

LLMゲートウェイを自分で作って、リクエスト単位ではなく、トークン単位の制限と会計を実装し、テナントごとのバジェット制御まで完成させます。

なぜ重要なのか

毎分60リクエストの制限は、LLMの前では、ほとんど無意味です。あるテナントが、毎回のリクエストに200,000トークンを入れると、リクエスト数は制限の範囲内ですが、コストとGPUの占有は、他のテナントの数百倍で、そのリクエストがKVキャッシュを独占して、他のリクエストをプリエンプションさせます。そのため、LLMゲートウェイの第1の設計原則は、トークンを数えることです。このラボで特に厄介なのが、ステップ4とステップ8です。出力トークンは、生成が終わって初めてわかるので、バジェットの検査を完了後にだけ行うと、同時に入ってきたリクエストが、すべて予算の範囲内だと判断して通過したあと、いっしょに超過します。max_tokensを上限としてあらかじめ予約し、完了後に精算する方式が、この問題を解きます。実際のサービスで、コストの事故が起きる地点が、まさにここです。

ステップ

/root/gw/gateway.pyを127.0.0.1:8171で起動して、POST /v1/generateを8170のバックエンドに転送してください。応答がバックエンドと同一である必要があります。テナントt1のAPIキーを/root/gw/key.txtに、t2のキーを/root/gw/key2.txtに、それぞれ1行で保存しておいてください。以降のステップの採点が、この2つのファイルを使います。
X-API-Keyヘッダーで、テナントを識別してください。キーとテナントのマッピングは、/opt/fixtures/llms/apikeys.jsonにあります。キーがないか、未知のキーなら、401です。
テナントごとの毎分のリクエスト制限10を適用してください。11番目のリクエストが429である必要があります。/root/gw/rpm.txtに、allowed=10 rejected=1を書いてください。
GET /v1/usageが、{"tenant":"...","prompt_tokens":<n>,"completion_tokens":<n>,"requests":<n>}を返すようにしてください。2つのトークンの値が、どちらも0より大きい必要があります。
毎分のトークン制限2000を、トークンバケットで適用してください。上限を超えるリクエストが、429を受け取る必要があります。/root/gw/tpm.txtに、limit=2000 consumed=<n> rejected=<n>を書き、rejectedは1以上である必要があります。
すべての429のレスポンスに、Retry-Afterヘッダーが、1以上60以下の整数で付く必要があります。/root/gw/retry.txtに、retry_after=<정수>(プレースホルダーは整数です)を書いてください。
/opt/fixtures/llms/pricing.jsonのモデルごとの単価でコストを計算して、/root/gw/cost.csvに、tenant,model,prompt_tokens,completion_tokens,cost_usdのヘッダーと2行以上を書いてください。
テナントt1の日次バジェットを0.01ドルにして、超過したら402を返してください。予約はmax_tokensを基準にあらかじめ確保し、完了後に実際の使用量で精算します。/root/gw/budget.txtに、budget_usd=0.01 spent_usd=<수> blocked=true(プレースホルダーは数値です)を書いてください。

参考

127.0.0.1:8170のバックエンドは、このPodにあらかじめ立ち上がってはいません。POST /generateで、{"prompt":..., "max_tokens":n}を受け取って、{"text":..., "usage":{"prompt_tokens":n,"completion_tokens":n}}を返せばよいので、前のトークンストリーミングのラボで作ったサーバーをそのまま使うか、同じ契約の最小のサーバーを自分で立ち上げてから始めてください。
トークンバケット: 容量と補充速度の2つの値で定義します。容量がバーストの最大値で、補充が持続可能な平均です。
ストリーミングで途中で切れたリクエストも、そこまでの出力トークンを計上する必要があります。
よくあるミス1: バジェットの検査を完了後にだけ行うことです。同時リクエストがすべて通過したあと、いっしょに超過します。
よくあるミス2: 入力と出力のトークンを合わせて数えることです。単価が違うので、コストが間違います。

ゲートウェイでバックエンドをプロキシする

/root/gw/gateway.pyを127.0.0.1:8171で起動して、POST /v1/generateを8170のバックエンドに転送してください。応答がバックエンドと同一である必要があります。テナントt1のAPIキーを/root/gw/key.txtに、t2のキーを/root/gw/key2.txtに、それぞれ1行で保存しておいてください。以降のステップの採点が、この2つのファイルを使います。

まず、そのまま通すプロキシを作って、機能を1つずつ載せていきます。

APIキーでテナントを識別する

X-API-Keyヘッダーで、テナントを識別してください。キーとテナントのマッピングは、/opt/fixtures/llms/apikeys.jsonにあります。キーがないか、未知のキーなら、401です。

キーがないか、未知のキーなら、401です。キーをテナントにマッピングしておいてください。

毎分のリクエスト制限をかける

テナントごとの毎分のリクエスト制限10を適用してください。11番目のリクエストが429である必要があります。/root/gw/rpm.txtに、allowed=10 rejected=1を書いてください。

固定ウィンドウが、最も単純です。超過したら、どのステータスコードになるかを考えてみてください。

入力・出力のトークンを会計する

GET /v1/usageが、{"tenant":"...","prompt_tokens":<n>,"completion_tokens":<n>,"requests":<n>}を返すようにしてください。2つのトークンの値が、どちらも0より大きい必要があります。

2つを別々に数える必要があります。単価が違うからです。

毎分のトークン制限をかける

毎分のトークン制限2000を、トークンバケットで適用してください。上限を超えるリクエストが、429を受け取る必要があります。/root/gw/tpm.txtに、limit=2000 consumed=<n> rejected=<n>を書き、rejectedは1以上である必要があります。

トークンバケットがよく合います。静かだったテナントのバーストを許容しながら、長期の平均を守ります。

429に再試行の時点を載せる

すべての429のレスポンスに、Retry-Afterヘッダーが、1以上60以下の整数で付く必要があります。/root/gw/retry.txtに、retry_after=<정수>(プレースホルダーは整数です)を書いてください。

残り時間を計算して、整数の秒で渡します。これがないと、クライアントたちが波を作ります。

モデルごとの単価でコストを計算する

/opt/fixtures/llms/pricing.jsonのモデルごとの単価でコストを計算して、/root/gw/cost.csvに、tenant,model,prompt_tokens,completion_tokens,cost_usdのヘッダーと2行以上を書いてください。

出力のほうが入力より高いのが普通です。単価表は、フィクスチャにあります。

バジェット超過を遮断して使用量を報告する

テナントt1の日次バジェットを0.01ドルにして、超過したら402を返してください。予約はmax_tokensを基準にあらかじめ確保し、完了後に実際の使用量で精算します。/root/gw/budget.txtに、budget_usd=0.01 spent_usd=<수> blocked=true(プレースホルダーは数値です)を書いてください。

遮断の代わりに、より安いモデルにダウングレードする選択肢もあります。ここでは、遮断で実装します。