fix model_uri usage issue

Author: Narrohag

src/sagemaker/estimator.py

@@ -2511,7 +2511,7 @@ def start_new(cls, estimator, inputs, experiment_config):
         train_args = cls._get_train_args(estimator, inputs, experiment_config)
 
         logger.debug("Train args after processing defaults: %s", train_args)
-        print("rohan debug: ", train_args)
+
         estimator.sagemaker_session.train(**train_args)
 
         return cls(estimator.sagemaker_session, estimator._current_job_name)

src/sagemaker/jumpstart/factory/estimator.py

@@ -56,6 +56,7 @@
     JUMPSTART_LOGGER,
     TRAINING_ENTRY_POINT_SCRIPT_NAME,
     SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
+    JUMPSTART_MODEL_HUB_NAME,
 )
 from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartModelType
 from sagemaker.jumpstart.factory import model
@@ -630,8 +631,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
+    is_private_hub = JUMPSTART_MODEL_HUB_NAME not in kwargs.hub_arn if kwargs.hub_arn else False
+    if (
+        _model_supports_training_model_uri(**get_model_info_default_kwargs(kwargs))
+        or is_private_hub
+    ):
         default_model_uri = model_uris.retrieve(
             model_scope=JumpStartScriptScope.TRAINING,
             instance_type=kwargs.instance_type,

src/sagemaker/jumpstart/types.py

@@ -1940,6 +1940,11 @@ def use_inference_script_uri(self) -> bool:
 
     def use_training_model_artifact(self) -> bool:
         """Returns True if the model should use a model uri when kicking off training job."""
+        # gated model never uses training model artifact
+        if self.gated_bucket:
+            return False
+
+        # otherwise, return true is a training model package is not set
         return len(self.training_model_package_artifact_uris or {}) == 0
 
     def is_gated_model(self) -> bool:

tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py

@@ -151,18 +151,18 @@ def test_jumpstart_hub_gated_estimator_with_eula(setup, add_model_references):
         },
     )
 
-    estimator = JumpStartEstimator.attach(
-        training_job_name=estimator.latest_training_job.name,
-        model_id=model_id,
-        model_version=model_version,
-    )
-
-    # uses ml.p3.2xlarge instance
     predictor = estimator.deploy(
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
     )
 
-    response = predictor.predict(["hello", "world"])
+    payload = {
+        "inputs": "some-payload",
+        "parameters": {"max_new_tokens": 256, "top_p": 0.9, "temperature": 0.6},
+    }
+
+    response = predictor.predict(payload, custom_attributes="accept_eula=true")
 
     assert response is not None
 

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

@@ -688,7 +688,6 @@ def test_gated_model_non_model_package_s3_uri(
             instance_count=1,
             image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pyt"
             "orch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04",
-            model_uri='s3://jumpstart-private-cache-prod-us-west-2/some/dummy/key',
             source_dir="s3://jumpstart-cache-prod-us-west-2/source-d"
             "irectory-tarballs/meta/transfer_learning/textgeneration/prepack/v1.0.1/sourcedir.tar.gz",
             entry_point="transfer_learning.py",

tests/unit/sagemaker/jumpstart/test_types.py

@@ -332,6 +332,9 @@ def test_jumpstart_model_header():
 def test_use_training_model_artifact():
     specs1 = JumpStartModelSpecs(BASE_SPEC)
     assert specs1.use_training_model_artifact()
+    specs1.gated_bucket = True
+    assert not specs1.use_training_model_artifact()
+    specs1.gated_bucket = False
     specs1.training_model_package_artifact_uris = {"region1": "blah", "region2": "blah2"}
     assert not specs1.use_training_model_artifact()