翻訳は機械翻訳により提供されています。提供された翻訳内容と英語版の間で齟齬、不一致または矛盾がある場合、英語版が優先します。

エージェントの準備

トレーニングジョブを起動する前に、トレーニングループに参加できるエージェントを設定する必要があります。このセクションでは、マネージドホスティングに Amazon Bedrock AgentCore を使用している場合でも、カスタムエージェントで独自のインフラストラクチャを使用する場合でも、エージェントを設定してデプロイする方法について説明します。

エージェントコード統合の概要

トレーニング中、SageMaker AI はトレーニングデータセットからエージェントにプロンプトを送信します。エージェントは各プロンプトを処理し、ポリシーモデルを呼び出して応答し、呼び出しツールなどのアクションを実行します。ポリシーモデルは、トレーニング中のモデルです。基本モデルとして開始され、トレーニングループを通じてフィードバックを受け取るにつれて重みが時間の経過とともに更新されます。その結果に基づいて、エージェントは報酬を SageMaker AI に報告し、トレーニングループを完了します。これはデータセット内のすべてのプロンプトに対して繰り返され、ポリシーモデルは収集された報酬に基づいて時間の経過とともに改善されます。

SageMaker AI ジョブランタイムサービス

トレーニング中、エージェントは SageMaker AI Job Runtime Service と通信して、推論のためにポリシーモデルを呼び出し、結果をトレーナーに報告します。SDK デコレータはこの統合を自動的に処理しますが、エージェントにカスタム要件がある場合は、ランタイム APIsを直接呼び出すことができます。

Endpoint

https://job-runtime.sagemaker.region.api.aws

API

認証

エージェントは、ベアラートークンを使用して Runtime Service で認証します。SDK を使用してトークンを生成します。

from sagemaker.core.token_generator import generate_token

token = generate_token(region="us-west-2")

Runtime Service を呼び出すときに、このトークンを API キーとして渡します。

直接 API 統合 (アドバンスド)

エージェントフレームワークが SDK デコレータを使用できない場合は、 APIsを直接呼び出します。

import requests
import openai
from sagemaker.core.token_generator import generate_token
from sagemaker.train.rft.headers import make_inference_headers

token = generate_token(region="us-west-2")
endpoint = "https://job-runtime.sagemaker.us-west-2.api.aws"
headers = {
    "Authorization": f"Bearer {token}",
    "Content-Type": "application/json",
    **make_inference_headers(metadata),  # injects jobArn, rolloutId tracking
}

# Sample (inference call)

# Option 1: Create OpenAI-compatible client pointing to the Runtime Service (recommended)
client = openai.OpenAI(
    base_url=endpoint + "/v1",
    api_key=token,
    default_headers=headers,
)

# Option 2: direct request call
response = requests.post(
    f"{endpoint}/v1/chat/completions",
    headers=headers,
    json={
        "model": "default",
        "messages": [{"role": "user", "content": "..."}],
        "max_tokens": 256,
        "temperature": 1.0,
    }
)

# CompleteRollout + UpdateReward (after all turns complete)
# The decorator handles this; for direct calls, use the SDK's
# RolloutFeedbackClient or call the APIs per the Smithy model.

エージェントデプロイオプション

SageMaker AI は、エージェントをトレーニングループに接続するための 2 つのオプションをサポートしています。

シナリオ 1: Bedrock AgentCore ランタイムのエージェント

フルマネージドホスティングのために、エージェントを Amazon Bedrock AgentCore にデプロイします。SageMaker AI は、モデルトレーニング中にエージェントを呼び出します。

前提条件

開始する前に、次の前提条件を完了します。

Bedrock AgentCore 実行ロール

トレーニング中、SageMaker AI は SageMaker 実行ロールを引き受けてエージェントを呼び出します。エージェントは、モデル推論と報酬レポートのために RFT ランタイムを呼び出すために、Bedrock AgentCore 実行ロールと呼ばれる独自の個別のロールも必要です。RFT ランタイムは、エージェントがトレーニング中に通信する SageMaker AI サービスエンドポイントです。推論中にポリシーモデルからのレスポンスを提供し、各ロールアウトの終了時にエージェントが報告する報酬を受け取るという 2 つのことを処理します。

次の信頼ポリシーを使用して Bedrock AgentCore 実行ロールを作成します。このポリシーは、トレーニング中にユーザーに代わってロールを引き受けるアクセス許可を Bedrock AgentCore に付与します。これがないと、Bedrock AgentCore はコンテナイメージにアクセスしたり、RFT ランタイムを呼び出すことができません。

{
    "Version": "2012-10-17",
    "Statement": [{
        "Effect": "Allow",
        "Principal": {"Service": "bedrock-agentcore.amazonaws.com"},
        "Action": "sts:AssumeRole"
    }]
}

次に、AmazonSageMakerJobRuntimeAccess 管理ポリシーをこのロールにアタッチします。これにより、エージェントが推論のためにポリシーモデルを呼び出し、トレーニング中に結果を報告するために必要なアクセス許可が付与されます。

エージェントコードを記述または更新する

エージェントは sagemaker-train SDK を使用し、デ@sagemaker_rft_handlerコレータをエージェントのエントリポイント関数に適用する必要があります。このデコレータは、AgentCore がトレーニング中に呼び出す HTTP サーバーをセットアップし、受信するロールアウトリクエストをリッスンし、プロンプトをエントリポイント関数にルーティングします。

SDK のインストール

次のいずれかの方法を使用して、エージェントの環境に SDK をインストールします。

オプション 1: 直接インストール

pip install sagemaker-train

オプション 2: ホイールを個別に構築してインストールします。直接インストールが環境のサイズ制約を超える場合に使用します。

# Clone the SDK repository
git clone https://github.com/aws/sagemaker-python-sdk.git

# Build the wheels
cd sagemaker-python-sdk/sagemaker-train
python -m build --wheel --outdir ./dist/

cd sagemaker-python-sdk/sagemaker-core
python -m build --wheel --outdir ./dist/

# Install
pip install ./sagemaker-train/dist/sagemaker_train-*.whl
pip install ./sagemaker-core/dist/sagemaker_core-*.whl

認証

エージェントは、ベアラートークンを使用して SageMaker AI で認証します。SDK の generate_token()メソッドを使用してトークンを生成し、推論のためにポリシーモデルを呼び出すときに API キーとして渡します。

エージェントの要件

エージェントは、ロールアウトごとに以下を実行する必要があります。

次の例は、Strands SDK を使用した基本的なエージェントテンプレートを示しています。

from sagemaker.train.rft import sagemaker_rft_handler, RolloutFeedbackClient
from sagemaker.train.rft.adapters.strands import wrap_model
from sagemaker.core.token_generator import generate_token
from strands import Agent, OpenAIModel
import os

@sagemaker_rft_handler
def handle_rollout(payload):
    metadata = payload.get("metadata", {})
    prompt = payload.get("prompt", "")
    endpoint = metadata.get("endpoint", os.environ.get("RFT_RUNTIME_ENDPOINT", ""))

    # Generate bearer token for authenticating with the RFT Runtime
    token = generate_token(region=os.environ.get("AWS_REGION", "us-west-2"))

    # Create client for model inference
    model = OpenAIModel(
        model_id="default",
        client_args={
            "api_key": token,
            "base_url": endpoint + "/v1",
        },
    )

    # Wrap model to auto-inject RFT tracking headers and inference parameters
    model = wrap_model(model)

    # Execute actions in your environment (tools, APIs, etc.)
    agent = Agent(model=model, tools=[..., ...])
    result = agent(prompt)

    # Return reward -   decorator handles CompleteRollout + UpdateReward
    return {"reward": compute_reward(result)}

Bedrock AgentCore にエージェントをデプロイする

AgentCore 開発ガイドに従って、エージェントを Amazon Bedrock AgentCore にデプロイします。

AgentCore CLI を使用してエージェントをデプロイしたら、出力のエージェントランタイム ARN を書き留めます。これは、トレーニングジョブを作成するときに必要になります。ランタイム ARN は の形式に従いますarn:aws:bedrock-agentcore:<region>:<account-id>:runtime/<agent-name>。

エージェントがデプロイされ、正常であることを確認します。

aws bedrock-agentcore-control list-agent-runtimes --region us-west-2

シナリオ 2: Lambda フォワーダーを使用するカスタムエージェント

Bedrock AgentCore に加えて、SageMaker AI はあらゆる環境でホストされているカスタムエージェントをサポートしています。エージェントは Lambda フォワーダーを介してトレーニングループに接続し、任意のエージェントフレームワークとホスティングプラットフォームを柔軟に使用できます。

エージェントは、Strands Agents SDK などの任意のフレームワークまたはプラットフォーム、または独自のカスタム実装を使用して構築できます。Amazon Bedrock AgentCore、Amazon EKS、Amazon EC2、 AWS Fargate、または独自のインフラストラクチャを含む任意のコンピューティング環境で実行できます。

Lambda 関数は SageMaker AI からロールアウトリクエストを受け取り、エージェントの HTTP エンドポイントに転送します。

以下のセクションでは、Lambda フォワーダーを使用してカスタムエージェントを設定する方法について詳しく説明し、Amazon EKS にデプロイされたカスタムエージェントを使用する例を示します。

前提条件

Lambda 実行ロール

トレーニング中、SageMaker AI は SageMaker 実行ロールを引き受けて Lambda フォワーダーを呼び出します。Lambda 関数は、 AWS Lambda が実行できるように独自の実行ロールを必要とします。このロールにより、Lambda はログを実行および書き込みできます。

aws iam create-role \
  --role-name RFTLambdaForwarderRole \
  --assume-role-policy-document '{
    "Version": "2012-10-17",
    "Statement": [{
      "Effect": "Allow",
      "Principal": {"Service": "lambda.amazonaws.com"},
      "Action": "sts:AssumeRole"
    }]
  }'

aws iam attach-role-policy \
  --role-name RFTLambdaForwarderRole \
  --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

エージェントコードを記述する

エージェントは、ロールアウトリクエストを受け入れて RFT ランタイムサービスとやり取りする HTTP エンドポイントを公開する必要があります。このエンドポイントを呼び出してモデル推論を行い、タスクが完了したら報酬を報告します。

ロールアウトリクエスト形式

エージェントは、Lambda フォワーダーから次の形式でリクエストを受け取ります。

{
    "prompt": "string",
    "metadata": {
        "jobArn": "string",
        "rolloutId": "string"
    },
    "inferenceParams": {
        "temperature": 1.0,
        "topP": 1.0,
        "maxTokens": 256
    }
}

エージェントの期待

Amazon SageMaker AI がロールアウトを正常に完了するには、エージェントが以下を行う必要があります。

次の例は、FastAPI エンドポイントを持つカスタムエージェントを示しています。

"""
Custom Lambda Agent for SageMaker RFT.

This is a minimal example showing the required integration points.
You can use any framework (Strands, LangChain, raw OpenAI, etc.)
or any HTTP framework (FastAPI, Flask, Django, etc.) for your agent.
"""
import os
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import Any, Dict
from sagemaker.train.rft import sagemaker_rft_handler, set_rollout_context
from sagemaker.train.rft.headers import make_inference_headers
from sagemaker.core.token_generator import generate_token

app = FastAPI()

@app.post("/rollout")
def rollout_endpoint(payload: dict):
    handle_rollout(payload)
    return {"status": "accepted"}

@sagemaker_rft_handler
def handle_rollout(payload):
    """Handle a rollout request from RFT."""
    prompt = payload.get("prompt", "")
    metadata = payload.get("metadata", {})
    inference_params = payload.get("inferenceParams", {})
    endpoint = metadata.get("endpoint", os.environ.get("RFT_RUNTIME_ENDPOINT", ""))

    # 1. Set rollout context and capture inference headers
    set_rollout_context(metadata)
    headers = make_inference_headers(metadata)

    # 2. Get bearer token for RFT Runtime authentication
    token = generate_token(region=os.environ.get("AWS_REGION", "us-west-2"))

    # 3. Run your agent logic
    result = run_agent(prompt, token, endpoint, inference_params, headers)

    # 4. Compute reward based on your task's success criteria
    reward = compute_reward(result)

    # 5. Report completion with reward to RFT Runtime
    # decorator handles CompleteRollout + UpdateReward
    return {"reward": reward}

def run_agent(prompt: str, token: str, endpoint: str, inference_params: dict, headers: dict) -> str:
    """Your agent logic goes here."""
    import openai
    client = openai.OpenAI(
        base_url=(endpoint or os.environ.get("RFT_RUNTIME_ENDPOINT", "")) + "/v1",
        api_key=token,
        default_headers=headers,
    )

    response = client.chat.completions.create(
        model="default",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=inference_params.get("maxTokens", 256),
        temperature=inference_params.get("temperature", 1.0),
        top_p=inference_params.get("topP", 1.0),
    )

    return response.choices[0].message.content

def compute_reward(result: str) -> float:
    """Implement your task-specific reward function here."""

@app.get("/health")
def health():
    return {"status": "ok"}

エージェント環境に必要な依存関係をインストールします。

pip install sagemaker-train fastapi uvicorn openai

エージェントをデプロイする

エージェントは、任意のコンピューティング環境で実行できます。推論と報酬レポートのために SageMaker AI ランタイムに到達するには、アウトバウンドネットワークアクセスが必要です。Lambda フォワーダーがエージェントを直接呼び出す場合、エージェントは Lambda から到達可能な HTTP エンドポイントも公開する必要があります。

次の手順では、カスタムエージェントを Amazon EKS にデプロイします。

EKS クラスターを作成する

EKS のセットアップ手順の詳細については、「EKS 入門ガイド」を参照してください。以下は、この例の最小限のセットアップです。

eksctl create cluster \
  --name external-agent \
  --region us-west-2 \
  --nodegroup-name agent-nodes \
  --node-type t3.medium \
  --nodes 1 \
  --nodes-min 1 \
  --nodes-max 2 \
  --managed

kubeconfig を更新してクラスターに接続します。

aws eks update-kubeconfig --name external-agent --region us-west-2

コンテナイメージを構築してプッシュします。

aws ecr create-repository --repository-name my-external-agent --region us-west-2

aws ecr get-login-password --region us-west-2 | \
  docker login --username AWS --password-stdin account-id.dkr.ecr.us-west-2.amazonaws.com

docker build -t my-external-agent .
docker tag my-external-agent:latest account-id.dkr.ecr.us-west-2.amazonaws.com/my-external-agent:latest
docker push account-id.dkr.ecr.us-west-2.amazonaws.com/my-external-agent:latest

EKS にデプロイする:

kubectl apply -f - <<EOF
apiVersion: apps/v1
kind: Deployment
metadata:
  name: external-agent
spec:
  replicas: 1
  selector:
    matchLabels:
      app: external-agent
  template:
    metadata:
      labels:
        app: external-agent
    spec:
      containers:
      - name: agent
        image: <account-id>.dkr.ecr.us-west-2.amazonaws.com/external-agent:latest
        ports:
        - containerPort: 8080
        env:
        - name: AWS_REGION
          value: us-west-2
        - name: RFT_RUNTIME_ENDPOINT
          value: https://job-runtime.sagemaker.us-west-2.api.aws
---
apiVersion: v1
kind: Service
metadata:
  name: external-agent
spec:
  selector:
    app: external-agent
  ports:
  - port: 80
    targetPort: 8080
  type: LoadBalancer
EOF

エージェントエンドポイントを取得します。

AGENT_ENDPOINT=$(kubectl get svc external-agent -o jsonpath='{.status.loadBalancer.ingress[0].hostname}')

echo "Agent endpoint: http://$AGENT_ENDPOINT"

デプロイを確認します。

curl -s "http://$AGENT_ENDPOINT/health"

# Expected: {"status": "ok"}

Lambda フォワーダーを作成する

Lambda フォワーダーは SageMaker AI からロールアウトリクエストを受け取り、エージェントに転送します。カスタマイズする主なものは、ロールアウトリクエストをエージェントの API 形式に変換する _call_agent()関数です。

エージェントにパブリック HTTP エンドポイントがない場合は、 _call_agent()の HTTP コールを SQS に置き換えsend_message、代わりにエージェントにキューをポーリングさせることができます。

"""
Lambda Template

Bridges SageMaker Job rollout requests to any agent platform with a public endpoint.
Implement _call_agent() with your platform-specific logic.

If your agent environment does not have a public endpoint, you can
replace the HTTP call with an SQS send_message to enqueue the request,
and have your agent poll the queue for work.

Env vars:
    AGENT_ENDPOINT  - target agent base URL
    AGENT_API_KEY   - API key for the target agent (prefer Secrets Manager)
"""

import json
import logging
import os
import re
import urllib.error
import urllib.request

logger = logging.getLogger()
logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))

AGENT_ENDPOINT = os.environ.get("AGENT_ENDPOINT", "")
AGENT_API_KEY = os.environ.get("AGENT_API_KEY", "")

_SAFE_ID = re.compile(r"^[\w\-.]+$")

# ---------------------------------------------------------------------------
# CUSTOMIZE THIS - translate rollout request to your platform's API
# ---------------------------------------------------------------------------
def _call_agent(prompt: str, metadata: dict, inference_params: dict):
    """
    Forward the prompt to your agent platform.
    Replace the body below with your platform's request format.
    """
    payload = json.dumps({
        "prompt": prompt,
        "metadata": metadata,
        "inferenceParams": inference_params,
    }).encode()

    req = urllib.request.Request(
        AGENT_ENDPOINT,
        data=payload,
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {AGENT_API_KEY}",
        },
        method="POST",
    )

    urllib.request.urlopen(req, timeout=120)


# ---------------------------------------------------------------------------
# Validation - no changes needed below
# ---------------------------------------------------------------------------
def _validate(event: dict) -> dict:
    body = json.loads(event["body"]) if isinstance(event.get("body"), str) else event

    prompt = body.get("prompt")
    if not isinstance(prompt, str) or not prompt.strip():
        raise ValueError("'prompt' is required and must be a non-empty string")

    meta = body.get("metadata")
    if not isinstance(meta, dict):
        raise ValueError("'metadata' is required")
    for key in ("jobArn", "rolloutId"):
        val = meta.get(key)
        if not isinstance(val, str) or not val.strip():
            raise ValueError(f"metadata.{key} must match [a-zA-Z0-9_\\-.]")

    params = body.get("inferenceParams") or {}
    if not isinstance(params, dict):
        raise ValueError("'inferenceParams' must be an object")

    return {
        "prompt": prompt.strip(),
        "metadata": meta,
        "inferenceParams": params,
    }


# ---------------------------------------------------------------------------
# CUSTOMIZE THIS - handle errors thrown from your agent environment
# ---------------------------------------------------------------------------
def _handle_agent_error(exc: Exception) -> dict:
    """
    Called when _call_agent() raises an exception. Customize this to map
    platform-specific errors to appropriate error types and messages.

    The return payload must follow this structure:
        {"errorType": "<type>", "errorMessage": "<description>"}

    Supported errorType values:
        ValidationError, InternalServerError, AccessDenied

    Examples:
          return {"errorType": "AccessDenied","errorMessage": "Agent denied access"}
          return {"errorType": "ValidationError","errorMessage": "Missing required field"}
          return {"errorType":"InternalServerError", "errorMessage": str(exc)}
    """

def handler(event, context):
    try:
        body = _validate(event)
    except ValueError as exc:
        logger.warning("Validation error: %s", exc)
        return {"errorType": "ValidationError", "errorMessage": str(exc)}

    try:
        result = _call_agent(body["prompt"], body["metadata"], body["inferenceParams"])
        logger.info("Rollout %s completed", body["metadata"]["rolloutId"])
        return {}
    except Exception as exc:
        return _handle_agent_error(exc)

Lambda 関数をデプロイします。

Lambda 関数をパッケージ化します。

zip lambda_forwarder.zip lambda_forwarder.py

関数を作成します。

aws lambda create-function \
  --function-name rft-agent-forwarder \
  --runtime python3.12 \
  --handler lambda_forwarder.handler \
  --role arn:aws:iam::account-id:role/RFTLambdaForwarderRole \
  --zip-file fileb://lambda_forwarder.zip \
  --timeout 600 \
  --environment "Variables={AGENT_ENDPOINT=http://$AGENT_ENDPOINT}" \
  --region us-west-2

Lambda をテストします。

aws lambda invoke \
  --function-name rft-agent-forwarder \
  --cli-binary-format raw-in-base64-out \
  --payload '{"prompt": "Plan a 3-day itinerary for a trip to Seattle", "metadata": {"jobArn": "arn:aws:sagemaker:us-west-2:123456789012:job/AgentRFT/test", "rolloutId": "roll-1"}}' \
  --region us-west-2 \
  /tmp/response.json && cat /tmp/response.json