chore: emit log when legacy fields used to get jumpstart image uri

Author: evakravi

src/sagemaker/image_uris.py

@@ -21,7 +21,7 @@
 from packaging.version import Version
 
 from sagemaker import utils
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION, JUMPSTART_LOGGER
 from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.jumpstart.utils import is_jumpstart_model_input
 from sagemaker.spark import defaults
@@ -154,6 +154,18 @@ def retrieve(
             )
 
     if is_jumpstart_model_input(model_id, model_version):
+        if non_none_fields := {
+            key: value
+            for key, value in args.items()
+            if key in {"version", "framework", "container_version", "py_version"}
+            and value is not None
+        }:
+            JUMPSTART_LOGGER.info(
+                "Ignoring the following fields when retriving image uri "
+                "for JumpStart model id '%s': %s",
+                model_id,
+                str(non_none_fields),
+            )
         return artifacts._retrieve_image_uri(
             model_id=model_id,
             model_version=model_version,

tests/unit/sagemaker/image_uris/jumpstart/test_common.py

@@ -186,3 +186,54 @@ def test_jumpstart_common_image_uri(
             model_id="pytorch-ic-mobilenet-v2",
             instance_type="ml.m5.xlarge",
         )
+
+
+@patch("sagemaker.image_uris.JUMPSTART_LOGGER.info")
+@patch("sagemaker.jumpstart.utils.validate_model_id_and_get_type")
+@patch("sagemaker.jumpstart.artifacts.image_uris.verify_model_region_and_return_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_image_uri_logging_extra_fields(
+    patched_get_model_specs,
+    patched_verify_model_region_and_return_specs,
+    patched_validate_model_id_and_get_type,
+    patched_info_log,
+):
+
+    patched_verify_model_region_and_return_specs.side_effect = verify_model_region_and_return_specs
+    patched_get_model_specs.side_effect = get_spec_from_base_spec
+    patched_validate_model_id_and_get_type.return_value = JumpStartModelType.OPEN_WEIGHTS
+
+    region = "us-west-2"
+    mock_client = boto3.client("s3")
+    mock_session = Mock(s3_client=mock_client, boto_region_name=region)
+
+    image_uris.retrieve(
+        framework=None,
+        region="us-west-2",
+        image_scope="training",
+        model_id="pytorch-ic-mobilenet-v2",
+        model_version="*",
+        instance_type="ml.m5.xlarge",
+        sagemaker_session=mock_session,
+    )
+
+    patched_info_log.assert_not_called()
+
+    image_uris.retrieve(
+        framework="framework",
+        container_version="1.2.3",
+        region="us-west-2",
+        image_scope="training",
+        model_id="pytorch-ic-mobilenet-v2",
+        model_version="*",
+        instance_type="ml.m5.xlarge",
+        sagemaker_session=mock_session,
+    )
+
+    patched_info_log.assert_called_once_with(
+        "Ignoring the following fields "
+        "when retriving image uri for "
+        "JumpStart model id '%s': %s",
+        "pytorch-ic-mobilenet-v2",
+        "{'framework': 'framework', 'container_version': '1.2.3'}",
+    )