はじめに
こんにちは！！！@HayatoHanaoka です。
せっかくのシルバーウィークなのに、台風襲来。外に出られない。
それなら記事を書けばいいじゃない！ということで、先月あたりに仕事でぶつかった壁についてまとめてみました。
私が所属するチームでは、Agentの実行結果をServer-Sent Events（SSE）で配信するアプリケーションを開発しています。
今回は、そのアプリケーションにGraceful Shutdownを導入したにも関わらず、デプロイ時にSSE接続が切れてエラーになったことについてです。
原因の調査と対応を通じて学びが多かったため、この記事で共有します。
目次
起きたこと
私たちが開発する「Speeda AI Agent」には、ユーザーの指示に応じてレポートを生成する機能があり、APIはAgentの生成結果をSSEで受け取り、ブラウザへSSEで配信しています。
図：今回の問題に関係するアプリケーション構成（簡略図）
今回起きたのは、 生成処理の途中でAPIをデプロイすると、なぜか生成完了前にSSE接続が切断され、ブラウザにはエラーが表示される といった現象です。
終了時の設定を確認したのですが、APIとAgentのどちらにもKubernetes側でPodの terminationGracePeriodSeconds を設定していました。
アプリケーション側の設定不備の可能性もあるため、API側に、SIGTERMを受信した際に、実行中の処理が完了するまで待つような終了処理を実装し、terminationGracePeriodSeconds を延長したのですが、それでも切断は解消しませんでした。
原因となる事実
詳しく調査を進めた結果、わかった原因は以下の2つです。
- デプロイ時のRollout戦略による挙動
- アプリケーションが依存するライブラリの仕様による挙動
1. デプロイ時のRollout戦略による挙動
APIとAgentのどちらも、Argo RolloutsのblueGreen方式でデプロイしていたのですが、ここに大きな見落としがありました。
BlueGreen方式とは、旧バージョンのPod群を稼働させたまま新バージョンのPod群を用意し、準備ができたら通信先を切り替える方式です。1
今回の構成では、旧Podが終了処理に入る前に、Serviceのセレクタが新Podを選択するように切り替わります。
これによって旧PodがServiceの接続先から外れ、 その変更を受けたIstio Sidecar（Envoy）が、旧Podへの既存接続を切断していました。
この切断は、旧Pod内のアプリケーションがSIGTERMを受け取る前に起きていたため、SIGTERMを受けてから処理完了を待つGraceful Shutdownでは、この切断を防げなかったことが分かりました。
2. アプリケーションが依存するライブラリの仕様による挙動
通信先の切り替えによる問題に加えて、Agentが終了処理に入ると、生成処理の完了を待たずにSSE接続が切断される現象もありました。
AI Agentは、SSEを返すためにsse-starlette（v3.4.6）というライブラリを使用しているのですが、このライブラリはUvicornのシグナル処理関数を独自の関数に差し替え、その内部で元々のUvicorn側にあるシグナル処理関数を呼ぶ仕様になっています。
そして、差し替えられる関数の実装を見ると、デフォルト挙動は SIGTERMを受け取ると、SSE配信を即時打ち切るような挙動 となっていました。2
class AppStatus:
"""Helper to capture a shutdown signal from Uvicorn so we can gracefully terminate SSE streams."""
should_exit = False
enable_automatic_graceful_drain = True
original_handler: Optional[Callable] = None
... # 省略
@staticmethod
def handle_exit(*args, **kwargs):
if AppStatus.enable_automatic_graceful_drain:
AppStatus.should_exit = True # ここのフラグを別の処理が検知し、SSE配信タスクを即時キャンセルする
if AppStatus.original_handler is not None:
AppStatus.original_handler(*args, **kwargs)
Uvicorn本来の終了処理も動いていましたが、sse-starletteがSSE配信タスクを即時キャンセルするため、Uvicorn側で処理完了を待つ設定をしていても、生成完了前にSSE接続が終了してしまっていました。
対応
デプロイのRollout戦略を、Canary方式にする
修正前の Kubernetes のマニフェストは以下の設定で、BlueGreen方式を指定していました。
※ 関連部分を一部抜粋。
apiVersion: argoproj.io/v1alpha1
kind: Rollout
metadata:
name: api
spec:
# その他の設定は省略
strategy:
blueGreen:
activeService: api-svc
previewService: api-preview-svc
autoPromotionEnabled: true
今回はこの設定を、一括でServiceの接続先を入れ替えるBlueGreen方式から、Serviceが新旧Podを共通のラベルで選択したまま新Podを増やし、旧Podを減らすCanary戦略に変更しました。
変更は以下の通りです。
※ 関連部分を一部抜粋。
apiVersion: argoproj.io/v1alpha1
kind: Rollout
metadata:
name: api
spec:
# その他の設定は省略
strategy:
canary:
maxSurge: "100%"
maxUnavailable: 0
steps:
- setWeight: 100
それぞれの変更内容を説明すると、
maxSurge -> 目標レプリカ数を超えて追加できるPod数の上限値です。3
maxUnavailable -> 更新中に、利用可能なPod数が目標を下回ることを許容する値です。4
steps[].setWeight -> 新バージョンへ振り分けるトラフィックの目標割合の設定値です。5
maxSurge: "100%" と maxUnavailable: 0 を組み合わせているので、通常必要な利用可能Pod数を維持しながら、目標レプリカ数と同じ数まで追加のPodを起動できる設定となっています。
setWeight: 100 は、新バージョンへのトラフィック配分を100%にする目標値です。
新Podを増やし、旧Podを減らすことで、最終的に新しいリクエストがすべて新Podへ向かう状態にします。5
この構成では、Serviceが新旧Podを共通のラベルで選択します。
処理を続けている旧Podは、Serviceの選択対象のまま終了処理に入り、エンドポイントも「終了中」として扱われるようになります。6
今回の環境では、その間もIstio／Envoyが既存接続を維持するため、Graceful Shutdownで終了猶予内の処理完了を待てるようになりました。
アプリケーションライブラリの設定を変更する
こちらについては、設定の変更を行うだけでOKでした。
具体的な対応としては、以下のコードをアプリケーションの起動処理（main.pyなど）に記述して、自動ドレインを無効化します。2
from sse_starlette.sse import AppStatus
AppStatus.disable_automatic_graceful_drain()
これにより、SIGTERMを契機にsse-starletteがSSE配信を自動終了する動作を止め、Uvicornの終了猶予内で生成処理とSSE配信の完了を待てるようにしました。
まとめ
これらの対応によって、デプロイ時にも終了猶予内で実行中の処理の完了を待ち、SSE配信を継続できるようになりました。
アジャイル開発のXPを取り入れている我々にとって、ユーザーに安全かつ迅速に価値を提供できる基盤を整えることは、とても意味のあるものだったので、このタイミングで対応できてよかったと思っております。
また、普段はアプリ開発に目が行きがちな私にとって、デプロイ時の通信経路の問題は完全に盲点でした。
今回の対応は、インフラ周りの知識を獲得する良い機会にもなったと思っています。
同じような問題に当たった方の参考になれば幸いです。