Skip to content

Commit b786a51

Browse files
laurenyulaurenyu
authored
breaking: rename session parameter to sagemaker_session in S3 utility classes (#1663)
This is for consistency with the sagemaker_session parameter that is used everywhere else in this library.
1 parent f55bc9d commit b786a51

File tree

11 files changed

+105
-107
lines changed

11 files changed

+105
-107
lines changed

‎src/sagemaker/estimator.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -404,7 +404,7 @@ def _prepare_rules(self):
404404
s3_uri = S3Uploader.upload(
405405
local_path=rule.rule_parameters["source_s3_uri"],
406406
desired_s3_uri=desired_s3_uri,
407-
session=self.sagemaker_session,
407+
sagemaker_session=self.sagemaker_session,
408408
)
409409
rule.rule_parameters["source_s3_uri"] = s3_uri
410410
# Save the request dictionary for the rule.

‎src/sagemaker/model_monitor/model_monitoring.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -871,7 +871,7 @@ def _normalize_baseline_inputs(self, baseline_inputs=None):
871871
S3Uploader.upload(
872872
local_path=file_input.source,
873873
desired_s3_uri=s3_uri,
874-
session=self.sagemaker_session,
874+
sagemaker_session=self.sagemaker_session,
875875
)
876876
file_input.source = s3_uri
877877
normalized_inputs.append(file_input)
@@ -944,7 +944,7 @@ def _s3_uri_from_local_path(self, path):
944944
str(uuid.uuid4()),
945945
)
946946
S3Uploader.upload(
947-
local_path=path, desired_s3_uri=s3_uri, session=self.sagemaker_session
947+
local_path=path, desired_s3_uri=s3_uri, sagemaker_session=self.sagemaker_session
948948
)
949949
path = os.path.join(s3_uri, os.path.basename(path))
950950
return path
@@ -1771,7 +1771,7 @@ def _upload_and_convert_to_processing_input(self, source, destination, name):
17711771
name,
17721772
)
17731773
S3Uploader.upload(
1774-
local_path=source, desired_s3_uri=s3_uri, session=self.sagemaker_session
1774+
local_path=source, desired_s3_uri=s3_uri, sagemaker_session=self.sagemaker_session
17751775
)
17761776
source = s3_uri
17771777

‎src/sagemaker/model_monitor/monitoring_files.py

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ def save(self, new_save_location_s3_uri=None):
7272
body=json.dumps(self.body_dict),
7373
desired_s3_uri=self.file_s3_uri,
7474
kms_key=self.kms_key,
75-
session=self.session,
75+
sagemaker_session=self.session,
7676
)
7777

7878

@@ -119,7 +119,9 @@ def from_s3_uri(cls, statistics_file_s3_uri, kms_key=None, sagemaker_session=Non
119119
"""
120120
try:
121121
body_dict = json.loads(
122-
S3Downloader.read_file(s3_uri=statistics_file_s3_uri, session=sagemaker_session)
122+
S3Downloader.read_file(
123+
s3_uri=statistics_file_s3_uri, sagemaker_session=sagemaker_session
124+
)
123125
)
124126
except ClientError as error:
125127
print(
@@ -163,7 +165,7 @@ def from_string(
163165
body=statistics_file_string,
164166
desired_s3_uri=desired_s3_uri,
165167
kms_key=kms_key,
166-
session=sagemaker_session,
168+
sagemaker_session=sagemaker_session,
167169
)
168170

169171
return Statistics.from_s3_uri(
@@ -243,7 +245,9 @@ def from_s3_uri(cls, constraints_file_s3_uri, kms_key=None, sagemaker_session=No
243245
"""
244246
try:
245247
body_dict = json.loads(
246-
S3Downloader.read_file(s3_uri=constraints_file_s3_uri, session=sagemaker_session)
248+
S3Downloader.read_file(
249+
s3_uri=constraints_file_s3_uri, sagemaker_session=sagemaker_session
250+
)
247251
)
248252
except ClientError as error:
249253
print(
@@ -290,7 +294,7 @@ def from_string(
290294
body=constraints_file_string,
291295
desired_s3_uri=desired_s3_uri,
292296
kms_key=kms_key,
293-
session=sagemaker_session,
297+
sagemaker_session=sagemaker_session,
294298
)
295299

296300
return Constraints.from_s3_uri(
@@ -396,7 +400,7 @@ def from_s3_uri(cls, constraint_violations_file_s3_uri, kms_key=None, sagemaker_
396400
try:
397401
body_dict = json.loads(
398402
S3Downloader.read_file(
399-
s3_uri=constraint_violations_file_s3_uri, session=sagemaker_session
403+
s3_uri=constraint_violations_file_s3_uri, sagemaker_session=sagemaker_session
400404
)
401405
)
402406
except ClientError as error:
@@ -445,7 +449,7 @@ def from_string(
445449
body=constraint_violations_file_string,
446450
desired_s3_uri=desired_s3_uri,
447451
kms_key=kms_key,
448-
session=sagemaker_session,
452+
sagemaker_session=sagemaker_session,
449453
)
450454

451455
return ConstraintViolations.from_s3_uri(

‎src/sagemaker/processing.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,7 @@ def _normalize_inputs(self, inputs=None):
222222
s3_uri = S3Uploader.upload(
223223
local_path=file_input.source,
224224
desired_s3_uri=desired_s3_uri,
225-
session=self.sagemaker_session,
225+
sagemaker_session=self.sagemaker_session,
226226
)
227227
file_input.source = s3_uri
228228
normalized_inputs.append(file_input)
@@ -480,7 +480,7 @@ def _upload_code(self, code):
480480
self._CODE_CONTAINER_INPUT_NAME,
481481
)
482482
return S3Uploader.upload(
483-
local_path=code, desired_s3_uri=desired_s3_uri, session=self.sagemaker_session
483+
local_path=code, desired_s3_uri=desired_s3_uri, sagemaker_session=self.sagemaker_session
484484
)
485485

486486
def _convert_code_and_add_to_inputs(self, inputs, s3_uri):

‎src/sagemaker/s3.py

Lines changed: 24 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -20,19 +20,6 @@
2020

2121
logger = logging.getLogger("sagemaker")
2222

23-
SESSION_V2_RENAME_MESSAGE = (
24-
"Parameter 'session' will be renamed to 'sagemaker_session' in SageMaker Python SDK v2."
25-
)
26-
27-
28-
def _session_v2_rename_warning(session):
29-
"""
30-
Args:
31-
session (sagemaker.session.Session):
32-
"""
33-
if session is not None:
34-
logger.warning(SESSION_V2_RENAME_MESSAGE)
35-
3623

3724
def parse_s3_url(url):
3825
"""Returns an (s3 bucket, key name/prefix) tuple from a url with an s3
@@ -53,15 +40,15 @@ class S3Uploader(object):
5340
"""Contains static methods for uploading directories or files to S3."""
5441

5542
@staticmethod
56-
def upload(local_path, desired_s3_uri, kms_key=None, session=None):
43+
def upload(local_path, desired_s3_uri, kms_key=None, sagemaker_session=None):
5744
"""Static method that uploads a given file or directory to S3.
5845
5946
Args:
6047
local_path (str): Path (absolute or relative) of local file or directory to upload.
6148
desired_s3_uri (str): The desired S3 location to upload to. It is the prefix to
6249
which the local filename will be added.
6350
kms_key (str): The KMS key to use to encrypt the files.
64-
session (sagemaker.session.Session): Session object which
51+
sagemaker_session (sagemaker.session.Session): Session object which
6552
manages interactions with Amazon SageMaker APIs and any other
6653
AWS services needed. If not specified, the estimator creates one
6754
using the default AWS configuration chain.
@@ -70,10 +57,7 @@ def upload(local_path, desired_s3_uri, kms_key=None, session=None):
7057
The S3 uri of the uploaded file(s).
7158
7259
"""
73-
if session is not None:
74-
_session_v2_rename_warning(session)
75-
76-
sagemaker_session = session or Session()
60+
sagemaker_session = sagemaker_session or Session()
7761
bucket, key_prefix = parse_s3_url(url=desired_s3_uri)
7862
if kms_key is not None:
7963
extra_args = {"SSEKMSKeyId": kms_key}
@@ -85,24 +69,23 @@ def upload(local_path, desired_s3_uri, kms_key=None, session=None):
8569
)
8670

8771
@staticmethod
88-
def upload_string_as_file_body(body, desired_s3_uri=None, kms_key=None, session=None):
72+
def upload_string_as_file_body(body, desired_s3_uri=None, kms_key=None, sagemaker_session=None):
8973
"""Static method that uploads a given file or directory to S3.
9074
9175
Args:
9276
body (str): String representing the body of the file.
9377
desired_s3_uri (str): The desired S3 uri to upload to.
9478
kms_key (str): The KMS key to use to encrypt the files.
95-
session (sagemaker.session.Session): AWS session to use. Automatically
96-
generates one if not provided.
79+
sagemaker_session (sagemaker.session.Session): Session object which
80+
manages interactions with Amazon SageMaker APIs and any other
81+
AWS services needed. If not specified, the estimator creates one
82+
using the default AWS configuration chain.
9783
9884
Returns:
9985
str: The S3 uri of the uploaded file(s).
10086
10187
"""
102-
if session is not None:
103-
_session_v2_rename_warning(session)
104-
105-
sagemaker_session = session or Session()
88+
sagemaker_session = sagemaker_session or Session()
10689
bucket, key = parse_s3_url(desired_s3_uri)
10790

10891
sagemaker_session.upload_string_as_file_body(
@@ -116,23 +99,19 @@ class S3Downloader(object):
11699
"""Contains static methods for downloading directories or files from S3."""
117100

118101
@staticmethod
119-
def download(s3_uri, local_path, kms_key=None, session=None):
102+
def download(s3_uri, local_path, kms_key=None, sagemaker_session=None):
120103
"""Static method that downloads a given S3 uri to the local machine.
121104
122105
Args:
123106
s3_uri (str): An S3 uri to download from.
124107
local_path (str): A local path to download the file(s) to.
125108
kms_key (str): The KMS key to use to decrypt the files.
126-
session (sagemaker.session.Session): Session object which
109+
sagemaker_session (sagemaker.session.Session): Session object which
127110
manages interactions with Amazon SageMaker APIs and any other
128111
AWS services needed. If not specified, the estimator creates one
129112
using the default AWS configuration chain.
130-
131113
"""
132-
if session is not None:
133-
_session_v2_rename_warning(session)
134-
135-
sagemaker_session = session or Session()
114+
sagemaker_session = sagemaker_session or Session()
136115
bucket, key_prefix = parse_s3_url(url=s3_uri)
137116
if kms_key is not None:
138117
extra_args = {"SSECustomerKey": kms_key}
@@ -144,43 +123,39 @@ def download(s3_uri, local_path, kms_key=None, session=None):
144123
)
145124

146125
@staticmethod
147-
def read_file(s3_uri, session=None):
126+
def read_file(s3_uri, sagemaker_session=None):
148127
"""Static method that returns the contents of an s3 uri file body as a string.
149128
150129
Args:
151130
s3_uri (str): An S3 uri that refers to a single file.
152-
session (sagemaker.session.Session): AWS session to use. Automatically
153-
generates one if not provided.
131+
sagemaker_session (sagemaker.session.Session): Session object which
132+
manages interactions with Amazon SageMaker APIs and any other
133+
AWS services needed. If not specified, the estimator creates one
134+
using the default AWS configuration chain.
154135
155136
Returns:
156137
str: The body of the file.
157-
158138
"""
159-
if session is not None:
160-
_session_v2_rename_warning(session)
161-
162-
sagemaker_session = session or Session()
139+
sagemaker_session = sagemaker_session or Session()
163140
bucket, key_prefix = parse_s3_url(url=s3_uri)
164141

165142
return sagemaker_session.read_s3_file(bucket=bucket, key_prefix=key_prefix)
166143

167144
@staticmethod
168-
def list(s3_uri, session=None):
145+
def list(s3_uri, sagemaker_session=None):
169146
"""Static method that lists the contents of an S3 uri.
170147
171148
Args:
172149
s3_uri (str): The S3 base uri to list objects in.
173-
session (sagemaker.session.Session): AWS session to use. Automatically
174-
generates one if not provided.
150+
sagemaker_session (sagemaker.session.Session): Session object which
151+
manages interactions with Amazon SageMaker APIs and any other
152+
AWS services needed. If not specified, the estimator creates one
153+
using the default AWS configuration chain.
175154
176155
Returns:
177156
[str]: The list of S3 URIs in the given S3 base uri.
178-
179157
"""
180-
if session is not None:
181-
_session_v2_rename_warning(session)
182-
183-
sagemaker_session = session or Session()
158+
sagemaker_session = sagemaker_session or Session()
184159
bucket, key_prefix = parse_s3_url(url=s3_uri)
185160

186161
file_keys = sagemaker_session.list_s3_files(bucket=bucket, key_prefix=key_prefix)

‎tests/integ/test_model_monitor.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1509,7 +1509,7 @@ def test_default_monitor_monitoring_execution_interactions(
15091509
desired_s3_uri = os.path.join(executions[-1].output.destination, file_name)
15101510

15111511
S3Uploader.upload_string_as_file_body(
1512-
body=file_body, desired_s3_uri=desired_s3_uri, session=sagemaker_session
1512+
body=file_body, desired_s3_uri=desired_s3_uri, sagemaker_session=sagemaker_session
15131513
)
15141514

15151515
statistics = my_attached_monitor.latest_monitoring_statistics()
@@ -1522,7 +1522,7 @@ def test_default_monitor_monitoring_execution_interactions(
15221522
desired_s3_uri = os.path.join(executions[-1].output.destination, file_name)
15231523

15241524
S3Uploader.upload_string_as_file_body(
1525-
body=file_body, desired_s3_uri=desired_s3_uri, session=sagemaker_session
1525+
body=file_body, desired_s3_uri=desired_s3_uri, sagemaker_session=sagemaker_session
15261526
)
15271527

15281528
constraint_violations = my_attached_monitor.latest_monitoring_constraint_violations()
@@ -2473,7 +2473,7 @@ def test_byoc_monitor_monitoring_execution_interactions(
24732473
desired_s3_uri = os.path.join(executions[-1].output.destination, file_name)
24742474

24752475
S3Uploader.upload_string_as_file_body(
2476-
body=file_body, desired_s3_uri=desired_s3_uri, session=sagemaker_session
2476+
body=file_body, desired_s3_uri=desired_s3_uri, sagemaker_session=sagemaker_session
24772477
)
24782478

24792479
statistics = my_attached_monitor.latest_monitoring_statistics()
@@ -2486,7 +2486,7 @@ def test_byoc_monitor_monitoring_execution_interactions(
24862486
desired_s3_uri = os.path.join(executions[-1].output.destination, file_name)
24872487

24882488
S3Uploader.upload_string_as_file_body(
2489-
body=file_body, desired_s3_uri=desired_s3_uri, session=sagemaker_session
2489+
body=file_body, desired_s3_uri=desired_s3_uri, sagemaker_session=sagemaker_session
24902490
)
24912491

24922492
constraint_violations = my_attached_monitor.latest_monitoring_constraint_violations()
@@ -2557,10 +2557,10 @@ def _upload_captured_data_to_endpoint(sagemaker_session, predictor):
25572557
S3Uploader.upload(
25582558
local_path=os.path.join(DATA_DIR, "monitor/captured-data.jsonl"),
25592559
desired_s3_uri=s3_uri_previous_hour,
2560-
session=sagemaker_session,
2560+
sagemaker_session=sagemaker_session,
25612561
)
25622562
S3Uploader.upload(
25632563
local_path=os.path.join(DATA_DIR, "monitor/captured-data.jsonl"),
25642564
desired_s3_uri=s3_uri_current_hour,
2565-
session=sagemaker_session,
2565+
sagemaker_session=sagemaker_session,
25662566
)

‎tests/integ/test_monitoring_files.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ def test_statistics_object_creation_from_s3_uri_with_customizations(
108108
body=file_body,
109109
desired_s3_uri=desired_s3_uri,
110110
kms_key=monitoring_files_kms_key,
111-
session=sagemaker_session,
111+
sagemaker_session=sagemaker_session,
112112
)
113113

114114
statistics = Statistics.from_s3_uri(
@@ -137,7 +137,7 @@ def test_statistics_object_creation_from_s3_uri_without_customizations(sagemaker
137137
)
138138

139139
s3_uri = S3Uploader.upload_string_as_file_body(
140-
body=file_body, desired_s3_uri=desired_s3_uri, session=sagemaker_session
140+
body=file_body, desired_s3_uri=desired_s3_uri, sagemaker_session=sagemaker_session
141141
)
142142

143143
statistics = Statistics.from_s3_uri(
@@ -259,7 +259,7 @@ def test_constraints_object_creation_from_s3_uri_with_customizations(
259259
body=file_body,
260260
desired_s3_uri=desired_s3_uri,
261261
kms_key=monitoring_files_kms_key,
262-
session=sagemaker_session,
262+
sagemaker_session=sagemaker_session,
263263
)
264264

265265
constraints = Constraints.from_s3_uri(
@@ -288,7 +288,7 @@ def test_constraints_object_creation_from_s3_uri_without_customizations(sagemake
288288
)
289289

290290
s3_uri = S3Uploader.upload_string_as_file_body(
291-
body=file_body, desired_s3_uri=desired_s3_uri, session=sagemaker_session
291+
body=file_body, desired_s3_uri=desired_s3_uri, sagemaker_session=sagemaker_session
292292
)
293293

294294
constraints = Constraints.from_s3_uri(
@@ -388,7 +388,7 @@ def test_constraint_violations_object_creation_from_s3_uri_with_customizations(
388388
body=file_body,
389389
desired_s3_uri=desired_s3_uri,
390390
kms_key=monitoring_files_kms_key,
391-
session=sagemaker_session,
391+
sagemaker_session=sagemaker_session,
392392
)
393393

394394
constraint_violations = ConstraintViolations.from_s3_uri(
@@ -419,7 +419,7 @@ def test_constraint_violations_object_creation_from_s3_uri_without_customization
419419
)
420420

421421
s3_uri = S3Uploader.upload_string_as_file_body(
422-
body=file_body, desired_s3_uri=desired_s3_uri, session=sagemaker_session
422+
body=file_body, desired_s3_uri=desired_s3_uri, sagemaker_session=sagemaker_session
423423
)
424424

425425
constraint_violations = ConstraintViolations.from_s3_uri(

‎tests/integ/test_multi_variant_endpoint.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -92,8 +92,8 @@ def multi_variant_endpoint(sagemaker_session):
9292
prefix = "sagemaker/DEMO-VariantTargeting"
9393
model_url = S3Uploader.upload(
9494
local_path=XG_BOOST_MODEL_LOCAL_PATH,
95-
desired_s3_uri="s3://" + bucket + "/" + prefix,
96-
session=sagemaker_session,
95+
desired_s3_uri="s3://{}/{}".format(bucket, prefix),
96+
sagemaker_session=sagemaker_session,
9797
)
9898

9999
image_uri = get_image_uri(sagemaker_session.boto_session.region_name, "xgboost", "0.90-1")

0 commit comments

Comments
 (0)