fix: resolve private hub Models and aliased references for modelTrainer

sagemaker-core/src/sagemaker/core/jumpstart/document.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """This module contains utilites for JumpStart model metadata."""
+
 from __future__ import absolute_import
 
 import json
@@ -47,26 +48,54 @@ def get_hub_content_and_document(
         logger.debug("No sagemaker session provided. Using default session.")
 
     hub_name = jumpstart_config.hub_name if jumpstart_config.hub_name else SAGEMAKER_PUBLIC_HUB
-    hub_content_type = "Model" if hub_name == SAGEMAKER_PUBLIC_HUB else "ModelReference"
 
     region = sagemaker_session.boto_region_name
 
-    try:
-        hub_content = HubContent.get(
-            hub_name=hub_name,
-            hub_content_name=jumpstart_config.model_id,
-            hub_content_version=jumpstart_config.model_version,
-            hub_content_type=hub_content_type,
-            session=sagemaker_session.boto_session,
-            region=region,
-        )
-    except ClientError as e:
-        if e.response["Error"]["Code"] == "ResourceNotFound":
-            logger.error(
-                f"Hub content {jumpstart_config.model_id} not found in {hub_name}.\n"
-                "Please check that the Model ID is availble in the specified hub."
+    # The hub content may be filed under an alias that differs from the public
+    # model_id, so honor hub_content_name when provided.
+    hub_content_name = (
+        jumpstart_config.hub_content_name
+        if getattr(jumpstart_config, "hub_content_name", None)
+        else jumpstart_config.model_id
+    )
+
+    # A private hub can contain either a ModelReference (a pointer to a public
+    # JumpStart model) or a privately-owned Model authored directly into the
+    # hub. We cannot tell which from the name alone, so probe: try
+    # ModelReference first, then fall back to Model. The public hub only holds
+    # Models. This mirrors ModelBuilder's resolution in accessors.py.
+    if hub_name == SAGEMAKER_PUBLIC_HUB:
+        content_types_to_try = ["Model"]
+    else:
+        content_types_to_try = ["ModelReference", "Model"]
+
+    hub_content = None
+    last_error: Optional[ClientError] = None
+    for content_type in content_types_to_try:
+        try:
+            hub_content = HubContent.get(
+                hub_name=hub_name,
+                hub_content_name=hub_content_name,
+                hub_content_version=jumpstart_config.model_version,
+                hub_content_type=content_type,
+                session=sagemaker_session.boto_session,
+                region=region,
             )
-        raise e
+            break
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "ResourceNotFound":
+                last_error = e
+                continue
+            raise e
+
+    if hub_content is None:
+        logger.error(
+            f"Hub content {hub_content_name} not found in {hub_name} as any of "
+            f"{content_types_to_try}.\n"
+            "Please check that the Model ID (or hub_content_name) is available "
+            "in the specified hub."
+        )
+        raise last_error
 
     logger.info(
         f"hub_content_name: {hub_content.hub_content_name}, "

sagemaker-core/tests/unit/jumpstart/test_document.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Test for JumpStart Document."""
+
 from __future__ import absolute_import
 
 import json
@@ -81,3 +82,145 @@ def test_get_hub_content_document_failure(jumpstart_session):
             get_hub_content_and_document(
                 jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session
             )
+
+
+# ---------------------------------------------------------------------------
+# Tests for private-hub content-type probing + hub_content_name alias support.
+#
+# A private hub can contain either a ModelReference (a pointer to a public
+# model) or a privately-owned Model. get_hub_content_and_document() must not
+# guess from the hub name; it probes ModelReference first, then falls back to
+# Model. The public hub only holds Models. It also honors hub_content_name when
+# the content is filed under an alias differing from model_id.
+#
+# Note: distinct model_id / hub_name values are used per test to avoid the
+# module-level lru_cache on get_hub_content_and_document returning a stale
+# result across tests.
+# ---------------------------------------------------------------------------
+
+
+def _hub_content(hub_name, name, content_type, doc):
+    return HubContent(
+        hub_name=hub_name,
+        hub_content_name=name,
+        hub_content_version="1.0.0",
+        hub_content_type=content_type,
+        hub_content_document=json.dumps(doc),
+    )
+
+
+def _not_found():
+    return ClientError(
+        error_response={"Error": {"Code": "ResourceNotFound"}},
+        operation_name="DescribeHubContent",
+    )
+
+
+def _load_doc():
+    cur_dir = os.path.dirname(os.path.abspath(__file__))
+    with open(os.path.join(cur_dir, "hub_content_document.json"), "r") as f:
+        return json.load(f)
+
+
+def test_public_hub_uses_model_type_only(jumpstart_session):
+    """Public hub: resolve as Model, and never probe ModelReference."""
+    doc = _load_doc()
+    jumpstart_config = JumpStartConfig(model_id="probe-public-model")
+
+    with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get:
+        mock_get.return_value = _hub_content(
+            "SageMakerPublicHub", "probe-public-model", "Model", doc
+        )
+        hub_content, _ = get_hub_content_and_document(
+            jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session
+        )
+
+    assert hub_content.hub_content_type == "Model"
+    # Public hub must be looked up exactly once, as Model.
+    assert mock_get.call_count == 1
+    assert mock_get.call_args.kwargs["hub_content_type"] == "Model"
+
+
+def test_private_hub_resolves_model_reference_first(jumpstart_session):
+    """Private hub holding a ModelReference: first probe (ModelReference) hits."""
+    doc = _load_doc()
+    jumpstart_config = JumpStartConfig(model_id="probe-ref-model", hub_name="my-private-hub-ref")
+
+    with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get:
+        mock_get.return_value = _hub_content(
+            "my-private-hub-ref", "probe-ref-model", "ModelReference", doc
+        )
+        hub_content, _ = get_hub_content_and_document(
+            jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session
+        )
+
+    assert hub_content.hub_content_type == "ModelReference"
+    # ModelReference is tried first and succeeds -> single call.
+    assert mock_get.call_count == 1
+    assert mock_get.call_args.kwargs["hub_content_type"] == "ModelReference"
+
+
+def test_private_hub_falls_back_to_model(jumpstart_session):
+    """Private hub holding a privately-owned Model: ModelReference misses, then
+    the Model fallback resolves it (the core of the fix)."""
+    doc = _load_doc()
+    jumpstart_config = JumpStartConfig(
+        model_id="probe-private-model", hub_name="my-private-hub-model"
+    )
+
+    with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get:
+        mock_get.side_effect = [
+            _not_found(),  # ModelReference lookup misses
+            _hub_content(  # Model fallback resolves
+                "my-private-hub-model", "probe-private-model", "Model", doc
+            ),
+        ]
+        hub_content, _ = get_hub_content_and_document(
+            jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session
+        )
+
+    assert hub_content.hub_content_type == "Model"
+    # Two probes: ModelReference (miss) then Model (hit).
+    assert mock_get.call_count == 2
+    assert [c.kwargs["hub_content_type"] for c in mock_get.call_args_list] == [
+        "ModelReference",
+        "Model",
+    ]
+
+
+def test_private_hub_honors_hub_content_name_alias(jumpstart_session):
+    """When hub_content_name is set (alias differs from model_id), the lookup
+    must use the alias, not the model_id."""
+    doc = _load_doc()
+    jumpstart_config = JumpStartConfig(
+        model_id="probe-alias-public-id",
+        hub_name="my-private-hub-alias",
+        hub_content_name="the-alias-name",
+    )
+
+    with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get:
+        mock_get.return_value = _hub_content(
+            "my-private-hub-alias", "the-alias-name", "ModelReference", doc
+        )
+        get_hub_content_and_document(
+            jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session
+        )
+
+    # Lookup used the alias, not the model_id.
+    assert mock_get.call_args.kwargs["hub_content_name"] == "the-alias-name"
+
+
+def test_private_hub_not_found_as_either_type_raises(jumpstart_session):
+    """Private hub where neither ModelReference nor Model exists: raise."""
+    jumpstart_config = JumpStartConfig(
+        model_id="probe-missing-model", hub_name="my-private-hub-missing"
+    )
+
+    with patch("sagemaker.core.jumpstart.document.HubContent.get") as mock_get:
+        mock_get.side_effect = [_not_found(), _not_found()]
+        with pytest.raises(ClientError):
+            get_hub_content_and_document(
+                jumpstart_config=jumpstart_config, sagemaker_session=jumpstart_session
+            )
+        # Both content types were attempted before giving up.
+        assert mock_get.call_count == 2