refactor: Handle exceptions and end traces in Graph class

This commit adds exception handling and trace ending functionality to the Graph class in the `base.py` file. It imports the `traceback` module and adds a try-except block to catch any exceptions that occur during graph execution. If an exception is caught, it logs the exception and ends all traces with the error message. Finally, it ensures that all traces are ended regardless of whether an exception occurred or not. This change improves the robustness and error handling of the Graph class.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-23 22:22:45 -03:00
commit a4c9ada178

View file

@ -1,4 +1,5 @@
import asyncio
import traceback
import uuid
from collections import defaultdict, deque
from datetime import datetime, timezone
@ -341,7 +342,11 @@ class Graph:
self.increment_run_count()
except Exception as exc:
logger.exception(exc)
tb = traceback.format_exc()
await self.end_all_traces(error=f"{exc.__class__.__name__}: {exc}\n\n{tb}")
raise ValueError(f"Error running graph: {exc}") from exc
finally:
await self.end_all_traces()
# Get the outputs
vertex_outputs = []
for vertex in self.vertices: