fix: account for "py36" and "py37" in image tag parsing (#1737)

Author: laurenyu

src/sagemaker/fw_utils.py

@@ -493,6 +493,7 @@ def _list_files_to_compress(script, directory):
 def framework_name_from_image(image_name):
     # noinspection LongLine
     """Extract the framework and Python version from the image name.
+
     Args:
         image_name (str): Image URI, which should be one of the following forms:
             legacy:
@@ -503,25 +504,32 @@ def framework_name_from_image(image_name):
             '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-<fw>:<fw_version>-<device>-<py_ver>'
             current:
             '<account>.dkr.ecr.<region>.amazonaws.com/sagemaker-rl-<fw>:<rl_toolkit><rl_version>-<device>-<py_ver>'
+            current:
+            '<account>.dkr.ecr.<region>.amazonaws.com/<fw>-<image_scope>:<fw_version>-<device>-<py_ver>'
+
     Returns:
         tuple: A tuple containing:
-            str: The framework name str: The Python version str: The image tag
-            str: If the image is script mode
-        """
+
+            - str: The framework name
+            - str: The Python version
+            - str: The image tag
+            - str: If the TensorFlow image is script mode
+    """
     sagemaker_pattern = re.compile(ECR_URI_PATTERN)
     sagemaker_match = sagemaker_pattern.match(image_name)
     if sagemaker_match is None:
         return None, None, None, None
+
     # extract framework, python version and image tag
     # We must support both the legacy and current image name format.
     name_pattern = re.compile(
-        r"^(?:sagemaker(?:-rl)?-)?(tensorflow|mxnet|chainer|pytorch|scikit-learn|xgboost)(?:-)?(scriptmode|training)?:(.*)-(.*?)-(py2|py3)$"  # noqa: E501 # pylint: disable=line-too-long
+        r"""^(?:sagemaker(?:-rl)?-)?
+        (tensorflow|mxnet|chainer|pytorch|scikit-learn|xgboost)(?:-)?
+        (scriptmode|training)?
+        :(.*)-(.*?)-(py2|py3[67]?)$""",
+        re.VERBOSE,
     )
-    legacy_name_pattern = re.compile(r"^sagemaker-(tensorflow|mxnet)-(py2|py3)-(cpu|gpu):(.*)$")
-
     name_match = name_pattern.match(sagemaker_match.group(9))
-    legacy_match = legacy_name_pattern.match(sagemaker_match.group(9))
-
     if name_match is not None:
         fw, scriptmode, ver, device, py = (
             name_match.group(1),
@@ -531,20 +539,25 @@ def framework_name_from_image(image_name):
             name_match.group(5),
         )
         return fw, py, "{}-{}-{}".format(ver, device, py), scriptmode
+
+    legacy_name_pattern = re.compile(r"^sagemaker-(tensorflow|mxnet)-(py2|py3)-(cpu|gpu):(.*)$")
+    legacy_match = legacy_name_pattern.match(sagemaker_match.group(9))
     if legacy_match is not None:
         return (legacy_match.group(1), legacy_match.group(2), legacy_match.group(4), None)
     return None, None, None, None
 
 
 def framework_version_from_tag(image_tag):
     """Extract the framework version from the image tag.
+
     Args:
         image_tag (str): Image tag, which should take the form
             '<framework_version>-<device>-<py_version>'
+
     Returns:
         str: The framework version.
     """
-    tag_pattern = re.compile("^(.*)-(cpu|gpu)-(py2|py3)$")
+    tag_pattern = re.compile("^(.*)-(cpu|gpu)-(py2|py3[67]?)$")
     tag_match = tag_pattern.match(image_tag)
     return None if tag_match is None else tag_match.group(1)
 

src/sagemaker/tensorflow/estimator.py

@@ -728,7 +728,7 @@ def _default_s3_path(self, directory, mpi=False):
 
     def _script_mode_enabled(self):
         """Placeholder docstring"""
-        return self.py_version == "py3" or self.script_mode
+        return self.py_version.startswith("py3") or self.script_mode
 
     def _validate_and_set_debugger_configs(self):
         """

tests/unit/test_fw_utils.py

@@ -1156,6 +1156,17 @@ def test_framework_name_from_image_rl():
     )
 
 
+def test_framework_name_from_image_python_versions():
+    image_name = "123.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:2.2-cpu-py37"
+    assert ("tensorflow", "py37", "2.2-cpu-py37", "training") == fw_utils.framework_name_from_image(
+        image_name
+    )
+
+    image_name = "123.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:1.15.2-cpu-py36"
+    expected_result = ("tensorflow", "py36", "1.15.2-cpu-py36", "training")
+    assert expected_result == fw_utils.framework_name_from_image(image_name)
+
+
 def test_legacy_name_from_framework_image():
     image_name = "123.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet-py3-gpu:2.5.6-gpu-py2"
     framework, py_ver, tag, _ = fw_utils.framework_name_from_image(image_name)
@@ -1200,8 +1211,17 @@ def test_legacy_name_from_image_any_tag():
 
 
 def test_framework_version_from_tag():
-    version = fw_utils.framework_version_from_tag("1.5rc-keras-gpu-py2")
-    assert version == "1.5rc-keras"
+    tags = (
+        "1.5rc-keras-cpu-py2",
+        "1.5rc-keras-gpu-py2",
+        "1.5rc-keras-cpu-py3",
+        "1.5rc-keras-gpu-py36",
+        "1.5rc-keras-gpu-py37",
+    )
+
+    for tag in tags:
+        version = fw_utils.framework_version_from_tag(tag)
+        assert "1.5rc-keras" == version
 
 
 def test_framework_version_from_tag_other():

tests/unit/test_tf_estimator.py

@@ -1126,8 +1126,9 @@ def test_legacy_mode_deprecated(sagemaker_session):
 
 
 def test_script_mode_enabled(sagemaker_session):
-    tf = _build_tf(sagemaker_session=sagemaker_session, py_version="py3")
-    assert tf._script_mode_enabled() is True
+    for py_version in ("py3", "py36", "py37"):
+        tf = _build_tf(sagemaker_session=sagemaker_session, py_version=py_version)
+        assert tf._script_mode_enabled() is True
 
     tf = _build_tf(sagemaker_session=sagemaker_session, script_mode=True)
     assert tf._script_mode_enabled() is True