feature: add support for TF 1.15.0, PyTorch 1.3.1 and MXNet 1.6rc0. (#280)

Author: nadiaya

src/sagemaker/fw_utils.py

@@ -60,6 +60,8 @@
 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 = "028651357192"
 
 MERGED_FRAMEWORKS_REPO_MAP = {
     "tensorflow-scriptmode": "tensorflow-training",
@@ -121,16 +123,18 @@ def _using_merged_images(region, framework, py_version, framework_version):
         framework_version:
     """
     is_gov_region = region in VALID_ACCOUNTS_BY_REGION
-    is_py3 = py_version == "py3" or py_version is None
+    not_py2 = py_version == "py3" or py_version is None
     is_merged_versions = _is_merged_versions(framework, framework_version)
 
     return (
         ((not is_gov_region) or region in ASIMOV_VALID_ACCOUNTS_BY_REGION)
         and is_merged_versions
+        # TODO: should be not mxnet-1.14.1-py2 instead?
         and (
-            is_py3
+            not_py2
             or _is_tf_14_or_later(framework, framework_version)
             or _is_pt_12_or_later(framework, framework_version)
+            or _is_mxnet_16_or_later(framework, framework_version)
         )
     )
 
@@ -162,6 +166,19 @@ def _is_pt_12_or_later(framework, framework_version):
     return is_pytorch and version >= asimov_lowest_pt[0 : len(version)]
 
 
+def _is_mxnet_16_or_later(framework, framework_version):
+    """
+    Args:
+        framework: Name of the frameowork
+        framework_version: framework version
+    """
+    # Asimov team now owns MXNet 1.6.0 py2 and py3
+    asimov_lowest_pt = [1, 6, 0]
+    version = [int(s) for s in framework_version.split(".")]
+    is_mxnet = framework in ("mxnet", "mxnet-serving")
+    return is_mxnet and version >= asimov_lowest_pt[0 : len(version)]
+
+
 def _registry_id(region, framework, py_version, account, framework_version):
     """
     Args:
@@ -177,7 +194,10 @@ def _registry_id(region, framework, py_version, account, framework_version):
             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 "763104351884"
+        # TODO: remove when ASIMOV images availabel in Prod
+        if framework in ("tensorflow-serving-eia", "mxnet-serving-eia"):
+            return ASIMOV_PROD_ACCOUNT
+        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)

src/sagemaker/mxnet/README.rst

@@ -4,9 +4,9 @@ Using MXNet with the SageMaker Python SDK
 
 With the SageMaker Python SDK, you can train and host MXNet models on Amazon SageMaker.
 
-Supported versions of MXNet: ``0.12.1``, ``1.0.0``, ``1.1.0``, ``1.2.1``, ``1.3.0``, ``1.4.0``, ``1.4.1``.
+Supported versions of MXNet: ``0.12.1``, ``1.0.0``, ``1.1.0``, ``1.2.1``, ``1.3.0``, ``1.4.0``, ``1.4.1``, ``1.6.0``.
 
-Supported versions of MXNet for Elastic Inference: ``1.3.0``, ``1.4.0``, ``1.4.1``.
+Supported versions of MXNet for Elastic Inference: ``1.3.0``, ``1.4.0``, ``1.4.1``, ``1.6.0``.
 
 For information about using MXNet with the SageMaker Python SDK, see https://sagemaker.readthedocs.io/en/stable/using_mxnet.html.
 

src/sagemaker/mxnet/estimator.py

@@ -36,7 +36,7 @@ class MXNet(Framework):
     __framework_name__ = "mxnet"
     _LOWEST_SCRIPT_MODE_VERSION = ["1", "3"]
 
-    LATEST_VERSION = "1.4.1"
+    LATEST_VERSION = "1.6.0"
     """The latest version of MXNet included in the SageMaker pre-built Docker images."""
 
     def __init__(

src/sagemaker/pytorch/README.rst

@@ -4,7 +4,7 @@ SageMaker PyTorch Estimators and Models
 
 With PyTorch Estimators and Models, you can train and host PyTorch models on Amazon SageMaker.
 
-Supported versions of PyTorch: ``0.4.0``, ``1.0.0``, ``1.1.0``, ``1.2.0``.
+Supported versions of PyTorch: ``0.4.0``, ``1.0.0``, ``1.1.0``, ``1.2.0``, ``1.3.1``.
 
 We recommend that you use the latest supported version, because that's where we focus most of our development efforts.
 

src/sagemaker/pytorch/estimator.py

@@ -35,7 +35,7 @@ class PyTorch(Framework):
 
     __framework_name__ = "pytorch"
 
-    LATEST_VERSION = "1.2.0"
+    LATEST_VERSION = "1.3.1"
     """The latest version of PyTorch included in the SageMaker pre-built Docker images."""
 
     def __init__(

src/sagemaker/tensorflow/README.rst

@@ -20,7 +20,7 @@ Documentation of the previous Legacy Mode versions: `1.4.1 <https://github.com/a
 | You can find the Legacy Mode README `here <https://github.com/aws/sagemaker-python-sdk/tree/v1.12.0/src/sagemaker/tensorflow#tensorflow-sagemaker-estimators-and-models>`_. |
 +-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 
-Supported versions of TensorFlow for Elastic Inference: ``1.11``, ``1.12``, ``1.13``.
+Supported versions of TensorFlow for Elastic Inference: ``1.11``, ``1.12``, ``1.13``, ``1.15``.
 
 For information about using TensorFlow with the SageMaker Python SDK, see https://sagemaker.readthedocs.io/en/stable/using_tf.html.
 
@@ -64,7 +64,7 @@ The containers include the following Python packages:
 +--------------------------------+---------------+
 | sklean                         | 0.20.3        |
 +--------------------------------+---------------+
-| tensorflow                     | 1.13.0        |
+| tensorflow                     | 1.15.0        |
 +--------------------------------+---------------+
 
 Script Mode TensorFlow Docker images support both Python 2.7 and Python 3.6. The Docker images extend Ubuntu 16.04.

src/sagemaker/tensorflow/estimator.py

@@ -197,13 +197,13 @@ class TensorFlow(Framework):
 
     __framework_name__ = "tensorflow"
 
-    LATEST_VERSION = "1.14"
+    LATEST_VERSION = "1.15"
     """The latest version of TensorFlow included in the SageMaker pre-built Docker images."""
 
     _LOWEST_SCRIPT_MODE_ONLY_VERSION = [1, 13]
-    # 1.14.0 now supports py2
-    # we will need to update this version number if future versions do not support py2 anymore
-    _LOWEST_PYTHON_2_ONLY_VERSION = [1, 14]
+    # 1.15.0 still supports py2
+    # we will need to update this version number if future versions still support py2
+    _HIGHEST_PYTHON_2_VERSION = [1, 15]
 
     def __init__(
         self,
@@ -371,9 +371,7 @@ def _only_script_mode_supported(self):
 
     def _only_python_3_supported(self):
         """Placeholder docstring"""
-        return [
-            int(s) for s in self.framework_version.split(".")
-        ] >= self._LOWEST_PYTHON_2_ONLY_VERSION
+        return [int(s) for s in self.framework_version.split(".")] >= self._HIGHEST_PYTHON_2_VERSION
 
     def _validate_requirements_file(self, requirements_file):
         """Placeholder docstring"""

tests/conftest.py

@@ -42,7 +42,7 @@ def pytest_addoption(parser):
     parser.addoption("--boto-config", action="store", default=None)
     parser.addoption("--chainer-full-version", action="store", default=Chainer.LATEST_VERSION)
     parser.addoption("--mxnet-full-version", action="store", default=MXNet.LATEST_VERSION)
-    parser.addoption("--ei-mxnet-full-version", action="store", default=MXNet.LATEST_VERSION)
+    parser.addoption("--ei-mxnet-full-version", action="store", default="1.4.1")
     parser.addoption("--pytorch-full-version", action="store", default=PyTorch.LATEST_VERSION)
     parser.addoption(
         "--rl-coach-mxnet-full-version",
@@ -125,6 +125,8 @@ def chainer_version(request):
     return request.param
 
 
+# TODO: current version fixtures are legacy fixtures that aren't useful
+# and no longer verify whether images are valid
 @pytest.fixture(
     scope="module",
     params=[