Fixing py2 support for latest TF version

Author: Sam Oshin

src/sagemaker/tensorflow/estimator.py

@@ -203,7 +203,7 @@ class TensorFlow(Framework):
     _LOWEST_SCRIPT_MODE_ONLY_VERSION = [1, 13]
     # 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]
+    _HIGHEST_PYTHON_2_VERSION = [1, 15, 0]
 
     def __init__(
         self,
@@ -371,7 +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._HIGHEST_PYTHON_2_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/unit/test_tf_estimator.py

@@ -987,6 +987,9 @@ def test_py2_version_deprecated(sagemaker_session):
     msg = "Python 2 containers are only available until January 1st, 2020. Please use a Python 3 container."
     assert msg in str(e.value)
 
+def test_py2_version_is_not_deprecated(sagemaker_session):
+    estimator = _build_tf(sagemaker_session=sagemaker_session, framework_version="1.15.0", py_version="py2")
+    assert estimator.py_version == "py2"
 
 def test_py3_is_default_version_before_tf1_14(sagemaker_session):
     estimator = _build_tf(sagemaker_session=sagemaker_session, framework_version="1.13")