Skip to content

Commit d8c62c4

Browse files
authored
Use custom user agent string at all times (#1)
* Use custom user agent string even if there is a user-supplied boto session or SageMaker client * Move user agent string construction to its own module
1 parent 2102374 commit d8c62c4

File tree

2 files changed

+37
-21
lines changed

2 files changed

+37
-21
lines changed

src/sagemaker/session.py

+5-21
Original file line numberDiff line numberDiff line change
@@ -16,21 +16,18 @@
1616
import re
1717

1818
import os
19-
import pkg_resources
2019
import sys
2120
import time
2221

2322
import boto3
24-
import botocore.session
2523
import json
2624
import six
2725
from botocore.exceptions import ClientError
2826

27+
from sagemaker.user_agent import prepend_user_agent
2928
from sagemaker.utils import name_from_image
3029
import sagemaker.logs
3130

32-
SDK_VERSION = pkg_resources.require('sagemaker')[0].version
33-
3431

3532
logging.basicConfig()
3633
LOGGER = logging.getLogger('sagemaker')
@@ -71,30 +68,17 @@ def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_c
7168
If not provided, one will be created using this instance's ``boto_session``.
7269
"""
7370
self._default_bucket = None
74-
75-
if boto_session is None:
76-
botocore_session = self._botocore_session_with_user_agent()
77-
self.boto_session = boto3.Session(botocore_session=botocore_session)
78-
else:
79-
self.boto_session = boto_session
71+
self.boto_session = boto_session or boto3.Session()
8072

8173
region = self.boto_session.region_name
8274
if region is None:
8375
raise ValueError('Must setup local AWS configuration with a region supported by SageMaker.')
8476

8577
self.sagemaker_client = sagemaker_client or self.boto_session.client('sagemaker')
86-
self.sagemaker_runtime_client = sagemaker_runtime_client or self.boto_session.client('runtime.sagemaker')
78+
prepend_user_agent(self.sagemaker_client)
8779

88-
def _botocore_session_with_user_agent(self):
89-
botocore_session = botocore.session.get_session()
90-
91-
# Some information is appended automatically, so this creates a user agent string that looks like
92-
# 'AWS-SageMaker-Python-SDK/1.0/1.7.45 Python/3.6.2 Darwin/15.6.0 Boto3/1.4.7 Botocore/1.7.45'
93-
botocore_session.user_agent_name = 'AWS-SageMaker-Python-SDK/{}'.format(SDK_VERSION)
94-
botocore_session.user_agent_extra = 'Boto3/{} Botocore/{}'.format(boto3.__version__,
95-
botocore_session.user_agent_version)
96-
97-
return botocore_session
80+
self.sagemaker_runtime_client = sagemaker_runtime_client or self.boto_session.client('runtime.sagemaker')
81+
prepend_user_agent(self.sagemaker_runtime_client)
9882

9983
@property
10084
def boto_region_name(self):

src/sagemaker/user_agent.py

+32
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
import pkg_resources
14+
import platform
15+
import sys
16+
17+
import boto3
18+
import botocore
19+
20+
SDK_VERSION = pkg_resources.require('sagemaker')[0].version
21+
OS_NAME = platform.system() or 'UnresolvedOS'
22+
OS_VERSION = platform.release() or 'UnresolvedOSVersion'
23+
PYTHON_VERSION = '{}.{}.{}'.format(sys.version_info.major, sys.version_info.minor, sys.version_info.micro)
24+
25+
26+
def prepend_user_agent(client):
27+
prefix = 'AWS-SageMaker-Python-SDK/{} Python/{} {}/{} Boto3/{} Botocore/{}'\
28+
.format(SDK_VERSION, PYTHON_VERSION, OS_NAME, OS_VERSION, boto3.__version__, botocore.__version__)
29+
if client._client_config.user_agent is None:
30+
client._client_config.user_agent = prefix
31+
else:
32+
client._client_config.user_agent = '{} {}'.format(prefix, client._client_config.user_agent)

0 commit comments

Comments
 (0)