Skip to content

feat(serve): add opt-in model source tag-based resource reuse#5993

Open
amazeAmazing wants to merge 2 commits into
aws:master-nova-follow-upsfrom
amazeAmazing:model-reuse-only
Open

feat(serve): add opt-in model source tag-based resource reuse#5993
amazeAmazing wants to merge 2 commits into
aws:master-nova-follow-upsfrom
amazeAmazing:model-reuse-only

Conversation

@amazeAmazing

@amazeAmazing amazeAmazing commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Model Source Tag-Based Resource Reuse

Summary

Adds opt-in resource reuse to the V3 serve builders so redeploying the same model artifacts doesn't create duplicate resources (or exhaust quota-limited Bedrock imported models). Each deploy tags the resource it creates with a model-source identifier; when reuse is requested, a later call for the same source discovers and returns the existing resource instead of creating a new one.

Controlled by reuse_resources: bool = False. On a hit we log a warning and return the existing resource

  • Nova checkpoint URI resolution surfaces both the manifest.json and output.tar.gz errors on failure, instead of masking the primary (raw-manifest) failure with a misleading tar.gz message.

Behavior

BedrockModelBuilder ModelBuilder (SageMaker)
Flag deploy(reuse_resources=) build(reuse_resources=) + deploy(reuse_resources=)
Discovery Bedrock list_custom_models + tags SageMaker list_endpoints + tags
Reuses custom model + active deployment endpoint (build also skips Model, sets built_model to the existing one)
Config re-validation no yes (env vars, image URI, instance type)
On match warn + return warn + return

reuse_resources is not inherited: set it on build() to skip Model creation and on deploy() to reuse the endpoint.

Inference components: IC deployments (inference_config is a ResourceRequirements, or a modelbuilder_list build) are not intercepted by reuse_resources. They manage their own reuse by IC name (a new component is added to the endpoint, or an existing one is updated in place), so passing reuse_resources=True never silently skips an IC create/update. deploy() logs a warning that the flag has no effect for IC deployments.

Nova inference-component deployment

Nova model-customization deploys supports inference components. When a ResourceRequirements inference_config is supplied, the Nova checkpoint is hosted as a single inference component referencing the built Model (which carries the Nova image, escrow artifacts, and env); without one, it keeps the direct model-on-variant path. All supported Nova checkpoints — full-rank custom models and LoRA-merged models — use this single-IC path.

  • _is_nova_model also identifies Nova from a package-less source (raw S3 checkpoint or trainer) via base_model_name, not just the model package recipe/hub-content name.
  • The IC endpoint config sets EnableNetworkIsolation to match the built Model (always True for Nova), fixing a CreateInferenceComponent rejection on mismatched network isolation.
  • Model-package-dependent logic (restricted-package path, PEFT/recipe metadata, lineage tracking) is guarded so package-less Nova checkpoints deploy cleanly.
  • Endpoints created on the shared IC path are tagged with the model-source tag so they stay discoverable.

Key changes

  • New sagemaker/serve/model_reuse.py: tag helpers + service-client discovery (find_existing_bedrock_model, find_active_bedrock_deployment_for_model,
    find_existing_sagemaker_endpoint) with status gating.
  • sagemaker/core/training/utils.py: shared Nova manifest/checkpoint helpers (consolidated from duplicated logic in both builders).
  • ModelBuilder: build()/deploy() reuse gates, config validation (PrimaryContainer with Containers[0] fallback for Nova), raw-S3 model
    input via model_metadata={"BASE_MODEL_NAME": ...}, model-source tagging.
    Reuse gates are skipped for inference-component builds/deploys so IC
    create/update is never intercepted. Discovery reuses the cached
    sagemaker_session.sagemaker_client / _get_bedrock_client().
  • BedrockModelBuilder: reuse of custom model + deployment; consolidated source resolution; response always includes modelArn.

Tests

  • Unit: test_model_reuse.py, additions to test_bedrock_model_builder.py / test_model_builder.py, and core test_training_utils.py.
  • Integ (this package): updated to verify tag usage
  • Manually tested Sagemaker and Bedrock deployments

"""
manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
try:
return read_nova_checkpoint_uri_from_manifest(s3_client, manifest_uri)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what if there's no manifest.json saved in the output folder and it is not a tar.gz file (for hyperpod)? wouldn't we in that case incorrectly throw an error that manifest file not found in the tar.gz output path?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the callout! I will adjust

Add reuse_resources to ModelBuilder.build/deploy and BedrockModelBuilder.deploy.
On a hit, discover an existing resource by the model-source tag and return it
instead of creating a duplicate (warn, do not raise). Honored per call.

- New sagemaker/serve/model_reuse.py: tag helpers + service-client discovery
- Consolidate Nova manifest/checkpoint reading into sagemaker/core/training/utils.py
- SageMaker: build() skips Model creation on reuse (sets built_model to the
  existing Model); deploy() reuses the endpoint after validating env vars/image/
  instance type (PrimaryContainer with Containers[0] fallback for Nova)
- Reuse gates are skipped for inference-component builds/deploys so IC
  create/update (via _deploy_for_ic) is never silently intercepted
- Bedrock: reuse custom model + active deployment; response includes modelArn
- Reuse discovery uses the cached session/bedrock clients
- Support raw S3 URI model input via model_metadata BASE_MODEL_NAME
- Unit tests + notebook examples
…euse

Route Nova model-customization deploys through the shared single-inference
-component path when a ResourceRequirements inference_config is supplied, so
each Nova checkpoint (full-rank or LoRA-merged) is hosted as one inference
component referencing the built Model. Nova without an inference_config keeps
the direct model-on-variant path.

- Broaden _is_nova_model to identify Nova from a package-less source (raw S3
  checkpoint or trainer) via base_model_name, in addition to the model
  package recipe/hub-content name.
- Set EnableNetworkIsolation on the IC endpoint config to match the built
  Model (always True for Nova), fixing CreateInferenceComponent rejection on
  mismatched network isolation.
- Guard model-package-dependent logic (restricted-package path, PEFT/recipe
  metadata, lineage tracking) so package-less Nova checkpoints deploy cleanly.
- Apply accumulated tags (including the model-source reuse tag) to endpoints
  created on the shared IC path so they remain discoverable.
- build(reuse_resources=True) only short-circuits when the backing Model can
  be resolved; IC endpoints and stale/deleted configs fall through and build
  a real Model, preventing a None built_model on later IC deploys.
- deploy() warns that reuse_resources has no effect for inference-component
  deployments, which manage their own reuse by component name.
- Surface both the manifest.json and output.tar.gz errors when Nova
  checkpoint URI resolution fails, instead of masking the primary failure.

Add unit tests covering the Nova IC path (routing, network isolation, IC
spec) and the model-on-variant fallback.
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