feat: Prepend SageMaker Studio App Type to boto3 User Agent string (#4440)

Author: knikure

src/sagemaker/user_agent.py

@@ -15,9 +15,18 @@
 
 import platform
 import sys
+import json
+import os
 
 import importlib_metadata
 
+SDK_PREFIX = "AWS-SageMaker-Python-SDK"
+STUDIO_PREFIX = "AWS-SageMaker-Studio"
+NOTEBOOK_PREFIX = "AWS-SageMaker-Notebook-Instance"
+
+NOTEBOOK_METADATA_FILE = "/etc/opt/ml/sagemaker-notebook-instance-version.txt"
+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"
@@ -27,30 +36,69 @@
 )
 
 
+def process_notebook_metadata_file():
+    """Check if the platform is SageMaker Notebook, if yes, return the InstanceType
+
+    Returns:
+        str: The InstanceType of the SageMaker Notebook if it exists, otherwise None
+    """
+    if os.path.exists(NOTEBOOK_METADATA_FILE):
+        with open(NOTEBOOK_METADATA_FILE, "r") as sagemaker_nbi_file:
+            return sagemaker_nbi_file.read().strip()
+
+    return None
+
+
+def process_studio_metadata_file():
+    """Check if the platform is SageMaker Studio, if yes, return the AppType
+
+    Returns:
+        str: The AppType of the SageMaker Studio if it exists, otherwise None
+    """
+    if os.path.exists(STUDIO_METADATA_FILE):
+        with open(STUDIO_METADATA_FILE, "r") as sagemaker_studio_file:
+            metadata = json.load(sagemaker_studio_file)
+            return metadata.get("AppType")
+
+    return None
+
+
 def determine_prefix(user_agent=""):
-    """Placeholder docstring"""
-    prefix = "AWS-SageMaker-Python-SDK/{}".format(SDK_VERSION)
+    """Determines the prefix for the user agent string.
+
+    Args:
+        user_agent (str): The user agent string to prepend the prefix to.
+
+    Returns:
+        str: The user agent string with the prefix prepended.
+    """
+    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)
 
-    try:
-        with open("/etc/opt/ml/sagemaker-notebook-instance-version.txt") as sagemaker_nbi_file:
-            prefix = "{} AWS-SageMaker-Notebook-Instance/{}".format(
-                prefix, sagemaker_nbi_file.read().strip()
-            )
-    except IOError:
-        # This file isn't expected to always exist, and we DO want to silently ignore failures.
-        pass
+    # 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)
+
+    # 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):
-    """Placeholder docstring"""
+    """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)
 
     if client._client_config.user_agent is None:

tests/unit/test_session.py

@@ -41,6 +41,11 @@
 from sagemaker.inputs import BatchDataCaptureConfig
 from sagemaker.config import MODEL_CONTAINERS_PATH
 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 (
     SAGEMAKER_CONFIG_MONITORING_SCHEDULE,
@@ -904,70 +909,64 @@ def test_delete_model(boto_session):
 
 
 def test_user_agent_injected(boto_session):
-    assert (
-        "AWS-SageMaker-Python-SDK" not in boto_session.client("sagemaker")._client_config.user_agent
-    )
+    assert SDK_PREFIX not in boto_session.client("sagemaker")._client_config.user_agent
 
     sess = Session(boto_session)
 
-    assert "AWS-SageMaker-Python-SDK" in sess.sagemaker_client._client_config.user_agent
-    assert "AWS-SageMaker-Python-SDK" in sess.sagemaker_runtime_client._client_config.user_agent
-    assert "AWS-SageMaker-Python-SDK" in sess.sagemaker_metrics_client._client_config.user_agent
-    assert "AWS-SageMaker-Notebook-Instance" not in sess.sagemaker_client._client_config.user_agent
-    assert (
-        "AWS-SageMaker-Notebook-Instance"
-        not in sess.sagemaker_runtime_client._client_config.user_agent
-    )
-    assert (
-        "AWS-SageMaker-Notebook-Instance"
-        not in sess.sagemaker_metrics_client._client_config.user_agent
-    )
+    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
 
 
-def test_user_agent_injected_with_nbi(boto_session):
-    assert (
-        "AWS-SageMaker-Python-SDK" not in boto_session.client("sagemaker")._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,
     )
 
-    with patch("six.moves.builtins.open", mock_open(read_data="120.0-0")) as mo:
-        sess = Session(boto_session)
+    for client in [
+        sess.sagemaker_client,
+        sess.sagemaker_runtime_client,
+        sess.sagemaker_metrics_client,
+    ]:
+        mock_process_notebook_metadata_file.assert_called()
 
-        mo.assert_called_with("/etc/opt/ml/sagemaker-notebook-instance-version.txt")
+        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 "AWS-SageMaker-Python-SDK" in sess.sagemaker_client._client_config.user_agent
-    assert "AWS-SageMaker-Python-SDK" in sess.sagemaker_runtime_client._client_config.user_agent
-    assert "AWS-SageMaker-Python-SDK" in sess.sagemaker_metrics_client._client_config.user_agent
-    assert "AWS-SageMaker-Notebook-Instance" in sess.sagemaker_client._client_config.user_agent
-    assert (
-        "AWS-SageMaker-Notebook-Instance" in sess.sagemaker_runtime_client._client_config.user_agent
-    )
-    assert (
-        "AWS-SageMaker-Notebook-Instance" in sess.sagemaker_metrics_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
 
-def test_user_agent_injected_with_nbi_ioerror(boto_session):
-    assert (
-        "AWS-SageMaker-Python-SDK" not in boto_session.client("sagemaker")._client_config.user_agent
+    sess = Session(
+        boto_session=boto_session,
     )
 
-    with patch("six.moves.builtins.open", MagicMock(side_effect=IOError("File not found"))) as mo:
-        sess = Session(boto_session)
+    for client in [
+        sess.sagemaker_client,
+        sess.sagemaker_runtime_client,
+        sess.sagemaker_metrics_client,
+    ]:
+        mock_process_studio_metadata_file.assert_called()
 
-        mo.assert_called_with("/etc/opt/ml/sagemaker-notebook-instance-version.txt")
-
-    assert "AWS-SageMaker-Python-SDK" in sess.sagemaker_client._client_config.user_agent
-    assert "AWS-SageMaker-Python-SDK" in sess.sagemaker_runtime_client._client_config.user_agent
-    assert "AWS-SageMaker-Python-SDK" in sess.sagemaker_metrics_client._client_config.user_agent
-    assert "AWS-SageMaker-Notebook-Instance" not in sess.sagemaker_client._client_config.user_agent
-    assert (
-        "AWS-SageMaker-Notebook-Instance"
-        not in sess.sagemaker_runtime_client._client_config.user_agent
-    )
-    assert (
-        "AWS-SageMaker-Notebook-Instance"
-        not in sess.sagemaker_metrics_client._client_config.user_agent
-    )
+        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
 
 
 def test_training_input_all_defaults():

tests/unit/test_user_agent.py

@@ -0,0 +1,104 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import json
+from mock import MagicMock, 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,
+)
+
+
+# Test process_notebook_metadata_file function
+def test_process_notebook_metadata_file_exists(tmp_path):
+    notebook_file = tmp_path / "sagemaker-notebook-instance-version.txt"
+    notebook_file.write_text("instance_type")
+
+    with patch("os.path.exists", return_value=True):
+        with patch("builtins.open", mock_open(read_data=notebook_file.read_text())):
+            assert process_notebook_metadata_file() == "instance_type"
+
+
+def test_process_notebook_metadata_file_not_exists(tmp_path):
+    with patch("os.path.exists", return_value=False):
+        assert process_notebook_metadata_file() is None
+
+
+# Test process_studio_metadata_file function
+def test_process_studio_metadata_file_exists(tmp_path):
+    studio_file = tmp_path / "resource-metadata.json"
+    studio_file.write_text(json.dumps({"AppType": "studio_type"}))
+
+    with patch("os.path.exists", return_value=True):
+        with patch("builtins.open", mock_open(read_data=studio_file.read_text())):
+            assert process_studio_metadata_file() == "studio_type"
+
+
+def test_process_studio_metadata_file_not_exists(tmp_path):
+    with patch("os.path.exists", return_value=False):
+        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"