CDKによるエンドポイント更新時のSageMaker自動スケーリング障害の修正

AWS CloudFormationやAWS Cloud Development Kit (CDK)を使用してAmazon SageMaker AIリアルタイムエンドポイントを更新すると、自動スケーリングの設定が破損することがよくあります。AWSから提供された新しいソリューションでは、スケーラブルターゲットの解除と再登録を自動的にオーケストレーションする再利用可能なCDKコンストラクトを導入しています。このアプローチにより、エンジニアは手動での介入やスタックの破損なしに、単一のデプロイでインスタンスタイプ、インスタンス数、環境変数を更新できます。

何が起きたか

ネイティブのCloudFormation自動スケーリングリソースを使ってSageMaker AIエンドポイントをデプロイする開発者は、エンドポイント設定を変更する際にデプロイ失敗に頻繁に見舞われます。根本的な問題は、インスタンスタイプやインスタンス数の変更など、エンドポイントを更新するには事前にスケーラブルターゲットを解除する必要があるというSageMaker AIの要件にあります。しかし、CloudFormationはこれらのリソースを宣言的に管理するため、自動スケーリングターゲットがまだアクティブな状態でサービスが更新をブロックするという競合が発生します。

エンジニアがAWS CLI経由でスケーラブルターゲットを手動で解除して回避しようとすると、2つの異なる失敗モードに直面します。エンドポイント設定の変更のみで再デプロイした場合、CloudFormationは自動スケーリングパラメータに変更がないと判断し、ターゲットの再登録をスキップします。エンドポイントは正常に更新されますが、自動スケーリングなしで実行されることになり、トラフィック急増に対してシステムが無防備なままとなり、チームに警告を出すエラーメッセージも表示されません。

手動で解除した後に自動スケーリングパラメータを変更した場合、CloudFormationは存在しないScalableTargetリソースの更新を試みます。これにより「Resource not found」エラーが発生します。さらに深刻なのは、ロールバックプロセスも失敗し、CloudFormationスタックがUPDATE_ROLLBACK_FAILED状態になることです。この状態からの回復にはコンソールでの面倒な手動介入が必要となり、自動化パイプラインが断絶し、運用リスクが増大します。

仕組み

提案されたソリューションでは、ネイティブのAWS::ApplicationAutoScaling::ScalableTargetおよびAWS::ApplicationAutoScaling::ScalingPolicyリソースを、CDKコンストラクトによって管理される3つのカスタムリソースに置き換えます。これらのカスタムリソースはnode.add_dependency呼び出しを使用して厳格な依存関係を強制し、エンドポイント更新ライフサイクル全体を通じて自動スケーリングが正しく処理されるようにします。プロセスは、エンドポイント更新開始前にスケーラブルターゲットを削除するDeregisterカスタムリソースから始まります。初回デプロイ時には、ターゲットが存在しない場合でもハンドラーはエラーを適切に無視します。

スケーラブルターゲットが削除されると、CDKはエンドポイントリソース自体の更新を進めます。エンドポイント更新が受け入れられた後、Register Scalable Targetカスタムリソースが引き継ぎます。このリソースは非同期完了パターンを使用しており、on_eventハンドラーは即座に返り、ウェーターが30秒ごとにis_completeハンドラーを再呼び出します。Lambdaの実行時間制限を避けるため、合計1時間のタイムアウトを設定し、エンドポイントがInServiceステータスに到達してからスケーラブルターゲットを登録します。

最後に、Attach Scaling Policyカスタムリソースがターゲット追跡ポリシーの適用を保証します。これも非同期パターンを使用し、登録後にスケーラブルターゲットがクエリ可能になるまで待機します（これは一瞬かかる場合があります）。ターゲットが見えるようになると、ポリシーを適用します。ConfigHashプロパティがこのシーケンス全体を駆動し、エンドポイントまたは自動スケーリング設定値のいずれかが変更されるたびにCloudFormationがLambda関数を再呼び出しするように強制することで、両システムの同期を保ちます。

主要な詳細

• このコンストラクトは、ネイティブのCloudFormation自動スケーリングリソースを、Deregister、Register Scalable Target、Attach Scaling Policyの3つのカスタムリソースに置き換えます。
• 正しい操作順序を確保するために、CDKのnode.add_dependency呼び出しを通じて依存関係の順序付けが強制されます。
• 非同期カスタムリソースは、AWS Step Functionsステートマシンでバックアップされたon_eventおよびis_completeハンドラーを持つプロバイダーフレームワークを使用します。
• Register Scalable Targetリソースの合計タイムアウトは1時間ですが、Attach Scaling Policyのタイムアウトは15分です。
• 単一のConfigHashプロパティが、エンドポイントまたは自動スケーリングパラメータが変更されたときにオーケストレーションシーケンスをトリガーします。
• デプロイされたスタックには12個のLambda関数と2つのStep Functionsステートマシンが含まれますが、その多くはCDKプロバイダーフレームワークによって生成されます。

なぜ重要なのか

本番AIインフラストラクチャを構築しているチームにとって、デプロイパイプラインの信頼性は譲れないものです。以前は、自動スケーリング設定が静かにドロップされたり、スタックが破損した状態のままになったりする挙動により、重大な運用負債が生じていました。エンジニアは、Infrastructure as Codeの原則を破る手動CLI介入を行うか、自動スケーリングなしでモデルを実行するリスクを受け入れるかのどちらかを選択する必要がありました。このソリューションは宣言型インフラストラクチャの約束事を取り戻し、チームがSageMakerエンドポイントを、安全かつ予測可能に更新できる他のクラウドリソースと同様に扱えるようにします。

さらに、非同期カスタムリソースの使用は、CloudFormationにおける一般的な課題である、Lambdaタイムアウトを超える長時間のオペレーションに対応します。待機ロジックをプロバイダーフレームワークとStep Functionsにオフロードすることで、このコンストラクトは一時的なタイミングの問題によるデプロイ失敗を防ぎます。この堅牢性は、エンドポイントの更新がフリート全体に伝播するのに数分かかる可能性がある大規模MLオペレーションにおいて極めて重要です。

実施できること

• GitHubからsample-sagemaker-endpoint-autoscalingリポジトリをクローンして、参照実装にアクセスしてください。
• uv syncを使用して依存関係をインストールし、uv run cdk bootstrapでAWS環境をブートストラップしてください。
• config/config.yamlでエンドポイントと自動スケーリング設定を定義し、ConfigHashを生成するために両セクションが含まれていることを確認してください。
• uv run cdk deploy InferenceStackを使用してサンプルスタックをデプロイし、AWS CLI経由でスケーラブルターゲットを検証してください。
• 設定ファイルでインスタンスタイプを変更して再デプロイすることで更新をテストし、自動スケーリングが維持されていることを確認してください。
• コンストラクトディレクトリをコピーし、エンドポイントリソースと設定ハッシュを渡すことで、既存のCDKスタックにSagemakerAutoscalingConstructを組み込んでください。