0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@hosonot(Takashi Hosono)

SageMaker Serverless の BYOC で OOM ─ ベースが同じイメージなのに worker 数が変わる理由を追跡した

0
Posted at

はじめに

SageMaker の推論環境を BYOC（Bring Your Own Container）化するとき、「ベースイメージが同じなら挙動も同じだろう」と考えることは少なくないと思います。
AWS 公式 DLC をベースに依存パッケージを足しただけなら、TorchServe や SageMaker Inference Toolkit といったサービングスタックは何も変わっていないのですから、自然な発想です。

ところが今回、公式 DLC（Deep Learning Containers） 構成と、それをベースにした BYOC 構成を比較したところ、同じメモリ設定（1024MB）なのに BYOC 側だけメモリ不足エラーが出るという挙動の差に遭遇しました。

この記事は、その原因の調査内容と対処方法をまとめたものです。

TL;DR

挙動の差の原因

  • 公式 DLC イメージのときだけ、SageMaker が worker 数を 1 に絞る環境変数を起動時に注入していた
  • 自前イメージ（BYOC）にするとこの注入が外れ、worker 数が2になり、メモリ使用量が増えていた

対処法

  • SageMaker Python SDK のPyTorchModel(model_server_workers=1)等でwoker数を明示する

前提となる構成

AIモデルを SageMaker Serverless Inference でホストしています。

Serverless Inference を簡単に説明すると、インスタンスタイプを選ばずメモリサイズ（1024〜6144MB）だけを指定する推論ホスティングです。
vCPU はメモリに比例して自動割当され、アイドル時はゼロまでスケールインし、リクエストが来ると実行環境が起動されます（そのためコールドスタートがあります）。
AWS Lambda に近い動作モデルです。

本記事での構成は以下のとおりです。

  • メモリ: 1024MB
  • ベースイメージ: AWS 公式 DLC の pytorch-inference（TorchServe + SageMaker Inference Toolkit 入り）

また、後の説明に関わるので、コンテナ内のサービングスタックについても簡単に記します。

  • TorchServe: リクエスト受付や worker 管理を行うフロントエンドは Java 製で、実際に推論コードを実行する model worker は独立した Python プロセスとして起動される。worker が 2 つあればモデルも 2 回ロードされる
  • SageMaker Inference Toolkit: SageMaker の起動規約と TorchServe の橋渡しをするアダプタ。コンテナ起動時に環境変数を読んで TorchServe の設定ファイルを生成し、TorchServe を起動する

開発当初は「公式 DLC をそのまま使い、requirements.txt をモデルに同梱してコンテナ起動時に pip install する」構成でしたが、PyPI 障害時にコールドスタートが失敗する等のリスクがあるため、依存パッケージを事前に焼き込んだイメージを以下の様な Dockerfile で作成し、自前の ECR に置く構成（BYOC）への移行を進めまていました。

# BYOC イメージ。ベースは公式 DLC そのもの
FROM <aws_account_id>.dkr.ecr.<region>.amazonaws.com/pytorch-inference:2.3-cpu-py311

COPY docker/inference-requirements.txt /tmp/inference-requirements.txt
RUN pip install --no-cache-dir -r /tmp/inference-requirements.txt

つまり、インストールするpipパッケージも、サービングスタック（TorchServe / toolkit / JVM）もBYOC化前後で同一なはずです。

何が起きたか

ところが、BYOC 構成に切り替えたところ、メモリ不足エラーが出るようになりました。

ModelError: Inference failed due to insufficient memory on the Endpoint.
Please add more memory to the endpoint.

公式 DLC 構成は同じモデル・同じ 1024MB でもメモリ不足になったことはありませんでした。
なので、 BYOC 化によりメモリ使用量が増えたと考えました。
とはいえ、BYOC 化前後で同じ構成はずなので、なにが変わっているかを調査する必要がありました。

調査

Step 1: ログ比較 ─ メモリ使用量が増えた要因は「2 つ目の worker」

公式 DLC 構成と BYOC 構成のCloudWatch Logs を以下のように確認したところ、wokerの数が1から2に増えていたことが分かりました。


# 公式 DLC 構成
aws logs filter-log-events --log-group-name <公式DLCロググループ> \
  --start-time <対象時刻>  \
  --filter-pattern '"Default workers per model"' \
  --query 'events[].message' --output text
# => Default workers per model: 1

# BYOC 構成
aws logs filter-log-events --log-group-name <BYOCロググループ> \
  --start-time <対象時刻> \
  --filter-pattern '"Default workers per model"' \
  --query 'events[].message' --output text
# => Default workers per model: 2

TorchServe は worker ごとに W-9000, W-9001... という ID を全リクエストログに付けるので、そこからも確認できます。
公式 DLC 構成のログには W-9001 が 1 件も存在せず、BYOC 構成では W-9001 が存在し、メモリ逼迫時に W-9001 が切断されて HTTP 507 を返していました。

org.pytorch.serve.wlm.WorkerThread - 9001 Worker disconnected. WORKER_MODEL_LOADED
W-9001-model_1.0 ACCESS_LOG - "POST /invocations HTTP/1.1" 507 3253

メモリ消費はざっくり以下のようなモデルで考えられます。

総メモリ ≈ worker 数 × (モデル + 推論スタックのロードに必要なメモリ) + 固定オーバーヘッド

worker ごとにモデルを独立にロードするので、worker が 2 つになるとメモリ消費が増えます。

ここまでで、woker 数の増加がメモリ不足の原因となっている事がわかりました。
ここからは、なぜ同じ構成なのに woker 数が変化したのかを調べていきます。

Step 2: イメージ内部の差分をレイヤー比較 ─ 差分はなかった

最初に疑ったのは DLC の mutable tag 問題です。
pytorch-inference:2.3-cpu-py311 の tag の実態が AWS により更新され、「ビルドした時点の中身」と「公式 DLC 構成が pull していた中身」が違っていた可能性があると思いました。

そこで、以下のコマンドで tag のビルド日時を調査しました。

# tag が指す中身のビルド日時
docker buildx imagetools inspect \
  <aws_account_id>.dkr.ecr.<region>.amazonaws.com/pytorch-inference:2.3-cpu-py311 \
  --format '{{json .Image}}' | jq -r '.created'
# => 2025-04-20T...

結果、ビルドされていたのはBYOC化を実施した1年以上前で、BYOC化前後でイメージが同一であることが分かりました。
（そもそも、これが原因であればBYOC化と DLC の更新がたまたま同時に起こったことになるので、可能性としてはかなり少なかったなと調査後に思いました。）

ということは、差が生まれているのはイメージそのものではなく、それをエンドポイントとして起動する過程にあるということになります。
そこで次は、実際に動いているエンドポイントのコンテナから状態を抜き出して比較することにしました。

Step 3: 動作中のエンドポイントの情報を抜き出して比較 ─ 環境変数に差分あり

条件をそろえるため、同じエンドポイント・同じモデル・同じコード・1024MB で、Model に登録するイメージだけを変えた 2 つのデプロイを用意し、起動後のコンテナの中身を比較します。

from sagemaker.pytorch.model import PyTorchModel

# 各変数はイメージです
common = dict(
    model_data="s3://my-bucket/path/to/model.tar.gz",
    role=ROLE_ARN,
    entry_point="inference.py",
    source_dir="inference",
    sagemaker_session=sm_session,
)

# A: BYOC（自前 ECR の URI。中身は公式 DLC + pip 数個）
model_a = PyTorchModel(
    image_uri=IMAGE_URI,
    **common,
)

# B: 公式 DLC （URI は SDK が解決する）
model_b = PyTorchModel(
    framework_version="2.3",
    py_version="py311",
    **common,
)

コンテナの中から実際に何が見えているかを観測するため、inference.py の先頭に診断コードを仕込みます。

import multiprocessing
import os

print(
    "DIAG ENV:",
    sorted((k, v) for k, v in os.environ.items()),
    flush=True,
)
print(
    f"DIAG cpu_count={multiprocessing.cpu_count()} "
    f"affinity={len(os.sched_getaffinity(0))}",
    flush=True,
)

print の出力は TorchServe が拾って CloudWatch Logs に MODEL_LOG として流してくれるので、デプロイ → 1 回 invoke → ログ回収、で両コンテナの全環境変数を突き合わせられます。

注意: 環境変数の全件ダンプは、Model の Environment に API キー等の秘匿情報を入れている環境ではそれがログに平文で残ります。真似する場合はキー名のみ出力する、検証後に該当ログを削除する、などの配慮をしてください。

全環境変数（A: 20 個 / B: 21 個）を突き合わせた結果、意味のある差は 1 個だけで、TS_DEFAULT_WORKERS_PER_MODEL の有無でした。

自前 ECR URI でデプロイしたものには TS_DEFAULT_WORKERS_PER_MODEL は未定義で、公式 DLC でデプロイしたものには TS_DEFAULT_WORKERS_PER_MODEL = 1 が定義されていました。

TorchServe は、TS_DEFAULT_WORKERS_PER_MODEL が定義されていればその数の worker を立て、未定義かつ設定ファイルでも指定がない場合は CPU 数だけ worker を立てます。
先ほど調査した際の CPU 数は2だったので、この環境変数定義の差が worker 数の差の原因だったことが分かりました。

SageMakerが環境変数を足す仕組を持つことが明記されたドキュメントは見つけれませんでしたが、イメージとモデルが同一であること、また、TS_DECODE_INPUT_REQUEST / TS_IPEX_ENABLE のような他の環境変数も起動時に注入されていそうだったことから、SageMaker は全コンテナに対して起動時に環境変数を足す仕組みを持っていて（docker run -e ... 相当をホスト側で実行）、その中で TS_DEFAULT_WORKERS_PER_MODEL だけを公式 DLC イメージかどうかで出し分けている、と判断しました。

対処

SageMaker Python SDK には worker 数を指定するためのパラメータがあるので、今回はそれを使って woker 数を明示することで解決しました。

model = PyTorchModel(
    image_uri=IMAGE_URI,
    model_data="s3://my-bucket/path/to/model.tar.gz,
    role=ROLE_ARN,
    entry_point="inference.py",
    source_dir="inference",
    sagemaker_session=sm_session,
    model_server_workers=1,  # ← これ
)

model_server_workers=1 は SAGEMAKER_MODEL_SERVER_WORKERS=1 として Model の環境変数に設定され、コンテナ内の toolkit が TorchServe の default_workers_per_model=1 を設定ファイルに書き込みます。
SDK を使わず boto3 や Terraform 等で Model を作っている場合は、Model の Environment に SAGEMAKER_MODEL_SERVER_WORKERS=1 を直接設定すれば同じ効果になります。

BYOC なら Dockerfile に ENV TS_DEFAULT_WORKERS_PER_MODEL=1 を焼き込む方法もあります。
ただ、worker 数は「serverless でホストするかどうか」というデプロイ形態に紐づく設定なので、イメージではなく Model 側に持たせる方が関心の分離として自然だと考え、SDK パラメータを選びました。

なお、優先順位上は TS_DEFAULT_WORKERS_PER_MODEL 環境変数の方が設定ファイルより強いため、厳密にはSageMaker が自前イメージにこの変数を 1 以外の値で設定したら上書きされます。
ただし現時点で SageMaker が設定する値は 1（しかも公式イメージのみ）なので競合せず、現実的なすべてのケースで worker 数は 1 に固定されます。

教訓

  1. BYOC は「中身が同じなら挙動も同じ」ではない。マネージドサービスは自社イメージを識別して特別扱いすることがあり、自前 ECR に置いた瞬間その外に出る。BYOC への移行は「プラットフォームの暗黙の最適化を剥ぎ取る」行為でもあるので、リソース系の設定はすべて自分で明示する前提で臨む
  2. 暗黙のデフォルトに依存している設定を洗い出して明示する。今回なら worker 数。公式 DLC 構成も「たまたま AWS が 1 にしてくれていた」だけで、明示していなかった点では同じ脆弱性を抱えていた
  3. コンテナ内の環境変数には「イメージ焼き込み」「Model 定義」「起動時に外から設定」の 3 つの由来がある。トラブルシュート時はこの分類で切り分けると速い

検証環境と注意

  • 2026 年 6 月時点での観測です
  • Serverless Inference / 1024MB / pytorch-inference:2.3-cpu-py311（TorchServe 0.11.0）/ SageMaker Python SDK 2.245 の組み合わせで検証しました
  • 本記事で扱った環境変数の設定は私が調べた限りドキュメントに記載のない挙動であり、AWS が予告なく変更する可能性があります
0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?