Skip to content

Commit 36d4503

Browse files
nadiayanadiaya
authored
fix: remove redundant information from the user_agent string. (#1746)
1 parent 8b2d5c8 commit 36d4503

11 files changed

+54
-47
lines changed

‎src/sagemaker/user_agent.py

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -16,28 +16,31 @@
1616
import platform
1717
import sys
1818

19-
import boto3
20-
import botocore
2119
import importlib_metadata
2220

2321
SDK_VERSION = importlib_metadata.version("sagemaker")
2422
OS_NAME = platform.system() or "UnresolvedOS"
2523
OS_VERSION = platform.release() or "UnresolvedOSVersion"
26-
PYTHON_VERSION = "{}.{}.{}".format(
24+
OS_NAME_VERSION = "{}/{}".format(OS_NAME, OS_VERSION)
25+
PYTHON_VERSION = "Python/{}.{}.{}".format(
2726
sys.version_info.major, sys.version_info.minor, sys.version_info.micro
2827
)
2928

3029

31-
def determine_prefix():
30+
def determine_prefix(user_agent=""):
3231
"""Placeholder docstring"""
33-
prefix = "AWS-SageMaker-Python-SDK/{} Python/{} {}/{} Boto3/{} Botocore/{}".format(
34-
SDK_VERSION, PYTHON_VERSION, OS_NAME, OS_VERSION, boto3.__version__, botocore.__version__
35-
)
32+
prefix = "AWS-SageMaker-Python-SDK/{}".format(SDK_VERSION)
33+
34+
if PYTHON_VERSION not in user_agent:
35+
prefix = "{} {}".format(prefix, PYTHON_VERSION)
36+
37+
if OS_NAME_VERSION not in user_agent:
38+
prefix = "{} {}".format(prefix, OS_NAME_VERSION)
3639

3740
try:
3841
with open("/etc/opt/ml/sagemaker-notebook-instance-version.txt") as sagemaker_nbi_file:
39-
prefix = "AWS-SageMaker-Notebook-Instance/{} {}".format(
40-
sagemaker_nbi_file.read().strip(), prefix
42+
prefix = "{} AWS-SageMaker-Notebook-Instance/{}".format(
43+
prefix, sagemaker_nbi_file.read().strip()
4144
)
4245
except IOError:
4346
# This file isn't expected to always exist, and we DO want to silently ignore failures.
@@ -51,7 +54,7 @@ def prepend_user_agent(client):
5154
Args:
5255
client:
5356
"""
54-
prefix = determine_prefix()
57+
prefix = determine_prefix(client._client_config.user_agent)
5558

5659
if client._client_config.user_agent is None:
5760
client._client_config.user_agent = prefix

‎tests/unit/test_create_deploy_entities.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from __future__ import absolute_import
1414

1515
import pytest
16-
from mock import Mock
16+
from mock import MagicMock, Mock
1717

1818
import sagemaker
1919

@@ -33,7 +33,7 @@
3333

3434
@pytest.fixture()
3535
def sagemaker_session():
36-
boto_mock = Mock(name="boto_session", region_name=REGION)
36+
boto_mock = MagicMock(name="boto_session", region_name=REGION)
3737
ims = sagemaker.Session(boto_session=boto_mock)
3838
ims.expand_role = Mock(return_value=EXPANDED_ROLE)
3939
return ims

‎tests/unit/test_default_bucket.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
import pytest
1616
from botocore.exceptions import ClientError
17-
from mock import Mock
17+
from mock import MagicMock
1818
import sagemaker
1919

2020
ACCOUNT_ID = "123"
@@ -24,7 +24,7 @@
2424

2525
@pytest.fixture()
2626
def sagemaker_session():
27-
boto_mock = Mock(name="boto_session", region_name=REGION)
27+
boto_mock = MagicMock(name="boto_session", region_name=REGION)
2828
boto_mock.client("sts").get_caller_identity.return_value = {"Account": ACCOUNT_ID}
2929
sagemaker_session = sagemaker.Session(boto_session=boto_mock)
3030
sagemaker_session.boto_session.resource("s3").Bucket().creation_date = None

‎tests/unit/test_endpoint_from_job.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from __future__ import absolute_import
1414

1515
import pytest
16-
from mock import Mock
16+
from mock import MagicMock, Mock
1717

1818
import sagemaker
1919

@@ -42,8 +42,10 @@
4242

4343
@pytest.fixture()
4444
def sagemaker_session():
45-
boto_mock = Mock(name="boto_session", region_name=REGION)
46-
ims = sagemaker.Session(sagemaker_client=Mock(name="sagemaker_client"), boto_session=boto_mock)
45+
boto_mock = MagicMock(name="boto_session", region_name=REGION)
46+
ims = sagemaker.Session(
47+
sagemaker_client=MagicMock(name="sagemaker_client"), boto_session=boto_mock
48+
)
4749
ims.sagemaker_client.describe_training_job = Mock(
4850
name="describe_training_job", return_value=TRAINING_JOB_RESPONSE
4951
)

‎tests/unit/test_endpoint_from_model_data.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
import pytest
1616
from botocore.exceptions import ClientError
17-
from mock import Mock
17+
from mock import MagicMock, Mock
1818
from mock import patch
1919

2020
import sagemaker
@@ -35,8 +35,10 @@
3535

3636
@pytest.fixture()
3737
def sagemaker_session():
38-
boto_mock = Mock(name="boto_session", region_name=REGION)
39-
ims = sagemaker.Session(sagemaker_client=Mock(name="sagemaker_client"), boto_session=boto_mock)
38+
boto_mock = MagicMock(name="boto_session", region_name=REGION)
39+
ims = sagemaker.Session(
40+
sagemaker_client=MagicMock(name="sagemaker_client"), boto_session=boto_mock
41+
)
4042
ims.sagemaker_client.describe_model = Mock(
4143
name="describe_model", side_effect=_raise_does_not_exist_client_error
4244
)

‎tests/unit/test_exception_on_bad_status.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@
2424

2525

2626
def get_sagemaker_session(returns_status):
27-
boto_mock = Mock(name="boto_session", region_name=REGION)
28-
client_mock = Mock()
27+
boto_mock = MagicMock(name="boto_session", region_name=REGION)
28+
client_mock = MagicMock()
2929
client_mock.describe_model_package = MagicMock(
3030
return_value={"ModelPackageStatus": returns_status}
3131
)

‎tests/unit/test_image.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,11 +76,11 @@
7676

7777
@pytest.fixture()
7878
def sagemaker_session():
79-
boto_mock = Mock(name="boto_session", region_name=REGION)
79+
boto_mock = MagicMock(name="boto_session", region_name=REGION)
8080
boto_mock.client("sts").get_caller_identity.return_value = {"Account": "123"}
8181
boto_mock.resource("s3").Bucket(BUCKET_NAME).objects.filter.return_value = []
8282

83-
sms = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
83+
sms = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
8484

8585
sms.default_bucket = Mock(name="default_bucket", return_value=BUCKET_NAME)
8686
sms.expand_role = Mock(return_value=EXPANDED_ROLE)

‎tests/unit/test_session.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -652,11 +652,11 @@ def test_s3_input_all_arguments():
652652

653653
@pytest.fixture()
654654
def sagemaker_session():
655-
boto_mock = Mock(name="boto_session")
655+
boto_mock = MagicMock(name="boto_session")
656656
boto_mock.client("sts", endpoint_url=STS_ENDPOINT).get_caller_identity.return_value = {
657657
"Account": "123"
658658
}
659-
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
659+
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
660660
ims.expand_role = Mock(return_value=EXPANDED_ROLE)
661661
return ims
662662

@@ -1352,10 +1352,10 @@ def __init__(self, code):
13521352

13531353
@pytest.fixture()
13541354
def sagemaker_session_complete():
1355-
boto_mock = Mock(name="boto_session")
1355+
boto_mock = MagicMock(name="boto_session")
13561356
boto_mock.client("logs").describe_log_streams.return_value = DEFAULT_LOG_STREAMS
13571357
boto_mock.client("logs").get_log_events.side_effect = DEFAULT_LOG_EVENTS
1358-
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
1358+
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
13591359
ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
13601360
ims.sagemaker_client.describe_transform_job.return_value = (
13611361
COMPLETED_DESCRIBE_TRANSFORM_JOB_RESULT
@@ -1365,21 +1365,21 @@ def sagemaker_session_complete():
13651365

13661366
@pytest.fixture()
13671367
def sagemaker_session_stopped():
1368-
boto_mock = Mock(name="boto_session")
1368+
boto_mock = MagicMock(name="boto_session")
13691369
boto_mock.client("logs").describe_log_streams.return_value = DEFAULT_LOG_STREAMS
13701370
boto_mock.client("logs").get_log_events.side_effect = DEFAULT_LOG_EVENTS
1371-
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
1371+
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
13721372
ims.sagemaker_client.describe_training_job.return_value = STOPPED_DESCRIBE_JOB_RESULT
13731373
ims.sagemaker_client.describe_transform_job.return_value = STOPPED_DESCRIBE_TRANSFORM_JOB_RESULT
13741374
return ims
13751375

13761376

13771377
@pytest.fixture()
13781378
def sagemaker_session_ready_lifecycle():
1379-
boto_mock = Mock(name="boto_session")
1379+
boto_mock = MagicMock(name="boto_session")
13801380
boto_mock.client("logs").describe_log_streams.return_value = DEFAULT_LOG_STREAMS
13811381
boto_mock.client("logs").get_log_events.side_effect = STREAM_LOG_EVENTS
1382-
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
1382+
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
13831383
ims.sagemaker_client.describe_training_job.side_effect = [
13841384
IN_PROGRESS_DESCRIBE_JOB_RESULT,
13851385
IN_PROGRESS_DESCRIBE_JOB_RESULT,
@@ -1395,10 +1395,10 @@ def sagemaker_session_ready_lifecycle():
13951395

13961396
@pytest.fixture()
13971397
def sagemaker_session_full_lifecycle():
1398-
boto_mock = Mock(name="boto_session")
1398+
boto_mock = MagicMock(name="boto_session")
13991399
boto_mock.client("logs").describe_log_streams.side_effect = LIFECYCLE_LOG_STREAMS
14001400
boto_mock.client("logs").get_log_events.side_effect = STREAM_LOG_EVENTS
1401-
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
1401+
ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
14021402
ims.sagemaker_client.describe_training_job.side_effect = [
14031403
IN_PROGRESS_DESCRIBE_JOB_RESULT,
14041404
IN_PROGRESS_DESCRIBE_JOB_RESULT,

‎tests/unit/test_upload_data.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
import os
1616

17-
from mock import Mock
17+
from mock import MagicMock, Mock
1818
import pytest
1919

2020
import sagemaker
@@ -30,7 +30,7 @@
3030

3131
@pytest.fixture()
3232
def sagemaker_session():
33-
boto_mock = Mock(name="boto_session")
33+
boto_mock = MagicMock(name="boto_session")
3434
ims = sagemaker.Session(boto_session=boto_mock)
3535
ims.default_bucket = Mock(name="default_bucket", return_value=BUCKET_NAME)
3636
return ims
@@ -39,9 +39,9 @@ def sagemaker_session():
3939
@pytest.fixture()
4040
def sagemaker_session_custom_endpoint():
4141

42-
boto_session = Mock("boto_session")
42+
boto_session = MagicMock("boto_session")
4343
resource_mock = Mock("resource")
44-
client_mock = Mock("client")
44+
client_mock = MagicMock("client")
4545
boto_attrs = {"region_name": "us-east-1"}
4646
boto_session.configure_mock(**boto_attrs)
4747
boto_session.resource = Mock(name="resource", return_value=resource_mock)

‎tests/unit/test_upload_string_as_file_body.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
import os
1616

17-
from mock import Mock
17+
from mock import MagicMock, Mock
1818
import pytest
1919

2020
import sagemaker
@@ -29,8 +29,8 @@
2929

3030
@pytest.fixture()
3131
def sagemaker_session():
32-
boto_mock = Mock(name="boto_session")
33-
client_mock = Mock()
32+
boto_mock = MagicMock(name="boto_session")
33+
client_mock = MagicMock()
3434
client_mock.get_caller_identity.return_value = {
3535
"UserId": "mock_user_id",
3636
"Account": "012345678910",

‎tests/unit/test_utils.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -309,8 +309,8 @@ def test_generate_tensorboard_url_domain_non_string():
309309

310310
@patch("os.makedirs")
311311
def test_download_folder(makedirs):
312-
boto_mock = Mock(name="boto_session")
313-
session = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
312+
boto_mock = MagicMock(name="boto_session")
313+
session = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
314314
s3_mock = boto_mock.resource("s3")
315315

316316
obj_mock = Mock()
@@ -363,10 +363,10 @@ def obj_mock_download(path):
363363

364364
@patch("os.makedirs")
365365
def test_download_folder_points_to_single_file(makedirs):
366-
boto_mock = Mock(name="boto_session")
366+
boto_mock = MagicMock(name="boto_session")
367367
boto_mock.client("sts").get_caller_identity.return_value = {"Account": "123"}
368368

369-
session = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
369+
session = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
370370

371371
train_data = Mock()
372372

@@ -390,11 +390,11 @@ def test_download_folder_points_to_single_file(makedirs):
390390

391391

392392
def test_download_file():
393-
boto_mock = Mock(name="boto_session")
393+
boto_mock = MagicMock(name="boto_session")
394394
boto_mock.client("sts").get_caller_identity.return_value = {"Account": "123"}
395395
bucket_mock = Mock()
396396
boto_mock.resource("s3").Bucket.return_value = bucket_mock
397-
session = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
397+
session = sagemaker.Session(boto_session=boto_mock, sagemaker_client=MagicMock())
398398

399399
sagemaker.utils.download_file(
400400
BUCKET_NAME, "/prefix/path/file.tar.gz", "/tmp/file.tar.gz", session

0 commit comments

Comments
 (0)