Skip to content

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

Closed
xiaoxshe wants to merge 2 commits intoaws:masterfrom
xiaoxshe:feat/aws-rft-sdk
Closed

feat: add aws-rft-sdk package for Reinforcement Fine-Tuning#5596
xiaoxshe wants to merge 2 commits intoaws:masterfrom
xiaoxshe:feat/aws-rft-sdk

Conversation

@xiaoxshe
Copy link
Collaborator

@xiaoxshe xiaoxshe commented Mar 3, 2026

Summary

  • Add standalone aws-rft-sdk package that integrates SageMaker RFT (Reinforcement Fine-Tuning) with Strands agent framework
  • RolloutFeedbackClient: calls real CompleteTrajectory and UpdateReward APIs via SigV4-signed requests
  • @rft_handler: decorator to extract rollout metadata from payloads and manage RFT context
  • RFTContext: thread-local context that produces X-Rft-Job-Arn, X-Trajectory-Id, X-Span-Id headers
  • wrap_model: Strands model adapter that injects these headers into every inference call via extra_headers

API Integration

Uses the AgenticRFTRuntimeService alpha endpoint (finetuning-job-runtime.alpha.sagemaker.us-west-2.api.aws):

  • POST /CompleteTrajectory — marks trajectory as complete (PENDING -> READY)
  • POST /UpdateReward — submits per-transition rewards (READY -> REWARD_RECEIVED)
  • POST /sample — OpenAI-compatible inference proxy with trajectory capture via X-Rft-Job-Arn, X-Trajectory-Id, X-Span-Id headers
  • SigV4 signing with service name sagemaker

Test Results

Against alpha endpoint:

  • CompleteTrajectory: 404 for non-existent trajectory (API reachable, auth working)
  • UpdateReward: 404 for non-existent trajectory (API reachable, auth working)
  • SampleInvoke: 400/500 with validation errors (API reachable, requires valid job ARN)

Unit/Integration tests (5/5 pass):

  • RFTContext produces correct X-Rft-Job-Arn, X-Trajectory-Id, X-Span-Id headers
  • wrap_model injects headers into inner model's stream() call
  • RolloutFeedbackClient sends SigV4-signed CompleteTrajectory request
  • RolloutFeedbackClient sends SigV4-signed UpdateReward request
  • End-to-end: @rft_handler -> RFTContext -> wrap_model -> correct headers injected

Strands SDK compatibility:

  • 72/72 OpenAI model tests pass (no regressions from extra_headers change)

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
- client.py: implement CompleteTrajectory and UpdateReward using SigV4-
  signed HTTP calls to finetuning-job-runtime.alpha.sagemaker endpoint
- context.py: align header names with service API (X-Rft-Job-Arn,
  X-Trajectory-Id, X-Span-Id); auto-generate span ID per inference call
- pyproject.toml: add requests dependency

Tested against alpha endpoint:
- CompleteTrajectory: 404 for non-existent trajectory (API reachable)
- UpdateReward: 404 for non-existent trajectory (API reachable)
- SigV4 signing with service name 'sagemaker' confirmed working
- End-to-end header injection via wrap_model verified
@xiaoxshe xiaoxshe closed this Mar 5, 2026
@xiaoxshe xiaoxshe deleted the feat/aws-rft-sdk branch March 5, 2026 00:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant