feat: add aws-rft-sdk package for Reinforcement Fine-Tuning

xiaoxshe · xiaoxshe · commit 9615fc90afae · 2026-03-03T22:00:51.000Z
Add a standalone SDK package that integrates SageMaker RFT (Reinforcement
Fine-Tuning) with Strands agent framework. Provides:

- RolloutFeedbackClient: report rewards back to the training service
- @rft_handler: decorator to extract rollout metadata from payloads
- RFTContext: thread-local context for propagating training metadata
- wrap_model: Strands model adapter that injects X-RFT-* headers
diff --git a/aws-rft-sdk/pyproject.toml b/aws-rft-sdk/pyproject.toml
@@ -0,0 +1,19 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "aws-rft-sdk"
+version = "0.1.0"
+description = "AWS Reinforcement Fine-Tuning SDK for online rollout-based training"
+readme = {text = "", content-type = "text/markdown"}
+requires-python = ">=3.9"
+dependencies = [
+    "boto3>=1.35.0",
+]
+
+[project.optional-dependencies]
+strands = ["strands-agents>=0.1.0"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/aws_rft_sdk"]
diff --git a/aws-rft-sdk/src/aws_rft_sdk/__init__.py b/aws-rft-sdk/src/aws_rft_sdk/__init__.py
@@ -0,0 +1,5 @@
+from aws_rft_sdk.client import RolloutFeedbackClient
+from aws_rft_sdk.handler import rft_handler
+from aws_rft_sdk.context import RFTContext
+
+__all__ = ["RolloutFeedbackClient", "rft_handler", "RFTContext"]
diff --git a/aws-rft-sdk/src/aws_rft_sdk/adapters/__init__.py b/aws-rft-sdk/src/aws_rft_sdk/adapters/__init__.py
diff --git a/aws-rft-sdk/src/aws_rft_sdk/adapters/strands.py b/aws-rft-sdk/src/aws_rft_sdk/adapters/strands.py
@@ -0,0 +1,78 @@
+"""Strands model adapter — wraps a Strands model to inject RFT headers.
+
+Usage::
+
+    from aws_rft_sdk.adapters.strands import wrap_model
+    from strands.models.openai import OpenAIModel
+
+    model = OpenAIModel(
+        client_args={"api_key": key, "base_url": endpoint},
+        model_id="my-model",
+    )
+    model = wrap_model(model)  # Now injects X-RFT-* headers on every call
+
+Requires the Strands OpenAIModel to pass through ``extra_headers`` kwarg
+to the underlying OpenAI client (supported since strands-agents >= X.Y.Z).
+"""
+
+import logging
+from typing import Any
+
+from aws_rft_sdk.context import RFTContext
+
+logger = logging.getLogger(__name__)
+
+
+def wrap_model(model: Any) -> Any:
+    """Wrap a Strands model to automatically inject RFT training headers.
+
+    The wrapper reads the current rollout context (populated by ``@rft_handler``)
+    and adds ``X-RFT-*`` headers to every inference request so the training
+    inference endpoint can correlate requests with rollouts.
+
+    Args:
+        model: A Strands model instance (e.g., ``OpenAIModel``).
+
+    Returns:
+        A wrapped model that transparently injects RFT headers.
+    """
+    return _RFTModelWrapper(model)
+
+
+class _RFTModelWrapper:
+    """Transparent proxy that injects RFT headers into Strands model calls.
+
+    Delegates all attribute access to the inner model so it quacks like
+    the original. Intercepts ``stream()`` to inject ``extra_headers``.
+    """
+
+    def __init__(self, inner_model: Any):
+        object.__setattr__(self, "_inner", inner_model)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._inner, name)
+
+    def __setattr__(self, name: str, value: Any):
+        if name == "_inner":
+            object.__setattr__(self, name, value)
+        else:
+            setattr(self._inner, name, value)
+
+    def stream(self, *args: Any, **kwargs: Any) -> Any:
+        """Intercept stream() to inject RFT headers via extra_headers kwarg."""
+        rft_headers = RFTContext.get_headers()
+        if rft_headers:
+            existing = kwargs.get("extra_headers") or {}
+            existing.update(rft_headers)
+            kwargs["extra_headers"] = existing
+            logger.debug("Injected RFT headers: %s", list(rft_headers.keys()))
+        return self._inner.stream(*args, **kwargs)
+
+    def update_config(self, **model_config: Any) -> None:
+        return self._inner.update_config(**model_config)
+
+    def get_config(self) -> Any:
+        return self._inner.get_config()
+
+    def structured_output(self, *args: Any, **kwargs: Any) -> Any:
+        return self._inner.structured_output(*args, **kwargs)
diff --git a/aws-rft-sdk/src/aws_rft_sdk/client.py b/aws-rft-sdk/src/aws_rft_sdk/client.py
@@ -0,0 +1,91 @@
+"""RolloutFeedbackClient — reports rewards and completion status back to the training service."""
+
+import logging
+from typing import Optional
+
+import boto3
+
+logger = logging.getLogger(__name__)
+
+
+class RolloutFeedbackClient:
+    """Client for reporting rollout feedback (rewards) to the RFT training service.
+
+    Typically used inside an @rft_handler-decorated entrypoint to report
+    the reward computed from the agent's rollout.
+
+    Example::
+
+        from aws_rft_sdk import RolloutFeedbackClient
+
+        client = RolloutFeedbackClient(payload.get("metadata"))
+        client.report_complete(reward=0.85)
+
+    Args:
+        metadata: The ``metadata`` dict from the rollout payload. Contains
+            training_job_arn, rollout_id, feedback_endpoint, etc.
+    """
+
+    def __init__(self, metadata: dict):
+        self._metadata = metadata or {}
+        self._training_job_arn = self._metadata.get("training_job_arn")
+        self._rollout_id = self._metadata.get("rollout_id")
+        self._feedback_endpoint = self._metadata.get("feedback_endpoint")
+        self._client = None
+
+    def _get_client(self):
+        if self._client is None:
+            kwargs = {}
+            if self._feedback_endpoint:
+                kwargs["endpoint_url"] = self._feedback_endpoint
+            self._client = boto3.client("sagemaker", **kwargs)
+        return self._client
+
+    def report_complete(self, reward: float):
+        """Report successful rollout completion with a reward score.
+
+        Args:
+            reward: The computed reward for this rollout (typically 0.0–1.0).
+        """
+        logger.info(
+            "Reporting rollout complete: training_job=%s rollout=%s reward=%s",
+            self._training_job_arn,
+            self._rollout_id,
+            reward,
+        )
+        # TODO: Replace with actual RFT feedback API call when available.
+        # The service API will accept:
+        #   - TrainingJobArn
+        #   - RolloutId
+        #   - Reward (float)
+        #   - Status (COMPLETED)
+        client = self._get_client()
+        # Placeholder — actual API TBD
+        # client.send_rollout_feedback(
+        #     TrainingJobArn=self._training_job_arn,
+        #     RolloutId=self._rollout_id,
+        #     Reward=reward,
+        #     Status="COMPLETED",
+        # )
+
+    def report_error(self, error: str, reward: Optional[float] = None):
+        """Report a rollout error.
+
+        Args:
+            error: Error description.
+            reward: Optional partial reward (defaults to 0.0).
+        """
+        logger.error(
+            "Reporting rollout error: training_job=%s rollout=%s error=%s",
+            self._training_job_arn,
+            self._rollout_id,
+            error,
+        )
+        # TODO: Replace with actual RFT feedback API call when available.
+        # client.send_rollout_feedback(
+        #     TrainingJobArn=self._training_job_arn,
+        #     RolloutId=self._rollout_id,
+        #     Reward=reward or 0.0,
+        #     Status="FAILED",
+        #     ErrorMessage=error,
+        # )
diff --git a/aws-rft-sdk/src/aws_rft_sdk/context.py b/aws-rft-sdk/src/aws_rft_sdk/context.py
@@ -0,0 +1,45 @@
+"""Thread-local context for RFT rollout metadata.
+
+The rft_handler decorator populates this context from the payload metadata.
+The Strands model wrapper reads it to inject per-request headers.
+"""
+
+import threading
+from typing import Optional
+
+_context = threading.local()
+
+
+class RFTContext:
+    """Access the current RFT rollout context.
+
+    Set by @rft_handler, read by wrap_model adapters to inject headers.
+    """
+
+    @staticmethod
+    def get_headers() -> dict:
+        """Return HTTP headers for the current rollout context."""
+        metadata = getattr(_context, "metadata", None)
+        if metadata is None:
+            return {}
+        headers = {}
+        if metadata.get("training_job_arn"):
+            headers["X-RFT-Training-Job-Arn"] = metadata["training_job_arn"]
+        if metadata.get("rollout_id"):
+            headers["X-RFT-Rollout-Id"] = metadata["rollout_id"]
+        if metadata.get("episode_id"):
+            headers["X-RFT-Episode-Id"] = metadata["episode_id"]
+        return headers
+
+    @staticmethod
+    def get_metadata() -> Optional[dict]:
+        """Return the raw metadata dict, or None if not in an RFT context."""
+        return getattr(_context, "metadata", None)
+
+
+def _set_metadata(metadata: dict):
+    _context.metadata = metadata
+
+
+def _clear_metadata():
+    _context.metadata = None
diff --git a/aws-rft-sdk/src/aws_rft_sdk/handler.py b/aws-rft-sdk/src/aws_rft_sdk/handler.py
@@ -0,0 +1,70 @@
+"""@rft_handler decorator — wraps an entrypoint to manage RFT rollout context."""
+
+import asyncio
+import functools
+import inspect
+import logging
+
+from aws_rft_sdk.client import RolloutFeedbackClient
+from aws_rft_sdk.context import _set_metadata, _clear_metadata
+
+logger = logging.getLogger(__name__)
+
+
+def rft_handler(func):
+    """Decorator that sets up RFT rollout context around an entrypoint.
+
+    Extracts ``metadata`` from the payload, makes it available via
+    ``RFTContext.get_headers()`` (used by ``wrap_model``), and auto-reports
+    errors if the function raises.
+
+    Works with both sync and async functions.
+
+    Example::
+
+        @app.entrypoint
+        @rft_handler
+        async def invoke_agent(payload):
+            user_input = payload.get("instance")
+            response = await agent.invoke_async(user_input)
+            return response.message["content"][0]["text"]
+    """
+
+    if asyncio.iscoroutinefunction(func):
+
+        @functools.wraps(func)
+        async def async_wrapper(payload, *args, **kwargs):
+            metadata = payload.get("metadata", {}) if isinstance(payload, dict) else {}
+            _set_metadata(metadata)
+            try:
+                return await func(payload, *args, **kwargs)
+            except Exception as e:
+                logger.error("RFT rollout failed: %s", e)
+                try:
+                    RolloutFeedbackClient(metadata).report_error(str(e))
+                except Exception:
+                    logger.exception("Failed to report rollout error")
+                raise
+            finally:
+                _clear_metadata()
+
+        return async_wrapper
+    else:
+
+        @functools.wraps(func)
+        def sync_wrapper(payload, *args, **kwargs):
+            metadata = payload.get("metadata", {}) if isinstance(payload, dict) else {}
+            _set_metadata(metadata)
+            try:
+                return func(payload, *args, **kwargs)
+            except Exception as e:
+                logger.error("RFT rollout failed: %s", e)
+                try:
+                    RolloutFeedbackClient(metadata).report_error(str(e))
+                except Exception:
+                    logger.exception("Failed to report rollout error")
+                raise
+            finally:
+                _clear_metadata()
+
+        return sync_wrapper
