Files
paddlepaddle--paddle/python/paddle/distributed/flex_checkpoint/aoa/traceback.py
T
2026-07-13 12:40:42 +08:00

134 lines
4.6 KiB
Python

# 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)