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_framework_model.py b/tests/unit/sagemaker/model/test_framework_model.py index c3498a839e..09cb5531ce 100644 --- a/tests/unit/sagemaker/model/test_framework_model.py +++ b/tests/unit/sagemaker/model/test_framework_model.py @@ -36,7 +36,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" @@ -50,11 +49,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): @@ -237,170 +231,6 @@ def test_deploy_update_endpoint_optional_args(sagemaker_session, tmpdir): sagemaker_session.create_endpoint.assert_not_called() -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..0d0492bf13 --- /dev/null +++ b/tests/unit/sagemaker/model/test_neo.py @@ -0,0 +1,156 @@ +# 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 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" + +REGION = "us-west-2" + +NEO_REGION_ACCOUNT = "301217895009" +DESCRIBE_COMPILATION_JOB_RESPONSE = { + "CompilationJobStatus": "Completed", + "ModelArtifacts": {"S3ModelArtifacts": "s3://output-path/model.tar.gz"}, +} + + +@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) + + 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)