breaking: deprecate fw_utils.create_image_uri()

Author: laurenyu

src/sagemaker/fw_utils.py

@@ -22,7 +22,6 @@
 
 import sagemaker.utils
 from sagemaker import s3
-from sagemaker.utils import get_ecr_image_uri_prefix, ECR_URI_PATTERN
 
 logger = logging.getLogger("sagemaker")
 
@@ -66,66 +65,8 @@
     "Please add framework_version={} to your constructor to avoid this error."
 )
 
-VALID_PY_VERSIONS = ["py2", "py3", "py37"]
-VALID_EIA_FRAMEWORKS = [
-    "tensorflow",
-    "tensorflow-serving",
-    "mxnet",
-    "mxnet-serving",
-    "pytorch-serving",
-]
-PY2_RESTRICTED_EIA_FRAMEWORKS = ["pytorch-serving"]
-PY37_SUPPORTED_FRAMEWORKS = ["tensorflow-scriptmode"]
-VALID_ACCOUNTS_BY_REGION = {
-    "us-gov-west-1": "246785580436",
-    "us-iso-east-1": "744548109606",
-    "cn-north-1": "422961961927",
-    "cn-northwest-1": "423003514399",
-}
-ASIMOV_VALID_ACCOUNTS_BY_REGION = {
-    "us-gov-west-1": "442386744353",
-    "us-iso-east-1": "886529160074",
-    "cn-north-1": "727897471807",
-    "cn-northwest-1": "727897471807",
-}
-OPT_IN_ACCOUNTS_BY_REGION = {"ap-east-1": "057415533634", "me-south-1": "724002660598"}
-ASIMOV_OPT_IN_ACCOUNTS_BY_REGION = {"ap-east-1": "871362719292", "me-south-1": "217643126080"}
-DEFAULT_ACCOUNT = "520713654638"
-ASIMOV_PROD_ACCOUNT = "763104351884"
-ASIMOV_DEFAULT_ACCOUNT = ASIMOV_PROD_ACCOUNT
 SINGLE_GPU_INSTANCE_TYPES = ("ml.p2.xlarge", "ml.p3.2xlarge")
 
-MERGED_FRAMEWORKS_REPO_MAP = {
-    "tensorflow-scriptmode": "tensorflow-training",
-    "tensorflow-serving": "tensorflow-inference",
-    "tensorflow-serving-eia": "tensorflow-inference-eia",
-    "mxnet": "mxnet-training",
-    "mxnet-serving": "mxnet-inference",
-    "mxnet-serving-eia": "mxnet-inference-eia",
-    "pytorch": "pytorch-training",
-    "pytorch-serving": "pytorch-inference",
-    "pytorch-serving-eia": "pytorch-inference-eia",
-}
-
-MERGED_FRAMEWORKS_LOWEST_VERSIONS = {
-    "tensorflow-scriptmode": {"py3": [1, 13, 1], "py2": [1, 14, 0], "py37": [1, 15, 2]},
-    "tensorflow-serving": [1, 13, 0],
-    "tensorflow-serving-eia": [1, 14, 0],
-    "mxnet": {"py3": [1, 4, 1], "py2": [1, 6, 0]},
-    "mxnet-serving": {"py3": [1, 4, 1], "py2": [1, 6, 0]},
-    "mxnet-serving-eia": [1, 4, 1],
-    "pytorch": [1, 2, 0],
-    "pytorch-serving": [1, 2, 0],
-    "pytorch-serving-eia": [1, 3, 1],
-}
-
-INFERENTIA_VERSION_RANGES = {
-    "neo-mxnet": [[1, 5, 1], [1, 5, 1]],
-    "neo-tensorflow": [[1, 15, 0], [1, 15, 0]],
-}
-
-INFERENTIA_SUPPORTED_REGIONS = ["us-east-1", "us-west-2"]
-
 DEBUGGER_UNSUPPORTED_REGIONS = ["us-gov-west-1", "us-iso-east-1"]
 
 
@@ -163,229 +104,6 @@ def is_version_equal_or_lower(highest_version, framework_version):
     return version_list <= highest_version[0 : len(version_list)]
 
 
-def _is_dlc_version(framework, framework_version, py_version):
-    """Return if the framework's version uses the corresponding DLC image.
-
-    Args:
-        framework (str): The framework name, e.g. "tensorflow-scriptmode"
-        framework_version (str): The framework version
-        py_version (str): The Python version, e.g. "py3"
-
-    Returns:
-        bool: Whether or not the framework's version uses the DLC image.
-    """
-    lowest_version_list = MERGED_FRAMEWORKS_LOWEST_VERSIONS.get(framework)
-    if isinstance(lowest_version_list, dict):
-        lowest_version_list = lowest_version_list[py_version]
-
-    if lowest_version_list:
-        return is_version_equal_or_higher(lowest_version_list, framework_version)
-    return False
-
-
-def _is_inferentia_supported(framework, framework_version):
-    """Return if Inferentia supports the framework and its version.
-
-    Args:
-        framework (str): The framework name, e.g. "tensorflow"
-        framework_version (str): The framework version
-
-    Returns:
-        bool: Whether or not Inferentia supports the framework and its version.
-    """
-    lowest_version_list = INFERENTIA_VERSION_RANGES.get(framework)[0]
-    highest_version_list = INFERENTIA_VERSION_RANGES.get(framework)[1]
-    return is_version_equal_or_higher(
-        lowest_version_list, framework_version
-    ) and is_version_equal_or_lower(highest_version_list, framework_version)
-
-
-def _registry_id(region, framework, py_version, account, framework_version):
-    """Return the Amazon ECR registry number (or AWS account ID) for
-    the given framework, framework version, Python version, and region.
-
-    Args:
-        region (str): The AWS region.
-        framework (str): The framework name, e.g. "tensorflow-scriptmode".
-        py_version (str): The Python version, e.g. "py3".
-        account (str): The AWS account ID to use as a default.
-        framework_version (str): The framework version.
-
-    Returns:
-        str: The appropriate Amazon ECR registry number. If there is no
-            specific one for the framework, framework version, Python version,
-            and region, then ``account`` is returned.
-    """
-    if _is_dlc_version(framework, framework_version, py_version):
-        if region in ASIMOV_OPT_IN_ACCOUNTS_BY_REGION:
-            return ASIMOV_OPT_IN_ACCOUNTS_BY_REGION.get(region)
-        if region in ASIMOV_VALID_ACCOUNTS_BY_REGION:
-            return ASIMOV_VALID_ACCOUNTS_BY_REGION.get(region)
-        return ASIMOV_DEFAULT_ACCOUNT
-    if region in OPT_IN_ACCOUNTS_BY_REGION:
-        return OPT_IN_ACCOUNTS_BY_REGION.get(region)
-    return VALID_ACCOUNTS_BY_REGION.get(region, account)
-
-
-def create_image_uri(
-    region,
-    framework,
-    instance_type,
-    framework_version,
-    py_version=None,
-    account=None,
-    accelerator_type=None,
-    optimized_families=None,
-):
-    """Return the ECR URI of an image.
-
-    Args:
-        region (str): AWS region where the image is uploaded.
-        framework (str): framework used by the image.
-        instance_type (str): SageMaker instance type. Used to determine device
-            type (cpu/gpu/family-specific optimized).
-        framework_version (str): The version of the framework.
-        py_version (str): Optional. Python version. If specified, should be one
-            of 'py2' or 'py3'. If not specified, image uri will not include a
-            python component.
-        account (str): AWS account that contains the image. (default:
-            '520713654638')
-        accelerator_type (str): SageMaker Elastic Inference accelerator type.
-        optimized_families (str): Instance families for which there exist
-            specific optimized images.
-
-    Returns:
-        str: The appropriate image URI based on the given parameters.
-    """
-    logger.warning(
-        "'create_image_uri' will be deprecated in favor of 'ImageURIProvider' class "
-        "in SageMaker Python SDK v2."
-    )
-
-    optimized_families = optimized_families or []
-
-    if py_version and py_version not in VALID_PY_VERSIONS:
-        raise ValueError("invalid py_version argument: {}".format(py_version))
-
-    if py_version == "py37" and framework not in PY37_SUPPORTED_FRAMEWORKS:
-        raise ValueError("{} does not support Python 3.7 at this time.".format(framework))
-
-    if _accelerator_type_valid_for_framework(
-        framework=framework,
-        py_version=py_version,
-        accelerator_type=accelerator_type,
-        optimized_families=optimized_families,
-    ):
-        framework += "-eia"
-
-    # Handle account number for specific cases (e.g. GovCloud, opt-in regions, DLC images etc.)
-    if account is None:
-        account = _registry_id(
-            region=region,
-            framework=framework,
-            py_version=py_version,
-            account=DEFAULT_ACCOUNT,
-            framework_version=framework_version,
-        )
-
-    # Handle Local Mode
-    if instance_type.startswith("local"):
-        device_type = "cpu" if instance_type == "local" else "gpu"
-    elif not instance_type.startswith("ml."):
-        raise ValueError(
-            "{} is not a valid SageMaker instance type. See: "
-            "https://aws.amazon.com/sagemaker/pricing/instance-types/".format(instance_type)
-        )
-    else:
-        family = instance_type.split(".")[1]
-
-        # For some frameworks, we have optimized images for specific families, e.g c5 or p3.
-        # In those cases, we use the family name in the image tag. In other cases, we use
-        # 'cpu' or 'gpu'.
-        if family in optimized_families:
-            device_type = family
-        elif family.startswith("inf"):
-            device_type = "inf"
-        elif family[0] in ["g", "p"]:
-            device_type = "gpu"
-        else:
-            device_type = "cpu"
-
-    if device_type == "inf":
-        if region not in INFERENTIA_SUPPORTED_REGIONS:
-            raise ValueError(
-                "Inferentia is not supported in region {}. Supported regions are {}".format(
-                    region, ", ".join(INFERENTIA_SUPPORTED_REGIONS)
-                )
-            )
-        if framework not in INFERENTIA_VERSION_RANGES:
-            raise ValueError(
-                "Inferentia does not support {}. Currently it supports "
-                "MXNet and TensorFlow with more frameworks coming soon.".format(
-                    framework.split("-")[-1]
-                )
-            )
-        if not _is_inferentia_supported(framework, framework_version):
-            raise ValueError(
-                "Inferentia is not supported with {} version {}.".format(
-                    framework.split("-")[-1], framework_version
-                )
-            )
-
-    use_dlc_image = _is_dlc_version(framework, framework_version, py_version)
-
-    if not py_version or (use_dlc_image and framework == "tensorflow-serving-eia"):
-        tag = "{}-{}".format(framework_version, device_type)
-    else:
-        tag = "{}-{}-{}".format(framework_version, device_type, py_version)
-
-    if use_dlc_image:
-        ecr_repo = MERGED_FRAMEWORKS_REPO_MAP[framework]
-    else:
-        ecr_repo = "sagemaker-{}".format(framework)
-
-    return "{}/{}:{}".format(get_ecr_image_uri_prefix(account, region), ecr_repo, tag)
-
-
-def _accelerator_type_valid_for_framework(
-    framework, py_version, accelerator_type=None, optimized_families=None
-):
-    """
-    Args:
-        framework:
-        py_version:
-        accelerator_type:
-        optimized_families:
-    """
-    if accelerator_type is None:
-        return False
-
-    if py_version == "py2" and framework in PY2_RESTRICTED_EIA_FRAMEWORKS:
-        raise ValueError(
-            "{} is not supported with Amazon Elastic Inference in Python 2.".format(framework)
-        )
-
-    if framework not in VALID_EIA_FRAMEWORKS:
-        raise ValueError(
-            "{} is not supported with Amazon Elastic Inference. Currently only "
-            "Python-based TensorFlow, MXNet, PyTorch are supported.".format(framework)
-        )
-
-    if optimized_families:
-        raise ValueError("Neo does not support Amazon Elastic Inference.")
-
-    if (
-        not accelerator_type.startswith("ml.eia")
-        and not accelerator_type == "local_sagemaker_notebook"
-    ):
-        raise ValueError(
-            "{} is not a valid SageMaker Elastic Inference accelerator type. "
-            "See: https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html".format(accelerator_type)
-        )
-
-    return True
-
-
 def validate_source_dir(script, directory):
     """Validate that the source directory exists and it contains the user script
     Args:
@@ -505,7 +223,7 @@ def framework_name_from_image(image_uri):
             str: The framework name str: The Python version str: The image tag
             str: If the image is script mode
         """
-    sagemaker_pattern = re.compile(ECR_URI_PATTERN)
+    sagemaker_pattern = re.compile(sagemaker.utils.ECR_URI_PATTERN)
     sagemaker_match = sagemaker_pattern.match(image_uri)
     if sagemaker_match is None:
         return None, None, None, None

src/sagemaker/image_uris.py

@@ -84,6 +84,8 @@ def _config_for_framework_and_scope(framework, image_scope, accelerator_type=Non
     config = config_for_framework(framework)
 
     if accelerator_type:
+        _validate_accelerator_type(accelerator_type)
+
         if image_scope not in ("eia", "inference"):
             logger.warning(
                 "Elastic inference is for inference only. Ignoring image scope: %s.", image_scope
@@ -118,6 +120,14 @@ def config_for_framework(framework):
         return json.load(f)
 
 
+def _validate_accelerator_type(accelerator_type):
+    if not accelerator_type.startswith("ml.eia") and accelerator_type != "local_sagemaker_notebook":
+        raise ValueError(
+            "Invalid SageMaker Elastic Inference accelerator type: {}. "
+            "See https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html".format(accelerator_type)
+        )
+
+
 def _validate_version_and_set_if_needed(version, config, framework):
     """Checks if the framework/algorithm version is one of the supported versions."""
     available_versions = list(config["versions"].keys())

tests/unit/sagemaker/image_uris/test_retrieve.py

@@ -143,14 +143,28 @@ def test_retrieve_eia(config_for_framework, caplog):
         version="1.0.0",
         py_version="py3",
         instance_type="ml.c4.xlarge",
-        accelerator_type="ml.eia1.medium",
+        accelerator_type="local_sagemaker_notebook",
         region="us-west-2",
         image_scope="training",
     )
     assert "123412341234.dkr.ecr.us-west-2.amazonaws.com/dummy-eia:1.0.0-cpu-py3" == uri
     assert "Ignoring image scope: training." in caplog.text
 
 
+@patch("sagemaker.image_uris.config_for_framework", return_value=BASE_CONFIG)
+def test_retrieve_invalid_accelerator(config_for_framework):
+    with pytest.raises(ValueError) as e:
+        image_uris.retrieve(
+            framework="useless-string",
+            version="1.0.0",
+            py_version="py3",
+            instance_type="ml.c4.xlarge",
+            accelerator_type="fake-accelerator",
+            region="us-west-2",
+        )
+    assert "Invalid SageMaker Elastic Inference accelerator type: fake-accelerator." in str(e.value)
+
+
 @patch("sagemaker.image_uris.config_for_framework")
 def test_retrieve_aliased_version(config_for_framework):
     version = "1.0.0-build123"