SageMaker上の自前モデルはトークン数が見えない — Bedrock AgentCoreとの組み合わせで可観測性を手動で補う実装パターン

8月15日、AWSが「Building agentic workflows with SageMaker AI and Bedrock AgentCore」と題した記事を公開した。この記事では、SageMaker AIにデプロイした自前モデルとBedrock AgentCoreを組み合わせてマルチエージェントワークフローを構築し、トークンレベルの可観測性を確保する方法について詳しく紹介されている。

何が面白いか：「SageMakerのトークン使用量が見えない」問題

Bedrock AgentCoreはOpenTelemetryによる自動計装を備えているが、Bedrockのモデル呼び出しにしか対応していない。SageMaker上でホストしたモデルをStrands Agents（AWSが公開しているオープンソースのエージェントフレームワーク）のOpenAIModelとして呼び出すと、トークン数がトレースに一切記録されない。コスト監視もレイテンシのデバッグもできない状態になる。

この記事の核心は、そのギャップを埋める具体的な実装方法だ。

なお元記事では、SageMakerにデプロイするモデルとして「Qwen3.5-9B」という型番が使用されている。Qwen系列の一般的な表記（Qwen2.5-7BやQwen3など）とは異なるが、本稿では元記事の記述に従って表記する。Claude側のバージョン表記（Claude Haiku 4.5 / Claude Sonnet 4.6）についても同様に元記事の記述をそのまま転記している。

アーキテクチャ概要

3つのエージェントが1つのAgentCoreコンテナ内で協調する構成だ：

オーケストレーターエージェント（Claude Haiku 4.5 / Bedrock）：ユーザーの意図を分類してルーティング
予算エージェント（Claude Sonnet 4.6 / Bedrock）：50/30/20予算配分をPydanticで構造化出力
財務分析エージェント（Qwen3.5-9B / SageMaker AI）：株式分析とポートフォリオ構築、ツール呼び出し対応

SageMakerエンドポイントはOpenAI互換APIとして公開されており、StrandsのOpenAIModelから透過的に呼び出せる。BedrockモデルはAnthropicのAPIをAWS経由で利用するため自動計装が効くが、SageMakerモデルはOpenAI互換エンドポイント経由となり、Strandsのトレース対象から外れる点がアーキテクチャ上の重要な非対称性だ。

Step 1：Qwen3.5-9BをSageMaker AIにデプロイ

vLLM DLC（vllm:0.22.1-gpu-py312-cu130）を使い、ml.g6e.2xlarge（L40S 48GB VRAM × 1）上にデプロイする。

region = "us-west-2"
model_id = "Qwen/Qwen3.5-9B"
instance_type = "ml.g6e.2xlarge"

inference_image = f"763104351884.dkr.ecr.{region}.amazonaws.com/vllm:0.22.1-gpu-py312-cu130-ubuntu22.04-sagemaker"

env = {
    "SM_VLLM_MODEL": model_id,
    "SM_VLLM_TENSOR_PARALLEL_SIZE": "1",
    "SM_VLLM_MAX_MODEL_LEN": "32768",
}

Step 2：トークンが切れる問題をhttpx.Authで解決

SageMakerのOpenAI互換エンドポイントはBearerトークン認証を要求する。トークンには有効期限があるため、リクエストごとに自動更新する仕組みが必要だ。httpx.Authのサブクラスとして実装する：

import httpx
from openai import AsyncOpenAI
from sagemaker.core.token_generator import generate_token

class SageMakerAuth(httpx.Auth):
    def __init__(self, region): self.region = region
    def auth_flow(self, request):
        request.headers["Authorization"] = f"Bearer {generate_token(region=self.region)}"
        yield request

strands_client = AsyncOpenAI(
    base_url=f"https://runtime.sagemaker.{REGION}.amazonaws.com/endpoints/{ENDPOINT_NAME}/openai/v1",
    api_key="sagemaker",
    http_client=httpx.AsyncClient(auth=SageMakerAuth(region=REGION)),
)

Step 3：SageMakerのトークン可観測性を手動で補う

問題の根本

StrandsはOpenAIModelプロバイダーに対してgen_ai.chatスパンを発行しない。結果として、Qwen3.5-9Bが消費したトークンはトレースに現れない。

解決策：カスタムOTELスパン

gen_ai.chatスパンを手動で発行し、StrandsのAgentResult.metrics.accumulated_usageからトークン数を取得して属性としてセットする：

from opentelemetry import trace

tracer = trace.get_tracer("financial_analysis_agent")

@tool
def financial_analysis_agent_tool(query: str) -> str:
    with tracer.start_as_current_span("gen_ai.chat", attributes={
        "gen_ai.system": "openai",
        "gen_ai.request.model": f"qwen3.5-9b ({SAGEMAKER_ENDPOINT_NAME})",
        "gen_ai.operation.name": "chat",
    }) as span:
        fa_agent = Agent(model=OpenAIModel(...), ...)
        result = fa_agent(query)
        usage = result.metrics.accumulated_usage
        span.set_attribute("gen_ai.usage.input_tokens", usage.get("inputTokens", 0))
        span.set_attribute("gen_ai.usage.output_tokens", usage.get("outputTokens", 0))
        span.set_attribute("gen_ai.usage.total_tokens", usage.get("totalTokens", 0))
        return str(result)

`stream_options`が必須な理由

vLLMはデフォルトでストリーミングレスポンスにusageチャンクを含めない。そのためaccumulated_usageの値がゼロのまま残る。以下のパラメータを必ず追加すること：

params={
    "temperature": 0.7,
    "max_tokens": 4096,
    "stream_options": {"include_usage": True},  # これがないとトークン数は常に0
},

この設定により、vLLMが最終チャンクにトークン数を付与するようになる。

実際のトレース出力例：

{
    "name": "gen_ai.chat",
    "attributes": {
        "gen_ai.request.model": "qwen3.5-9b (qwen35-9b-260612-082732)",
        "gen_ai.usage.input_tokens": 1391,
        "gen_ai.usage.output_tokens": 1432,
        "gen_ai.usage.total_tokens": 2823
    },
    "durationNano": 37237386894
}

実装上の注意点まとめ

項目	内容
BedrockモデルのOTEL	自動計装済み。追加作業不要
SageMaker OpenAIエンドポイントのOTEL	手動でgen_ai.chatスパンが必要
トークン取得キー	`inputTokens` / `outputTokens` / `totalTokens`
X-Rayサンプリングレート	デフォルト1%。開発中は100%に変更すること
エージェントインスタンス	リクエストごとに新規生成。シングルトンは並行実行エラーの原因になる

応用パターン

記事では以下の発展的な使い方も紹介されている：

ファインチューニング済みモデルへの差し替え：SM_VLLM_MODELをS3上のチェックポイントに変更するだけで認証・OTEL・AgentCoreデプロイはそのまま再利用できる
A/Bテスト：同一エンドポイントにベースモデルとファインチューニング済みモデルをバリアントとしてデプロイし、OTELスパンのバリアント属性でトレース上の品質を比較する
コスト考慮ルーティング：クエリの複雑さに応じて、単純な照会はBedrock上のHaikuへ、多段階推論はSageMaker GPUエンドポイントへ振り分ける

ソースコードはGitHubリポジトリで公開されており（元記事内にリンクが掲載されている）、可観測性の詳細はOBSERVABILITY.mdを参照する形になっている。

詳細はBuilding agentic workflows with SageMaker AI and Bedrock AgentCoreを参照していただきたい。