|
16 | 16 | import re |
17 | 17 |
|
18 | 18 | import os |
19 | | -import pkg_resources |
20 | 19 | import sys |
21 | 20 | import time |
22 | 21 |
|
23 | 22 | import boto3 |
24 | | -import botocore.session |
25 | 23 | import json |
26 | 24 | import six |
27 | 25 | from botocore.exceptions import ClientError |
28 | 26 |
|
| 27 | +from sagemaker.user_agent import prepend_user_agent |
29 | 28 | from sagemaker.utils import name_from_image |
30 | 29 | import sagemaker.logs |
31 | 30 |
|
32 | | -SDK_VERSION = pkg_resources.require('sagemaker')[0].version |
33 | | - |
34 | 31 |
|
35 | 32 | logging.basicConfig() |
36 | 33 | LOGGER = logging.getLogger('sagemaker') |
@@ -71,30 +68,17 @@ def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_c |
71 | 68 | If not provided, one will be created using this instance's ``boto_session``. |
72 | 69 | """ |
73 | 70 | 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() |
80 | 72 |
|
81 | 73 | region = self.boto_session.region_name |
82 | 74 | if region is None: |
83 | 75 | raise ValueError('Must setup local AWS configuration with a region supported by SageMaker.') |
84 | 76 |
|
85 | 77 | 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) |
87 | 79 |
|
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) |
98 | 82 |
|
99 | 83 | @property |
100 | 84 | def boto_region_name(self): |
|
0 commit comments