Prepend SageMaker Notebook Instance version ID to boto3 User Agent (#600)

SageMaker Notebook Instances provide the version ID as a file on a
Notebook Instance.  In order to enhance version adoption from within
SageMaker, provide an extra segment to the user_agent supplied to the
boto3 ClientConfig.

Added new unit tests:

* Basic test to verify User Agent propagation
* Test to ensure that notebook instance is injected
* Test to ensure handling of the IOError when reading the file.

Author: mckev-amazon

CHANGELOG.rst

@@ -3,6 +3,10 @@ CHANGELOG
 =========
 
 
+1.18.2.dev
+======
+* enhancement: Include SageMaker Notebook Instance version number in boto3 user agent, if available.
+
 1.18.1
 ======
 

src/sagemaker/user_agent.py

@@ -25,9 +25,23 @@
 PYTHON_VERSION = '{}.{}.{}'.format(sys.version_info.major, sys.version_info.minor, sys.version_info.micro)
 
 
-def prepend_user_agent(client):
+def determine_prefix():
     prefix = 'AWS-SageMaker-Python-SDK/{} Python/{} {}/{} Boto3/{} Botocore/{}'\
         .format(SDK_VERSION, PYTHON_VERSION, OS_NAME, OS_VERSION, boto3.__version__, botocore.__version__)
+
+    try:
+        with open('/etc/opt/ml/sagemaker-notebook-instance-version.txt') as sagemaker_nbi_file:
+            prefix = 'AWS-SageMaker-Notebook-Instance/{} {}'.format(sagemaker_nbi_file.read().strip(), prefix)
+    except IOError:
+        # This file isn't expected to always exist, and we DO want to silently ignore failures.
+        pass
+
+    return prefix
+
+
+def prepend_user_agent(client):
+    prefix = determine_prefix()
+
     if client._client_config.user_agent is None:
         client._client_config.user_agent = prefix
     else:

tests/unit/test_session.py

@@ -19,7 +19,7 @@
 import pytest
 import six
 from botocore.exceptions import ClientError
-from mock import Mock, patch, call
+from mock import MagicMock, Mock, patch, call, mock_open
 
 import sagemaker
 from sagemaker import s3_input, Session, get_execution_role
@@ -36,6 +36,12 @@
 @pytest.fixture()
 def boto_session():
     boto_session = Mock(region_name=REGION)
+
+    mock_client = Mock()
+    mock_client._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'
+
+    boto_session.client.return_value = mock_client
     return boto_session
 
 
@@ -139,6 +145,45 @@ def test_delete_endpoint(boto_session):
     boto_session.client().delete_endpoint.assert_called_with(EndpointName='my_endpoint')
 
 
+def test_user_agent_injected(boto_session):
+    assert 'AWS-SageMaker-Python-SDK' 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-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
+
+
+def test_user_agent_injected_with_nbi(boto_session):
+    assert 'AWS-SageMaker-Python-SDK' not in boto_session.client('sagemaker')._client_config.user_agent
+
+    with patch('six.moves.builtins.open', mock_open(read_data='120.0-0')) as mo:
+        sess = Session(boto_session)
+
+        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-Notebook-Instance' in sess.sagemaker_client._client_config.user_agent
+    assert 'AWS-SageMaker-Notebook-Instance' in sess.sagemaker_runtime_client._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
+
+    with patch('six.moves.builtins.open', MagicMock(side_effect=IOError('File not found'))) as mo:
+        sess = Session(boto_session)
+
+        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-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
+
+
 def test_s3_input_all_defaults():
     prefix = 'pre'
     actual = s3_input(s3_data=prefix)