Address comments: only add is_hub_content to user aggent suffix if its available

Author: Malav Shastri

src/sagemaker/jumpstart/hub/hub.py

@@ -78,7 +78,10 @@ def __init__(
         """
         self.hub_name = hub_name
         self.region = sagemaker_session.boto_region_name
-        self._sagemaker_session = sagemaker_session or utils.get_default_jumpstart_session_with_user_agent_suffix(is_hub_content=True)
+        self._sagemaker_session = (
+            sagemaker_session
+            or utils.get_default_jumpstart_session_with_user_agent_suffix(is_hub_content=True)
+        )
         self.hub_storage_location = self._generate_hub_storage_location(bucket_name)
 
     def _fetch_hub_bucket_name(self) -> str:

src/sagemaker/jumpstart/utils.py

@@ -1022,12 +1022,15 @@ def get_jumpstart_user_agent_extra_suffix(
 
     if os.getenv(constants.ENV_VARIABLE_DISABLE_JUMPSTART_TELEMETRY, None):
         headers = sagemaker_python_sdk_headers
-    elif model_id is None and model_version is None:
-        headers = f"{sagemaker_python_sdk_headers} {hub_specific_suffix}"
+    elif is_hub_content is True:
+        if model_id is None and model_version is None:
+            headers = f"{sagemaker_python_sdk_headers} {hub_specific_suffix}"
+        else:
+            headers = (
+                f"{sagemaker_python_sdk_headers} {jumpstart_specific_suffix} {hub_specific_suffix}"
+            )
     else:
-        headers = (
-            f"{sagemaker_python_sdk_headers} {jumpstart_specific_suffix} {hub_specific_suffix}"
-        )
+        headers = f"{sagemaker_python_sdk_headers} {jumpstart_specific_suffix}"
 
     return headers
 

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -1716,11 +1716,11 @@ def test_get_jumpstart_user_agent_extra_suffix(self, mock_getenv):
         mock_getenv.return_value = False
         assert utils.get_jumpstart_user_agent_extra_suffix(
             "some-id", "some-version", "False"
-        ).endswith("md/js_model_id#some-id md/js_model_ver#some-version md/js_is_hub_content#False")
+        ).endswith("md/js_model_id#some-id md/js_model_ver#some-version")
         mock_getenv.return_value = None
         assert utils.get_jumpstart_user_agent_extra_suffix(
             "some-id", "some-version", "False"
-        ).endswith("md/js_model_id#some-id md/js_model_ver#some-version md/js_is_hub_content#False")
+        ).endswith("md/js_model_id#some-id md/js_model_ver#some-version")
         mock_getenv.return_value = "True"
         assert not utils.get_jumpstart_user_agent_extra_suffix(
             "some-id", "some-version", "True"