change: use recommended inference image uri from Neo API (#2923)

Author: HappyAmazonian

src/sagemaker/model.py

@@ -23,7 +23,6 @@
 import sagemaker
 from sagemaker import (
     fw_utils,
-    image_uris,
     local,
     s3,
     session,
@@ -657,34 +656,6 @@ def _compilation_job_config(
             "job_name": job_name,
         }
 
-    def _compilation_image_uri(self, region, target_instance_type, framework, framework_version):
-        """Retrieve the Neo or Inferentia image URI.
-
-        Args:
-            region (str): The AWS region.
-            target_instance_type (str): Identifies the device on which you want to run
-                your model after compilation, for example: ml_c5. For valid values, see
-                https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
-            framework (str): The framework name.
-            framework_version (str): The framework version.
-        """
-        framework_prefix = ""
-        framework_suffix = ""
-
-        if framework == "xgboost":
-            framework_suffix = "-neo"
-        elif target_instance_type.startswith("ml_inf"):
-            framework_prefix = "inferentia-"
-        else:
-            framework_prefix = "neo-"
-
-        return image_uris.retrieve(
-            "{}{}{}".format(framework_prefix, framework, framework_suffix),
-            region,
-            instance_type=target_instance_type,
-            version=framework_version,
-        )
-
     def package_for_edge(
         self,
         output_path,
@@ -849,12 +820,7 @@ def compile(
             if target_instance_family == "ml_eia2":
                 pass
             elif target_instance_family.startswith("ml_"):
-                self.image_uri = self._compilation_image_uri(
-                    self.sagemaker_session.boto_region_name,
-                    target_instance_family,
-                    framework,
-                    framework_version,
-                )
+                self.image_uri = job_status.get("InferenceImage", None)
                 self._is_compiled_model = True
             else:
                 LOGGER.warning(

tests/unit/sagemaker/model/test_neo.py

@@ -20,12 +20,15 @@
 MODEL_DATA = "s3://bucket/model.tar.gz"
 MODEL_IMAGE = "mi"
 
+IMAGE_URI = "inference-container-uri"
+
 REGION = "us-west-2"
 
 NEO_REGION_ACCOUNT = "301217895009"
 DESCRIBE_COMPILATION_JOB_RESPONSE = {
     "CompilationJobStatus": "Completed",
     "ModelArtifacts": {"S3ModelArtifacts": "s3://output-path/model.tar.gz"},
+    "InferenceImage": IMAGE_URI,
 }
 
 
@@ -52,12 +55,7 @@ def test_compile_model_for_inferentia(sagemaker_session):
         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_uri
-    )
+    assert DESCRIBE_COMPILATION_JOB_RESPONSE["InferenceImage"] == model.image_uri
     assert model._is_compiled_model is True
 
 
@@ -271,11 +269,12 @@ def test_deploy_add_compiled_model_suffix_to_endpoint_name_from_model_name(sagem
     assert model.endpoint_name.startswith("{}-ml-c4".format(model_name))
 
 
-@patch("sagemaker.session.Session")
-def test_compile_with_framework_version_15(session):
-    session.return_value.boto_region_name = REGION
+def test_compile_with_framework_version_15(sagemaker_session):
+    sagemaker_session.wait_for_compilation_job = Mock(
+        return_value=DESCRIBE_COMPILATION_JOB_RESPONSE
+    )
 
-    model = _create_model()
+    model = _create_model(sagemaker_session)
     model.compile(
         target_instance_family="ml_c4",
         input_shape={"data": [1, 3, 1024, 1024]},
@@ -286,14 +285,15 @@ def test_compile_with_framework_version_15(session):
         job_name="compile-model",
     )
 
-    assert "1.5" in model.image_uri
+    assert IMAGE_URI == model.image_uri
 
 
-@patch("sagemaker.session.Session")
-def test_compile_with_framework_version_16(session):
-    session.return_value.boto_region_name = REGION
+def test_compile_with_framework_version_16(sagemaker_session):
+    sagemaker_session.wait_for_compilation_job = Mock(
+        return_value=DESCRIBE_COMPILATION_JOB_RESPONSE
+    )
 
-    model = _create_model()
+    model = _create_model(sagemaker_session)
     model.compile(
         target_instance_family="ml_c4",
         input_shape={"data": [1, 3, 1024, 1024]},
@@ -304,26 +304,7 @@ def test_compile_with_framework_version_16(session):
         job_name="compile-model",
     )
 
-    assert "1.6" in model.image_uri
-
-
-@patch("sagemaker.session.Session")
-def test_compile_validates_framework_version(session):
-    session.return_value.boto_region_name = REGION
-
-    model = _create_model()
-    with pytest.raises(ValueError) as e:
-        model.compile(
-            target_instance_family="ml_c4",
-            input_shape={"data": [1, 3, 1024, 1024]},
-            output_path="s3://output",
-            role="role",
-            framework="pytorch",
-            framework_version="1.6.1",
-            job_name="compile-model",
-        )
-
-    assert "Unsupported neo-pytorch version: 1.6.1." in str(e)
+    assert IMAGE_URI == model.image_uri
 
 
 @patch("sagemaker.session.Session")
@@ -347,3 +328,25 @@ def test_compile_with_pytorch_neo_in_ml_inf(session):
         )
         != model.image_uri
     )
+
+
+def test_compile_validates_framework_version(sagemaker_session):
+    sagemaker_session.wait_for_compilation_job = Mock(
+        return_value={
+            "CompilationJobStatus": "Completed",
+            "ModelArtifacts": {"S3ModelArtifacts": "s3://output-path/model.tar.gz"},
+            "InferenceImage": None,
+        }
+    )
+    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="pytorch",
+        framework_version="1.6.1",
+        job_name="compile-model",
+    )
+
+    assert model.image_uri is None

tests/unit/test_mxnet.py

@@ -68,6 +68,8 @@
 
 ENV_INPUT = {"env_key1": "env_val1", "env_key2": "env_val2", "env_key3": "env_val3"}
 
+INFERENCE_IMAGE_URI = "inference-uri"
+
 
 @pytest.fixture()
 def sagemaker_session():
@@ -83,7 +85,10 @@ def sagemaker_session():
     )
 
     describe = {"ModelArtifacts": {"S3ModelArtifacts": "s3://m/m.tar.gz"}}
-    describe_compilation = {"ModelArtifacts": {"S3ModelArtifacts": "s3://m/model_c5.tar.gz"}}
+    describe_compilation = {
+        "ModelArtifacts": {"S3ModelArtifacts": "s3://m/model_c5.tar.gz"},
+        "InferenceImage": INFERENCE_IMAGE_URI,
+    }
     session.sagemaker_client.create_model_package.side_effect = MODEL_PKG_RESPONSE
     session.sagemaker_client.describe_training_job = Mock(return_value=describe)
     session.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
@@ -195,12 +200,6 @@ def _create_compilation_job(input_shape, output_location):
     }
 
 
-def _neo_inference_image(mxnet_version):
-    return "301217895009.dkr.ecr.us-west-2.amazonaws.com/sagemaker-inference-{}:{}-cpu-py3".format(
-        FRAMEWORK.lower(), mxnet_version
-    )
-
-
 @patch("sagemaker.estimator.name_from_base")
 @patch("sagemaker.utils.create_tar_file", MagicMock())
 def test_create_model(
@@ -422,7 +421,7 @@ def test_mxnet_neo(time, strftime, sagemaker_session, neo_mxnet_version):
     actual_compile_model_args = sagemaker_session.method_calls[3][2]
     assert expected_compile_model_args == actual_compile_model_args
 
-    assert compiled_model.image_uri == _neo_inference_image(neo_mxnet_version)
+    assert compiled_model.image_uri == INFERENCE_IMAGE_URI
 
     predictor = mx.deploy(1, CPU, use_compiled_model=True)
     assert isinstance(predictor, MXNetPredictor)