【SageMaker AI】非同期推論エンドポイントのオートスケーリングについて

記事タイトルとURLをコピーする

はじめに

こんにちは、アプリケーションサービス部 ディベロップメントサービス1課の北出です。
以前の記事でSagemaker AIの非同期推論エンドポイントの作成方法の流れについて作成しました。
今回はそれに関連して、非同期推論エンドポイントのオートスケーリングについて膨らませた内容についてと、利用に際して自分がぶつかった壁について紹介しようと思います。

非同期推論とは

そもそも非同期推論とはという内容はこちらの記事で記載していますので参照ください。

非同期推論のインスタンスについて

非同期推論では、推論を行うサーバーとしてのインスタンスを起動し、そこで推論を行います。
使用するインスタンスのインスタンスタイプはエンドポイント設定を作成する際に設定します。
使用可能なインスタンスタイプと料金はこちらの公式ドキュメントを参照してください。
見てもらえればわかりますが、最低でも「large」サイズからとなっており、不要な時間は停止させておくなどの対策が求められるかと思います。

非同期推論のオートスケーリングについて

SageMaker AI 非同期推論エンドポイントのオートスケーリングについてのAWS公式のドキュメントはこちらになります。
このドキュメントを見ると、基本的なスケーリングはターゲット追跡スケーリングポリシーで行い、インスタンス数がゼロからスケールアウトする場合のみ、ステップスケーリングポリシーで行うようです。

ステップスケーリングポリシーの必要性

なぜターゲット追跡スケーリングポリシーだけじゃだめなのかと疑問に思うかもしれません。AWS公式ドキュメントにも記述されていますが、こちらでも説明します。

まず、ターゲット追跡スケーリングポリシーの内容についてですが、公式ドキュメントの例をそのまま実装したとすると、ApproximateBacklogSizePerInstanceのメトリクスが「5」になるようにスケーリングするようです。
この ApproximateBacklogSizePerInstanceこちらのドキュメントに記述されており、「キュー内の項目数をエンドポイントの背後にあるインスタンス数で割った値」とのことです。
「キュー内の項目数」とは、インスタンスで推論される前の順番待ちしている推論リクエストの数のことです。
インスタンス数がゼロの時は単純に順番待ちしている推論リクエストの数になります。

ターゲット追跡スケーリングポリシーで多少前後しますが、基本的にはこれが目標値である5を上回るとスケールアウトして、下回るとスケールインするという挙動になります。

この時、例えばインスタンス数がゼロで推論リクエストが1件だけあった場合、スケールアウトされず、リクエストが5件になるまでずっと放置されることになります。
それを防ぐために、インスタンス数が0→1にするステップスケーリングポリシーが追加で必要になります。

非同期推論のオートスケーリングの困りごと

AWS公式ドキュメントをそのまま実装して検証していたところ、以下のような困りごとがあったのでどう対処したのかを記述します。

  • 1件の推論リクエストに対してインスタンスが2個起動されることがある
  • 推論完了してからインスタンスがスケールインするまでの時間が長い
  • リアルタイム性を上げるために日中は最低1台インスタンスを稼働させ続けたい

インスタンスが2個起動されることがある

インスタンス数がゼロの状態で1件だけ推論リクエストを送ったときに、CloudShellでスケーリングの履歴を調べると以下のようになっていました。

~ $ aws application-autoscaling describe-scaling-activities \
     --service-namespace sagemaker \
     --resource-id "endpoint/async-inference-endpoint/variant/async-inference-variant" \
     --query "ScalingActivities[*].[StatusCode,Description,StartTime]" \
     --output table
--------------------------------------------------------------------------------------------
|                                 DescribeScalingActivities                                |
+------------+----------------------------------------+------------------------------------+
|  Successful|  Setting desired instance count to 0.  |  2025-02-21T02:38:56.256000+00:00  |
|  Successful|  Setting desired instance count to 1.  |  2025-02-21T02:26:56.271000+00:00  |
|  Successful|  Setting desired instance count to 2.  |  2025-02-21T02:02:14.346000+00:00  |
|  Successful|  Setting desired instance count to 1.  |  2025-02-21T01:26:14.403000+00:00  |
+------------+----------------------------------------+------------------------------------+

この原因はスケーリングのクールダウン時間が関係しているのではないかと推測しました。
クールダウン時間が300秒と設定していましたが、インスタンスの起動に300秒以上かかっていた場合、推論リクエストは処理されず、再度スケーリングが起動したのではないかと思います。
そのため、クールダウンをインスタンスの起動にかかる時間以上にすることで対処しました。
インスタンス数の時系列の推移を調べる方法は見つけられませんでした。

推論完了してからインスタンスがスケールインするまでの時間が長い

上記クールダウンの対応後、インスタンス数がゼロの状態で1件だけ推論リクエストを送ったときに、CloudShellでスケーリングの履歴を調べると以下のようになっていました。

~ $ aws application-autoscaling describe-scaling-activities \
>     --service-namespace sagemaker \
>     --resource-id "endpoint/async-inference-endpoint-small/variant/async-inference-variant-small" \
>     --query "ScalingActivities[*].[StatusCode,Description,StartTime]" \
>     --output table
--------------------------------------------------------------------------------------------
|                                 DescribeScalingActivities                                |
+------------+----------------------------------------+------------------------------------+
|  Successful|  Setting desired instance count to 0.  |  2025-02-25T02:27:56.231000+00:00  |
|  Successful|  Setting desired instance count to 1.  |  2025-02-25T01:56:46.722000+00:00  |
+------------+----------------------------------------+------------------------------------+

これを見ると、日本時間で10:56にスケールアウトの指示を出し、11:27にスケールインの指示を出しており、一つのリクエストのために30分インスタンスが起動していることになります。

おおよその内訳はスケールアウト指示~インスタンス起動準備~推論完了までで15分で、推論完了~スケールイン指示で15分となっており、推論完了してからスケールインするまでの時間を減らしたいと思いました。

まず、スケールインまでで15分かかっている原因を調べていたのですが、どうやらターゲット追跡スケーリングポリシーに原因があるようでした。
ターゲット追跡スケーリングポリシーで作成されたCloudWatchアラームを確認すると画像のようになっていました。

注目は説明としきい値の部分になります。

しきい値は「15分内の15データデータポイントのApproximateBacklogSizePerInstance < 2.7」となっています。(2.7は目標値3で設定したためです)
これは、15分間のデータがすべてしきい値以下だったらアラームが発火しますという意味になります。つまり、アラームが発火するまでに最低でも15分必要ということになります。
この設定ですが、ターゲット追跡スケーリングポリシーのデプロイ時に自動で設定されるもので、15分というのはユーザー側で設定することはできませんでした。
デプロイ後にCloudWatchAlarmの設定を変更することも考えたのですが、先ほどのアラーム設定の画像の「説明」の部分で DO NOT EDIT OR DELETEと書かれていたり、こちらの公式ドキュメントにも「ターゲット追跡スケーリングポリシーで使用される CloudWatch アラームを作成、編集、削除しないでください。」と書かれていたので変更しないほうがいいと判断しました。

じゃあどうしたらいいのかということで、ターゲット追跡スケーリングポリシーをやめて、ステップスケーリングポリシーのみにするのがベターなのかと思います。

# スケーラブルターゲットを登録
response = autoscaling.register_scalable_target(
    ServiceNamespace="sagemaker",
    ResourceId=resource_id,
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    MinCapacity=0,
    MaxCapacity=3,
)

# スケールアウトポリシー 
scale_out_policy = autoscaling.put_scaling_policy(
    PolicyName=f"{endpoint_name}-ScaleOutPolicy",
    ServiceNamespace="sagemaker",
    ResourceId=resource_id,
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    PolicyType="StepScaling",
    StepScalingPolicyConfiguration={
        "AdjustmentType": "ChangeInCapacity",
        "MetricAggregationType": "Average",
        "Cooldown": 300,
        "StepAdjustments": [
            {
                "MetricIntervalLowerBound": 0.5,  # バックログが 1 以上
                "ScalingAdjustment": 1,  # インスタンス +1
            },
        ],
    },
)

# スケールインポリシー
scale_in_policy_backlog = autoscaling.put_scaling_policy(
    PolicyName=f"{endpoint_name}-ScaleInPolicy-Backlog",
    ServiceNamespace="sagemaker",
    ResourceId=resource_id,
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    PolicyType="StepScaling",
    StepScalingPolicyConfiguration={
        "AdjustmentType": "ChangeInCapacity",
        "MetricAggregationType": "Average",
        "Cooldown": 300,
        "StepAdjustments": [
            {
                "MetricIntervalUpperBound": 0,  # バックログが0以下
                "ScalingAdjustment": -1,  # インスタンス -1
            }
        ],
    },
)

# スケールアウトアラーム 
cw_client.put_metric_alarm(
    AlarmName=f"{endpoint_name}-ScaleOut-Backlog-Alarm",
    MetricName="ApproximateBacklogSize",
    Namespace="AWS/SageMaker",
    Statistic="Average",
    EvaluationPeriods=1,
    DatapointsToAlarm=1,
    Threshold=0.5,  # バックログ 0.5 以上
    ComparisonOperator="GreaterThanOrEqualToThreshold",
    TreatMissingData="breaching",
    Dimensions=[{"Name": "EndpointName", "Value": endpoint_name}],
    Period=60,
    AlarmActions=[scale_out_policy["PolicyARN"]],
)


# スケールインアラーム 
cw_client.put_metric_alarm(
    AlarmName=f"{endpoint_name}-ScaleIn-Backlog-Alarm",
    MetricName="ApproximateBacklogSize",
    Namespace="AWS/SageMaker",
    Statistic="Average",
    EvaluationPeriods=1,
    DatapointsToAlarm=1,
    Threshold=0.1,  # バックログが 0 の場合
    ComparisonOperator="LessThanOrEqualToThreshold",
    TreatMissingData="notBreaching",
    Dimensions=[{"Name": "EndpointName", "Value": endpoint_name}],
    Period=60,
    AlarmActions=[scale_in_policy_backlog["PolicyARN"]],
)

print("SageMaker Auto Scaling with Corrected Step Adjustments configured successfully.")

このようにすることで、順番待ちしている推論リクエストの数に応じたシンプルなスケーリングポリシーになります。
評価期間も最短にしているので、過敏に反応するデメリットはありますが、インスタンスの停止までの時間を早くできます。

ただ、ここまでするなら、処理ジョブやFargateなどのサービスを使って推論を使うべきではと自分は思っています。

日中は最低1台インスタンスを稼働させ続けたい

最後はスケジューリングについてです。
ターゲット追跡スケーリングポリシーを使った場合で、日中(例えば9:00~18:00)は最低一台はインスタンスを起動させ、それ以外の時間は0台を最低にするという設定にします。

やることはシンプルです。こちらのドキュメントに従って下記のように設定しているとします。

# Common class representing application autoscaling for SageMaker 
client = boto3.client('application-autoscaling') 

# This is the format in which application autoscaling references the endpoint
resource_id='endpoint/' + <endpoint_name> + '/variant/' + <'variant1'> 

# Define and register your endpoint variant
response = client.register_scalable_target(
    ServiceNamespace='sagemaker', 
    ResourceId=resource_id,
    ScalableDimension='sagemaker:variant:DesiredInstanceCount', # The number of EC2 instances for your Amazon SageMaker model endpoint variant.
    MinCapacity=0,
    MaxCapacity=5
)

この MinCapacity = 0 の設定を時間帯に応じて変更するという対応をします。
本記事ではそこまで詳細には説明しませんが、EventBridgeを使って、RegisterScalableTarget のAPIを指定した時間に実行するようにします。

この画像のように ApplicationAutoScalingを選択し、RegisterScalableTarget を選択します。

実行するAPIのJSONフォームは以下のようにします。

{
    "ServiceNamespace": "sagemaker",
    "ResourceId": "endpoint/async-inference-endpoint/variant/async-inference-variant",
    "ScalableDimension": "sagemaker:variant:DesiredInstanceCount",
    "MinCapacity": 1,
    "MaxCapacity": 5
}

これを毎日9:00に実行することで最低1台のインスタンスは稼働することになります。

反対に MinCapacity = 0 のスケジュールを18:00に設定すると、18:00以降はリクエストがなければインスタンスは0台になります。

邪道かもしれないですが、このようにすることで日中だけ最低1台を稼働させることができるようになります。

おわりに

オートスケーリングについて不慣れですが、SageMaker AIでのオートスケーリングに関する資料があまりなかったので、今回作成してみました。
一通り検証するなかで、推論だからSageMaker AI を使おうではなく、FargateやEC2を使う方法も視野にいれて考えることも必要になるかなと感じました。

北出 宏紀(執筆記事の一覧)

アプリケーションサービス本部ディベロップメントサービス3課

2024年9月中途入社です。