Skip to content

feat: bedrock-oss-provisioned-throughput-polling#5893

Open
lucasjia-aws wants to merge 2 commits into
aws:masterfrom
lucasjia-aws:bedrock-oss-provisioned-throughput-polling
Open

feat: bedrock-oss-provisioned-throughput-polling#5893
lucasjia-aws wants to merge 2 commits into
aws:masterfrom
lucasjia-aws:bedrock-oss-provisioned-throughput-polling

Conversation

@lucasjia-aws
Copy link
Copy Markdown
Collaborator

@lucasjia-aws lucasjia-aws commented May 26, 2026

Description

Add provisioned throughput polling for Bedrock OSS (non-Nova) model deployments.

Previously, BedrockModelBuilder.deploy() for OSS models would fire off a create_model_import_job call and return immediately, leaving the user to manually wait for the import to complete and then create provisioned throughput themselves. This PR makes the non-Nova path synchronous end-to-end, matching the behavior of the Nova deployment path.

Changes

Implementation (bedrock_model_builder.py):

  • Non-Nova deploy() now waits for the import job to complete, then creates provisioned model throughput and waits for it to become InService before returning.
  • New public method: create_provisioned_throughput() — creates provisioned throughput for an imported model and polls until ready.
  • New private methods: _wait_for_import_job_complete(), _wait_for_provisioned_throughput_in_service().
  • New deploy() parameters for non-Nova path: provisioned_model_name, model_units, commitment_duration, provisioned_model_tags.

Unit tests (test_bedrock_model_builder.py):

  • Updated existing non-Nova deploy tests to reflect the new synchronous behavior.
  • Added test classes: TestWaitForImportJobComplete, TestCreateProvisionedThroughput, TestWaitForProvisionedThroughputInService.

Integration test (test_bedrock_provisioned_throughput.py):

  • New integ test file covering the full deploy flow and standalone create_provisioned_throughput() usage.

Testing

  • Unit tests cover all new methods (happy path, polling, failure, timeout).
  • Integration test validated locally against real Bedrock APIs — import job creation, polling to completion, and CreateProvisionedModelThroughput API call all succeeded (blocked only by account-level service quota, not code issues).

@lucasjia-aws lucasjia-aws force-pushed the bedrock-oss-provisioned-throughput-polling branch from 07f6443 to 456b788 Compare May 27, 2026 06:07
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