fix: resolve private hub Models and aliased references for modelTrainer by tanvikab4 · Pull Request #6201 · aws/sagemaker-python-sdk

tanvikab4 · 2026-08-21T16:55:31Z

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

Narrohag · 2026-08-25T16:05:40Z

+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        try:
+            resp = sm.list_hub_contents(HubName=hub_name, HubContentType=content_type)


can we use describe_hub_content directly?

switched _wait_for_content to use describe_hub_content directly

Narrohag · 2026-08-25T16:06:32Z

+
+def _default_training_dataset(region, model_id):
+    """Resolve the model's default training dataset S3 URI from JS metadata."""
+    from sagemaker.core.jumpstart.accessors import JumpStartModelsAccessor


let's avoid the lazy imports, claude loves to add them for some reason lol

Narrohag · 2026-08-25T16:08:49Z

+        compute=Compute(instance_type="ml.m5.xlarge"),
+        sagemaker_session=sagemaker_session,
+    )
+    model_trainer.train(input_data_config=[InputData(channel_name="training", data_source=dataset)])


is there anything to assert or just validating nothing is thrown?

This test's purpose is to verify the aliased-reference resolution path so when hub_content_name differs from the public model_id, from_jumpstart_config resolves the reference by its alias. I added these assertions to strengthen the test:

_jumpstart_config.hub_content_name == ALIASED_REFERENCE_NAME (the alias was threaded through resolution)

training_image is set

the model channel's S3 source carries a HubAccessConfig with a hub_content_arn

Narrohag · 2026-08-25T16:10:22Z

    model_trainer.train()
+
+
+def test_jumpstart_train_from_private_hub_reference(private_hub, sagemaker_session):


can you add a unit test for training a model reference to a gated model? there should have been a similar test in v2 so you can use that same model. There's some ModelAccessConfig/accept_eula stuff that we should verify works

Added a unit test class TestJumpStartTrainDefaultsGatedModelReferenceEula in test_defaults.py using the same gated model as v2 (mocks the resolver seam and verifies the
ModelAccessConfig/accept_eula behavior)

Sorry, meant integ test. My bad

might have missed it but do we have this as an integ test?

I looked into the v2 tests and there exists a gated private hub test on the ModelBuilder/deploy side, but not one for the training path. I'll add a corresponding integ test for ModelTrainer rn

sahilper · 2026-08-25T17:50:49Z

+            HubDescription="SDK integ test JumpStart training private hub",
+        )
+    except ClientError as e:
+        pytest.skip(f"Cannot create private hub (missing permissions?): {e}")


Skipping gracefully without hub permissions is the right intent, but the implementation is wider than that: every setup step here is a bare except ClientError → skip (create_hub, create_hub_content_reference, import_hub_content, describe_hub_content, and both _wait_for_content timeouts).

The assertion itself fails loudly — from_jumpstart_config isn't wrapped. But if import_hub_content for a private-hub Model breaks service-side, test_jumpstart_train_from_private_owned_model skips and CI stays green, so the test for the core case in this fix silently stops running.

Suggest matching only the specific authorization codes you expect in a restricted account and letting everything else fail; _wait_for_content returning False should be a pytest.fail once the hub was creatable.

Added a _skip_if_unauthorized helper that skips only on an authorization allowlist (AccessDeniedException/AccessForbiddenException/UnauthorizedOperation) and re-raises everything else

sahilper · 2026-08-25T17:52:40Z

+        compute=Compute(instance_type="ml.m5.xlarge"),
+        sagemaker_session=sagemaker_session,
+    )
+    model_trainer.train(input_data_config=[InputData(channel_name="training", data_source=dataset)])


This explicit training channel bypasses the hub-aware channel construction the fix enables. _create_training_job_args merges with "method parameter taking precedence" (existing_channels[new_input.channel_name] = new_input), so this bare InputData replaces the channel built by JumpStartTrainDefaults.get_training_dataset_input — the only place ModelAccessConfig and HubAccessConfig(hub_content_arn=...) are attached (defaults.py, hub_content_type == "ModelReference" branch).

So the reference and alias tests never assert that a HubAccessConfig was derived from the resolved hub content, and with a non-gated model like catboost-regression-model they'd pass even if that plumbing were wrong.

The existing test_jumpstart_train calls train() with no arguments and relies on SDK-resolved channels — suggest the same here and dropping _default_training_dataset, which also removes the hardcoded jumpstart-cache-prod-{region} bucket.

Dropped explicit channel and _default_training_dataset. Train() now uses SDK-resolved channels like test_jumpstart_train. Also added assertions on the resolved channels so the plumbing is guarded

tanvikab4 had a problem deploying to manual-approval August 21, 2026 16:55 — with GitHub Actions Error

tanvikab4 had a problem deploying to manual-approval August 24, 2026 14:18 — with GitHub Actions Error

tanvikab4 had a problem deploying to manual-approval August 24, 2026 14:19 — with GitHub Actions Error

tanvikab4 requested a deployment to manual-approval August 25, 2026 14:23 — with GitHub Actions Waiting

tanvikab4 had a problem deploying to manual-approval August 25, 2026 14:23 — with GitHub Actions Error

Narrohag reviewed Aug 25, 2026

View reviewed changes

sahilper reviewed Aug 25, 2026

View reviewed changes

tanvikab4 had a problem deploying to manual-approval August 25, 2026 18:47 — with GitHub Actions Error

tanvikab4 requested a review from sahilper August 25, 2026 19:21

tanvikab4 had a problem deploying to manual-approval August 27, 2026 16:07 — with GitHub Actions Error

tanvikab4 force-pushed the fix/jumpstart-private-hub-modeltrainer branch from b96a38b to b0a2b7b Compare August 27, 2026 16:08

tanvikab4 had a problem deploying to manual-approval August 27, 2026 16:08 — with GitHub Actions Error

sahilper approved these changes Aug 27, 2026

View reviewed changes

tanvikab4 requested a review from Narrohag August 27, 2026 17:05

fix: resolve private hub Models and aliased references for ModelTrainer

f66b654

tanvikab4 force-pushed the fix/jumpstart-private-hub-modeltrainer branch from b0a2b7b to f66b654 Compare August 27, 2026 19:54

tanvikab4 had a problem deploying to manual-approval August 27, 2026 19:54 — with GitHub Actions Error

Narrohag approved these changes Aug 27, 2026

View reviewed changes

Merge branch 'master' into fix/jumpstart-private-hub-modeltrainer

b5afc72

tanvikab4 requested a deployment to manual-approval August 27, 2026 20:06 — with GitHub Actions Waiting

		model_trainer.train()


		def test_jumpstart_train_from_private_hub_reference(private_hub, sagemaker_session):

Uh oh!

Conversation

tanvikab4 commented Aug 21, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants