Skip to content

fix(serve): BedrockModelBuilder accepts BaseTrainer as model input - #6104

Open
amazeAmazing wants to merge 1 commit into
aws:master-nova-follow-upsfrom
amazeAmazing:fix/bedrock-model-builder-trainer-input
Open

fix(serve): BedrockModelBuilder accepts BaseTrainer as model input#6104
amazeAmazing wants to merge 1 commit into
aws:master-nova-follow-upsfrom
amazeAmazing:fix/bedrock-model-builder-trainer-input

Conversation

@amazeAmazing

Copy link
Copy Markdown
Contributor

BedrockModelBuilder(model=trainer) now works for SFTTrainer, DPOTrainer, RLVRTrainer, etc. Previously only ModelTrainer and S3 URIs were supported.

Issue #, if available:

Description of changes:

fix(serve): BedrockModelBuilder accepts BaseTrainer as model input

Problem

BedrockModelBuilder(model=trainer) failed when trainer was an SFTTrainer,
DPOTrainer, RLVRTrainer, or any BaseTrainer subclass. Users had to extract the
S3 checkpoint URI manually:

# Worked
builder = BedrockModelBuilder(model="s3://bucket/checkpoint/step_10")

# Didn't work
builder = BedrockModelBuilder(model=trainer)

Changes

  • Added BaseTrainer branch in _fetch_model_package() to resolve
    output_model_package_arn from _latest_training_job
  • Added manifest fallback in _get_s3_artifacts() for HyperPod jobs where
    model_artifacts is not populated
  • Updated type annotation and docstrings to include BaseTrainer

Testing

Unit tests:

Integration tests by resolution path:

Category Model Compute Trainer Resolution path
Nova (SMTJ) nova-micro (DPO) TrainingJobCompute DPOTrainer model_artifacts (S3 direct)
Nova (HyperPod) nova-lite-v2 (RLVR) HyperPodCompute RLVRTrainer manifest fallback
Nova (Bedrock deploy) nova-lite-v2 (RLVR) RLVRTrainer manifest fallback → create_custom_model
OSS (SMTJ) gpt-oss-20b (SFT) TrainingJobCompute SFTTrainer model_package → hf_merged probe
OSS (TrainingJob direct) gpt-oss-20b (SFT) TrainingJobCompute model_artifacts (S3 direct)

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

BedrockModelBuilder(model=trainer) now works for SFTTrainer, DPOTrainer,
RLVRTrainer, etc. Previously only ModelTrainer and S3 URIs were supported.
@amazeAmazing
amazeAmazing marked this pull request as ready for review July 27, 2026 17:33
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.

2 participants