chore: always use model channel for private hub models, add unit tests

Author: evakravi

src/sagemaker/jumpstart/factory/estimator.py

@@ -54,6 +54,7 @@
 from sagemaker.jumpstart.constants import (
     JUMPSTART_DEFAULT_REGION_NAME,
     JUMPSTART_LOGGER,
+    JUMPSTART_MODEL_HUB_NAME,
     TRAINING_ENTRY_POINT_SCRIPT_NAME,
     SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
@@ -631,7 +632,13 @@ def _add_model_reference_arn_to_kwargs(
 
 def _add_model_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstimatorInitKwargs:
     """Sets model uri in kwargs based on default or override, returns full kwargs."""
-    if _model_supports_training_model_uri(**get_model_info_default_kwargs(kwargs)):
+    # hub_arn is by default None unless the user specifies the hub_name
+    # If no hub_name is specified, it is assumed the public hub
+    # Training platform enforces that private hub models must use model channel
+    is_private_hub = JUMPSTART_MODEL_HUB_NAME not in kwargs.hub_arn if kwargs.hub_arn else False
+    if is_private_hub or _model_supports_training_model_uri(
+        **get_model_info_default_kwargs(kwargs)
+    ):
         default_model_uri = model_uris.retrieve(
             model_scope=JumpStartScriptScope.TRAINING,
             instance_type=kwargs.instance_type,

tests/unit/sagemaker/jumpstart/factory/test_estimator.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 import pytest
 from unittest.mock import patch
+from sagemaker.jumpstart.constants import JUMPSTART_MODEL_HUB_NAME
 from sagemaker.jumpstart.factory.estimator import (
     _add_model_uri_to_kwargs,
     get_model_info_default_kwargs,
@@ -119,3 +120,43 @@ def test_add_model_uri_to_kwargs_no_training_support(self, mock_supports_trainin
 
         mock_supports_training.assert_called_once()
         assert result.model_uri is None
+
+    @patch(
+        "sagemaker.jumpstart.factory.estimator._model_supports_training_model_uri",
+        return_value=False,
+    )
+    @patch("sagemaker.jumpstart.factory.estimator.model_uris.retrieve")
+    def test_add_model_uri_to_kwargs_private_hub(
+        self, mock_retrieve, mock_supports_training, mock_kwargs
+    ):
+        """Test when model is from a private hub."""
+        default_uri = "s3://jumpstart-models/training/test-model/1.0.0"
+        mock_retrieve.return_value = default_uri
+        mock_kwargs.hub_arn = "arn:aws:sagemaker:us-west-2:123456789012:hub/private-hub"
+
+        result = _add_model_uri_to_kwargs(mock_kwargs)
+
+        # Should not check if model supports training model URI for private hub
+        mock_supports_training.assert_not_called()
+        mock_retrieve.assert_called_once()
+        assert result.model_uri == default_uri
+
+    @patch(
+        "sagemaker.jumpstart.factory.estimator._model_supports_training_model_uri",
+        return_value=False,
+    )
+    @patch("sagemaker.jumpstart.factory.estimator.model_uris.retrieve")
+    def test_add_model_uri_to_kwargs_public_hub(
+        self, mock_retrieve, mock_supports_training, mock_kwargs
+    ):
+        """Test when model is from the public hub."""
+        mock_kwargs.hub_arn = (
+            f"arn:aws:sagemaker:us-west-2:123456789012:hub/{JUMPSTART_MODEL_HUB_NAME}"
+        )
+
+        result = _add_model_uri_to_kwargs(mock_kwargs)
+
+        # Should check if model supports training model URI for public hub
+        mock_supports_training.assert_called_once()
+        mock_retrieve.assert_not_called()
+        assert result.model_uri is None