infra: refactor toolkit tests. (#332)

* Separate toolkit tests from container tests;
* Remove test logic for other libraries installed in the container;
* Provide options to build and push images as part of running tests.
* Add PR build logic specific to just toolkit.

Author: nadiaya

.flake8

@@ -1,3 +1,3 @@
 [flake8]
-application_import_names = sagemaker_tensorflow_container, test, test-toolkit, timeout, utils
+application_import_names = image_utils, integration, sagemaker_tensorflow_container, test, test-toolkit, timeout, utils
 import-order-style = google

buildspec-toolkit.yml

@@ -0,0 +1,93 @@
+version: 0.2
+
+env:
+  variables:
+    FRAMEWORK_VERSION: '1.15.2'
+    CPU_INSTANCE_TYPE: 'ml.c4.xlarge'
+    GPU_INSTANCE_TYPE: 'ml.p2.xlarge'
+    ECR_REPO: 'sagemaker-test'
+    GITHUB_REPO: 'sagemaker-tensorflow-container'
+    DLC_ACCOUNT: '763104351884'
+    SETUP_FILE: 'setup_cmds.sh'
+    SETUP_CMDS: '#!/bin/bash\npip install --upgrade pip\npip install -U -e .\npip install -U -e .[test]'
+
+phases:
+  pre_build:
+    commands:
+      - start-dockerd
+      - ACCOUNT=$(aws --region $AWS_DEFAULT_REGION sts --endpoint-url https://sts.$AWS_DEFAULT_REGION.amazonaws.com get-caller-identity --query 'Account' --output text)
+      - PREPROD_IMAGE="$ACCOUNT.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/$ECR_REPO"
+      - PR_NUM=$(echo $CODEBUILD_SOURCE_VERSION | grep -o '[0-9]\+')
+      - BUILD_ID="$(echo $CODEBUILD_BUILD_ID | sed -e 's/:/-/g')"
+      - echo 'Pull request number:' $PR_NUM '. No value means this build is not from pull request.'
+
+  build:
+    commands:
+      - TOX_PARALLEL_NO_SPINNER=1
+      - PY_COLORS=0
+
+      # install
+      - pip3 install -U -e .[test]
+
+      # run linters
+      - tox -e flake8,twine
+
+      # run unit tests
+      - tox -e py36,py27 test-toolkit/unit
+
+      # define tags
+      - GENERIC_TAG="$FRAMEWORK_VERSION-tensorflow-$BUILD_ID"
+      - DLC_CPU_TAG="$FRAMEWORK_VERSION-dlc-cpu-$BUILD_ID"
+      - DLC_GPU_TAG="$FRAMEWORK_VERSION-dlc-gpu-$BUILD_ID"
+
+      # run local CPU integration tests (build and push the image to ECR repo)
+      - test_cmd="pytest test-toolkit/integration/local --build-image --push-image --dockerfile-type tf --region $AWS_DEFAULT_REGION --docker-base-name $ECR_REPO --account-id $ACCOUNT --framework-version $FRAMEWORK_VERSION --processor cpu --tag $GENERIC_TAG"
+      - execute-command-if-has-matching-changes "$test_cmd" "test-toolkit/" "src/*.py" "setup.py" "setup.cfg" "buildspec-toolkit.yml"
+      - test_cmd="pytest test-toolkit/integration/local --build-image --push-image --dockerfile-type dlc.cpu --region $AWS_DEFAULT_REGION --docker-base-name $ECR_REPO --account-id $ACCOUNT --framework-version $FRAMEWORK_VERSION --processor cpu --tag $DLC_CPU_TAG"
+      - execute-command-if-has-matching-changes "$test_cmd" "test-toolkit/" "src/*.py" "setup.py" "setup.cfg" "buildspec-toolkit.yml"
+
+      # launch remote GPU instance
+      - prefix='ml.'
+      - instance_type=${GPU_INSTANCE_TYPE#"$prefix"}
+      - create-key-pair
+      - launch-ec2-instance --instance-type $instance_type --ami-name dlami-ubuntu-latest
+
+      # build DLC GPU image because the base DLC image is too big and takes too long to build as part of the test
+      - python3 setup.py sdist
+      - build_dir="test-toolkit/docker/$FRAMEWORK_VERSION"
+      - $(aws ecr get-login --registry-ids $DLC_ACCOUNT --no-include-email --region $AWS_DEFAULT_REGION)
+      - docker build -f "$build_dir/Dockerfile.dlc.gpu" -t $PREPROD_IMAGE:$DLC_GPU_TAG --build-arg region=$AWS_DEFAULT_REGION .
+      # push DLC GPU image to ECR
+      - $(aws ecr get-login --registry-ids $ACCOUNT --no-include-email --region $AWS_DEFAULT_REGION)
+      - docker push $PREPROD_IMAGE:$DLC_GPU_TAG
+
+      # run GPU local integration tests
+      - printf "$SETUP_CMDS" > $SETUP_FILE
+      # no reason to rebuild the image again since it was already built and pushed to ECR during CPU tests
+      - generic_cmd="pytest test-toolkit/integration/local --region $AWS_DEFAULT_REGION --docker-base-name $ECR_REPO --account-id $ACCOUNT --framework-version $FRAMEWORK_VERSION --processor gpu --tag $GENERIC_TAG"
+      - test_cmd="remote-test --github-repo $GITHUB_REPO --test-cmd \"$generic_cmd\" --setup-file $SETUP_FILE --pr-number \"$PR_NUM\""
+      - execute-command-if-has-matching-changes "$test_cmd" "test-toolkit/" "src/*.py" "setup.py" "setup.cfg" "buildspec-toolkit.yml"
+      - dlc_cmd="pytest test-toolkit/integration/local --region $AWS_DEFAULT_REGION --docker-base-name $ECR_REPO --account-id $ACCOUNT --framework-version $FRAMEWORK_VERSION --processor gpu --tag $DLC_GPU_TAG"
+      - test_cmd="remote-test --github-repo $GITHUB_REPO --test-cmd \"$dlc_cmd\" --setup-file $SETUP_FILE --pr-number \"$PR_NUM\" --skip-setup"
+      - execute-command-if-has-matching-changes "$test_cmd" "test-toolkit/" "src/*.py" "setup.py" "setup.cfg" "buildspec-toolkit.yml"
+
+      # run CPU sagemaker integration tests
+      - test_cmd="pytest test-toolkit/integration/sagemaker --region $AWS_DEFAULT_REGION --docker-base-name $ECR_REPO --account-id $ACCOUNT --framework-version $FRAMEWORK_VERSION --processor cpu --instance-type $CPU_INSTANCE_TYPE --tag $GENERIC_TAG"
+      - execute-command-if-has-matching-changes "$test_cmd" "test-toolkit/" "src/*.py" "setup.py" "setup.cfg" "buildspec-toolkit.yml"
+      - test_cmd="pytest test-toolkit/integration/sagemaker --region $AWS_DEFAULT_REGION --docker-base-name $ECR_REPO --account-id $ACCOUNT --framework-version $FRAMEWORK_VERSION --processor cpu --instance-type $CPU_INSTANCE_TYPE --tag $DLC_CPU_TAG"
+      - execute-command-if-has-matching-changes "$test_cmd" "test-toolkit/" "src/*.py" "setup.py" "setup.cfg" "buildspec-toolkit.yml"
+
+      # run GPU sagemaker integration tests
+      - test_cmd="pytest test-toolkit/integration/sagemaker --region $AWS_DEFAULT_REGION --docker-base-name $ECR_REPO --account-id $ACCOUNT --framework-version $FRAMEWORK_VERSION --processor gpu --instance-type $GPU_INSTANCE_TYPE --tag $GENERIC_TAG"
+      - execute-command-if-has-matching-changes "$test_cmd" "test-toolkit/" "src/*.py" "setup.py" "setup.cfg" "buildspec-toolkit.yml"
+      - test_cmd="pytest test-toolkit/integration/sagemaker --region $AWS_DEFAULT_REGION --docker-base-name $ECR_REPO --account-id $ACCOUNT --framework-version $FRAMEWORK_VERSION --processor gpu --instance-type $GPU_INSTANCE_TYPE --tag $DLC_GPU_TAG"
+      - execute-command-if-has-matching-changes "$test_cmd" "test-toolkit/" "src/*.py" "setup.py" "setup.cfg" "buildspec-toolkit.yml"
+    finally:
+      # shut down remote GPU instance
+      - cleanup-gpu-instances
+      - cleanup-key-pairs
+
+      # remove ECR image
+      - aws ecr batch-delete-image --repository-name $ECR_REPO --region $AWS_DEFAULT_REGION --image-ids imageTag=$GENERIC_TAG
+      - aws ecr batch-delete-image --repository-name $ECR_REPO --region $AWS_DEFAULT_REGION --image-ids imageTag=$DLC_CPU_TAG
+      - aws ecr batch-delete-image --repository-name $ECR_REPO --region $AWS_DEFAULT_REGION --image-ids imageTag=$DLC_GPU_TAG

test-toolkit/docker/1.15.2/Dockerfile.dlc.cpu

@@ -0,0 +1,6 @@
+ARG region
+FROM 763104351884.dkr.ecr.$region.amazonaws.com/tensorflow-training:1.15.2-cpu-py2
+
+COPY dist/sagemaker_tensorflow_training-*.tar.gz /sagemaker_tensorflow_training.tar.gz
+RUN pip install --upgrade --no-cache-dir /sagemaker_tensorflow_training.tar.gz && \
+    rm /sagemaker_tensorflow_training.tar.gz

test-toolkit/docker/1.15.2/Dockerfile.dlc.gpu

@@ -0,0 +1,6 @@
+ARG region
+FROM 763104351884.dkr.ecr.$region.amazonaws.com/tensorflow-training:1.15.2-gpu-py3
+
+COPY dist/sagemaker_tensorflow_training-*.tar.gz /sagemaker_tensorflow_training.tar.gz
+RUN pip install --upgrade --no-cache-dir /sagemaker_tensorflow_training.tar.gz && \
+    rm /sagemaker_tensorflow_training.tar.gz

test-toolkit/docker/1.15.2/Dockerfile.tf

@@ -0,0 +1,7 @@
+FROM tensorflow/tensorflow:1.15.2-gpu-py3
+
+ENV SAGEMAKER_TRAINING_MODULE sagemaker_tensorflow_container.training:main
+
+COPY dist/sagemaker_tensorflow_training-*.tar.gz /sagemaker_tensorflow_training.tar.gz
+RUN pip install --upgrade --no-cache-dir /sagemaker_tensorflow_training.tar.gz && \
+    rm /sagemaker_tensorflow_training.tar.gz

test-toolkit/integration/conftest.py

@@ -18,9 +18,10 @@
 import boto3
 import pytest
 from sagemaker import LocalSession, Session
-from sagemaker.tensorflow import TensorFlow
 
-from test.integration import NO_P2_REGIONS, NO_P3_REGIONS
+from integration import image_utils
+from integration import NO_P2_REGIONS, NO_P3_REGIONS
+
 
 logger = logging.getLogger(__name__)
 logging.getLogger('boto').setLevel(logging.INFO)
@@ -29,14 +30,19 @@
 logging.getLogger('auth.py').setLevel(logging.INFO)
 logging.getLogger('connectionpool.py').setLevel(logging.INFO)
 
-SCRIPT_PATH = os.path.dirname(os.path.realpath(__file__))
+DIR_PATH = os.path.dirname(os.path.realpath(__file__))
 
 
 def pytest_addoption(parser):
-    parser.addoption('--docker-base-name', default='sagemaker-tensorflow-scriptmode')
+    parser.addoption('--build-image', '-B', action='store_true')
+    parser.addoption('--push-image', '-P', action='store_true')
+    parser.addoption('--dockerfile-type', '-T', choices=['dlc.cpu', 'dlc.gpu', 'tf'],
+                     default='tf')
+    parser.addoption('--dockerfile', '-D', default=None)
+    parser.addoption('--docker-base-name', default='sagemaker-tensorflow-training')
     parser.addoption('--tag', default=None)
     parser.addoption('--region', default='us-west-2')
-    parser.addoption('--framework-version', default=TensorFlow.LATEST_VERSION)
+    parser.addoption('--framework-version', default='1.15.2')
     parser.addoption('--processor', default='cpu', choices=['cpu', 'gpu', 'cpu,gpu'])
     parser.addoption('--py-version', default='3', choices=['2', '3', '2,3'])
     parser.addoption('--account-id', default='142577830533')
@@ -48,6 +54,38 @@ def pytest_configure(config):
     os.environ['TEST_PROCESSORS'] = config.getoption('--processor')
 
 
+@pytest.fixture(scope='session', name='dockerfile_type')
+def fixture_dockerfile_type(request):
+    return request.config.getoption('--dockerfile-type')
+
+
+@pytest.fixture(scope='session', name='dockerfile')
+def fixture_dockerfile(request, dockerfile_type):
+    dockerfile = request.config.getoption('--dockerfile')
+    return dockerfile if dockerfile else 'Dockerfile.{}'.format(dockerfile_type)
+
+
+@pytest.fixture(scope='session', name='build_image', autouse=True)
+def fixture_build_image(request, framework_version, dockerfile, image_uri, region):
+    build_image = request.config.getoption('--build-image')
+    if build_image:
+        return image_utils.build_image(framework_version=framework_version,
+                                       dockerfile=dockerfile,
+                                       image_uri=image_uri,
+                                       region=region,
+                                       cwd=os.path.join(DIR_PATH, '..', '..'))
+
+    return image_uri
+
+
+@pytest.fixture(scope='session', name='push_image', autouse=True)
+def fixture_push_image(request, image_uri, region, account_id):
+    push_image = request.config.getoption('--push-image')
+    if push_image:
+        return image_utils.push_image(image_uri, region, account_id)
+    return None
+
+
 @pytest.fixture(scope='session')
 def docker_base_name(request):
     return request.config.getoption('--docker-base-name')
@@ -63,7 +101,7 @@ def framework_version(request):
     return request.config.getoption('--framework-version')
 
 
-@pytest.fixture
+@pytest.fixture(scope='session')
 def tag(request, framework_version, processor, py_version):
     provided_tag = request.config.getoption('--tag')
     default_tag = '{}-{}-py{}'.format(framework_version, processor, py_version)
@@ -107,12 +145,20 @@ def skip_gpu_instance_restricted_regions(region, instance_type):
         pytest.skip('Skipping GPU test in region {}'.format(region))
 
 
-@pytest.fixture
-def docker_image(docker_base_name, tag):
-    return '{}:{}'.format(docker_base_name, tag)
+@pytest.fixture(autouse=True)
+def skip_by_dockerfile_type(request, dockerfile_type):
+    is_generic = (dockerfile_type == 'tf')
+    if request.node.get_closest_marker('skip_generic') and is_generic:
+        pytest.skip('Skipping because running generic image without mpi and horovod')
 
 
-@pytest.fixture
-def ecr_image(account_id, docker_base_name, tag, region):
-    return '{}.dkr.ecr.{}.amazonaws.com/{}:{}'.format(
-        account_id, region, docker_base_name, tag)
+@pytest.fixture(name='docker_registry', scope='session')
+def fixture_docker_registry(account_id, region):
+    return '{}.dkr.ecr.{}.amazonaws.com'.format(account_id, region) if account_id else None
+
+
+@pytest.fixture(name='image_uri', scope='session')
+def fixture_image_uri(docker_registry, docker_base_name, tag):
+    if docker_registry:
+        return '{}/{}:{}'.format(docker_registry, docker_base_name, tag)
+    return '{}:{}'.format(docker_base_name, tag)

test-toolkit/integration/image_utils.py

@@ -0,0 +1,59 @@
+# 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 os
+import subprocess
+import sys
+
+CYAN_COLOR = '\033[36m'
+END_COLOR = '\033[0m'
+DLC_AWS_ID = '763104351884'
+
+
+def build_image(framework_version, dockerfile, image_uri, region, cwd='.'):
+    _check_call('python setup.py sdist')
+
+    if 'dlc' in dockerfile:
+        ecr_login(region, DLC_AWS_ID)
+
+    dockerfile_location = os.path.join('test-toolkit', 'docker', framework_version, dockerfile)
+
+    subprocess.check_call(
+        ['docker', 'build', '-t', image_uri, '-f', dockerfile_location, '--build-arg',
+         'region={}'.format(region), cwd], cwd=cwd)
+    print('created image {}'.format(image_uri))
+    return image_uri
+
+
+def push_image(ecr_image, region, aws_id):
+    ecr_login(region, aws_id)
+    _check_call('docker push {}'.format(ecr_image))
+
+
+def ecr_login(region, aws_id):
+    login = _check_call('aws ecr get-login --registry-ids {} '.format(aws_id)
+                        + '--no-include-email --region {}'.format(region))
+    _check_call(login.decode('utf-8').rstrip('\n'))
+
+
+def _check_call(cmd, *popenargs, **kwargs):
+    if isinstance(cmd, str):
+        cmd = cmd.split(" ")
+    _print_cmd(cmd)
+    return subprocess.check_output(cmd, *popenargs, **kwargs)
+
+
+def _print_cmd(cmd):
+    print('executing docker command: {}{}{}'.format(CYAN_COLOR, ' '.join(cmd), END_COLOR))
+    sys.stdout.flush()

test-toolkit/integration/local/test_horovod.py

@@ -19,12 +19,13 @@
 import pytest
 from sagemaker.tensorflow import TensorFlow
 
-from test.integration.utils import processor, py_version  # noqa: F401
+from integration.utils import processor, py_version  # noqa: F401
 
 RESOURCE_PATH = os.path.join(os.path.dirname(__file__), '..', '..', 'resources')
 
 
 @pytest.mark.skip_gpu
+@pytest.mark.skip_generic
 @pytest.mark.parametrize('instances, processes', [
     [1, 2],
     (2, 1),
@@ -33,7 +34,7 @@
 def test_distributed_training_horovod_basic(instances,
                                             processes,
                                             sagemaker_local_session,
-                                            docker_image,
+                                            image_uri,
                                             tmpdir,
                                             framework_version):
     output_path = 'file://%s' % tmpdir
@@ -43,7 +44,7 @@ def test_distributed_training_horovod_basic(instances,
         train_instance_type='local',
         sagemaker_session=sagemaker_local_session,
         train_instance_count=instances,
-        image_name=docker_image,
+        image_name=image_uri,
         output_path=output_path,
         framework_version=framework_version,
         hyperparameters={'sagemaker_mpi_enabled': True,