sagemaker.train.rft.decorators

sagemaker.train.rft.decorators#

Decorators for RFT integration.

Provides sagemaker_rft_handler decorator for AgentCore Runtime entrypoints.

Functions

sagemaker_rft_handler(func)

Decorator for AgentCore Runtime entrypoints to handle RFT rollout lifecycle.
sagemaker.train.rft.decorators.sagemaker_rft_handler(func: Callable) → Callable[source]#

Decorator for AgentCore Runtime entrypoints to handle RFT rollout lifecycle.

Automatically: 1. Sets rollout context (metadata + inference_params) for header injection 2. On success, calls CompleteTrajectory + UpdateReward if “reward” in result 3. On error, calls report_error 4. Clears context when done

Works with both sync and async functions.

Example:

from sagemaker.train.rft import rft_handler
from sagemaker.train.rft.adapters.strands import wrap_model
from strands import Agent
from strands.models.openai import OpenAIModel

model = wrap_model(OpenAIModel(client_args={...}, model_id="my-model"))
agent = Agent(model=model, tools=[...])

@app.entrypoint
@rft_handler
async def invoke_agent(payload):
    result = await agent.invoke_async(payload["instance"]["prompt"])
    return {"reward": compute_reward(result)}
Contents