From 2c195a563bbd1aca6030bcc7f00079d006c76e4e Mon Sep 17 00:00:00 2001 From: Bryan Lu <55512809+luyuzhe111@users.noreply.github.com> Date: Fri, 6 Nov 2026 23:37:45 +0000 Subject: [PATCH] fix: preserve AgentCore context in RFT handlers --- .../src/sagemaker/train/rft/decorators.py | 16 +-- .../tests/unit/train/test_rft_decorators.py | 98 +++++++++++++++++++ 2 files changed, 107 insertions(+), 7 deletions(-) create mode 100644 sagemaker-train/tests/unit/train/test_rft_decorators.py diff --git a/sagemaker-train/src/sagemaker/train/rft/decorators.py b/sagemaker-train/src/sagemaker/train/rft/decorators.py index 4b5db6b69d..1dde58b903 100644 --- a/sagemaker-train/src/sagemaker/train/rft/decorators.py +++ b/sagemaker-train/src/sagemaker/train/rft/decorators.py @@ -25,11 +25,13 @@ def sagemaker_rft_handler(func: Callable) -> Callable: 3. On error, calls report_error 4. Clears context when done - Works with both sync and async functions. + Works with both sync and async functions accepting a ``payload`` dict. + Any additional positional and keyword arguments are forwarded unchanged + to the handler. Example:: - from sagemaker.train.rft import rft_handler + from sagemaker.train.rft import sagemaker_rft_handler from sagemaker.train.rft.adapters.strands import wrap_model from strands import Agent from strands.models.openai import OpenAIModel @@ -38,7 +40,7 @@ def sagemaker_rft_handler(func: Callable) -> Callable: agent = Agent(model=model, tools=[...]) @app.entrypoint - @rft_handler + @sagemaker_rft_handler async def invoke_agent(payload): result = await agent.invoke_async(payload["instance"]["prompt"]) return {"reward": compute_reward(result)} @@ -46,13 +48,13 @@ async def invoke_agent(payload): if asyncio.iscoroutinefunction(func): @wraps(func) - async def async_wrapper(payload: dict) -> Any: + async def async_wrapper(payload: dict, *args, **kwargs) -> Any: metadata = payload.get("metadata") or {} inference_params = payload.get("inferenceParams") or payload.get("inference_params") set_rollout_context(metadata, inference_params) feedback = RolloutFeedbackClient(metadata) try: - result = await func(payload) + result = await func(payload, *args, **kwargs) except Exception as e: error_str = str(e) if _is_trajectory_already_processed(error_str): @@ -77,13 +79,13 @@ async def async_wrapper(payload: dict) -> Any: else: @wraps(func) - def sync_wrapper(payload: dict) -> Any: + def sync_wrapper(payload: dict, *args, **kwargs) -> Any: metadata = payload.get("metadata") or {} inference_params = payload.get("inferenceParams") or payload.get("inference_params") set_rollout_context(metadata, inference_params) feedback = RolloutFeedbackClient(metadata) try: - result = func(payload) + result = func(payload, *args, **kwargs) except Exception as e: error_str = str(e) if _is_trajectory_already_processed(error_str): diff --git a/sagemaker-train/tests/unit/train/test_rft_decorators.py b/sagemaker-train/tests/unit/train/test_rft_decorators.py new file mode 100644 index 0000000000..f542fd4d32 --- /dev/null +++ b/sagemaker-train/tests/unit/train/test_rft_decorators.py @@ -0,0 +1,98 @@ +"""Unit tests for RFT handler argument forwarding.""" + +import asyncio +import json +from functools import wraps + +import pytest + +from sagemaker.train.rft import RolloutFeedbackClient, sagemaker_rft_handler +from sagemaker.train.rft.context import get_inference_params, get_rollout_context + + +@pytest.fixture(params=["sync", "async"]) +def decorate(request): + def decorator(handler): + if request.param == "async": + + @wraps(handler) + async def async_handler(*args, **kwargs): + return handler(*args, **kwargs) + + return sagemaker_rft_handler(async_handler) + return sagemaker_rft_handler(handler) + + return decorator + + +@pytest.fixture +def feedback(monkeypatch): + reports = {} + + def record(self, path, body): + reports[path] = json.loads(body) + + monkeypatch.setattr(RolloutFeedbackClient, "_bearer_post", record) + return reports + + +@pytest.mark.parametrize("with_context", [False, True]) +def test_handler_invocation(decorate, feedback, with_context): + def handler(payload): + return {"reward": payload["score"]} + + def handler_with_context(payload, context): + return {**handler(payload), "context": context} + + decorated = decorate(handler_with_context if with_context else handler) + payload = {"score": 0.75, "metadata": {"jobArn": "test-job", "rolloutId": "test-rollout"}} + context = object() + + async def invoke(): + result = decorated(payload, context) if with_context else decorated(payload) + return await result if asyncio.iscoroutine(result) else result + + result = asyncio.run(invoke()) + assert result["reward"] == 0.75 + if with_context: + assert result["context"] is context + assert feedback["/complete-rollout"]["Status"] == "ready" + assert feedback["/update-reward"]["Rewards"] == [0.75] + + +@pytest.mark.parametrize("fail", [False, True]) +def test_keyword_context_and_rollout_lifecycle(decorate, feedback, fail): + context = object() + metadata = {"jobArn": "test-job", "rolloutId": "test-rollout"} + inference_params = {"temperature": 0.5} + + def handler(payload, context=None): + assert get_rollout_context() == metadata + assert get_inference_params() == inference_params + if fail: + raise ValueError("agent failed") + return {"reward": payload["score"], "context": context} + + decorated = decorate(handler) + payload = {"score": 0.75, "metadata": metadata, "inferenceParams": inference_params} + + async def invoke(): + try: + result = decorated(payload, context=context) + if asyncio.iscoroutine(result): + result = await result + return result + finally: + assert get_rollout_context() is None + assert get_inference_params() is None + + if fail: + with pytest.raises(ValueError, match="agent failed"): + asyncio.run(invoke()) + else: + result = asyncio.run(invoke()) + assert result["context"] is context + assert result["reward"] == 0.75 + + assert feedback["/complete-rollout"]["Status"] == ("failed" if fail else "ready") + assert feedback["/update-reward"]["Rewards"] == [0.0 if fail else 0.75]