Skip to content

Commit 05be636

Browse files
laurenyulaurenyu
committed
infra: clean up pickle.load logic in integ tests
Because we no longer run our tests with Python 2, we no longer need the branched logic for pickle.load args.
1 parent c0134e5 commit 05be636

11 files changed

+27
-60
lines changed

‎tests/integ/test_airflow_config.py

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import gzip
1616
import os
1717
import pickle
18-
import sys
1918
import pytest
2019
import tests.integ
2120

@@ -102,11 +101,10 @@ def test_byo_airflow_config_uploads_data_source_to_s3_when_inputs_provided(
102101
def test_kmeans_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_instance_type):
103102
with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
104103
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
105-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
106104

107105
# Load the data into memory as numpy arrays
108106
with gzip.open(data_path, "rb") as f:
109-
train_set, _, _ = pickle.load(f, **pickle_args)
107+
train_set, _, _ = pickle.load(f, encoding="latin1")
110108

111109
kmeans = KMeans(
112110
role=ROLE,
@@ -141,11 +139,10 @@ def test_kmeans_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_
141139
def test_fm_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_instance_type):
142140
with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
143141
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
144-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
145142

146143
# Load the data into memory as numpy arrays
147144
with gzip.open(data_path, "rb") as f:
148-
train_set, _, _ = pickle.load(f, **pickle_args)
145+
train_set, _, _ = pickle.load(f, encoding="latin1")
149146

150147
fm = FactorizationMachines(
151148
role=ROLE,
@@ -207,11 +204,10 @@ def test_ipinsights_airflow_config_uploads_data_source_to_s3(sagemaker_session,
207204
def test_knn_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_instance_type):
208205
with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
209206
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
210-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
211207

212208
# Load the data into memory as numpy arrays
213209
with gzip.open(data_path, "rb") as f:
214-
train_set, _, _ = pickle.load(f, **pickle_args)
210+
train_set, _, _ = pickle.load(f, encoding="latin1")
215211

216212
knn = KNN(
217213
role=ROLE,
@@ -278,11 +274,10 @@ def test_linearlearner_airflow_config_uploads_data_source_to_s3(
278274
):
279275
with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
280276
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
281-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
282277

283278
# Load the data into memory as numpy arrays
284279
with gzip.open(data_path, "rb") as f:
285-
train_set, _, _ = pickle.load(f, **pickle_args)
280+
train_set, _, _ = pickle.load(f, encoding="latin1")
286281

287282
train_set[1][:100] = 1
288283
train_set[1][100:200] = 0
@@ -381,11 +376,10 @@ def test_ntm_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_ins
381376
def test_pca_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_instance_type):
382377
with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
383378
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
384-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
385379

386380
# Load the data into memory as numpy arrays
387381
with gzip.open(data_path, "rb") as f:
388-
train_set, _, _ = pickle.load(f, **pickle_args)
382+
train_set, _, _ = pickle.load(f, encoding="latin1")
389383

390384
pca = PCA(
391385
role=ROLE,

‎tests/integ/test_byo_estimator.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import json
1717
import os
1818
import pickle
19-
import sys
2019

2120
import pytest
2221

@@ -58,10 +57,9 @@ def test_byo_estimator(sagemaker_session, region, cpu_instance_type):
5857

5958
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
6059
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
61-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
6260

6361
with gzip.open(data_path, "rb") as f:
64-
train_set, _, _ = pickle.load(f, **pickle_args)
62+
train_set, _, _ = pickle.load(f, encoding="latin1")
6563

6664
prefix = "test_byo_estimator"
6765
key = "recordio-pb-data"
@@ -107,10 +105,9 @@ def test_async_byo_estimator(sagemaker_session, region, cpu_instance_type):
107105

108106
with timeout(minutes=5):
109107
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
110-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
111108

112109
with gzip.open(data_path, "rb") as f:
113-
train_set, _, _ = pickle.load(f, **pickle_args)
110+
train_set, _, _ = pickle.load(f, encoding="latin1")
114111

115112
prefix = "test_byo_estimator"
116113
key = "recordio-pb-data"

‎tests/integ/test_factorization_machines.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,11 +29,10 @@ def test_factorization_machines(sagemaker_session, cpu_instance_type):
2929

3030
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
3131
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
32-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
3332

3433
# Load the data into memory as numpy arrays
3534
with gzip.open(data_path, "rb") as f:
36-
train_set, _, _ = pickle.load(f, **pickle_args)
35+
train_set, _, _ = pickle.load(f, encoding="latin1")
3736

3837
fm = FactorizationMachines(
3938
role="SageMakerRole",
@@ -71,11 +70,10 @@ def test_async_factorization_machines(sagemaker_session, cpu_instance_type):
7170

7271
with timeout(minutes=5):
7372
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
74-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
7573

7674
# Load the data into memory as numpy arrays
7775
with gzip.open(data_path, "rb") as f:
78-
train_set, _, _ = pickle.load(f, **pickle_args)
76+
train_set, _, _ = pickle.load(f, encoding="latin1")
7977

8078
fm = FactorizationMachines(
8179
role="SageMakerRole",

‎tests/integ/test_kmeans.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import json
1717
import os
1818
import pickle
19-
import sys
2019
import time
2120

2221
import pytest
@@ -31,11 +30,10 @@ def test_kmeans(sagemaker_session, cpu_instance_type):
3130
job_name = unique_name_from_base("kmeans")
3231
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
3332
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
34-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
3533

3634
# Load the data into memory as numpy arrays
3735
with gzip.open(data_path, "rb") as f:
38-
train_set, _, _ = pickle.load(f, **pickle_args)
36+
train_set, _, _ = pickle.load(f, encoding="latin1")
3937

4038
kmeans = KMeans(
4139
role="SageMakerRole",
@@ -94,11 +92,10 @@ def test_async_kmeans(sagemaker_session, cpu_instance_type):
9492

9593
with timeout(minutes=5):
9694
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
97-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
9895

9996
# Load the data into memory as numpy arrays
10097
with gzip.open(data_path, "rb") as f:
101-
train_set, _, _ = pickle.load(f, **pickle_args)
98+
train_set, _, _ = pickle.load(f, encoding="latin1")
10299

103100
kmeans = KMeans(
104101
role="SageMakerRole",

‎tests/integ/test_knn.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import gzip
1616
import os
1717
import pickle
18-
import sys
1918
import time
2019

2120
from sagemaker import KNN, KNNModel
@@ -29,11 +28,10 @@ def test_knn_regressor(sagemaker_session, cpu_instance_type):
2928

3029
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
3130
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
32-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
3331

3432
# Load the data into memory as numpy arrays
3533
with gzip.open(data_path, "rb") as f:
36-
train_set, _, _ = pickle.load(f, **pickle_args)
34+
train_set, _, _ = pickle.load(f, encoding="latin1")
3735

3836
knn = KNN(
3937
role="SageMakerRole",
@@ -66,11 +64,10 @@ def test_async_knn_classifier(sagemaker_session, cpu_instance_type):
6664

6765
with timeout(minutes=5):
6866
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
69-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
7067

7168
# Load the data into memory as numpy arrays
7269
with gzip.open(data_path, "rb") as f:
73-
train_set, _, _ = pickle.load(f, **pickle_args)
70+
train_set, _, _ = pickle.load(f, encoding="latin1")
7471

7572
knn = KNN(
7673
role="SageMakerRole",

‎tests/integ/test_linear_learner.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import gzip
1616
import os
1717
import pickle
18-
import sys
1918
import time
2019

2120
import numpy as np
@@ -33,11 +32,10 @@ def test_linear_learner(sagemaker_session, cpu_instance_type):
3332

3433
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
3534
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
36-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
3735

3836
# Load the data into memory as numpy arrays
3937
with gzip.open(data_path, "rb") as f:
40-
train_set, _, _ = pickle.load(f, **pickle_args)
38+
train_set, _, _ = pickle.load(f, encoding="latin1")
4139

4240
train_set[1][:100] = 1
4341
train_set[1][100:200] = 0
@@ -102,11 +100,10 @@ def test_linear_learner_multiclass(sagemaker_session, cpu_instance_type):
102100

103101
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
104102
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
105-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
106103

107104
# Load the data into memory as numpy arrays
108105
with gzip.open(data_path, "rb") as f:
109-
train_set, _, _ = pickle.load(f, **pickle_args)
106+
train_set, _, _ = pickle.load(f, encoding="latin1")
110107

111108
train_set = train_set[0], train_set[1].astype(np.dtype("float32"))
112109

@@ -137,11 +134,10 @@ def test_async_linear_learner(sagemaker_session, cpu_instance_type):
137134

138135
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
139136
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
140-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
141137

142138
# Load the data into memory as numpy arrays
143139
with gzip.open(data_path, "rb") as f:
144-
train_set, _, _ = pickle.load(f, **pickle_args)
140+
train_set, _, _ = pickle.load(f, encoding="latin1")
145141

146142
train_set[1][:100] = 1
147143
train_set[1][100:200] = 0

‎tests/integ/test_pca.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import gzip
1616
import os
1717
import pickle
18-
import sys
1918
import time
2019

2120
import sagemaker.amazon.pca
@@ -29,11 +28,10 @@ def test_pca(sagemaker_session, cpu_instance_type):
2928

3029
with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
3130
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
32-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
3331

3432
# Load the data into memory as numpy arrays
3533
with gzip.open(data_path, "rb") as f:
36-
train_set, _, _ = pickle.load(f, **pickle_args)
34+
train_set, _, _ = pickle.load(f, encoding="latin1")
3735

3836
pca = sagemaker.amazon.pca.PCA(
3937
role="SageMakerRole",
@@ -72,11 +70,10 @@ def test_async_pca(sagemaker_session, cpu_instance_type):
7270

7371
with timeout(minutes=5):
7472
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
75-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
7673

7774
# Load the data into memory as numpy arrays
7875
with gzip.open(data_path, "rb") as f:
79-
train_set, _, _ = pickle.load(f, **pickle_args)
76+
train_set, _, _ = pickle.load(f, encoding="latin1")
8077

8178
pca = sagemaker.amazon.pca.PCA(
8279
role="SageMakerRole",

‎tests/integ/test_record_set.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import gzip
1616
import os
1717
import pickle
18-
import sys
1918

2019
from six.moves.urllib.parse import urlparse
2120

@@ -29,9 +28,8 @@ def test_record_set(sagemaker_session, cpu_instance_type):
2928
In particular, test that the objects uploaded to the S3 bucket are encrypted.
3029
"""
3130
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
32-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
3331
with gzip.open(data_path, "rb") as file_object:
34-
train_set, _, _ = pickle.load(file_object, **pickle_args)
32+
train_set, _, _ = pickle.load(file_object, encoding="latin1")
3533
kmeans = KMeans(
3634
role="SageMakerRole",
3735
train_instance_count=1,

‎tests/integ/test_transformer.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import json
1717
import os
1818
import pickle
19-
import sys
2019
import time
2120

2221
import pytest
@@ -109,12 +108,11 @@ def test_transform_mxnet(
109108
@pytest.mark.canary_quick
110109
def test_attach_transform_kmeans(sagemaker_session, cpu_instance_type):
111110
data_path = os.path.join(DATA_DIR, "one_p_mnist")
112-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
113111

114112
# Load the data into memory as numpy arrays
115113
train_set_path = os.path.join(data_path, "mnist.pkl.gz")
116114
with gzip.open(train_set_path, "rb") as f:
117-
train_set, _, _ = pickle.load(f, **pickle_args)
115+
train_set, _, _ = pickle.load(f, encoding="latin1")
118116

119117
kmeans = KMeans(
120118
role="SageMakerRole",
@@ -177,7 +175,7 @@ def test_transform_pytorch_vpc_custom_model_bucket(
177175
entry_point=os.path.join(data_dir, "mnist.py"),
178176
role="SageMakerRole",
179177
framework_version=pytorch_full_version,
180-
py_version="py3",
178+
py_version=PYTHON_VERSION,
181179
sagemaker_session=sagemaker_session,
182180
vpc_config={"Subnets": subnet_ids, "SecurityGroupIds": [security_group_id]},
183181
code_location="s3://{}".format(custom_bucket_name),
@@ -232,13 +230,12 @@ def test_transform_mxnet_tags(
232230

233231
def test_transform_byo_estimator(sagemaker_session, cpu_instance_type):
234232
data_path = os.path.join(DATA_DIR, "one_p_mnist")
235-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
236233
tags = [{"Key": "some-tag", "Value": "value-for-tag"}]
237234

238235
# Load the data into memory as numpy arrays
239236
train_set_path = os.path.join(data_path, "mnist.pkl.gz")
240237
with gzip.open(train_set_path, "rb") as f:
241-
train_set, _, _ = pickle.load(f, **pickle_args)
238+
train_set, _, _ = pickle.load(f, encoding="latin1")
242239

243240
kmeans = KMeans(
244241
role="SageMakerRole",

‎tests/integ/test_tuner.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import json
1717
import os
1818
import pickle
19-
import sys
2019
import time
2120

2221
import numpy as np
@@ -55,10 +54,10 @@
5554
@pytest.fixture(scope="module")
5655
def kmeans_train_set(sagemaker_session):
5756
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
58-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
57+
5958
# Load the data into memory as numpy arrays
6059
with gzip.open(data_path, "rb") as f:
61-
train_set, _, _ = pickle.load(f, **pickle_args)
60+
train_set, _, _ = pickle.load(f, encoding="latin1")
6261

6362
return train_set
6463

@@ -847,10 +846,9 @@ def test_tuning_byo_estimator(sagemaker_session, cpu_instance_type):
847846

848847
with timeout(minutes=TUNING_DEFAULT_TIMEOUT_MINUTES):
849848
data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
850-
pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
851849

852850
with gzip.open(data_path, "rb") as f:
853-
train_set, _, _ = pickle.load(f, **pickle_args)
851+
train_set, _, _ = pickle.load(f, encoding="latin1")
854852

855853
prefix = "test_byo_estimator"
856854
key = "recordio-pb-data"

0 commit comments

Comments
 (0)