Skip to content

breaking: deprecate fw_utils.create_image_uri() #1770

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 29, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
284 changes: 1 addition & 283 deletions src/sagemaker/fw_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down Expand Up @@ -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"]


Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
11 changes: 11 additions & 0 deletions src/sagemaker/image_uris.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,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
Expand Down Expand Up @@ -121,6 +123,15 @@ def config_for_framework(framework):
return json.load(f)


def _validate_accelerator_type(accelerator_type):
"""Raises a ``ValueError`` if ``accelerator_type`` is invalid."""
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())
Expand Down
16 changes: 15 additions & 1 deletion tests/unit/sagemaker/image_uris/test_retrieve.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Loading