Add SerDe compatability

Author: Balaji Veeramani

src/sagemaker/cli/compatibility/v2/ast_transformer.py

@@ -35,14 +35,24 @@
     modifiers.renamed_params.SessionCreateEndpointImageURIRenamer(),
     modifiers.training_params.TrainPrefixRemover(),
     modifiers.training_input.TrainingInputConstructorRefactor(),
+    modifiers.serde.SerdeConstructorRenamer(),
 ]
 
 IMPORT_MODIFIERS = [modifiers.tfs.TensorFlowServingImportRenamer()]
 
+NAME_MODIFIERS = [modifiers.serde.SerdeObjectRenamer()]
+
+MODULE_MODIFIERS = [
+    modifiers.serde.SerializerImportInserter(),
+    modifiers.serde.DeserializerImportInserter(),
+]
+
 IMPORT_FROM_MODIFIERS = [
     modifiers.predictors.PredictorImportFromRenamer(),
     modifiers.tfs.TensorFlowServingImportFromRenamer(),
     modifiers.training_input.TrainingInputImportFromRenamer(),
+    modifiers.serde.SerdeImportFromAmazonCommonRenamer(),
+    modifiers.serde.SerdeImportFromPredictorRenamer(),
 ]
 
 
@@ -52,52 +62,88 @@ class ASTTransformer(ast.NodeTransformer):
     """
 
     def visit_Call(self, node):
-        """Visits an ``ast.Call`` node and returns a modified node, if needed.
+        """Visits an ``ast.Call`` node and returns a modified node or None.
         See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
 
         Args:
             node (ast.Call): a node that represents a function call.
 
         Returns:
-            ast.Call: a node that represents a function call, which has
-                potentially been modified from the original input.
+            ast.AST: if the returned node is None, the original node is removed
+            from its location. Otherwise, the original node is replaced with the
+            returned node.
         """
         for function_checker in FUNCTION_CALL_MODIFIERS:
-            function_checker.check_and_modify_node(node)
+            node = function_checker.check_and_modify_node(node)
+        return ast.fix_missing_locations(node) if node else None
+
+    def visit_Name(self, node):
+        """Visits an ``ast.Name`` node and returns a modified node or None.
+        See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
 
-        ast.fix_missing_locations(node)
+        Args:
+            node (ast.Name): a node that represents an identifier.
+
+        Returns:
+            ast.AST: if the returned node is None, the original node is removed
+            from its location. Otherwise, the original node is replaced with the
+            returned node.
+        """
+        for name_checker in NAME_MODIFIERS:
+            node = name_checker.check_and_modify_node(node)
+            if node is None:
+                return None
+            node = ast.fix_missing_locations(node)
         return node
 
     def visit_Import(self, node):
-        """Visits an ``ast.Import`` node and returns a modified node, if needed.
+        """Visits an ``ast.Import`` node and returns a modified node or None.
         See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
 
         Args:
             node (ast.Import): a node that represents an import statement.
 
         Returns:
-            ast.Import: a node that represents an import statement, which has
-                potentially been modified from the original input.
+            ast.AST: if the returned node is None, the original node is removed
+            from its location. Otherwise, the original node is replaced with the
+            returned node.
         """
         for import_checker in IMPORT_MODIFIERS:
-            import_checker.check_and_modify_node(node)
+            node = import_checker.check_and_modify_node(node)
+        return ast.fix_missing_locations(node) if node else None
 
-        ast.fix_missing_locations(node)
-        return node
+    def visit_Module(self, node):
+        """Visits an ``ast.Module`` node and returns a modified node or None.
+        See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
+
+        The ``ast.NodeTransformer`` walks the abstract syntax tree and modifies
+        all other nodes before modifying the ``ast.Module`` node.
+
+        Args:
+            node (ast.Module): a node that represents a Python module.
+
+        Returns:
+            ast.AST: if the returned node is None, the original node is removed
+            from its location. Otherwise, the original node is replaced with the
+            returned node.
+        """
+        self.generic_visit(node)
+        for module_checker in MODULE_MODIFIERS:
+            node = module_checker.check_and_modify_node(node)
+        return ast.fix_missing_locations(node) if node else None
 
     def visit_ImportFrom(self, node):
-        """Visits an ``ast.ImportFrom`` node and returns a modified node, if needed.
+        """Visits an ``ast.ImportFrom`` node and returns a modified node or None.
         See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
 
         Args:
             node (ast.ImportFrom): a node that represents an import statement.
 
         Returns:
-            ast.ImportFrom: a node that represents an import statement, which has
-                potentially been modified from the original input.
+            ast.AST: if the returned node is None, the original node is removed
+            from its location. Otherwise, the original node is replaced with the
+            returned node.
         """
         for import_checker in IMPORT_FROM_MODIFIERS:
-            import_checker.check_and_modify_node(node)
-
-        ast.fix_missing_locations(node)
-        return node
+            node = import_checker.check_and_modify_node(node)
+        return ast.fix_missing_locations(node) if node else None

src/sagemaker/cli/compatibility/v2/modifiers/__init__.py

@@ -19,6 +19,7 @@
     framework_version,
     predictors,
     renamed_params,
+    serde,
     tf_legacy_mode,
     tfs,
     training_params,

src/sagemaker/cli/compatibility/v2/modifiers/airflow.py

@@ -58,9 +58,13 @@ def modify_node(self, node):
         Args:
             node (ast.Call): a node that represents either a ``model_config`` call or
                 a ``model_config_from_estimator`` call.
+
+        Returns:
+            ast.AST: the original node, which has been potentially modified.
         """
         instance_type = node.args.pop(0)
         node.keywords.append(ast.keyword(arg="instance_type", value=instance_type))
+        return node
 
 
 class ModelConfigImageURIRenamer(renamed_params.ParamRenamer):

src/sagemaker/cli/compatibility/v2/modifiers/deprecated_params.py

@@ -54,7 +54,11 @@ def modify_node(self, node):
 
         Args:
             node (ast.Call): a node that represents a TensorFlow constructor.
+
+        Returns:
+            ast.AST: the original node, which has been potentially modified.
         """
         for kw in node.keywords:
             if kw.arg == "script_mode":
                 node.keywords.remove(kw)
+        return node

src/sagemaker/cli/compatibility/v2/modifiers/framework_version.py

@@ -94,6 +94,9 @@ def modify_node(self, node):
 
         Args:
             node (ast.Call): a node that represents the constructor of a framework class.
+
+        Returns:
+            ast.AST: the original node, which has been potentially modified.
         """
         framework, is_model = _framework_from_node(node)
 
@@ -109,6 +112,7 @@ def modify_node(self, node):
             py_version = _py_version_defaults(framework, framework_version, is_model)
             if py_version:
                 node.keywords.append(ast.keyword(arg=PY_ARG, value=ast.Str(s=py_version)))
+        return node
 
 
 def _py_version_defaults(framework, framework_version, is_model=False):

src/sagemaker/cli/compatibility/v2/modifiers/modifier.py

@@ -22,9 +22,10 @@ class Modifier(object):
     """
 
     def check_and_modify_node(self, node):
-        """Check an AST node, and modify it if applicable."""
+        """Check an AST node, and modify, replace, or remove it if applicable."""
         if self.node_should_be_modified(node):
-            self.modify_node(node)
+            node = self.modify_node(node)
+        return node
 
     @abstractmethod
     def node_should_be_modified(self, node):

src/sagemaker/cli/compatibility/v2/modifiers/predictors.py

@@ -62,9 +62,13 @@ def modify_node(self, node):
 
         Args:
             node (ast.Call): a node that represents a *Predictor constructor.
+
+        Returns:
+            ast.AST: the original node, which has been potentially modified.
         """
         _rename_class(node)
         _rename_endpoint(node)
+        return node
 
 
 def _rename_class(node):
@@ -106,7 +110,11 @@ def modify_node(self, node):
         Args:
             node (ast.ImportFrom): a node that represents a ``from ... import ... `` statement.
                 For more, see https://docs.python.org/3/library/ast.html#abstract-grammar.
+
+        Returns:
+            ast.AST: the original node, which has been potentially modified.
         """
         for name in node.names:
             if name.name == BASE_PREDICTOR:
                 name.name = "Predictor"
+        return node

src/sagemaker/cli/compatibility/v2/modifiers/renamed_params.py

@@ -61,9 +61,13 @@ def modify_node(self, node):
 
         Args:
             node (ast.Call): a node that represents the relevant function call.
+
+        Returns:
+            ast.AST: the original node, which has been potentially modified.
         """
         keyword = parsing.arg_from_keywords(node, self.old_param_name)
         keyword.arg = self.new_param_name
+        return node
 
 
 class MethodParamRenamer(ParamRenamer):