|
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