fix: Fix UserAgent logging in Python SDK (#4647)

Author: knikure

src/sagemaker/session.py

@@ -121,7 +121,7 @@
 from sagemaker.deprecations import deprecated_class
 from sagemaker.enums import EndpointType
 from sagemaker.inputs import ShuffleConfig, TrainingInput, BatchDataCaptureConfig
-from sagemaker.user_agent import prepend_user_agent
+from sagemaker.user_agent import get_user_agent_extra_suffix
 from sagemaker.utils import (
     name_from_image,
     secondary_training_status_changed,
@@ -285,6 +285,7 @@ def _initialize(
         Creates or uses a boto_session, sagemaker_client and sagemaker_runtime_client.
         Sets the region_name.
         """
+
         self.boto_session = boto_session or boto3.DEFAULT_SESSION or boto3.Session()
 
         self._region_name = self.boto_session.region_name
@@ -293,19 +294,30 @@ def _initialize(
                 "Must setup local AWS configuration with a region supported by SageMaker."
             )
 
-        self.sagemaker_client = sagemaker_client or self.boto_session.client("sagemaker")
-        prepend_user_agent(self.sagemaker_client)
+        # Make use of user_agent_extra field of the botocore_config object
+        # to append SageMaker Python SDK specific user_agent suffix
+        # to the current User-Agent header value from boto3
+        # This config will also make sure that user_agent never fails to log the User-Agent string
+        # even if boto User-Agent header format is updated in the future
+        # Ref: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+        botocore_config = botocore.config.Config(user_agent_extra=get_user_agent_extra_suffix())
+
+        # Create sagemaker_client with the botocore_config object
+        # This config is customized to append SageMaker Python SDK specific user_agent suffix
+        self.sagemaker_client = sagemaker_client or self.boto_session.client(
+            "sagemaker", config=botocore_config
+        )
 
         if sagemaker_runtime_client is not None:
             self.sagemaker_runtime_client = sagemaker_runtime_client
         else:
-            config = botocore.config.Config(read_timeout=80)
+            config = botocore.config.Config(
+                read_timeout=80, user_agent_extra=get_user_agent_extra_suffix()
+            )
             self.sagemaker_runtime_client = self.boto_session.client(
                 "runtime.sagemaker", config=config
             )
 
-        prepend_user_agent(self.sagemaker_runtime_client)
-
         if sagemaker_featurestore_runtime_client:
             self.sagemaker_featurestore_runtime_client = sagemaker_featurestore_runtime_client
         else:
@@ -316,8 +328,9 @@ def _initialize(
         if sagemaker_metrics_client:
             self.sagemaker_metrics_client = sagemaker_metrics_client
         else:
-            self.sagemaker_metrics_client = self.boto_session.client("sagemaker-metrics")
-        prepend_user_agent(self.sagemaker_metrics_client)
+            self.sagemaker_metrics_client = self.boto_session.client(
+                "sagemaker-metrics", config=botocore_config
+            )
 
         self.s3_client = self.boto_session.client("s3", region_name=self.boto_region_name)
         self.s3_resource = self.boto_session.resource("s3", region_name=self.boto_region_name)

src/sagemaker/user_agent.py

@@ -13,8 +13,6 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-import platform
-import sys
 import json
 import os
 
@@ -28,12 +26,6 @@
 STUDIO_METADATA_FILE = "/opt/ml/metadata/resource-metadata.json"
 
 SDK_VERSION = importlib_metadata.version("sagemaker")
-OS_NAME = platform.system() or "UnresolvedOS"
-OS_VERSION = platform.release() or "UnresolvedOSVersion"
-OS_NAME_VERSION = "{}/{}".format(OS_NAME, OS_VERSION)
-PYTHON_VERSION = "Python/{}.{}.{}".format(
-    sys.version_info.major, sys.version_info.minor, sys.version_info.micro
-)
 
 
 def process_notebook_metadata_file():
@@ -63,45 +55,24 @@ def process_studio_metadata_file():
     return None
 
 
-def determine_prefix(user_agent=""):
-    """Determines the prefix for the user agent string.
+def get_user_agent_extra_suffix():
+    """Get the user agent extra suffix string specific to SageMaker Python SDK
 
-    Args:
-        user_agent (str): The user agent string to prepend the prefix to.
+    Adhers to new boto recommended User-Agent 2.0 header format
 
     Returns:
-        str: The user agent string with the prefix prepended.
+        str: The user agent extra suffix string to be appended
     """
-    prefix = "{}/{}".format(SDK_PREFIX, SDK_VERSION)
-
-    if PYTHON_VERSION not in user_agent:
-        prefix = "{} {}".format(prefix, PYTHON_VERSION)
-
-    if OS_NAME_VERSION not in user_agent:
-        prefix = "{} {}".format(prefix, OS_NAME_VERSION)
+    suffix = "lib/{}#{}".format(SDK_PREFIX, SDK_VERSION)
 
     # Get the notebook instance type and prepend it to the user agent string if exists
     notebook_instance_type = process_notebook_metadata_file()
     if notebook_instance_type:
-        prefix = "{} {}/{}".format(prefix, NOTEBOOK_PREFIX, notebook_instance_type)
+        suffix = "{} md/{}#{}".format(suffix, NOTEBOOK_PREFIX, notebook_instance_type)
 
     # Get the studio app type and prepend it to the user agent string if exists
     studio_app_type = process_studio_metadata_file()
     if studio_app_type:
-        prefix = "{} {}/{}".format(prefix, STUDIO_PREFIX, studio_app_type)
-
-    return prefix
-
-
-def prepend_user_agent(client):
-    """Prepends the user agent string with the SageMaker Python SDK version.
-
-    Args:
-        client (botocore.client.BaseClient): The client to prepend the user agent string for.
-    """
-    prefix = determine_prefix(client._client_config.user_agent)
+        suffix = "{} md/{}#{}".format(suffix, STUDIO_PREFIX, studio_app_type)
 
-    if client._client_config.user_agent is None:
-        client._client_config.user_agent = prefix
-    else:
-        client._client_config.user_agent = "{} {}".format(prefix, client._client_config.user_agent)
+    return suffix

tests/unit/test_session.py

@@ -43,8 +43,6 @@
 from sagemaker.utils import update_list_of_dicts_with_values_from_config
 from sagemaker.user_agent import (
     SDK_PREFIX,
-    STUDIO_PREFIX,
-    NOTEBOOK_PREFIX,
 )
 from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
 from tests.unit import (
@@ -87,15 +85,20 @@
     limits={},
 )
 
+SDK_DEFAULT_SUFFIX = f"lib/{SDK_PREFIX}#2.218.0"
+NOTEBOOK_SUFFIX = f"{SDK_DEFAULT_SUFFIX} md/AWS-SageMaker-Notebook-Instance#instance_type"
+STUDIO_SUFFIX = f"{SDK_DEFAULT_SUFFIX} md/AWS-SageMaker-Studio#app_type"
 
-@pytest.fixture()
-def boto_session():
-    boto_mock = Mock(name="boto_session", region_name=REGION)
 
+@pytest.fixture
+def boto_session(request):
+    boto_user_agent = "Boto3/1.33.9 md/Botocore#1.33.9 ua/2.0 os/linux#linux-ver md/arch#x86_64 lang/python#3.10.6"
+    user_agent_suffix = getattr(request, "param", "")
+    boto_mock = Mock(name="boto_session", region_name=REGION)
     client_mock = Mock()
-    client_mock._client_config.user_agent = (
-        "Boto3/1.9.69 Python/3.6.5 Linux/4.14.77-70.82.amzn1.x86_64 Botocore/1.12.69 Resource"
-    )
+    user_agent = f"{boto_user_agent} {SDK_DEFAULT_SUFFIX} {user_agent_suffix}"
+    with patch("sagemaker.user_agent.get_user_agent_extra_suffix", return_value=user_agent_suffix):
+        client_mock._client_config.user_agent = user_agent
     boto_mock.client.return_value = client_mock
     return boto_mock
 
@@ -887,65 +890,42 @@ def test_delete_model(boto_session):
     boto_session.client().delete_model.assert_called_with(ModelName=model_name)
 
 
+@pytest.mark.parametrize("boto_session", [""], indirect=True)
 def test_user_agent_injected(boto_session):
-    assert SDK_PREFIX not in boto_session.client("sagemaker")._client_config.user_agent
-
     sess = Session(boto_session)
-
+    expected_user_agent_suffix = "lib/AWS-SageMaker-Python-SDK#2.218.0"
     for client in [
         sess.sagemaker_client,
         sess.sagemaker_runtime_client,
         sess.sagemaker_metrics_client,
     ]:
-        assert SDK_PREFIX in client._client_config.user_agent
-        assert NOTEBOOK_PREFIX not in client._client_config.user_agent
-        assert STUDIO_PREFIX not in client._client_config.user_agent
+        assert expected_user_agent_suffix in client._client_config.user_agent
 
 
-@patch("sagemaker.user_agent.process_notebook_metadata_file", return_value="ml.t3.medium")
-def test_user_agent_injected_with_nbi(
-    mock_process_notebook_metadata_file,
-    boto_session,
-):
-    assert SDK_PREFIX not in boto_session.client("sagemaker")._client_config.user_agent
-
-    sess = Session(
-        boto_session=boto_session,
+@pytest.mark.parametrize("boto_session", [f"{NOTEBOOK_SUFFIX}"], indirect=True)
+def test_user_agent_with_notebook_instance_type(boto_session):
+    sess = Session(boto_session)
+    expected_user_agent_suffix = (
+        "lib/AWS-SageMaker-Python-SDK#2.218.0 md/AWS-SageMaker-Notebook-Instance#instance_type"
     )
-
     for client in [
         sess.sagemaker_client,
         sess.sagemaker_runtime_client,
         sess.sagemaker_metrics_client,
     ]:
-        mock_process_notebook_metadata_file.assert_called()
-
-        assert SDK_PREFIX in client._client_config.user_agent
-        assert NOTEBOOK_PREFIX in client._client_config.user_agent
-        assert STUDIO_PREFIX not in client._client_config.user_agent
+        assert expected_user_agent_suffix in client._client_config.user_agent
 
 
-@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="dymmy-app-type")
-def test_user_agent_injected_with_studio_app_type(
-    mock_process_studio_metadata_file,
-    boto_session,
-):
-    assert SDK_PREFIX not in boto_session.client("sagemaker")._client_config.user_agent
-
-    sess = Session(
-        boto_session=boto_session,
-    )
-
+@pytest.mark.parametrize("boto_session", [f"{STUDIO_SUFFIX}"], indirect=True)
+def test_user_agent_with_studio_app_type(boto_session):
+    sess = Session(boto_session)
+    expected_user_agent = "lib/AWS-SageMaker-Python-SDK#2.218.0 md/AWS-SageMaker-Studio#app_type"
     for client in [
         sess.sagemaker_client,
         sess.sagemaker_runtime_client,
         sess.sagemaker_metrics_client,
     ]:
-        mock_process_studio_metadata_file.assert_called()
-
-        assert SDK_PREFIX in client._client_config.user_agent
-        assert NOTEBOOK_PREFIX not in client._client_config.user_agent
-        assert STUDIO_PREFIX in client._client_config.user_agent
+        assert expected_user_agent in client._client_config.user_agent
 
 
 def test_training_input_all_defaults():

tests/unit/test_user_agent.py

@@ -13,20 +13,17 @@
 from __future__ import absolute_import
 
 import json
-from mock import MagicMock, patch, mock_open
+from mock import patch, mock_open
 
 
 from sagemaker.user_agent import (
     SDK_PREFIX,
     SDK_VERSION,
-    PYTHON_VERSION,
-    OS_NAME_VERSION,
     NOTEBOOK_PREFIX,
     STUDIO_PREFIX,
     process_notebook_metadata_file,
     process_studio_metadata_file,
-    determine_prefix,
-    prepend_user_agent,
+    get_user_agent_extra_suffix,
 )
 
 
@@ -60,45 +57,18 @@ def test_process_studio_metadata_file_not_exists(tmp_path):
         assert process_studio_metadata_file() is None
 
 
-# Test determine_prefix function
-def test_determine_prefix_notebook_instance_type(monkeypatch):
-    monkeypatch.setattr(
-        "sagemaker.user_agent.process_notebook_metadata_file", lambda: "instance_type"
-    )
-    assert (
-        determine_prefix()
-        == f"{SDK_PREFIX}/{SDK_VERSION} {PYTHON_VERSION} {OS_NAME_VERSION} {NOTEBOOK_PREFIX}/instance_type"
-    )
-
-
-def test_determine_prefix_studio_app_type(monkeypatch):
-    monkeypatch.setattr(
-        "sagemaker.user_agent.process_studio_metadata_file", lambda: "studio_app_type"
-    )
-    assert (
-        determine_prefix()
-        == f"{SDK_PREFIX}/{SDK_VERSION} {PYTHON_VERSION} {OS_NAME_VERSION} {STUDIO_PREFIX}/studio_app_type"
-    )
-
-
-def test_determine_prefix_no_metadata(monkeypatch):
-    monkeypatch.setattr("sagemaker.user_agent.process_notebook_metadata_file", lambda: None)
-    monkeypatch.setattr("sagemaker.user_agent.process_studio_metadata_file", lambda: None)
-    assert determine_prefix() == f"{SDK_PREFIX}/{SDK_VERSION} {PYTHON_VERSION} {OS_NAME_VERSION}"
-
-
-# Test prepend_user_agent function
-def test_prepend_user_agent_existing_user_agent(monkeypatch):
-    client = MagicMock()
-    client._client_config.user_agent = "existing_user_agent"
-    monkeypatch.setattr("sagemaker.user_agent.determine_prefix", lambda _: "prefix")
-    prepend_user_agent(client)
-    assert client._client_config.user_agent == "prefix existing_user_agent"
-
-
-def test_prepend_user_agent_no_user_agent(monkeypatch):
-    client = MagicMock()
-    client._client_config.user_agent = None
-    monkeypatch.setattr("sagemaker.user_agent.determine_prefix", lambda _: "prefix")
-    prepend_user_agent(client)
-    assert client._client_config.user_agent == "prefix"
+# Test get_user_agent_extra_suffix function
+def test_get_user_agent_extra_suffix():
+    assert get_user_agent_extra_suffix() == f"lib/{SDK_PREFIX}#{SDK_VERSION}"
+
+    with patch("sagemaker.user_agent.process_notebook_metadata_file", return_value="instance_type"):
+        assert (
+            get_user_agent_extra_suffix()
+            == f"lib/{SDK_PREFIX}#{SDK_VERSION} md/{NOTEBOOK_PREFIX}#instance_type"
+        )
+
+    with patch("sagemaker.user_agent.process_studio_metadata_file", return_value="studio_type"):
+        assert (
+            get_user_agent_extra_suffix()
+            == f"lib/{SDK_PREFIX}#{SDK_VERSION} md/{STUDIO_PREFIX}#studio_type"
+        )