[SOT] Add JSON dumping and Base64 encoding for log output (#71525)

Author: DrRyanHuang

python/paddle/jit/sot/symbolic/compile_cache.py

@@ -223,7 +223,7 @@ def collect_subgraph_info(self, program: Program):
 
         InfoCollector().attach(
             SubGraphInfo,
-            program,
+            str(program),
             self.graph_size(),
             self.SIR.name,
         )

python/paddle/jit/sot/utils/envs.py

@@ -141,6 +141,7 @@ def parse_parameterized_key(input_str: str) -> dict[str, list[str]]:
     "SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE", False
 )
 ENV_SOT_COLLECT_INFO = PEP508LikeEnvironmentVariable("SOT_COLLECT_INFO", {})
+ENV_SOT_SERIALIZE_INFO = BooleanEnvironmentVariable("SOT_SERIALIZE_INFO", False)
 ENV_SOT_FORCE_FALLBACK_SIR_IDS = StringEnvironmentVariable(
     "SOT_FORCE_FALLBACK_SIR_IDS", ""
 )

python/paddle/jit/sot/utils/info_collector.py

@@ -15,22 +15,28 @@
 from __future__ import annotations
 
 import atexit
+import base64
+import json
 import sys
 from abc import ABC, abstractmethod
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, ClassVar, NamedTuple
+from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple
 
 from typing_extensions import Self
 
-from .envs import ENV_SOT_COLLECT_INFO
+from .envs import ENV_SOT_COLLECT_INFO, ENV_SOT_SERIALIZE_INFO
 from .utils import Singleton
 
 if TYPE_CHECKING:
     import types
 
     from .exceptions import BreakGraphReasonBase
 
+PREFIX = "<sot>"
+SUFFIX = "</sot>"
+ENCODING = "utf-8"
+
 
 def try_import_graphviz():
     try:
@@ -102,7 +108,10 @@ def generate_report(self, info_dict: dict[str, list[InfoBase]]) -> str:
         for info_class_name, info_list in info_dict.items():
             cls = info_list[0].__class__
             report += f"{info_class_name} ({cls.SHORT_NAME}):\n"
-            report += cls.summary(info_list)
+            if ENV_SOT_SERIALIZE_INFO.get():
+                report += cls.json_report(info_list)
+            else:
+                report += cls.summary(info_list)
             report += "\n"
         return report
 
@@ -120,6 +129,22 @@ def __init__(self): ...
     @abstractmethod
     def summary(cls, history: list[Self]) -> str: ...
 
+    @classmethod
+    def serialize(cls, obj: dict[str:Any]) -> str:
+
+        json_data = json.dumps(obj)
+        b64_bytes = base64.b64encode(json_data.encode(ENCODING))
+
+        return b64_bytes.decode(ENCODING)
+
+    @classmethod
+    def deserialize(cls, data: bytes | str) -> dict:
+        if isinstance(data, str):
+            data = data.encode(ENCODING)
+        json_str = base64.b64decode(data).decode(ENCODING)
+
+        return json.loads(json_str)
+
 
 class NewSymbolHitRateInfo(InfoBase):
     SHORT_NAME = "new_symbol_hit_rate"
@@ -154,6 +179,11 @@ def summary(cls, history: list[Self]) -> str:
         summary += f"Hit rate: {hit_count / all_count:.2f}"
         return summary
 
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+        # TODO: need to support serialize the output
+        return cls.summary(history)
+
 
 class SubGraphRelationInfo(InfoBase):
     SHORT_NAME = "subgraph_relation"
@@ -241,6 +271,11 @@ def to_tensor_node_name(
         dot.render(directory / filename, format="svg", cleanup=True)
         return f"Please check {directory / filename}.svg for subgraph relation"
 
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+        # TODO: need to support serialize the output
+        return cls.summary(history)
+
 
 class CompileCountInfo(InfoBase):
     SHORT_NAME = "compile_count"
@@ -268,6 +303,11 @@ def summary(cls, history: list[Self]) -> str:
         summary = "\n".join(summary_lines)
         return summary
 
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+        # TODO: need to support serialize the output
+        return cls.summary(history)
+
 
 class BreakGraphReasonInfo(InfoBase):
     SHORT_NAME = "breakgraph_reason"
@@ -278,17 +318,24 @@ def __init__(self, reason: BreakGraphReasonBase):
         self.reason = reason
 
     @classmethod
-    def summary(cls, history: list[Self]) -> str:
-        reason_dict = {}
+    def classify(cls, history: list[Self]) -> str:
+        reasons_dict = {}
 
         for info in history:
             name = info.reason.__class__.__name__
-            if name not in reason_dict:
-                reason_dict[name] = []
-            reason_dict[name].append(str(info.reason))
+            if name not in reasons_dict:
+                reasons_dict[name] = []
+            reasons_dict[name].append(str(info.reason))
+
+        sorted_reasons = list(reasons_dict.items())
+        sorted_reasons.sort(key=lambda x: len(x[1]), reverse=True)
 
-        reason_list = list(reason_dict.items())
-        reason_list.sort(key=lambda x: len(x[1]), reverse=True)
+        return reasons_dict, sorted_reasons
+
+    @classmethod
+    def summary(cls, history: list[Self]) -> str:
+
+        reason_dict, reason_list = cls.classify(history)
 
         return "\n".join(
             [
@@ -297,6 +344,33 @@ def summary(cls, history: list[Self]) -> str:
             ]
         )
 
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+
+        reason_dict, sorted_reasons = cls.classify(history)
+        reason_dict["count"] = {k: len(v) for k, v in sorted_reasons}
+        serialized = cls.serialize({cls.SHORT_NAME: reason_dict})
+
+        return f"{PREFIX}{serialized}{SUFFIX}"
+
+    @classmethod
+    def restore_from_string(cls, serialized: str) -> list[Self]:
+        # This method is the inverse of json_report
+
+        from paddle.jit.sot.utils import exceptions
+
+        history = []
+        obj = cls.deserialize(serialized)[cls.SHORT_NAME]
+        obj.pop("count")
+
+        for classname in obj:
+
+            ReasonClass = getattr(exceptions, classname, None)
+            for reason in obj[classname]:
+                history.append(cls(ReasonClass(reason_str=reason)))
+
+        return history
+
     @staticmethod
     def collect_break_graph_reason(reason: BreakGraphReasonBase):
         if not InfoCollector().need_collect(BreakGraphReasonInfo):
@@ -309,7 +383,8 @@ class SubGraphInfo(InfoBase):
     SHORT_NAME = "subgraph_info"
     TYPE = InfoType.STEP_INFO
 
-    def __init__(self, graph, op_num, sir_name):
+    def __init__(self, graph: str, op_num: int, sir_name: str):
+        # NOTE: All data should be serializable
         super().__init__()
         self.graph = graph
         self.op_num = op_num
@@ -320,11 +395,12 @@ def __str__(self):
 
     @classmethod
     def summary(cls, history: list[Self]) -> str:
-
         num_of_subgraph = len(history)
         sum_of_op_num = sum(item.op_num for item in history)
 
-        need_details = "details" in ENV_SOT_COLLECT_INFO.get()[cls.SHORT_NAME]
+        need_details = "details" in ENV_SOT_COLLECT_INFO.get().get(
+            cls.SHORT_NAME, []
+        )
 
         details = ""
         if need_details:
@@ -338,3 +414,59 @@ def summary(cls, history: list[Self]) -> str:
         summary = f"[Number of subgraph]: {num_of_subgraph} [Sum of opnum]: {sum_of_op_num}"
 
         return f"{summary}\n{details}"
+
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+        need_details = "details" in ENV_SOT_COLLECT_INFO.get().get(
+            cls.SHORT_NAME, []
+        )
+
+        aggregated_info_list = []
+        for idx, record in enumerate(history):
+            entry_data = {}
+
+            entry_data["SIR_name"] = record.sir_name
+            entry_data["OpNum"] = record.op_num
+            entry_data["Graph"] = ""
+            if need_details:
+                entry_data["Graph"] = str(record.graph)
+            aggregated_info_list.append(entry_data)
+
+        serialized = cls.serialize({cls.SHORT_NAME: aggregated_info_list})
+
+        return f"{PREFIX}{serialized}{SUFFIX}"
+
+    @classmethod
+    def restore_from_string(cls, serialized: str) -> list[Self]:
+        # This method is the inverse of json_report
+
+        history = []
+        obj = cls.deserialize(serialized)[cls.SHORT_NAME]
+
+        for entry in obj:
+
+            history.append(
+                SubGraphInfo(
+                    graph=entry["Graph"],
+                    op_num=entry["OpNum"],
+                    sir_name=entry["SIR_name"],
+                )
+            )
+
+        return history
+
+    def __eq__(self, other):
+
+        need_graph_equal = "details" in ENV_SOT_COLLECT_INFO.get().get(
+            self.SHORT_NAME, []
+        )
+
+        graph_equal_or_not = True
+        if need_graph_equal:
+            graph_equal_or_not = self.graph == other.graph
+
+        return (
+            graph_equal_or_not
+            and self.op_num == other.op_num
+            and self.sir_name == other.sir_name
+        )

test/sot/test_info_collect.py

@@ -0,0 +1,121 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import random
+import string
+import unittest
+
+from test_case_base import TestCaseBase
+
+from paddle.jit.sot.utils.exceptions import (
+    BuiltinFunctionBreak,
+    DataDependencyControlFlowBreak,
+    DataDependencyDynamicShapeBreak,
+    DataDependencyOperationBreak,
+    DygraphInconsistentWithStaticBreak,
+    FallbackInlineCallBreak,
+    InferMetaBreak,
+    InlineCallBreak,
+    OtherInlineCallBreak,
+    PsdbBreakReason,
+    SideEffectBreak,
+    UnsupportedIteratorBreak,
+    UnsupportedPaddleAPIBreak,
+)
+from paddle.jit.sot.utils.info_collector import (
+    BreakGraphReasonInfo,
+    InfoBase,
+    SubGraphInfo,
+)
+
+generate_random_string = lambda N: ''.join(
+    random.choices(string.ascii_uppercase + string.digits, k=N)
+)
+
+
+class TestSerialize(TestCaseBase):
+    def test_case(self):
+        x_dict = {
+            'a': 'b',
+            'c': 'd',
+        }
+
+        x_str = InfoBase.serialize(x_dict)
+        y_dict = InfoBase.deserialize(x_str)
+
+        self.assertEqual(x_dict, y_dict)
+
+
+class TestBreakGraphReasonInfo(TestCaseBase):
+    def test_case(self):
+        history = [
+            BreakGraphReasonInfo(
+                BreakReasonClass(
+                    reason_str=generate_random_string(random.randint(1, 5))
+                )
+            )
+            for BreakReasonClass in [
+                FallbackInlineCallBreak,
+                DataDependencyControlFlowBreak,
+                DataDependencyDynamicShapeBreak,
+                DataDependencyOperationBreak,
+                UnsupportedPaddleAPIBreak,
+                BuiltinFunctionBreak,
+                SideEffectBreak,
+                UnsupportedIteratorBreak,
+                InlineCallBreak,
+                OtherInlineCallBreak,
+                DygraphInconsistentWithStaticBreak,
+                PsdbBreakReason,
+                InferMetaBreak,
+            ]
+        ]
+
+        serialized = BreakGraphReasonInfo.json_report(history)
+        deserialized = BreakGraphReasonInfo.restore_from_string(
+            serialized[5:-6]  # remove `<sot>` and `</sot>`
+        )  # `removeprefix` & `removesuffix` are only available from python3.9
+
+        origin_reasons_dict, _ = BreakGraphReasonInfo.classify(history)
+        origin_reasons2count = {
+            k: len(v) for k, v in origin_reasons_dict.items()
+        }
+
+        new_reasons_dict, _ = BreakGraphReasonInfo.classify(deserialized)
+        new_reasons2count = {k: len(v) for k, v in new_reasons_dict.items()}
+
+        self.assertEqual(origin_reasons2count, new_reasons2count)
+
+
+class TestSubGraphInfo(TestCaseBase):
+    def test_case(self):
+        history = [
+            SubGraphInfo(
+                generate_random_string(random.randint(1, 5)),
+                random.randint(0, 20),
+                generate_random_string(random.randint(1, 5)),
+            ),
+        ] * 10
+
+        serialized = SubGraphInfo.json_report(history)
+        deserialized = SubGraphInfo.restore_from_string(
+            serialized[5:-6]  # remove `<sot>` and `</sot>`
+        )  # `removeprefix` & `removesuffix` are only available from python3.9
+
+        self.assertEqual(history, deserialized)
+
+
+if __name__ == "__main__":
+    unittest.main()