From 74bc1ed7009047af6ed8947e1a9ab327ce1bb37b Mon Sep 17 00:00:00 2001 From: knikure Date: Wed, 21 Feb 2024 00:02:49 +0000 Subject: [PATCH] feat: Prepend SageMaker Studio App Type to boto3 User Agent string --- src/sagemaker/user_agent.py | 70 +++++++++++++++++++---- tests/unit/test_session.py | 99 ++++++++++++++++---------------- tests/unit/test_user_agent.py | 104 ++++++++++++++++++++++++++++++++++ 3 files changed, 212 insertions(+), 61 deletions(-) create mode 100644 tests/unit/test_user_agent.py diff --git a/src/sagemaker/user_agent.py b/src/sagemaker/user_agent.py index 6a6c6b5204..8af89696c2 100644 --- a/src/sagemaker/user_agent.py +++ b/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,9 +36,43 @@ ) +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) @@ -37,20 +80,25 @@ def determine_prefix(user_agent=""): 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: diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index de543b6f53..6661422e4f 100644 --- a/tests/unit/test_session.py +++ b/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(): diff --git a/tests/unit/test_user_agent.py b/tests/unit/test_user_agent.py new file mode 100644 index 0000000000..c116fef951 --- /dev/null +++ b/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"