Skip to content

fix: resolve private hub Models and aliased references for modelTrainer - #6201

Open
tanvikab4 wants to merge 1 commit into
aws:masterfrom
tanvikab4:fix/jumpstart-private-hub-modeltrainer
Open

fix: resolve private hub Models and aliased references for modelTrainer#6201
tanvikab4 wants to merge 1 commit into
aws:masterfrom
tanvikab4:fix/jumpstart-private-hub-modeltrainer

Conversation

@tanvikab4

Copy link
Copy Markdown

Description of changes

Summary

ModelTrainer.from_jumpstart_config(...) could not resolve models that live in a private hub — only public
JumpStart models and (by a fragile assumption) plain private-hub references worked. This change fixes hub-content
resolution so ModelTrainer reaches parity with ModelBuilder, supporting:

  • Public JumpStart models (unchanged)
  • Model References in a private hub (pointer to a public model)
  • Aliased references — filed under a hub content name that differs from the public model_id
  • Privately-owned Models authored directly into a private hub

Root cause

sagemaker.core.jumpstart.document.get_hub_content_and_document() guessed the hub content type from the hub name:

hub_content_type = "Model" if hub_name == SAGEMAKER_PUBLIC_HUB else "ModelReference"

A private hub can hold either a Model or a ModelReference. This guess meant:

  • Privately-owned Models were looked up as ModelReference → ResourceNotFound → resolution failed.
  • The lookup used model_id and ignored hub_content_name, so aliased references were never found.

Fix (sagemaker-core/src/sagemaker/core/jumpstart/document.py)

  • Replace the guess with a probe: for a private hub, try ModelReference first, then fall back to Model; the
    public hub uses Model only. This mirrors ModelBuilder's resolution in accessors.py.
  • Honor hub_content_name (falling back to model_id) so aliased references resolve.
  • On miss, raise a combined error naming both content types attempted.

No changes were needed elsewhere: defaults.py already attaches HubAccessConfig based on hub_content_type, and
model_trainer.py / JumpStartConfig already support hub_name/hub_content_name — they become correct automatically
once the content type is resolved honestly.

Testing

Unit (sagemaker-core/tests/unit/jumpstart/test_document.py) — 5 new tests, all passing:

  • public hub resolves as Model (single lookup, no probe)
  • private-hub reference resolves on the first probe
  • private-hub Model resolves via the fallback (asserts probe order ["ModelReference", "Model"])
  • hub_content_name alias is used for lookup
  • neither type present → raises after attempting both

Integration (sagemaker-train/tests/integ/jumpstart/test_jumpstart_train.py) — 3 new tests, each creates a
temporary private hub, runs a real training job, and tears down (skips gracefully without hub permissions). All
verified passing end-to-end against AWS:

  • test_jumpstart_train_from_private_hub_reference
  • test_jumpstart_train_from_aliased_reference
  • test_jumpstart_train_from_private_owned_model

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