chore: address PR comments

Author: evakravi

src/sagemaker/jumpstart/factory/model.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module stores JumpStart Model factory methods."""
 from __future__ import absolute_import
+import json
 
 
 from typing import Any, Dict, List, Optional, Union
@@ -217,19 +218,21 @@ def _add_model_data_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartMode
     )
 
     if isinstance(model_data, str) and model_data.startswith("s3://") and model_data.endswith("/"):
-        if kwargs.model_data:
-            JUMPSTART_LOGGER.info(
-                "S3 prefix model_data detected for JumpStartModel: '%s'. "
-                "Converting to S3DataSource dictionary.",
-                model_data,
-            )
+        old_model_data_str = model_data
         model_data = {
             "S3DataSource": {
                 "S3Uri": model_data,
                 "S3DataType": "S3Prefix",
                 "CompressionType": "None",
             }
         }
+        if kwargs.model_data:
+            JUMPSTART_LOGGER.info(
+                "S3 prefix model_data detected for JumpStartModel: '%s'. "
+                "Converting to S3DataSource dictionary: '%s'.",
+                old_model_data_str,
+                json.dumps(model_data),
+            )
 
     kwargs.model_data = model_data
 

tests/unit/sagemaker/jumpstart/constants.py

@@ -1721,7 +1721,7 @@
             "py_version": "py38",
             "huggingface_transformers_version": "4.17.0",
         },
-        "hosting_artifact_key": "huggingface-infer/infer-huggingface-text2text-flan-t5-xxl-fp16.tar.gz",
+        "hosting_artifact_key": "huggingface-infer/",
         "hosting_script_key": "source-directory-tarballs/huggingface/inference/text2text/v1.0.3/sourcedir.tar.gz",
         "hosting_prepacked_artifact_key": "huggingface-infer/prepack/v1.0.1/",
         "hosting_prepacked_artifact_version": "1.0.1",

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -737,8 +737,11 @@ def test_model_data_s3_prefix_override(
         )
 
         mock_js_info_logger.assert_called_with(
-            "S3 prefix model_data detected for JumpStartModel: '%s'. Converting to S3DataSource dictionary.",
+            "S3 prefix model_data detected for JumpStartModel: '%s'. "
+            "Converting to S3DataSource dictionary: '%s'.",
             "s3://some-bucket/path/to/prefix/",
+            '{"S3DataSource": {"S3Uri": "s3://some-bucket/path/to/prefix/", '
+            '"S3DataType": "S3Prefix", "CompressionType": "None"}}',
         )
 
     @mock.patch("sagemaker.jumpstart.model.is_valid_model_id")