change: use recommended inference image URI from Neo API

src/sagemaker/model.py

@@ -658,12 +658,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["InferenceImage"]
                 self._is_compiled_model = True
             else:
                 LOGGER.warning(

tests/unit/sagemaker/model/test_neo.py

@@ -26,6 +26,7 @@
 DESCRIBE_COMPILATION_JOB_RESPONSE = {
     "CompilationJobStatus": "Completed",
     "ModelArtifacts": {"S3ModelArtifacts": "s3://output-path/model.tar.gz"},
+    "InferenceImage": "inference-container-uri",
 }
 
 
@@ -52,12 +53,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
 
 
@@ -286,7 +282,7 @@ def test_compile_with_framework_version_15(session):
         job_name="compile-model",
     )
 
-    assert "1.5" in model.image_uri
+    assert model.image_uri is not None
 
 
 @patch("sagemaker.session.Session")
@@ -304,26 +300,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 model.image_uri is not None
 
 
 @patch("sagemaker.session.Session")

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)