Fix the tests.

buildspec_integ.yaml

@@ -4,6 +4,7 @@ phases:
   pre_build:
     commands:
       - python setup.py develop
+      - pip install -e .[test]
   build:
     commands:
       - tox tests/integ/

buildspec_unit.yaml

@@ -4,6 +4,7 @@ phases:
   pre_build:
     commands:
       - python setup.py develop
+      - pip install -e .[test]
   build:
     commands:
       - tox tests/unit/

setup.py

@@ -61,13 +61,13 @@ def read_version():
     install_requires=required_packages,
     extras_require={
         "test": [
-            "tox==3.13.1",
-            "pytest==4.4.1",
+            "tox>=3.13.1",
+            "pytest>=4.4.1",
             "stopit==1.1.2",
             "tensorflow>=1.3.0",
             "mock>=2.0.0",
             "contextlib2==0.5.5",
-            "IPython==7.8.0"
+            "IPython>=7.8.0"
         ]
     }
 )

tests/integ/test_inference_pipeline.py

@@ -46,7 +46,7 @@ def sklearn_preprocessor(sagemaker_role_arn, sagemaker_session):
     sklearn_preprocessor = SKLearn(
         entry_point=script_path,
         role=sagemaker_role_arn,
-        train_instance_type="ml.m5.large",
+        train_instance_type="ml.m5.xlarge",
         sagemaker_session=sagemaker_session,
         hyperparameters={"epochs": 1},
     )
@@ -61,7 +61,7 @@ def sklearn_estimator(sagemaker_role_arn, sagemaker_session):
     sklearn_estimator = SKLearn(
         entry_point=script_path,
         role=sagemaker_role_arn,
-        train_instance_type="ml.m5.large",
+        train_instance_type="ml.m5.xlarge",
         sagemaker_session=sagemaker_session,
         hyperparameters={"epochs": 1},
         input_mode='File'

tests/integ/test_sagemaker_steps.py

@@ -43,7 +43,7 @@
 )
 
 INSTANCE_COUNT = 1
-INSTANCE_TYPE = "ml.m5.large"
+INSTANCE_TYPE = "ml.m5.xlarge"
 
 CREATE_ENDPOINT_TIMEOUT_MINUTES = 20
 
@@ -255,7 +255,7 @@ def test_tuning_step(sfn_client, record_set_for_hyperparameter_tuning, sagemaker
     kmeans = KMeans(
         role=sagemaker_role_arn,
         train_instance_count=1,
-        train_instance_type='ml.m5.large',
+        train_instance_type=INSTANCE_TYPE,
         k=10
     )
 

tests/integ/test_state_machine_definition.py

@@ -36,7 +36,7 @@ def training_job_parameters(sagemaker_session, sagemaker_role_arn, record_set_fi
         },
         "ResourceConfig": { 
             "InstanceCount": 1,
-            "InstanceType": "ml.m5.large",
+            "InstanceType": "ml.m5.xlarge",
             "VolumeSizeInGB": 30
         },
         "RoleArn": sagemaker_role_arn,

tests/integ/test_training_pipeline_estimators.py

@@ -52,7 +52,7 @@ def pca_estimator(sagemaker_role_arn):
         role=sagemaker_role_arn,
         num_components=1,
         train_instance_count=1,
-        train_instance_type='ml.m5.large',
+        train_instance_type='ml.m5.xlarge',
         )
 
     pca_estimator.feature_dim=500
@@ -125,7 +125,7 @@ def test_pca_estimator(sfn_client, sagemaker_session, sagemaker_role_arn, sfn_ro
            'OutputDataConfig': {'S3OutputPath': s3_output_path},
            'StoppingCondition': {'MaxRuntimeInSeconds': 86400},
            'ResourceConfig': {'InstanceCount': 1,
-            'InstanceType': 'ml.m5.large',
+            'InstanceType': 'ml.m5.xlarge',
             'VolumeSizeInGB': 30},
            'RoleArn': sagemaker_role_arn,
            'InputDataConfig': [{'DataSource': {'S3DataSource': {'S3DataDistributionType': 'ShardedByS3Key',
@@ -141,7 +141,7 @@ def test_pca_estimator(sfn_client, sagemaker_session, sagemaker_role_arn, sfn_ro
            'ExecutionRoleArn': sagemaker_role_arn},
           'Configure Endpoint': {'EndpointConfigName': job_name,
            'ProductionVariants': [{'ModelName': job_name,
-             'InstanceType': 'ml.m5.large',
+             'InstanceType': 'ml.m5.xlarge',
              'InitialInstanceCount': 1,
              'VariantName': 'AllTraffic'}]},
           'Deploy': {'EndpointName': job_name,

tests/integ/test_training_pipeline_framework_estimator.py

@@ -39,7 +39,7 @@ def torch_estimator(sagemaker_role_arn):
         role=sagemaker_role_arn,
         framework_version='1.1.0',
         train_instance_count=1,
-        train_instance_type='ml.m5.large',
+        train_instance_type='ml.m5.xlarge',
         hyperparameters={
             'epochs': 6,
             'backend': 'gloo'
@@ -53,7 +53,7 @@ def sklearn_estimator(sagemaker_role_arn):
         entry_point=script_path,
         role=sagemaker_role_arn,
         train_instance_count=1,
-        train_instance_type='ml.m5.large',
+        train_instance_type='ml.m5.xlarge',
         framework_version='0.20.0',
         hyperparameters={
             "epochs": 1

tests/unit/test_pipeline.py

@@ -13,29 +13,33 @@
 from __future__ import absolute_import
 
 import pytest
-import sagemaker
+from unittest.mock import MagicMock, patch
 
+import sagemaker
 from sagemaker.sklearn.estimator import SKLearn
-from unittest.mock import MagicMock, patch
+from sagemaker.amazon.amazon_estimator import get_image_uri
 from stepfunctions.template import TrainingPipeline, InferencePipeline
 
 from tests.unit.utils import mock_boto_api_call
 
 SAGEMAKER_EXECUTION_ROLE = 'SageMakerExecutionRole'
 STEPFUNCTIONS_EXECUTION_ROLE = 'StepFunctionsExecutionRole'
-PCA_IMAGE = '382416733822.dkr.ecr.us-east-1.amazonaws.com/pca:1'
-LINEAR_LEARNER_IMAGE = '382416733822.dkr.ecr.us-east-1.amazonaws.com/linear-learner:1'
+PCA_IMAGE = get_image_uri('us-east-1', 'pca')
+LINEAR_LEARNER_IMAGE = get_image_uri('us-east-1', 'linear-learner')
 
 @pytest.fixture
 def pca_estimator():
     s3_output_location = 's3://sagemaker/models'
+    sagemaker_session = MagicMock()
+    sagemaker_session.boto_region_name = 'us-east-1'
 
     pca = sagemaker.estimator.Estimator(
         PCA_IMAGE,
         role=SAGEMAKER_EXECUTION_ROLE,
         train_instance_count=1,
         train_instance_type='ml.c4.xlarge',
-        output_path=s3_output_location
+        output_path=s3_output_location,
+        sagemaker_session=sagemaker_session
     )
 
     pca.set_hyperparameters(
@@ -50,23 +54,26 @@ def pca_estimator():
 
 @pytest.fixture
 def sklearn_preprocessor():
-    sagemaker_session = MagicMock()
     script_path = 'sklearn_abalone_featurizer.py'
     source_dir = 's3://sagemaker/source'
+    sagemaker_session = MagicMock()
+    sagemaker_session.boto_region_name = 'us-east-1'
 
     sklearn_preprocessor = SKLearn(
         entry_point=script_path,
         role=SAGEMAKER_EXECUTION_ROLE,
         train_instance_type="ml.c4.xlarge",
-        source_dir=source_dir
+        source_dir=source_dir,
+        sagemaker_session=sagemaker_session
     )
 
     return sklearn_preprocessor
 
 @pytest.fixture
 def linear_learner_estimator():
-    sagemaker_session = MagicMock()
     s3_output_location = 's3://sagemaker/models'
+    sagemaker_session = MagicMock()
+    sagemaker_session.boto_region_name = 'us-east-1'
 
     ll_estimator = sagemaker.estimator.Estimator(
         LINEAR_LEARNER_IMAGE,
@@ -76,7 +83,8 @@ def linear_learner_estimator():
         train_volume_size=20,
         train_max_run=3600,
         input_mode='File',
-        output_path=s3_output_location
+        output_path=s3_output_location,
+        sagemaker_session=sagemaker_session
     )
 
     ll_estimator.set_hyperparameters(feature_dim=10, predictor_type='regressor', mini_batch_size=32)
@@ -218,7 +226,7 @@ def test_inference_pipeline(sklearn_preprocessor, linear_learner_estimator):
     s3_bucket = 'sagemaker-us-east-1'
 
     pipeline = InferencePipeline(
-        preprocessor=sklearn_preprocessor, 
+        preprocessor=sklearn_preprocessor,
         estimator=linear_learner_estimator,
         data=s3_inputs,
         s3_bucket=s3_bucket,