import json
import os
import platform
import sys
import tarfile
import tempfile
from dataclasses import asdict, dataclass
from pathlib import Path
from traceback import print_exc
from types import TracebackType
from typing import List, Optional, Type
import click
from latch_cli.constants import latch_constants
from latch_cli.exceptions.traceback import _Traceback
from latch_cli.utils import get_local_package_version
@dataclass(frozen=True)
class _Metadata:
os_name: str = os.name
os_version: str = platform.version()
latch_version: str = get_local_package_version()
py_version: str = sys.version
platform: str = platform.platform()
def print(self):
click.secho("Crash info:", fg="red", bold=True)
click.echo(
" ".join(
[
click.style("Latch SDK version:", fg="red"),
self.latch_version,
]
)
)
click.echo(
" ".join(
[
click.style("Python version:", fg="red"),
self.py_version.replace("\n", ";"),
]
)
)
click.echo(" ".join([click.style("Platform:", fg="red"), self.platform]))
click.echo(
" ".join(
[
click.style("OS:", fg="red"),
f"{self.os_name}; {self.os_version}",
]
)
)
[docs]class CrashHandler:
"""Display and store useful information when the program fails
* Display tracebacks
* Parse and display opaque flytekit serialization error messages
* Write necessary information to reproduce failure to a tarball
"""
def __init__(self):
self.metadata: _Metadata = _Metadata()
self.message: Optional[str] = None
self.pkg_root: Optional[str] = None
def _write_state_to_tarball(self):
"""Bundles files needed to reproduce failed state into a tarball.
Tarball contains:
* JSON holding platform + package version metadata
* logs directory holding docker build logs
* Text file holding traceback
* (If register) workflow package (python files, Dockerfile, etc.)
"""
tarball_path = Path(".latch_report.tar.gz").resolve()
tarball_path.unlink(missing_ok=True)
with tarfile.open(tarball_path, mode="x:gz") as tf:
# If calling stack frame is handling an exception, we want to store
# the traceback in a log file.
if sys.exc_info()[0] is not None:
with tempfile.NamedTemporaryFile("w+") as ntf:
print_exc(file=ntf)
ntf.seek(0)
tf.add(ntf.name, arcname="traceback.txt")
if self.pkg_root is not None:
logs_path = Path(self.pkg_root) / ".latch" / ".logs"
if logs_path.exists():
tf.add(logs_path, arcname="logs")
pkg_files: List[Path] = []
for dp, _, fnames in os.walk(logs_path):
for f in fnames:
p = (Path(dp) / f).resolve()
if (
latch_constants.ignore_regex.search(str(p))
or p.is_symlink()
or p.stat().st_size > latch_constants.file_max_size
):
continue
pkg_files.append(p)
for file_path in pkg_files:
tf.add(file_path)
with tempfile.NamedTemporaryFile("w+") as ntf:
json.dump(asdict(self.metadata), ntf)
ntf.seek(0)
tf.add(ntf.name, arcname="metadata.json")
print(f"Crash report written to {tarball_path}")
[docs] def init(self):
"""Custom error handling.
When an exception is thrown:
- Display an optional message
- Pretty print the traceback
- Parse unintuitive errors and redisplay
- Store useful system information in a tarball for debugging
"""
def _excepthook(
type_: Type[BaseException],
value: BaseException,
traceback: Optional[TracebackType],
) -> None:
click.secho(f"\n{self.message}:\n", fg="red", bold=True)
_Traceback(type_, value, traceback).pretty_print()
self.metadata.print()
if os.environ.get("LATCH_NO_CRASH_REPORT") == "1":
click.secho(
"Not generating crash report due to $LATCH_NO_CRASH_REPORT",
bold=True,
)
return
if not click.confirm("Generate a crash report?"):
return
print("Generating...")
self._write_state_to_tarball()
sys.excepthook = _excepthook