breaking: deprecate sagemaker.amazon.amazon_estimator.get_image_uri()

doc/overview.rst

@@ -993,12 +993,17 @@ the ML Pipeline.
 
 .. code:: python
 
-   xgb_image = get_image_uri(sess.boto_region_name, 'xgboost', repo_version="latest")
-   xgb_model = Model(model_data='s3://path/to/model.tar.gz', image_uri=xgb_image)
-   sparkml_model = SparkMLModel(model_data='s3://path/to/model.tar.gz', env={'SAGEMAKER_SPARKML_SCHEMA': schema})
+   from sagemaker import image_uris, session
+   from sagemaker.model import Model
+   from sagemaker.pipeline import PipelineModel
+   from sagemaker.sparkml import SparkMLModel
 
-   model_name = 'inference-pipeline-model'
-   endpoint_name = 'inference-pipeline-endpoint'
+   xgb_image = image_uris.retrieve("xgboost", session.Session().boto_region_name, repo_version="latest")
+   xgb_model = Model(model_data="s3://path/to/model.tar.gz", image_uri=xgb_image)
+   sparkml_model = SparkMLModel(model_data="s3://path/to/model.tar.gz", env={"SAGEMAKER_SPARKML_SCHEMA": schema})
+
+   model_name = "inference-pipeline-model"
+   endpoint_name = "inference-pipeline-endpoint"
    sm_model = PipelineModel(name=model_name, role=sagemaker_role, models=[sparkml_model, xgb_model])
 
 This defines a ``PipelineModel`` consisting of SparkML model and an XGBoost model stacked sequentially.

src/sagemaker/amazon/amazon_estimator.py

@@ -19,13 +19,13 @@
 
 from six.moves.urllib.parse import urlparse
 
+from sagemaker import image_uris
 from sagemaker.amazon import validation
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.common import write_numpy_to_dense_tensor
 from sagemaker.estimator import EstimatorBase, _TrainingJob
 from sagemaker.inputs import FileSystemInput, TrainingInput
-from sagemaker.model import NEO_IMAGE_ACCOUNT
-from sagemaker.utils import sagemaker_timestamp, get_ecr_image_uri_prefix
+from sagemaker.utils import sagemaker_timestamp
 
 logger = logging.getLogger(__name__)
 
@@ -93,8 +93,8 @@ def __init__(
 
     def train_image(self):
         """Placeholder docstring"""
-        return get_image_uri(
-            self.sagemaker_session.boto_region_name, type(self).repo_name, type(self).repo_version
+        return image_uris.retrieve(
+            self.repo_name, self.sagemaker_session.boto_region_name, version=self.repo_version,
         )
 
     def hyperparameters(self):
@@ -452,167 +452,3 @@ def upload_numpy_to_s3_shards(
                 s3.Object(bucket, key_prefix + file).delete()
         finally:
             raise ex
-
-
-def registry(region_name, algorithm=None):
-    """Return docker registry for the given AWS region
-
-    Note: Not all the algorithms listed below have an Amazon Estimator
-    implemented. For full list of pre-implemented Estimators, look at:
-
-    https://github.com/aws/sagemaker-python-sdk/tree/master/src/sagemaker/amazon
-
-    Args:
-        region_name (str): The region name for the account.
-        algorithm (str): The algorithm for the account.
-
-    Raises:
-        ValueError: If invalid algorithm passed in or if mapping does not exist for given algorithm
-            and region.
-    """
-    region_to_accounts = {}
-    if algorithm in [
-        None,
-        "pca",
-        "kmeans",
-        "linear-learner",
-        "factorization-machines",
-        "ntm",
-        "randomcutforest",
-        "knn",
-        "object2vec",
-        "ipinsights",
-    ]:
-        region_to_accounts = {
-            "us-east-1": "382416733822",
-            "us-east-2": "404615174143",
-            "us-west-2": "174872318107",
-            "eu-west-1": "438346466558",
-            "eu-central-1": "664544806723",
-            "ap-northeast-1": "351501993468",
-            "ap-northeast-2": "835164637446",
-            "ap-southeast-2": "712309505854",
-            "us-gov-west-1": "226302683700",
-            "ap-southeast-1": "475088953585",
-            "ap-south-1": "991648021394",
-            "ca-central-1": "469771592824",
-            "eu-west-2": "644912444149",
-            "us-west-1": "632365934929",
-            "us-iso-east-1": "490574956308",
-            "ap-east-1": "286214385809",
-            "eu-north-1": "669576153137",
-            "eu-west-3": "749696950732",
-            "sa-east-1": "855470959533",
-            "me-south-1": "249704162688",
-            "cn-north-1": "390948362332",
-            "cn-northwest-1": "387376663083",
-        }
-    elif algorithm in ["lda"]:
-        region_to_accounts = {
-            "us-east-1": "766337827248",
-            "us-east-2": "999911452149",
-            "us-west-2": "266724342769",
-            "eu-west-1": "999678624901",
-            "eu-central-1": "353608530281",
-            "ap-northeast-1": "258307448986",
-            "ap-northeast-2": "293181348795",
-            "ap-southeast-2": "297031611018",
-            "us-gov-west-1": "226302683700",
-            "ap-southeast-1": "475088953585",
-            "ap-south-1": "991648021394",
-            "ca-central-1": "469771592824",
-            "eu-west-2": "644912444149",
-            "us-west-1": "632365934929",
-            "us-iso-east-1": "490574956308",
-        }
-    elif algorithm in ["forecasting-deepar"]:
-        region_to_accounts = {
-            "us-east-1": "522234722520",
-            "us-east-2": "566113047672",
-            "us-west-2": "156387875391",
-            "eu-west-1": "224300973850",
-            "eu-central-1": "495149712605",
-            "ap-northeast-1": "633353088612",
-            "ap-northeast-2": "204372634319",
-            "ap-southeast-2": "514117268639",
-            "us-gov-west-1": "226302683700",
-            "ap-southeast-1": "475088953585",
-            "ap-south-1": "991648021394",
-            "ca-central-1": "469771592824",
-            "eu-west-2": "644912444149",
-            "us-west-1": "632365934929",
-            "us-iso-east-1": "490574956308",
-            "ap-east-1": "286214385809",
-            "eu-north-1": "669576153137",
-            "eu-west-3": "749696950732",
-            "sa-east-1": "855470959533",
-            "me-south-1": "249704162688",
-            "cn-north-1": "390948362332",
-            "cn-northwest-1": "387376663083",
-        }
-    elif algorithm in [
-        "xgboost",
-        "seq2seq",
-        "image-classification",
-        "blazingtext",
-        "object-detection",
-        "semantic-segmentation",
-    ]:
-        region_to_accounts = {
-            "us-east-1": "811284229777",
-            "us-east-2": "825641698319",
-            "us-west-2": "433757028032",
-            "eu-west-1": "685385470294",
-            "eu-central-1": "813361260812",
-            "ap-northeast-1": "501404015308",
-            "ap-northeast-2": "306986355934",
-            "ap-southeast-2": "544295431143",
-            "us-gov-west-1": "226302683700",
-            "ap-southeast-1": "475088953585",
-            "ap-south-1": "991648021394",
-            "ca-central-1": "469771592824",
-            "eu-west-2": "644912444149",
-            "us-west-1": "632365934929",
-            "us-iso-east-1": "490574956308",
-            "ap-east-1": "286214385809",
-            "eu-north-1": "669576153137",
-            "eu-west-3": "749696950732",
-            "sa-east-1": "855470959533",
-            "me-south-1": "249704162688",
-            "cn-north-1": "390948362332",
-            "cn-northwest-1": "387376663083",
-        }
-    elif algorithm in ["image-classification-neo", "xgboost-neo"]:
-        region_to_accounts = NEO_IMAGE_ACCOUNT
-    else:
-        raise ValueError(
-            "Algorithm class:{} does not have mapping to account_id with images".format(algorithm)
-        )
-
-    if region_name in region_to_accounts:
-        account_id = region_to_accounts[region_name]
-        return get_ecr_image_uri_prefix(account_id, region_name)
-
-    raise ValueError(
-        "Algorithm ({algorithm}) is unsupported for region ({region_name}).".format(
-            algorithm=algorithm, region_name=region_name
-        )
-    )
-
-
-def get_image_uri(region_name, repo_name, repo_version=1):
-    """Return algorithm image URI for the given AWS region, repository name, and
-    repository version
-
-    Args:
-        region_name:
-        repo_name:
-        repo_version:
-    """
-    logger.warning(
-        "'get_image_uri' method will be deprecated in favor of 'ImageURIProvider' class "
-        "in SageMaker Python SDK v2."
-    )
-
-    repo = "{}:{}".format(repo_name, repo_version)
-    return "{}/{}".format(registry(region_name, repo_name), repo)

src/sagemaker/amazon/factorization_machines.py

@@ -13,7 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
+from sagemaker import image_uris
+from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin, ge
@@ -309,8 +310,11 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             **kwargs:
         """
         sagemaker_session = sagemaker_session or Session()
-        repo = "{}:{}".format(FactorizationMachines.repo_name, FactorizationMachines.repo_version)
-        image_uri = "{}/{}".format(registry(sagemaker_session.boto_session.region_name), repo)
+        image_uri = image_uris.retrieve(
+            FactorizationMachines.repo_name,
+            sagemaker_session.boto_region_name,
+            version=FactorizationMachines.repo_version,
+        )
         super(FactorizationMachinesModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/ipinsights.py

@@ -13,7 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
+from sagemaker import image_uris
+from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, le
 from sagemaker.deserializers import JSONDeserializer
@@ -219,11 +220,11 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             **kwargs:
         """
         sagemaker_session = sagemaker_session or Session()
-        repo = "{}:{}".format(IPInsights.repo_name, IPInsights.repo_version)
-        image_uri = "{}/{}".format(
-            registry(sagemaker_session.boto_session.region_name, IPInsights.repo_name), repo
+        image_uri = image_uris.retrieve(
+            IPInsights.repo_name,
+            sagemaker_session.boto_region_name,
+            version=IPInsights.repo_version,
         )
-
         super(IPInsightsModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/kmeans.py

@@ -13,7 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
+from sagemaker import image_uris
+from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin, ge, le
@@ -242,8 +243,9 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             **kwargs:
         """
         sagemaker_session = sagemaker_session or Session()
-        repo = "{}:{}".format(KMeans.repo_name, KMeans.repo_version)
-        image_uri = "{}/{}".format(registry(sagemaker_session.boto_session.region_name), repo)
+        image_uri = image_uris.retrieve(
+            KMeans.repo_name, sagemaker_session.boto_region_name, version=KMeans.repo_version,
+        )
         super(KMeansModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/knn.py

@@ -13,7 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
+from sagemaker import image_uris
+from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, isin
@@ -230,12 +231,11 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             **kwargs:
         """
         sagemaker_session = sagemaker_session or Session()
-        repo = "{}:{}".format(KNN.repo_name, KNN.repo_version)
-        image = "{}/{}".format(
-            registry(sagemaker_session.boto_session.region_name, KNN.repo_name), repo
+        image_uri = image_uris.retrieve(
+            KNN.repo_name, sagemaker_session.boto_region_name, version=KNN.repo_version,
         )
         super(KNNModel, self).__init__(
-            image,
+            image_uri,
             model_data,
             role,
             predictor_cls=KNNPredictor,

src/sagemaker/amazon/lda.py

@@ -13,7 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
+from sagemaker import image_uris
+from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt
@@ -214,9 +215,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             **kwargs:
         """
         sagemaker_session = sagemaker_session or Session()
-        repo = "{}:{}".format(LDA.repo_name, LDA.repo_version)
-        image_uri = "{}/{}".format(
-            registry(sagemaker_session.boto_session.region_name, LDA.repo_name), repo
+        image_uri = image_uris.retrieve(
+            LDA.repo_name, sagemaker_session.boto_region_name, version=LDA.repo_version,
         )
         super(LDAModel, self).__init__(
             image_uri,

src/sagemaker/amazon/linear_learner.py

@@ -13,7 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
+from sagemaker import image_uris
+from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import isin, gt, lt, ge, le
@@ -473,8 +474,11 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             **kwargs:
         """
         sagemaker_session = sagemaker_session or Session()
-        repo = "{}:{}".format(LinearLearner.repo_name, LinearLearner.repo_version)
-        image_uri = "{}/{}".format(registry(sagemaker_session.boto_session.region_name), repo)
+        image_uri = image_uris.retrieve(
+            LinearLearner.repo_name,
+            sagemaker_session.boto_region_name,
+            version=LinearLearner.repo_version,
+        )
         super(LinearLearnerModel, self).__init__(
             image_uri,
             model_data,

src/sagemaker/amazon/ntm.py

@@ -13,7 +13,8 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
+from sagemaker import image_uris
+from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase
 from sagemaker.amazon.common import RecordSerializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, le, isin
@@ -244,9 +245,8 @@ def __init__(self, model_data, role, sagemaker_session=None, **kwargs):
             **kwargs:
         """
         sagemaker_session = sagemaker_session or Session()
-        repo = "{}:{}".format(NTM.repo_name, NTM.repo_version)
-        image_uri = "{}/{}".format(
-            registry(sagemaker_session.boto_session.region_name, NTM.repo_name), repo
+        image_uri = image_uris.retrieve(
+            NTM.repo_name, sagemaker_session.boto_region_name, version=NTM.repo_version,
         )
         super(NTMModel, self).__init__(
             image_uri,