From 875f4a1c783270b385c0263251090b01fb6a8589 Mon Sep 17 00:00:00 2001 From: Lauren Yu <6631887+laurenyu@users.noreply.github.com> Date: Thu, 16 Apr 2020 09:27:24 -0700 Subject: [PATCH 1/2] infra: move Neo unit tests to a new file --- .../sagemaker/model/test_framework_model.py | 170 --------------- tests/unit/sagemaker/model/test_neo.py | 203 ++++++++++++++++++ 2 files changed, 203 insertions(+), 170 deletions(-) create mode 100644 tests/unit/sagemaker/model/test_neo.py diff --git a/tests/unit/sagemaker/model/test_framework_model.py b/tests/unit/sagemaker/model/test_framework_model.py index 635a59ae79..f972bdfe04 100644 --- a/tests/unit/sagemaker/model/test_framework_model.py +++ b/tests/unit/sagemaker/model/test_framework_model.py @@ -37,7 +37,6 @@ ACCELERATOR_TYPE = "ml.eia.medium" IMAGE_NAME = "fakeimage" REGION = "us-west-2" -NEO_REGION_ACCOUNT = "301217895009" MODEL_NAME = "{}-{}".format(MODEL_IMAGE, TIMESTAMP) GIT_REPO = "https://github.com/aws/sagemaker-python-sdk.git" BRANCH = "test-branch-git-config" @@ -51,11 +50,6 @@ CODECOMMIT_BRANCH = "master" REPO_DIR = "/tmp/repo_dir" -DESCRIBE_COMPILATION_JOB_RESPONSE = { - "CompilationJobStatus": "Completed", - "ModelArtifacts": {"S3ModelArtifacts": "s3://output-path/model.tar.gz"}, -} - class DummyFrameworkModel(FrameworkModel): def __init__(self, sagemaker_session, **kwargs): @@ -400,170 +394,6 @@ def test_delete_non_deployed_model(sagemaker_session): model.delete_model() -def test_compile_model_for_inferentia(sagemaker_session, tmpdir): - sagemaker_session.wait_for_compilation_job = Mock( - return_value=DESCRIBE_COMPILATION_JOB_RESPONSE - ) - model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir)) - model.compile( - target_instance_family="ml_inf", - input_shape={"data": [1, 3, 1024, 1024]}, - output_path="s3://output", - role="role", - framework="tensorflow", - framework_version="1.15.0", - job_name="compile-model", - ) - assert ( - "{}.dkr.ecr.{}.amazonaws.com/sagemaker-neo-tensorflow:1.15.0-inf-py3".format( - NEO_REGION_ACCOUNT, REGION - ) - == model.image - ) - assert model._is_compiled_model is True - - -def test_compile_model_for_edge_device(sagemaker_session, tmpdir): - sagemaker_session.wait_for_compilation_job = Mock( - return_value=DESCRIBE_COMPILATION_JOB_RESPONSE - ) - model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir)) - model.compile( - target_instance_family="deeplens", - input_shape={"data": [1, 3, 1024, 1024]}, - output_path="s3://output", - role="role", - framework="tensorflow", - job_name="compile-model", - ) - assert model._is_compiled_model is False - - -def test_compile_model_for_edge_device_tflite(sagemaker_session, tmpdir): - sagemaker_session.wait_for_compilation_job = Mock( - return_value=DESCRIBE_COMPILATION_JOB_RESPONSE - ) - model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir)) - model.compile( - target_instance_family="deeplens", - input_shape={"data": [1, 3, 1024, 1024]}, - output_path="s3://output", - role="role", - framework="tflite", - job_name="tflite-compile-model", - ) - assert model._is_compiled_model is False - - -def test_compile_model_for_cloud(sagemaker_session, tmpdir): - sagemaker_session.wait_for_compilation_job = Mock( - return_value=DESCRIBE_COMPILATION_JOB_RESPONSE - ) - model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir)) - model.compile( - target_instance_family="ml_c4", - input_shape={"data": [1, 3, 1024, 1024]}, - output_path="s3://output", - role="role", - framework="tensorflow", - job_name="compile-model", - ) - assert model._is_compiled_model is True - - -def test_compile_model_for_cloud_tflite(sagemaker_session, tmpdir): - sagemaker_session.wait_for_compilation_job = Mock( - return_value=DESCRIBE_COMPILATION_JOB_RESPONSE - ) - model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir)) - model.compile( - target_instance_family="ml_c4", - input_shape={"data": [1, 3, 1024, 1024]}, - output_path="s3://output", - role="role", - framework="tflite", - job_name="tflite-compile-model", - ) - assert model._is_compiled_model is True - - -@patch("sagemaker.session.Session") -@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock()) -def test_compile_creates_session(session): - session.return_value.boto_region_name = "us-west-2" - - model = DummyFrameworkModel(sagemaker_session=None) - model.compile( - target_instance_family="ml_c4", - input_shape={"data": [1, 3, 1024, 1024]}, - output_path="s3://output", - role="role", - framework="tensorflow", - job_name="compile-model", - ) - - assert model.sagemaker_session == session.return_value - - -def test_check_neo_region(sagemaker_session, tmpdir): - sagemaker_session.wait_for_compilation_job = Mock( - return_value=DESCRIBE_COMPILATION_JOB_RESPONSE - ) - model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir)) - ec2_region_list = [ - "us-east-2", - "us-east-1", - "us-west-1", - "us-west-2", - "ap-east-1", - "ap-south-1", - "ap-northeast-3", - "ap-northeast-2", - "ap-southeast-1", - "ap-southeast-2", - "ap-northeast-1", - "ca-central-1", - "cn-north-1", - "cn-northwest-1", - "eu-central-1", - "eu-west-1", - "eu-west-2", - "eu-west-3", - "eu-north-1", - "sa-east-1", - "us-gov-east-1", - "us-gov-west-1", - ] - neo_support_region = [ - "us-west-1", - "us-west-2", - "us-east-1", - "us-east-2", - "eu-west-1", - "eu-west-2", - "eu-west-3", - "eu-central-1", - "eu-north-1", - "ap-northeast-1", - "ap-northeast-2", - "ap-east-1", - "ap-south-1", - "ap-southeast-1", - "ap-southeast-2", - "sa-east-1", - "ca-central-1", - "me-south-1", - "cn-north-1", - "cn-northwest-1", - "us-gov-west-1", - ] - for region_name in ec2_region_list: - if region_name in neo_support_region: - assert model.check_neo_region(region_name) is True - else: - assert model.check_neo_region(region_name) is False - - @patch("sagemaker.git_utils.git_clone_repo") @patch("sagemaker.model.fw_utils.tar_and_upload_dir") def test_git_support_succeed(tar_and_upload_dir, git_clone_repo, sagemaker_session): diff --git a/tests/unit/sagemaker/model/test_neo.py b/tests/unit/sagemaker/model/test_neo.py new file mode 100644 index 0000000000..54c0ef5590 --- /dev/null +++ b/tests/unit/sagemaker/model/test_neo.py @@ -0,0 +1,203 @@ +# Copyright 2017-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 pytest +from mock import Mock, patch + +from sagemaker.model import Model + +MODEL_DATA = "s3://bucket/model.tar.gz" +MODEL_IMAGE = "mi" + +REGION = "us-west-2" + +NEO_REGION_ACCOUNT = "301217895009" +DESCRIBE_COMPILATION_JOB_RESPONSE = { + "CompilationJobStatus": "Completed", + "ModelArtifacts": {"S3ModelArtifacts": "s3://output-path/model.tar.gz"}, +} + +EC2_REGION_LIST = [ + "us-east-2", + "us-east-1", + "us-west-1", + "us-west-2", + "ap-east-1", + "ap-south-1", + "ap-northeast-3", + "ap-northeast-2", + "ap-southeast-1", + "ap-southeast-2", + "ap-northeast-1", + "ca-central-1", + "cn-north-1", + "cn-northwest-1", + "eu-central-1", + "eu-west-1", + "eu-west-2", + "eu-west-3", + "eu-north-1", + "sa-east-1", + "us-gov-east-1", + "us-gov-west-1", +] + +NEO_REGION_LIST = [ + "us-west-1", + "us-west-2", + "us-east-1", + "us-east-2", + "eu-west-1", + "eu-west-2", + "eu-west-3", + "eu-central-1", + "eu-north-1", + "ap-northeast-1", + "ap-northeast-2", + "ap-east-1", + "ap-south-1", + "ap-southeast-1", + "ap-southeast-2", + "sa-east-1", + "ca-central-1", + "me-south-1", + "cn-north-1", + "cn-northwest-1", + "us-gov-west-1", +] + + +@pytest.fixture +def sagemaker_session(): + return Mock(boto_region_name=REGION) + + +def _create_model(sagemaker_session=None): + return Model(MODEL_DATA, MODEL_IMAGE, sagemaker_session=sagemaker_session) + + +def test_compile_model_for_inferentia(sagemaker_session): + sagemaker_session.wait_for_compilation_job = Mock( + return_value=DESCRIBE_COMPILATION_JOB_RESPONSE + ) + model = _create_model(sagemaker_session) + model.compile( + target_instance_family="ml_inf", + input_shape={"data": [1, 3, 1024, 1024]}, + output_path="s3://output", + role="role", + framework="tensorflow", + framework_version="1.15.0", + job_name="compile-model", + ) + assert ( + "{}.dkr.ecr.{}.amazonaws.com/sagemaker-neo-tensorflow:1.15.0-inf-py3".format( + NEO_REGION_ACCOUNT, REGION + ) + == model.image + ) + assert model._is_compiled_model is True + + +def test_compile_model_for_edge_device(sagemaker_session): + sagemaker_session.wait_for_compilation_job = Mock( + return_value=DESCRIBE_COMPILATION_JOB_RESPONSE + ) + model = _create_model(sagemaker_session) + model.compile( + target_instance_family="deeplens", + input_shape={"data": [1, 3, 1024, 1024]}, + output_path="s3://output", + role="role", + framework="tensorflow", + job_name="compile-model", + ) + assert model._is_compiled_model is False + + +def test_compile_model_for_edge_device_tflite(sagemaker_session): + sagemaker_session.wait_for_compilation_job = Mock( + return_value=DESCRIBE_COMPILATION_JOB_RESPONSE + ) + model = _create_model(sagemaker_session) + model.compile( + target_instance_family="deeplens", + input_shape={"data": [1, 3, 1024, 1024]}, + output_path="s3://output", + role="role", + framework="tflite", + job_name="tflite-compile-model", + ) + assert model._is_compiled_model is False + + +def test_compile_model_for_cloud(sagemaker_session): + sagemaker_session.wait_for_compilation_job = Mock( + return_value=DESCRIBE_COMPILATION_JOB_RESPONSE + ) + model = _create_model(sagemaker_session) + model.compile( + target_instance_family="ml_c4", + input_shape={"data": [1, 3, 1024, 1024]}, + output_path="s3://output", + role="role", + framework="tensorflow", + job_name="compile-model", + ) + assert model._is_compiled_model is True + + +def test_compile_model_for_cloud_tflite(sagemaker_session): + sagemaker_session.wait_for_compilation_job = Mock( + return_value=DESCRIBE_COMPILATION_JOB_RESPONSE + ) + model = _create_model(sagemaker_session) + model.compile( + target_instance_family="ml_c4", + input_shape={"data": [1, 3, 1024, 1024]}, + output_path="s3://output", + role="role", + framework="tflite", + job_name="tflite-compile-model", + ) + assert model._is_compiled_model is True + + +@patch("sagemaker.session.Session") +def test_compile_creates_session(session): + session.return_value.boto_region_name = REGION + + model = _create_model() + model.compile( + target_instance_family="ml_c4", + input_shape={"data": [1, 3, 1024, 1024]}, + output_path="s3://output", + role="role", + framework="tensorflow", + job_name="compile-model", + ) + + assert session.return_value == model.sagemaker_session + + +def test_check_neo_region(sagemaker_session): + sagemaker_session.wait_for_compilation_job = Mock( + return_value=DESCRIBE_COMPILATION_JOB_RESPONSE + ) + model = _create_model(sagemaker_session) + for region_name in EC2_REGION_LIST: + if region_name in NEO_REGION_LIST: + assert model.check_neo_region(region_name) is True + else: + assert model.check_neo_region(region_name) is False From 831ab6305a0fe2c31b308b90ce0d416da09d221c Mon Sep 17 00:00:00 2001 From: Lauren Yu <6631887+laurenyu@users.noreply.github.com> Date: Thu, 16 Apr 2020 12:39:31 -0700 Subject: [PATCH 2/2] move hardcoded region list --- tests/unit/__init__.py | 24 ++++++++++ tests/unit/sagemaker/model/test_neo.py | 61 +++----------------------- 2 files changed, 31 insertions(+), 54 deletions(-) diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py index db50764966..a94934a1bc 100644 --- a/tests/unit/__init__.py +++ b/tests/unit/__init__.py @@ -15,3 +15,27 @@ import os DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data") + +NEO_REGION_LIST = [ + "us-west-1", + "us-west-2", + "us-east-1", + "us-east-2", + "eu-west-1", + "eu-west-2", + "eu-west-3", + "eu-central-1", + "eu-north-1", + "ap-northeast-1", + "ap-northeast-2", + "ap-east-1", + "ap-south-1", + "ap-southeast-1", + "ap-southeast-2", + "sa-east-1", + "ca-central-1", + "me-south-1", + "cn-north-1", + "cn-northwest-1", + "us-gov-west-1", +] diff --git a/tests/unit/sagemaker/model/test_neo.py b/tests/unit/sagemaker/model/test_neo.py index 54c0ef5590..0d0492bf13 100644 --- a/tests/unit/sagemaker/model/test_neo.py +++ b/tests/unit/sagemaker/model/test_neo.py @@ -12,10 +12,12 @@ # language governing permissions and limitations under the License. from __future__ import absolute_import +import boto3 import pytest from mock import Mock, patch from sagemaker.model import Model +from tests.unit import NEO_REGION_LIST MODEL_DATA = "s3://bucket/model.tar.gz" MODEL_IMAGE = "mi" @@ -28,55 +30,6 @@ "ModelArtifacts": {"S3ModelArtifacts": "s3://output-path/model.tar.gz"}, } -EC2_REGION_LIST = [ - "us-east-2", - "us-east-1", - "us-west-1", - "us-west-2", - "ap-east-1", - "ap-south-1", - "ap-northeast-3", - "ap-northeast-2", - "ap-southeast-1", - "ap-southeast-2", - "ap-northeast-1", - "ca-central-1", - "cn-north-1", - "cn-northwest-1", - "eu-central-1", - "eu-west-1", - "eu-west-2", - "eu-west-3", - "eu-north-1", - "sa-east-1", - "us-gov-east-1", - "us-gov-west-1", -] - -NEO_REGION_LIST = [ - "us-west-1", - "us-west-2", - "us-east-1", - "us-east-2", - "eu-west-1", - "eu-west-2", - "eu-west-3", - "eu-central-1", - "eu-north-1", - "ap-northeast-1", - "ap-northeast-2", - "ap-east-1", - "ap-south-1", - "ap-southeast-1", - "ap-southeast-2", - "sa-east-1", - "ca-central-1", - "me-south-1", - "cn-north-1", - "cn-northwest-1", - "us-gov-west-1", -] - @pytest.fixture def sagemaker_session(): @@ -196,8 +149,8 @@ def test_check_neo_region(sagemaker_session): return_value=DESCRIBE_COMPILATION_JOB_RESPONSE ) model = _create_model(sagemaker_session) - for region_name in EC2_REGION_LIST: - if region_name in NEO_REGION_LIST: - assert model.check_neo_region(region_name) is True - else: - assert model.check_neo_region(region_name) is False + + boto_session = boto3.Session() + for partition in boto_session.get_available_partitions(): + for region_name in boto_session.get_available_regions("ec2", partition_name=partition): + assert (region_name in NEO_REGION_LIST) is model.check_neo_region(region_name)