infra: clean up pickle.load logic in integ tests

tests/integ/datasets.py

@@ -0,0 +1,27 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import gzip
+import os
+import pickle
+
+from tests.integ import DATA_DIR
+
+
+def one_p_mnist():
+ data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
+ with gzip.open(data_path, "rb") as f:
+ training_set, _, _ = pickle.load(f, encoding="latin1")
+
+ return training_set

tests/integ/test_airflow_config.py

@@ -12,15 +12,17 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import gzip
 import os
-import pickle
-import sys
-import pytest
-import tests.integ
 
+import airflow
+import pytest
 import numpy as np
+from airflow import DAG
+from airflow.contrib.operators.sagemaker_training_operator import SageMakerTrainingOperator
+from airflow.contrib.operators.sagemaker_transform_operator import SageMakerTransformOperator
+from six.moves.urllib.parse import urlparse
 
+import tests.integ
 from sagemaker import (
  KMeans,
  FactorizationMachines,
@@ -40,21 +42,13 @@
 from sagemaker.pytorch.estimator import PyTorch
 from sagemaker.sklearn import SKLearn
 from sagemaker.tensorflow import TensorFlow
-from sagemaker.workflow import airflow as sm_airflow
 from sagemaker.utils import sagemaker_timestamp
-
-import airflow
-from airflow import DAG
-from airflow.contrib.operators.sagemaker_training_operator import SageMakerTrainingOperator
-from airflow.contrib.operators.sagemaker_transform_operator import SageMakerTransformOperator
-
+from sagemaker.workflow import airflow as sm_airflow
 from sagemaker.xgboost import XGBoost
-from tests.integ import DATA_DIR, PYTHON_VERSION
+from tests.integ import datasets, DATA_DIR, PYTHON_VERSION
 from tests.integ.record_set import prepare_record_set_from_local_files
 from tests.integ.timeout import timeout
 
-from six.moves.urllib.parse import urlparse
-
 PYTORCH_MNIST_DIR = os.path.join(DATA_DIR, "pytorch_mnist")
 PYTORCH_MNIST_SCRIPT = os.path.join(PYTORCH_MNIST_DIR, "mnist.py")
 AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS = 10
@@ -101,13 +95,6 @@ def test_byo_airflow_config_uploads_data_source_to_s3_when_inputs_provided(
 @pytest.mark.canary_quick
 def test_kmeans_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_instance_type):
  with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
- data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
- pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
-
- # Load the data into memory as numpy arrays
- with gzip.open(data_path, "rb") as f:
- train_set, _, _ = pickle.load(f, **pickle_args)
-
  kmeans = KMeans(
  role=ROLE,
  train_instance_count=SINGLE_INSTANCE_COUNT,
@@ -126,7 +113,7 @@ def test_kmeans_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_
  kmeans.center_factor = 1
  kmeans.eval_metrics = ["ssd", "msd"]
 
- records = kmeans.record_set(train_set[0][:100])
+ records = kmeans.record_set(datasets.one_p_mnist()[0][:100])
 
  training_config = _build_airflow_workflow(
  estimator=kmeans, instance_type=cpu_instance_type, inputs=records
@@ -140,13 +127,6 @@ def test_kmeans_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_
 
 def test_fm_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_instance_type):
  with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
- data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
- pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
-
- # Load the data into memory as numpy arrays
- with gzip.open(data_path, "rb") as f:
- train_set, _, _ = pickle.load(f, **pickle_args)
-
  fm = FactorizationMachines(
  role=ROLE,
  train_instance_count=SINGLE_INSTANCE_COUNT,
@@ -160,7 +140,8 @@ def test_fm_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_inst
  sagemaker_session=sagemaker_session,
  )
 
- records = fm.record_set(train_set[0][:200], train_set[1][:200].astype("float32"))
+ training_set = datasets.one_p_mnist()
+ records = fm.record_set(training_set[0][:200], training_set[1][:200].astype("float32"))
 
  training_config = _build_airflow_workflow(
  estimator=fm, instance_type=cpu_instance_type, inputs=records
@@ -206,13 +187,6 @@ def test_ipinsights_airflow_config_uploads_data_source_to_s3(sagemaker_session,
 
 def test_knn_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_instance_type):
  with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
- data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
- pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
-
- # Load the data into memory as numpy arrays
- with gzip.open(data_path, "rb") as f:
- train_set, _, _ = pickle.load(f, **pickle_args)
-
  knn = KNN(
  role=ROLE,
  train_instance_count=SINGLE_INSTANCE_COUNT,
@@ -223,7 +197,8 @@ def test_knn_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_ins
  sagemaker_session=sagemaker_session,
  )
 
- records = knn.record_set(train_set[0][:200], train_set[1][:200].astype("float32"))
+ training_set = datasets.one_p_mnist()
+ records = knn.record_set(training_set[0][:200], training_set[1][:200].astype("float32"))
 
  training_config = _build_airflow_workflow(
  estimator=knn, instance_type=cpu_instance_type, inputs=records
@@ -277,16 +252,10 @@ def test_linearlearner_airflow_config_uploads_data_source_to_s3(
  sagemaker_session, cpu_instance_type
 ):
  with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
- data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
- pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
-
- # Load the data into memory as numpy arrays
- with gzip.open(data_path, "rb") as f:
- train_set, _, _ = pickle.load(f, **pickle_args)
-
- train_set[1][:100] = 1
- train_set[1][100:200] = 0
- train_set = train_set[0], train_set[1].astype(np.dtype("float32"))
+ training_set = datasets.one_p_mnist()
+ training_set[1][:100] = 1
+ training_set[1][100:200] = 0
+ training_set = training_set[0], training_set[1].astype(np.dtype("float32"))
 
  ll = LinearLearner(
  ROLE,
@@ -331,7 +300,7 @@ def test_linearlearner_airflow_config_uploads_data_source_to_s3(
  ll.early_stopping_tolerance = 0.0001
  ll.early_stopping_patience = 3
 
- records = ll.record_set(train_set[0][:200], train_set[1][:200])
+ records = ll.record_set(training_set[0][:200], training_set[1][:200])
 
  training_config = _build_airflow_workflow(
  estimator=ll, instance_type=cpu_instance_type, inputs=records
@@ -380,13 +349,6 @@ def test_ntm_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_ins
 @pytest.mark.canary_quick
 def test_pca_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_instance_type):
  with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
- data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
- pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
-
- # Load the data into memory as numpy arrays
- with gzip.open(data_path, "rb") as f:
- train_set, _, _ = pickle.load(f, **pickle_args)
-
  pca = PCA(
  role=ROLE,
  train_instance_count=SINGLE_INSTANCE_COUNT,
@@ -399,7 +361,7 @@ def test_pca_airflow_config_uploads_data_source_to_s3(sagemaker_session, cpu_ins
  pca.subtract_mean = True
  pca.extra_components = 5
 
- records = pca.record_set(train_set[0][:100])
+ records = pca.record_set(datasets.one_p_mnist()[0][:100])
 
  training_config = _build_airflow_workflow(
  estimator=pca, instance_type=cpu_instance_type, inputs=records

tests/integ/test_byo_estimator.py

@@ -12,19 +12,16 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import gzip
 import json
 import os
-import pickle
-import sys
 
 import pytest
 
 import sagemaker
 from sagemaker.amazon.amazon_estimator import get_image_uri
 from sagemaker.estimator import Estimator
 from sagemaker.utils import unique_name_from_base
-from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
+from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES, datasets
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
@@ -33,6 +30,11 @@ def region(sagemaker_session):
  return sagemaker_session.boto_session.region_name
 
 
+@pytest.fixture
+def training_set():
+ return datasets.one_p_mnist()
+
+
 def fm_serializer(data):
  js = {"instances": []}
  for row in data:
@@ -41,7 +43,7 @@ def fm_serializer(data):
 
 
 @pytest.mark.canary_quick
-def test_byo_estimator(sagemaker_session, region, cpu_instance_type):
+def test_byo_estimator(sagemaker_session, region, cpu_instance_type, training_set):
  """Use Factorization Machines algorithm as an example here.
 
  First we need to prepare data for training. We take standard data set, convert it to the
@@ -57,12 +59,6 @@ def test_byo_estimator(sagemaker_session, region, cpu_instance_type):
  job_name = unique_name_from_base("byo")
 
  with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
- data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
- pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
-
- with gzip.open(data_path, "rb") as f:
- train_set, _, _ = pickle.load(f, **pickle_args)
-
  prefix = "test_byo_estimator"
  key = "recordio-pb-data"
 
@@ -92,26 +88,20 @@ def test_byo_estimator(sagemaker_session, region, cpu_instance_type):
  predictor.content_type = "application/json"
  predictor.deserializer = sagemaker.predictor.json_deserializer
 
- result = predictor.predict(train_set[0][:10])
+ result = predictor.predict(training_set[0][:10])
 
  assert len(result["predictions"]) == 10
  for prediction in result["predictions"]:
  assert prediction["score"] is not None
 
 
-def test_async_byo_estimator(sagemaker_session, region, cpu_instance_type):
+def test_async_byo_estimator(sagemaker_session, region, cpu_instance_type, training_set):
  image_name = get_image_uri(region, "factorization-machines")
  endpoint_name = unique_name_from_base("byo")
  training_data_path = os.path.join(DATA_DIR, "dummy_tensor")
  job_name = unique_name_from_base("byo")
 
  with timeout(minutes=5):
- data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
- pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
-
- with gzip.open(data_path, "rb") as f:
- train_set, _, _ = pickle.load(f, **pickle_args)
-
  prefix = "test_byo_estimator"
  key = "recordio-pb-data"
 
@@ -144,7 +134,7 @@ def test_async_byo_estimator(sagemaker_session, region, cpu_instance_type):
  predictor.content_type = "application/json"
  predictor.deserializer = sagemaker.predictor.json_deserializer
 
- result = predictor.predict(train_set[0][:10])
+ result = predictor.predict(training_set[0][:10])
 
  assert len(result["predictions"]) == 10
  for prediction in result["predictions"]:

tests/integ/test_factorization_machines.py

@@ -12,29 +12,25 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import gzip
-import os
-import pickle
-import sys
 import time
 
+import pytest
+
 from sagemaker import FactorizationMachines, FactorizationMachinesModel
 from sagemaker.utils import unique_name_from_base
-from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
+from tests.integ import datasets, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
-def test_factorization_machines(sagemaker_session, cpu_instance_type):
- job_name = unique_name_from_base("fm")
+@pytest.fixture
+def training_set():
+ return datasets.one_p_mnist()
 
- with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
- data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
- pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
 
- # Load the data into memory as numpy arrays
- with gzip.open(data_path, "rb") as f:
- train_set, _, _ = pickle.load(f, **pickle_args)
+def test_factorization_machines(sagemaker_session, cpu_instance_type, training_set):
+ job_name = unique_name_from_base("fm")
 
+ with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
  fm = FactorizationMachines(
  role="SageMakerRole",
  train_instance_count=1,
@@ -50,7 +46,7 @@ def test_factorization_machines(sagemaker_session, cpu_instance_type):
 
  # training labels must be 'float32'
  fm.fit(
- fm.record_set(train_set[0][:200], train_set[1][:200].astype("float32")),
+ fm.record_set(training_set[0][:200], training_set[1][:200].astype("float32")),
  job_name=job_name,
  )
 
@@ -59,24 +55,17 @@ def test_factorization_machines(sagemaker_session, cpu_instance_type):
  fm.model_data, role="SageMakerRole", sagemaker_session=sagemaker_session
  )
  predictor = model.deploy(1, cpu_instance_type, endpoint_name=job_name)
- result = predictor.predict(train_set[0][:10])
+ result = predictor.predict(training_set[0][:10])
 
  assert len(result) == 10
  for record in result:
  assert record.label["score"] is not None
 
 
-def test_async_factorization_machines(sagemaker_session, cpu_instance_type):
+def test_async_factorization_machines(sagemaker_session, cpu_instance_type, training_set):
  job_name = unique_name_from_base("fm")
 
  with timeout(minutes=5):
- data_path = os.path.join(DATA_DIR, "one_p_mnist", "mnist.pkl.gz")
- pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}
-
- # Load the data into memory as numpy arrays
- with gzip.open(data_path, "rb") as f:
- train_set, _, _ = pickle.load(f, **pickle_args)
-
  fm = FactorizationMachines(
  role="SageMakerRole",
  train_instance_count=1,
@@ -92,7 +81,7 @@ def test_async_factorization_machines(sagemaker_session, cpu_instance_type):
 
  # training labels must be 'float32'
  fm.fit(
- fm.record_set(train_set[0][:200], train_set[1][:200].astype("float32")),
+ fm.record_set(training_set[0][:200], training_set[1][:200].astype("float32")),
  job_name=job_name,
  wait=False,
  )
@@ -109,7 +98,7 @@ def test_async_factorization_machines(sagemaker_session, cpu_instance_type):
  estimator.model_data, role="SageMakerRole", sagemaker_session=sagemaker_session
  )
  predictor = model.deploy(1, cpu_instance_type, endpoint_name=job_name)
- result = predictor.predict(train_set[0][:10])
+ result = predictor.predict(training_set[0][:10])
 
  assert len(result) == 10
  for record in result: