infra: use ECR-hosted image for ubuntu:16.04 (#1987)

Author: ajaykarpur

tests/data/multimodel/container/Dockerfile

@@ -1,4 +1,4 @@
-FROM ubuntu:16.04
+FROM 142577830533.dkr.ecr.us-east-2.amazonaws.com/ubuntu:16.04
 
 # Set a docker label to advertise multi-model support on the container
 LABEL com.amazonaws.sagemaker.capabilities.multi-models=true

tests/integ/test_multidatamodel.py

@@ -22,6 +22,8 @@
 import pytest
 from botocore.exceptions import ClientError
 
+import tests
+
 from sagemaker import utils
 from sagemaker.amazon.randomcutforest import RandomCutForest
 from sagemaker.deserializers import StringDeserializer
@@ -46,18 +48,26 @@ def container_image(sagemaker_session):
     algorithm_name = unique_name_from_base("sagemaker-multimodel-integ-test")
     ecr_image = _ecr_image_uri(sagemaker_session, algorithm_name)
 
-    # Build and tag docker image locally
+    ecr_client = sagemaker_session.boto_session.client("ecr")
+    username, password = _ecr_login(ecr_client)
+
     docker_client = docker.from_env()
+
+    # Base image pull
+    base_image = "142577830533.dkr.ecr.us-east-2.amazonaws.com/ubuntu:16.04"
+    docker_client.images.pull(base_image, auth_config={"username": username, "password": password})
+
+    # Build and tag docker image locally
     image, build_log = docker_client.images.build(
-        path=os.path.join(DATA_DIR, "multimodel", "container"), tag=algorithm_name, rm=True
+        path=os.path.join(DATA_DIR, "multimodel", "container"),
+        tag=algorithm_name,
+        rm=True,
     )
     image.tag(ecr_image, tag="latest")
 
     # Create AWS ECR and push the local docker image to it
-    ecr_client = sagemaker_session.boto_session.client("ecr")
     _create_repository(ecr_client, algorithm_name)
 
-    username, password = _ecr_login(ecr_client)
     # Retry docker image push
     for _ in retries(3, "Upload docker image to ECR repo", seconds_to_sleep=10):
         try:
@@ -124,6 +134,10 @@ def _ecr_login(ecr_client):
     return username, password
 
 
+@pytest.mark.skipif(
+    tests.integ.test_region() != "us-east-2",
+    reason="Pulling the base image is currently limited to us-east-2.",
+)
 def test_multi_data_model_deploy_pretrained_models(
     container_image, sagemaker_session, cpu_instance_type
 ):
@@ -184,6 +198,10 @@ def test_multi_data_model_deploy_pretrained_models(
 
 
 @pytest.mark.local_mode
+@pytest.mark.skipif(
+    tests.integ.test_region() != "us-east-2",
+    reason="Pulling the base image is currently limited to us-east-2.",
+)
 def test_multi_data_model_deploy_pretrained_models_local_mode(container_image, sagemaker_session):
     timestamp = sagemaker_timestamp()
     endpoint_name = "test-multimodel-endpoint-{}".format(timestamp)
@@ -244,6 +262,10 @@ def test_multi_data_model_deploy_pretrained_models_local_mode(container_image, s
         assert "Could not find endpoint" in str(exception.value)
 
 
+@pytest.mark.skipif(
+    tests.integ.test_region() != "us-east-2",
+    reason="Pulling the base image is currently limited to us-east-2.",
+)
 def test_multi_data_model_deploy_trained_model_from_framework_estimator(
     container_image,
     sagemaker_session,
@@ -358,6 +380,10 @@ def _mxnet_training_job(
         return mx.create_model(image_uri=container_image)
 
 
+@pytest.mark.skipif(
+    tests.integ.test_region() != "us-east-2",
+    reason="Pulling the base image is currently limited to us-east-2.",
+)
 def test_multi_data_model_deploy_train_model_from_amazon_first_party_estimator(
     container_image, sagemaker_session, cpu_instance_type
 ):
@@ -455,6 +481,10 @@ def __rcf_training_job(
         return rcf_model
 
 
+@pytest.mark.skipif(
+    tests.integ.test_region() != "us-east-2",
+    reason="Pulling the base image is currently limited to us-east-2.",
+)
 def test_multi_data_model_deploy_pretrained_models_update_endpoint(
     container_image, sagemaker_session, cpu_instance_type, alternative_cpu_instance_type
 ):