feat: add aws-rft-sdk package for Reinforcement Fine-Tuning#5596
Closed
xiaoxshe wants to merge 2 commits intoaws:masterfrom
Closed
feat: add aws-rft-sdk package for Reinforcement Fine-Tuning#5596xiaoxshe wants to merge 2 commits intoaws:masterfrom
xiaoxshe wants to merge 2 commits intoaws:masterfrom
Conversation
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
7a8081c to
9615fc9
Compare
- 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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
aws-rft-sdkpackage that integrates SageMaker RFT (Reinforcement Fine-Tuning) with Strands agent frameworkRolloutFeedbackClient: calls real CompleteTrajectory and UpdateReward APIs via SigV4-signed requests@rft_handler: decorator to extract rollout metadata from payloads and manage RFT contextRFTContext: thread-local context that producesX-Rft-Job-Arn,X-Trajectory-Id,X-Span-Idheaderswrap_model: Strands model adapter that injects these headers into every inference call viaextra_headersAPI 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 viaX-Rft-Job-Arn,X-Trajectory-Id,X-Span-IdheaderssagemakerTest Results
Against alpha endpoint:
Unit/Integration tests (5/5 pass):
Strands SDK compatibility: