chore: import upstream snapshot with attribution

python/paddle/distributed/flex_checkpoint/aoa/traceback.py

@@ -0,0 +1,133 @@
+# 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.
+
+from __future__ import annotations
+
+
+class AOATraceback:
+    """
+    When error occurs, print the chain of "original aoa_statement -> ... -> current expression".
+    """
+
+    def __init__(self) -> None:
+        self.records: list[dict] = []
+        self.last_error_chain: list[str] = []
+        self.last_error_message: str = ""
+        self.last_error_stage: str = ""
+        self.last_error_type: str = ""
+        self.parent_map: dict[str, str | None] = {}
+        self.child_macro_map: dict[str, str] = {}
+
+    def register_roots(self, expressions: list[str]) -> None:
+        """Register the original aoa_statements as the root nodes of the chain."""
+        for expr in expressions:
+            self.parent_map.setdefault(expr, None)
+
+    def record_children(
+        self, parent: str, children: list[str], macro_name: str | None = None
+    ) -> None:
+        """Record the children expressions obtained by the parent expression, and mark the macro name used."""
+        macro = macro_name or "Expanded"
+        for child in children:
+            if child == parent:
+                continue
+            self.parent_map[child] = parent
+            self.child_macro_map[child] = macro
+
+    def build_chain(self, expr: str) -> list[str]:
+        """Build the chain from the root to expr by tracing back from the current expression."""
+        chain: list[str] = []
+        visited = set()
+        cur = expr
+        while cur is not None and cur not in visited:
+            chain.append(cur)
+            visited.add(cur)
+            cur = self.parent_map.get(cur)
+        chain.reverse()
+        return chain
+
+    def add_error(
+        self,
+        error_message: str,
+        stage: str,
+        chain: list[str],
+        error_type: str = "",
+    ) -> None:
+        """Record the error chain and information."""
+        self.last_error_chain = chain
+        self.last_error_message = error_message
+        self.last_error_stage = stage
+        self.last_error_type = error_type or ""
+        self.records.append(
+            {
+                "type": "error",
+                "stage": stage,
+                "message": error_message,
+                "error_type": self.last_error_type,
+                "chain": chain,
+            }
+        )
+
+    def format_traceback(self) -> str:
+        lines: list[str] = []
+        header_text = " AOA Traceback (related chain) "
+        header = f"===={header_text}===="
+        footer = "=" * len(header)
+
+        if self.last_error_chain:
+            lines.append(header)
+            indent_unit = "    "
+
+            lines.append("| Origin AOA Statement")
+            origin_expr = self.last_error_chain[0].replace("\n", " ")
+            lines.append(f"|-> {origin_expr}")
+
+            for level, expr in enumerate(self.last_error_chain[1:], start=1):
+                indent = indent_unit * level
+                single_line_expr = expr.replace("\n", " ")
+                macro = self.child_macro_map.get(
+                    expr, self.last_error_stage or "Expanded"
+                )
+                lines.append(f"{indent}| {macro}")
+                lines.append(f"{indent}|-> {single_line_expr}")
+
+            if self.last_error_message:
+                err_title = self.last_error_type or "Error"
+                stage_str = (
+                    f" [{self.last_error_stage}]"
+                    if self.last_error_stage
+                    else ""
+                )
+                err_level = len(self.last_error_chain)
+                indent = indent_unit * err_level
+                single_line_msg = self.last_error_message.replace("\n", " ")
+                lines.append(f"{indent}| Error")
+                lines.append(
+                    f"{indent}|-> ({err_title}{stage_str}) {single_line_msg}"
+                )
+
+            lines.append(footer)
+        else:
+            lines.append(header)
+            lines.append("(No trace records)")
+            lines.append(footer)
+
+        return "\n".join(lines)
+
+    def print(self, logger=None) -> None:
+        text = self.format_traceback()
+        if logger:
+            logger.error(text)
+        else:
+            print(text)