|
|
import traceback as tb |
|
|
from typing import Any |
|
|
|
|
|
|
|
|
WRAPPED_EXCEPTION = tuple[BaseException, tb.StackSummary] |
|
|
|
|
|
__all__ = ["CheckpointException"] |
|
|
|
|
|
|
|
|
def _wrap_exception(exc: BaseException) -> WRAPPED_EXCEPTION: |
|
|
return (exc, tb.extract_tb(exc.__traceback__)) |
|
|
|
|
|
|
|
|
def _is_wrapped_exception(obj: Any) -> bool: |
|
|
if not isinstance(obj, tuple): |
|
|
return False |
|
|
if len(obj) != 2: |
|
|
return False |
|
|
return isinstance(obj[0], BaseException) and isinstance(obj[1], tb.StackSummary) |
|
|
|
|
|
|
|
|
class CheckpointException(BaseException): |
|
|
"""Exception raised if failure was detected as part of a checkpoint load or save.""" |
|
|
|
|
|
def __init__(self, msg: str, failures: dict[int, WRAPPED_EXCEPTION]): |
|
|
super().__init__(msg, failures) |
|
|
self._failures = failures |
|
|
|
|
|
@property |
|
|
def failures(self) -> dict[int, WRAPPED_EXCEPTION]: |
|
|
"""Return a dictionary mapping node ranks to their associated exceptions in case of failure.""" |
|
|
return self._failures |
|
|
|
|
|
def __str__(self) -> str: |
|
|
str = f"CheckpointException ranks:{self._failures.keys()}\n" |
|
|
for rank, exc_pair in self._failures.items(): |
|
|
exc, trace = exc_pair |
|
|
str += f"Traceback (most recent call last): (RANK {rank})\n" |
|
|
if trace is not None: |
|
|
str += "".join(tb.format_list(trace)) |
|
|
str += "".join(tb.format_exception_only(type(exc), value=exc)) |
|
|
return str |
|
|
|
