[Pass] Remove metadata_props and doc_string from the model (#2182)

Fix #2163

---------

Co-authored-by: Justin Chu <[email protected]>

Author: titaiwangms

onnxscript/ir/_core.py

@@ -2583,16 +2583,14 @@ class Function(_protocols.FunctionProtocol, Sequence[Node], _display.PrettyPrint
         outputs: The output values of the function.
         opset_imports: Opsets imported by the function.
         doc_string: Documentation string.
-        metadata_props: Metadata that will be serialized to the ONNX file.
         meta: Metadata store for graph transform passes.
+        metadata_props: Metadata that will be serialized to the ONNX file.
     """
 
     __slots__ = (
         "_attributes",
         "_domain",
         "_graph",
-        "_metadata",
-        "_metadata_props",
         "_name",
         "_overload",
     )
@@ -2607,15 +2605,12 @@ def __init__(
         # and not from an outer scope
         graph: Graph,
         attributes: Sequence[Attr],
-        metadata_props: dict[str, str] | None = None,
     ) -> None:
         self._domain = domain
         self._name = name
         self._overload = overload
         self._graph = graph
         self._attributes = OrderedDict((attr.name, attr) for attr in attributes)
-        self._metadata: _metadata.MetadataStore | None = None
-        self._metadata_props: dict[str, str] | None = metadata_props
 
     def identifier(self) -> _protocols.OperatorIdentifier:
         return self.domain, self.name, self.overload
@@ -2687,15 +2682,11 @@ def meta(self) -> _metadata.MetadataStore:
         Write to the :attr:`metadata_props` if you would like the metadata to be serialized
         to the ONNX proto.
         """
-        if self._metadata is None:
-            self._metadata = _metadata.MetadataStore()
-        return self._metadata
+        return self._graph.meta
 
     @property
     def metadata_props(self) -> dict[str, str]:
-        if self._metadata_props is None:
-            self._metadata_props = {}
-        return self._metadata_props
+        return self._graph.metadata_props
 
     # Mutation methods
     def append(self, node: Node, /) -> None:

onnxscript/ir/passes/common/clear_metadata_and_docstring.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Clear all metadata and docstring from the model, graphs, nodes, and functions."""
+
+from __future__ import annotations
+
+__all__ = [
+    "ClearMetadataAndDocStringPass",
+]
+
+import logging
+
+from onnxscript import ir
+
+logger = logging.getLogger(__name__)
+
+
+class ClearMetadataAndDocStringPass(ir.passes.InPlacePass):
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        # 0. TODO: Should we clean model metadata and docstring?
+
+        # 1. Clean up the graph and the belonged nodes metadata properties
+        modified = self._clear_graph_or_function_metadata_and_docstring(model.graph)
+
+        # 2. Clean up all of the functions metadata properties
+        for function in model.functions.values():
+            modified = (
+                self._clear_graph_or_function_metadata_and_docstring(function) or modified
+            )
+        return ir.passes.PassResult(model, modified=modified)
+
+    def _clear_graph_or_function_metadata_and_docstring(
+        self,
+        graph_or_function: ir.Graph | ir.Function,
+    ) -> bool:
+        """Clear metadata and docstring from the graph or function."""
+        checked_graphs_or_functions: set[ir.Graph | ir.Function] = set()
+        modified = False
+        # Clean up all of the nodes metadata properties
+        for node in ir.traversal.RecursiveGraphIterator(graph_or_function):
+            if node.metadata_props:
+                modified = True
+                logger.debug("Removed metadata from %s nodes", node.name)
+            node.metadata_props.clear()
+            node.doc_string = None
+
+            # Clean up the owning graph/function metadata properties
+            # and doc_string if the graph/function is not already checked
+            assert node.graph is not None
+            if node.graph not in checked_graphs_or_functions and (
+                node.graph.metadata_props or node.graph.doc_string
+            ):
+                modified = True
+                logger.debug("Removed metadata from %s graph/function", node.graph.name)
+                node.graph.metadata_props.clear()
+                node.graph.doc_string = None
+                checked_graphs_or_functions.add(node.graph)
+        return modified

onnxscript/ir/passes/common/clear_metadata_and_docstring_test.py

@@ -0,0 +1,95 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+
+from onnxscript import ir
+from onnxscript.ir.passes.common import clear_metadata_and_docstring
+
+
+class TestClearMetadataAndDocStringPass(unittest.TestCase):
+    def test_pass_with_clear_metadata_and_docstring(self):
+        # Create a model (node, graph, function) with metadata and docstring
+        inputs = [
+            ir.Value(
+                name="input_a", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3))
+            ),
+            ir.Value(
+                name="input_b", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3))
+            ),
+        ]
+        add_node = ir.node(
+            "Add",
+            inputs=inputs,
+            num_outputs=1,
+            metadata_props={"add_key": "add_value"},
+            doc_string="This is an Add node",
+        )
+        mul_node = ir.node(
+            "Mul",
+            inputs=[add_node.outputs[0], inputs[1]],
+            num_outputs=1,
+            metadata_props={"mul_key": "mul_value"},
+            doc_string="This is a Mul node",
+        )
+        function = ir.Function(
+            graph=ir.Graph(
+                name="my_function",
+                inputs=inputs,
+                outputs=mul_node.outputs,
+                nodes=[add_node, mul_node],
+                opset_imports={"": 20},
+                doc_string="This is a function docstring",
+                metadata_props={"function_key": "function_value"},
+            ),
+            name="my_function",
+            domain="my_domain",
+            attributes=[],
+        )
+        # Create a model with the graph and function
+        constant_tensor = ir.tensor(np.random.rand(2, 3).astype(ir.DataType.FLOAT.numpy()))
+        const_node = ir.node(
+            "Constant",
+            inputs=[],
+            attributes={"value": constant_tensor},
+            num_outputs=1,
+            metadata_props={"const_key": "const_value"},
+            doc_string="This is a Constant node",
+        )
+        sub_node = ir.node(
+            "Sub",
+            inputs=[function.outputs[0], const_node.outputs[0]],
+            num_outputs=1,
+            metadata_props={"sub_key": "sub_value"},
+            doc_string="This is a Sub node",
+        )
+        model = ir.Model(
+            graph=ir.Graph(
+                inputs=inputs,
+                outputs=sub_node.outputs,
+                nodes=[const_node, sub_node],
+                opset_imports={"": 20},
+                doc_string="This is a graph docstring",
+                metadata_props={"graph_key": "graph_value"},
+            ),
+            ir_version=10,
+            functions=[function],
+        )
+        # Create a pass to clear metadata and docstring
+        clear_pass = clear_metadata_and_docstring.ClearMetadataAndDocStringPass()
+        # Apply the pass
+        result = clear_pass(model)
+        # Check that the pass was applied
+        self.assertTrue(result.modified)
+        # Check that the metadata and docstring were cleared
+        self.assertEqual(model.graph.doc_string, None)
+        self.assertEqual(model.graph.metadata_props, {})
+        for node in model.graph:
+            self.assertEqual(node.metadata_props, {})
+            self.assertEqual(node.doc_string, None)
+        # Check that the function docstring and metadata were cleared
+        self.assertEqual(function.doc_string, None)
+        self.assertEqual(function.metadata_props, {})

onnxscript/ir/serde.py

@@ -699,6 +699,7 @@ def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
             if hasattr(proto, "overload") and proto.overload
             else ""
         ),
+        metadata_props=deserialize_metadata_props(proto.metadata_props),
     )
     attributes = [_deserialize_attribute(attr, []) for attr in proto.attribute_proto]
     # Attributes without defaults
@@ -711,7 +712,6 @@ def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
         overload=getattr(proto, "overload", ""),
         graph=graph,
         attributes=typing.cast(List[_core.Attr], attributes),
-        metadata_props=deserialize_metadata_props(proto.metadata_props),
     )