breaking: deprecate fw_registry module and use image_uris.retrieve() for SparkML

Author: laurenyu

src/sagemaker/fw_registry.py

@@ -0,0 +1,32 @@
+{
+    "scope": ["inference"],
+    "versions": {
+        "2.2": {
+            "registries": {
+                "ap-east-1": "651117190479",
+                "ap-northeast-1": "354813040037",
+                "ap-northeast-2": "366743142698",
+                "ap-south-1": "720646828776",
+                "ap-southeast-1": "121021644041",
+                "ap-southeast-2": "783357654285",
+                "ca-central-1": "341280168497",
+                "cn-north-1": "450853457545",
+                "cn-northwest-1": "451049120500",
+                "eu-central-1": "492215442770",
+                "eu-north-1": "662702820516",
+                "eu-west-1": "141502667606",
+                "eu-west-2": "764974769150",
+                "eu-west-3": "659782779980",
+                "me-south-1": "801668240914",
+                "sa-east-1": "737474898029",
+                "us-east-1": "683313688378",
+                "us-east-2": "257758044811",
+                "us-gov-west-1": "414596584902",
+                "us-iso-east-1": "833128469047",
+                "us-west-1": "746614075791",
+                "us-west-2": "246618743249"
+            },
+            "repository": "sagemaker-sparkml-serving"
+        }
+    }
+}

src/sagemaker/image_uri_config/sparkml-serving.json

@@ -13,13 +13,11 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from sagemaker import Model, Predictor, Session
+from sagemaker import Model, Predictor, Session, image_uris
 from sagemaker.content_types import CONTENT_TYPE_CSV
-from sagemaker.fw_registry import registry
 from sagemaker.serializers import CSVSerializer
 
 framework_name = "sparkml-serving"
-repo_name = "sagemaker-sparkml-serving"
 
 
 class SparkMLPredictor(Predictor):
@@ -94,9 +92,7 @@ def __init__(self, model_data, role=None, spark_version=2.2, sagemaker_session=N
         # For local mode, sagemaker_session should be passed as None but we need a session to get
         # boto_region_name
         region_name = (sagemaker_session or Session()).boto_region_name
-        image_uri = "{}/{}:{}".format(
-            registry(region_name, framework_name), repo_name, spark_version
-        )
+        image_uri = image_uris.retrieve(framework_name, region_name, version=spark_version)
         super(SparkMLModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/sparkml/model.py

@@ -0,0 +1,52 @@
+# Copyright 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
+
+from sagemaker import image_uris
+from tests.unit.sagemaker.image_uris import expected_uris, regions
+
+ACCOUNTS = {
+    "ap-east-1": "651117190479",
+    "ap-northeast-1": "354813040037",
+    "ap-northeast-2": "366743142698",
+    "ap-south-1": "720646828776",
+    "ap-southeast-1": "121021644041",
+    "ap-southeast-2": "783357654285",
+    "ca-central-1": "341280168497",
+    "cn-north-1": "450853457545",
+    "cn-northwest-1": "451049120500",
+    "eu-central-1": "492215442770",
+    "eu-north-1": "662702820516",
+    "eu-west-1": "141502667606",
+    "eu-west-2": "764974769150",
+    "eu-west-3": "659782779980",
+    "me-south-1": "801668240914",
+    "sa-east-1": "737474898029",
+    "us-east-1": "683313688378",
+    "us-east-2": "257758044811",
+    "us-gov-west-1": "414596584902",
+    "us-iso-east-1": "833128469047",
+    "us-west-1": "746614075791",
+    "us-west-2": "246618743249",
+}
+VERSION = "2.2"
+
+
+def test_sparkml():
+    for region in regions.regions():
+        uri = image_uris.retrieve("sparkml-serving", region=region, version=VERSION)
+
+        expected = expected_uris.algo_uri(
+            "sagemaker-sparkml-serving", ACCOUNTS[region], region, version=VERSION
+        )
+        assert expected == uri