Merge branch 'zwei' into tf-docs

laurenyu · laurenyu · commit b584b7ef106a · 2020-06-08T10:12:43.000-07:00
diff --git a/doc/conf.py b/doc/conf.py
@@ -14,34 +14,7 @@
 from __future__ import absolute_import
 
 import pkg_resources
-import sys
 from datetime import datetime
-from unittest.mock import MagicMock
-
-
-class Mock(MagicMock):
-    @classmethod
-    def __getattr__(cls, name):
-        """
-        Args:
-            name:
-        """
-        if name == "__version__":
-            return "1.4.0"
-        else:
-            return MagicMock()
-
-
-MOCK_MODULES = [
-    "tensorflow",
-    "tensorflow.core",
-    "tensorflow.core.framework",
-    "tensorflow.python",
-    "tensorflow.python.framework",
-    "tensorflow_serving",
-    "tensorflow_serving.apis",
-]
-sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES)
 
 project = u"sagemaker"
 version = pkg_resources.require(project)[0].version
diff --git a/setup.py b/setup.py
@@ -35,6 +35,7 @@ def read_version():
 # Declare minimal set for installation
 required_packages = [
     "boto3>=1.13.6",
+    "google-pasta",
     "numpy>=1.9.0",
     "protobuf>=3.1",
     "protobuf3-to-dict>=0.1.5",
@@ -51,7 +52,6 @@ def read_version():
         "docker-compose>=1.25.2",
         "PyYAML>=5.3, <6",  # PyYAML version has to match docker-compose requirements
     ],
-    "tensorflow": ["tensorflow>=1.3.0"],
     "scipy": ["scipy>=0.19.0"],
 }
 # Meta dependency groups
diff --git a/src/sagemaker/cli/compatibility/v2/modifiers/framework_version.py b/src/sagemaker/cli/compatibility/v2/modifiers/framework_version.py
@@ -29,6 +29,7 @@
 # TODO: check for sagemaker.tensorflow.serving.Model
 FRAMEWORK_CLASSES = FRAMEWORKS + ["{}Model".format(fw) for fw in FRAMEWORKS]
 FRAMEWORK_MODULES = [fw.lower() for fw in FRAMEWORKS]
+FRAMEWORK_SUBMODULES = ("model", "estimator")
 
 
 class FrameworkVersionEnforcer(Modifier):
@@ -68,19 +69,30 @@ def _is_framework_constructor(self, node):
         if isinstance(node.func, ast.Name):
             return node.func.id in FRAMEWORK_CLASSES
 
-        # Check for sagemaker.<framework>.<Framework> call
-        ends_with_framework_constructor = (
-            isinstance(node.func, ast.Attribute) and node.func.attr in FRAMEWORK_CLASSES
-        )
+        # Check for something.that.ends.with.<framework>.<Framework> call
+        if not (isinstance(node.func, ast.Attribute) and node.func.attr in FRAMEWORK_CLASSES):
+            return False
 
-        is_in_framework_module = (
+        # Check for sagemaker.<frameworks>.<estimator/model>.<Framework> call
+        if (
             isinstance(node.func.value, ast.Attribute)
-            and node.func.value.attr in FRAMEWORK_MODULES
-            and isinstance(node.func.value.value, ast.Name)
-            and node.func.value.value.id == "sagemaker"
-        )
+            and node.func.value.attr in FRAMEWORK_SUBMODULES
+        ):
+            return self._is_in_framework_module(node.func.value)
 
-        return ends_with_framework_constructor and is_in_framework_module
+        # Check for sagemaker.<framework>.<Framework> call
+        return self._is_in_framework_module(node.func)
+
+    def _is_in_framework_module(self, node):
+        """Checks if the node is an ``ast.Attribute`` that represents a
+        ``sagemaker.<framework>`` module.
+        """
+        return (
+            isinstance(node.value, ast.Attribute)
+            and node.value.attr in FRAMEWORK_MODULES
+            and isinstance(node.value.value, ast.Name)
+            and node.value.value.id == "sagemaker"
+        )
 
     def _fw_version_in_keywords(self, node):
         """Checks if the ``ast.Call`` node's keywords contain ``framework_version``."""
diff --git a/src/sagemaker/cli/compatibility/v2/modifiers/tf_legacy_mode.py b/src/sagemaker/cli/compatibility/v2/modifiers/tf_legacy_mode.py
@@ -54,6 +54,7 @@ def node_should_be_modified(self, node):
 
         - ``TensorFlow``
         - ``sagemaker.tensorflow.TensorFlow``
+        - ``sagemaker.tensorflow.estimator.TensorFlow``
 
         Legacy mode is enabled if (1) ``script_mode`` is ``False``, ``None``, or not specified,
         and (2) if ``py_version`` is ``py2`` or not specified.
@@ -68,27 +69,35 @@ def node_should_be_modified(self, node):
         return self._is_tf_constructor(node) and self._is_legacy_mode(node)
 
     def _is_tf_constructor(self, node):
-        """Checks if the ``ast.Call`` node represents a call of the form
-        ``TensorFlow`` or ``sagemaker.tensorflow.TensorFlow``.
+        """Checks if the ``ast.Call`` node represents a call of the form ``TensorFlow``,
+        ``sagemaker.tensorflow.TensorFlow``, or ``sagemaker.tensorflow.estimator.TensorFlow``.
         """
         # Check for TensorFlow()
         if isinstance(node.func, ast.Name):
             return node.func.id == "TensorFlow"
 
+        # Check for something.that.ends.with.TensorFlow()
+        if not (isinstance(node.func, ast.Attribute) and node.func.attr == "TensorFlow"):
+            return False
+
+        # Check for sagemaker.tensorflow.estimator.TensorFlow()
+        if isinstance(node.func.value, ast.Attribute) and node.func.value.attr == "estimator":
+            return self._is_in_tensorflow_module(node.func.value)
+
         # Check for sagemaker.tensorflow.TensorFlow()
-        ends_with_tensorflow_constructor = (
-            isinstance(node.func, ast.Attribute) and node.func.attr == "TensorFlow"
-        )
+        return self._is_in_tensorflow_module(node.func)
 
-        is_in_tensorflow_module = (
-            isinstance(node.func.value, ast.Attribute)
-            and node.func.value.attr == "tensorflow"
-            and isinstance(node.func.value.value, ast.Name)
-            and node.func.value.value.id == "sagemaker"
+    def _is_in_tensorflow_module(self, node):
+        """Checks if the node is an ``ast.Attribute`` that represents the
+        ``sagemaker.tensorflow`` module.
+        """
+        return (
+            isinstance(node.value, ast.Attribute)
+            and node.value.attr == "tensorflow"
+            and isinstance(node.value.value, ast.Name)
+            and node.value.value.id == "sagemaker"
         )
 
-        return ends_with_tensorflow_constructor and is_in_tensorflow_module
-
     def _is_legacy_mode(self, node):
         """Checks if the ``ast.Call`` node's keywords signal using legacy mode."""
         script_mode = False
diff --git a/tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_framework_version.py b/tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_framework_version.py
@@ -32,24 +32,34 @@ def test_node_should_be_modified_fw_constructor_no_fw_version():
     fw_constructors = (
         "TensorFlow()",
         "sagemaker.tensorflow.TensorFlow()",
+        "sagemaker.tensorflow.estimator.TensorFlow()",
         "TensorFlowModel()",
         "sagemaker.tensorflow.TensorFlowModel()",
+        "sagemaker.tensorflow.model.TensorFlowModel()",
         "MXNet()",
         "sagemaker.mxnet.MXNet()",
+        "sagemaker.mxnet.estimator.MXNet()",
         "MXNetModel()",
         "sagemaker.mxnet.MXNetModel()",
+        "sagemaker.mxnet.model.MXNetModel()",
         "Chainer()",
         "sagemaker.chainer.Chainer()",
+        "sagemaker.chainer.estimator.Chainer()",
         "ChainerModel()",
         "sagemaker.chainer.ChainerModel()",
+        "sagemaker.chainer.model.ChainerModel()",
         "PyTorch()",
         "sagemaker.pytorch.PyTorch()",
+        "sagemaker.pytorch.estimator.PyTorch()",
         "PyTorchModel()",
         "sagemaker.pytorch.PyTorchModel()",
+        "sagemaker.pytorch.model.PyTorchModel()",
         "SKLearn()",
         "sagemaker.sklearn.SKLearn()",
+        "sagemaker.sklearn.estimator.SKLearn()",
         "SKLearnModel()",
         "sagemaker.sklearn.SKLearnModel()",
+        "sagemaker.sklearn.model.SKLearnModel()",
     )
 
     modifier = framework_version.FrameworkVersionEnforcer()
@@ -63,24 +73,34 @@ def test_node_should_be_modified_fw_constructor_with_fw_version():
     fw_constructors = (
         "TensorFlow(framework_version='2.2')",
         "sagemaker.tensorflow.TensorFlow(framework_version='2.2')",
+        "sagemaker.tensorflow.estimator.TensorFlow(framework_version='2.2')",
         "TensorFlowModel(framework_version='1.10')",
         "sagemaker.tensorflow.TensorFlowModel(framework_version='1.10')",
+        "sagemaker.tensorflow.model.TensorFlowModel(framework_version='1.10')",
         "MXNet(framework_version='1.6')",
         "sagemaker.mxnet.MXNet(framework_version='1.6')",
+        "sagemaker.mxnet.estimator.MXNet(framework_version='1.6')",
         "MXNetModel(framework_version='1.6')",
         "sagemaker.mxnet.MXNetModel(framework_version='1.6')",
+        "sagemaker.mxnet.model.MXNetModel(framework_version='1.6')",
         "PyTorch(framework_version='1.4')",
         "sagemaker.pytorch.PyTorch(framework_version='1.4')",
+        "sagemaker.pytorch.estimator.PyTorch(framework_version='1.4')",
         "PyTorchModel(framework_version='1.4')",
         "sagemaker.pytorch.PyTorchModel(framework_version='1.4')",
+        "sagemaker.pytorch.model.PyTorchModel(framework_version='1.4')",
         "Chainer(framework_version='5.0')",
         "sagemaker.chainer.Chainer(framework_version='5.0')",
+        "sagemaker.chainer.estimator.Chainer(framework_version='5.0')",
         "ChainerModel(framework_version='5.0')",
         "sagemaker.chainer.ChainerModel(framework_version='5.0')",
+        "sagemaker.chainer.model.ChainerModel(framework_version='5.0')",
         "SKLearn(framework_version='0.20.0')",
         "sagemaker.sklearn.SKLearn(framework_version='0.20.0')",
+        "sagemaker.sklearn.estimator.SKLearn(framework_version='0.20.0')",
         "SKLearnModel(framework_version='0.20.0')",
         "sagemaker.sklearn.SKLearnModel(framework_version='0.20.0')",
+        "sagemaker.sklearn.model.SKLearnModel(framework_version='0.20.0')",
     )
 
     modifier = framework_version.FrameworkVersionEnforcer()
@@ -97,51 +117,36 @@ def test_node_should_be_modified_random_function_call():
 
 
 def test_modify_node_tf():
-    classes = (
-        "TensorFlow" "sagemaker.tensorflow.TensorFlow",
-        "TensorFlowModel",
-        "sagemaker.tensorflow.TensorFlowModel",
-    )
-    _test_modify_node(classes, "1.11.0")
+    _test_modify_node("TensorFlow", "1.11.0")
 
 
 def test_modify_node_mx():
-    classes = ("MXNet", "sagemaker.mxnet.MXNet", "MXNetModel", "sagemaker.mxnet.MXNetModel")
-    _test_modify_node(classes, "1.2.0")
+    _test_modify_node("MXNet", "1.2.0")
 
 
 def test_modify_node_chainer():
-    classes = (
-        "Chainer",
-        "sagemaker.chainer.Chainer",
-        "ChainerModel",
-        "sagemaker.chainer.ChainerModel",
-    )
-    _test_modify_node(classes, "4.1.0")
+    _test_modify_node("Chainer", "4.1.0")
 
 
 def test_modify_node_pt():
-    classes = (
-        "PyTorch",
-        "sagemaker.pytorch.PyTorch",
-        "PyTorchModel",
-        "sagemaker.pytorch.PyTorchModel",
-    )
-    _test_modify_node(classes, "0.4.0")
+    _test_modify_node("PyTorch", "0.4.0")
 
 
 def test_modify_node_sklearn():
-    classes = (
-        "SKLearn",
-        "sagemaker.sklearn.SKLearn",
-        "SKLearnModel",
-        "sagemaker.sklearn.SKLearnModel",
-    )
-    _test_modify_node(classes, "0.20.0")
+    _test_modify_node("SKLearn", "0.20.0")
 
 
-def _test_modify_node(classes, default_version):
+def _test_modify_node(framework, default_version):
     modifier = framework_version.FrameworkVersionEnforcer()
+
+    classes = (
+        "{}".format(framework),
+        "sagemaker.{}.{}".format(framework.lower(), framework),
+        "sagemaker.{}.estimator.{}".format(framework.lower(), framework),
+        "{}Model".format(framework),
+        "sagemaker.{}.{}Model".format(framework.lower(), framework),
+        "sagemaker.{}.model.{}Model".format(framework.lower(), framework),
+    )
     for cls in classes:
         node = ast_call("{}()".format(cls))
         modifier.modify_node(node)
diff --git a/tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_tf_legacy_mode.py b/tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_tf_legacy_mode.py
@@ -42,6 +42,10 @@ def test_node_should_be_modified_tf_constructor_legacy_mode():
         "sagemaker.tensorflow.TensorFlow(script_mode=None)",
         "sagemaker.tensorflow.TensorFlow(py_version='py2')",
         "sagemaker.tensorflow.TensorFlow()",
+        "sagemaker.tensorflow.estimator.TensorFlow(script_mode=False)",
+        "sagemaker.tensorflow.estimator.TensorFlow(script_mode=None)",
+        "sagemaker.tensorflow.estimator.TensorFlow(py_version='py2')",
+        "sagemaker.tensorflow.estimator.TensorFlow()",
     )
 
     modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
@@ -61,6 +65,10 @@ def test_node_should_be_modified_tf_constructor_script_mode():
         "sagemaker.tensorflow.TensorFlow(py_version='py3')",
         "sagemaker.tensorflow.TensorFlow(py_version='py37')",
         "sagemaker.tensorflow.TensorFlow(py_version='py3', script_mode=False)",
+        "sagemaker.tensorflow.estimator.TensorFlow(script_mode=True)",
+        "sagemaker.tensorflow.estimator.TensorFlow(py_version='py3')",
+        "sagemaker.tensorflow.estimator.TensorFlow(py_version='py37')",
+        "sagemaker.tensorflow.estimator.TensorFlow(py_version='py3', script_mode=False)",
     )
 
     modifier = tf_legacy_mode.TensorFlowLegacyModeConstructorUpgrader()
diff --git a/tox.ini b/tox.ini
@@ -63,7 +63,7 @@ passenv =
 # Can be used to specify which tests to run, e.g.: tox -- -s
 commands =
     coverage run --source sagemaker -m pytest {posargs}
-    {env:IGNORE_COVERAGE:} coverage report --fail-under=86 --omit */tensorflow/tensorflow_serving/*
+    {env:IGNORE_COVERAGE:} coverage report --fail-under=86
 extras = test
 
 [testenv:flake8]
