add except hook override to write traceback to file

Signed-off-by: Ayush Kamat <[email protected]>

src/latch/__init__.py

@@ -4,6 +4,11 @@
 Latch platform.
 """
 
+import os
+import sys
+import traceback
+from pathlib import Path
+
 from latch.functions.messages import message
 from latch.functions.operators import (
     combine,
@@ -20,17 +25,30 @@
 from latch.resources.tasks import (
     custom_memory_optimized_task,
     custom_task,
-    g6e_xlarge_task,
     g6e_2xlarge_task,
     g6e_4xlarge_task,
     g6e_8xlarge_task,
     g6e_12xlarge_task,
     g6e_16xlarge_task,
     g6e_24xlarge_task,
+    g6e_xlarge_task,
     large_gpu_task,
     large_task,
     medium_task,
     small_gpu_task,
     small_task,
 )
 from latch.resources.workflow import workflow
+
+
+def _except_hook(type, value, tb):
+    traceback_path = Path(os.environ.get("latch_traceback_path", "~/traceback.txt"))
+    traceback_path.parent.mkdir(exist_ok=True)
+
+    with traceback_path.open("w") as f:
+        traceback.print_tb(tb, file=f)
+
+    return sys.__excepthook__(type, value, tb)
+
+
+sys.excepthook = _except_hook