torchx - add exception_type, exception_message, and exception_source_…

…location to torchx event (#966) (#966) Summary: Add exception type, exception message, exception source location to torchx event. This allows for better logging of exception details for further analysis. Differential Revision: D64406552 Co-authored-by: Tony Kao <[email protected]>
pytorch · Oct 16, 2024 · af5acd6 · af5acd6
1 parent 3855ae4
commit af5acd6
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/torchx/runner/events/__init__.py b/torchx/runner/events/__init__.py
@@ -20,7 +20,9 @@
 
 """
 
+import json
 import logging
+import sys
 import time
 import traceback
 from types import TracebackType
@@ -123,6 +125,20 @@ def __exit__(
         ) // 1000
         if traceback_type:
             self._torchx_event.raw_exception = traceback.format_exc()
+            typ, value, tb = sys.exc_info()
+            if tb:
+                last_frame = traceback.extract_tb(tb)[-1]
+                self._torchx_event.exception_source_location = json.dumps(
+                    {
+                        "filename": last_frame.filename,
+                        "lineno": last_frame.lineno,
+                        "name": last_frame.name,
+                    }
+                )
+        if exec_type:
+            self._torchx_event.exception_type = exec_type.__name__
+        if exec_value:
+            self._torchx_event.exception_message = str(exec_value)
         record(self._torchx_event)
 
     def _generate_torchx_event(

diff --git a/torchx/runner/events/api.py b/torchx/runner/events/api.py
@@ -52,6 +52,9 @@ class TorchxEvent:
     wall_time_usec: Optional[int] = None
     start_epoch_time_usec: Optional[int] = None
     workspace: Optional[str] = None
+    exception_type: Optional[str] = None
+    exception_message: Optional[str] = None
+    exception_source_location: Optional[str] = None
 
     def __str__(self) -> str:
         return self.serialize()