add compile model warning and integ test

Author: Deng

src/sagemaker/model.py

@@ -464,6 +464,11 @@ def deploy(
         if self.role is None:
             raise ValueError("Role can not be null for deploying a model")
 
+        if instance_type.startswith("ml.inf") and not self._is_compiled_model:
+            LOGGER.warning(
+                "Your model is not compiled, please compile your model before using Inferentia."
+            )
+
         compiled_model_suffix = "-".join(instance_type.split(".")[:-1])
         if self._is_compiled_model:
             name_prefix = self.name or utils.name_from_image(self.image)

tests/conftest.py

@@ -269,6 +269,11 @@ def cpu_instance_type(sagemaker_session, request):
         return "ml.m4.xlarge"
 
 
+@pytest.fixture(scope="session")
+def inf_instance_type(sagemaker_session, request):
+    return "ml.inf1.xlarge"
+
+
 @pytest.fixture(scope="session")
 def ec2_instance_type(cpu_instance_type):
     return cpu_instance_type[3:]
@@ -289,6 +294,11 @@ def cpu_instance_family(cpu_instance_type):
     return "_".join(cpu_instance_type.split(".")[0:2])
 
 
+@pytest.fixture(scope="session")
+def inf_instance_family(inf_instance_type):
+    return "_".join(inf_instance_type.split(".")[0:2])
+
+
 def pytest_generate_tests(metafunc):
     if "instance_type" in metafunc.fixturenames:
         boto_config = metafunc.config.getoption("--boto-config")

tests/integ/test_neo_mxnet.py

@@ -24,6 +24,7 @@
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 NEO_MXNET_VERSION = "1.4.1"  # Neo doesn't support MXNet 1.6 yet.
+INF_MXNET_VERSION = "1.5.0"
 
 
 @pytest.fixture(scope="module")
@@ -110,3 +111,38 @@ def test_deploy_model(
         predictor.content_type = "application/vnd+python.numpy+binary"
         data = numpy.zeros(shape=(1, 1, 28, 28))
         predictor.predict(data)
+
+
+@pytest.mark.skip(reason="Inferentia is not supported yet.")
+def test_inferentia_deploy_model(
+    mxnet_training_job, sagemaker_session, inf_instance_type, inf_instance_family
+):
+    endpoint_name = unique_name_from_base("test-neo-deploy-model")
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+        desc = sagemaker_session.sagemaker_client.describe_training_job(
+            TrainingJobName=mxnet_training_job
+        )
+        model_data = desc["ModelArtifacts"]["S3ModelArtifacts"]
+        script_path = os.path.join(DATA_DIR, "mxnet_mnist", "mnist_neo.py")
+        role = "SageMakerRole"
+        model = MXNetModel(
+            model_data,
+            role,
+            entry_point=script_path,
+            framework_version=NEO_MXNET_VERSION,
+            sagemaker_session=sagemaker_session,
+        )
+
+        model.compile(
+            target_instance_family=inf_instance_family,
+            input_shape={"data": [1, 1, 28, 28]},
+            role=role,
+            job_name=unique_name_from_base("test-deploy-model-compilation-job"),
+            output_path="/".join(model_data.split("/")[:-1]),
+        )
+        predictor = model.deploy(1, inf_instance_type, endpoint_name=endpoint_name)
+
+        predictor.content_type = "application/vnd+python.numpy+binary"
+        data = numpy.zeros(shape=(1, 1, 28, 28))
+        predictor.predict(data)

tests/unit/test_model.py

@@ -39,6 +39,7 @@
 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"
@@ -544,6 +545,29 @@ 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