BFD-4841: Improve IDR Pipeline error handling

apps/bfd-pipeline-idr/pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
     "click>=8.3.3",
     "anyio",
     "loguru",
+    "tblib>=3.2.2",
 ]
 
 [dependency-groups]

apps/bfd-pipeline-idr/src/idr_pipeline/__init__.py

@@ -1,5 +1,6 @@
 import atexit
 import multiprocessing
+import sys
 from datetime import UTC, datetime
 
 import anyio
@@ -149,7 +150,7 @@ async def run_worker_and_stages() -> None:
                 failure_time=resolve_test_date(load_mode),
             )
         logger.opt(exception=True).error("Unrecoverable exception raised during pipeline load:")
-        raise
+        sys.exit(1)
     finally:
         if idr_job_events:
             update_completion_times(

apps/bfd-pipeline-idr/src/idr_pipeline/batch_worker.py

@@ -3,6 +3,7 @@
 import random
 import signal
 import string
+import sys
 import threading
 import uuid
 from collections import Counter, OrderedDict
@@ -13,6 +14,7 @@
 from math import ceil
 from multiprocessing import Manager, Process
 from multiprocessing.managers import SyncManager
+from pathlib import Path
 from queue import Empty, Queue
 from threading import Event
 from types import FrameType
@@ -28,6 +30,11 @@
 from psycopg.errors import DeadlockDetected, InFailedSqlTransaction
 from psycopg_pool.abc import ACT
 
+from .exception_utils import (
+    SerializedExceptionChain,
+    rebuild_exception_chain,
+    serialize_exception_chain,
+)
 from .load_partition import LoadPartition
 from .model.base_model import DbType, IdrBaseModel
 from .model.load_progress import LoadProgress
@@ -183,7 +190,7 @@ class _LoadingBatchWorker(Process):
     def __init__(
         self,
         task_queue: Queue[_TaskSequence],
-        errors_queue: Queue[BaseException],
+        errors_queue: Queue[SerializedExceptionChain],
         started_signal: Event,
         cancel_signal: Event,
         root_logger: Logger,
@@ -200,24 +207,31 @@ def __init__(
         self._running_tasks: set[_Task] = set()
 
     def run(self) -> None:
-        self._logger.reinstall()
+        sys.unraisablehook = lambda _: None
+        sys.excepthook = lambda _, __, ___: None
+        with Path(os.devnull).open("w") as devnull:
+            sys.stderr = devnull
 
-        def _watch_for_parent_cancel(stop_signal: Event) -> None:
-            stop_signal.wait()
-            os.kill(os.getpid(), signal.SIGUSR1)
+            self._logger.reinstall()
 
-        def sigusr1_handler(signum: int, frame: FrameType | None) -> Never:  # noqa: ARG001
-            raise ExternallyCanceled("Externally canceled, interrupting")
+            def _watch_for_parent_cancel(stop_signal: Event) -> None:
+                stop_signal.wait()
+                os.kill(os.getpid(), signal.SIGUSR1)
 
-        signal.signal(signal.SIGUSR1, sigusr1_handler)
-        threading.Thread(
-            target=lambda: _watch_for_parent_cancel(self._cancel_signal), daemon=True
-        ).start()
+            def sigusr1_handler(signum: int, frame: FrameType | None) -> Never:  # noqa: ARG001
+                raise ExternallyCanceled("Externally canceled, interrupting")
 
-        try:
-            anyio.run(self._worker_main)
-        except BaseException as ex:
-            self.errors_queue.put(ex)
+            signal.signal(signal.SIGUSR1, sigusr1_handler)
+            threading.Thread(
+                target=lambda: _watch_for_parent_cancel(self._cancel_signal), daemon=True
+            ).start()
+
+            try:
+                anyio.run(self._worker_main)
+            except ExternallyCanceled:
+                pass
+            except BaseException as ex:
+                self.errors_queue.put(serialize_exception_chain(ex))
 
     async def _worker_main(self) -> None:
         task_send, task_receive = anyio.create_memory_object_stream[_TaskSequence](
@@ -508,7 +522,7 @@ async def start(self, stop: anyio.Event, task_status: TaskStatus | None = None)
             return
 
         self._started_signal.clear()
-        errors_queue: Queue[BaseException] = self._manager.Queue()
+        errors_queue: Queue[SerializedExceptionChain] = self._manager.Queue()
         cancel_signal = self._manager.Event()
 
         self._worker = _LoadingBatchWorker(
@@ -521,28 +535,29 @@ async def start(self, stop: anyio.Event, task_status: TaskStatus | None = None)
         )
         self._worker.start()
 
-        # Block until the worker signals it has started
-        self._started_signal.wait()
-
-        logger.info("LoadingBatchWorker signaled startup")
+        # Block until the worker signals it has started, or 10 seconds have passed
+        self._started_signal.wait(10)
+        if self._started_signal.is_set():
+            logger.info("LoadingBatchWorker signaled startup")
+        else:
+            logger.error("LoadingBatchWorker start signal never set. See exception for detail")
 
         if task_status:
             task_status.started()
 
         async def watch_queue() -> None:
             while not stop.is_set():
                 with contextlib.suppress(Empty):
-                    errors = errors_queue.get_nowait()
-                    raise errors
+                    raise rebuild_exception_chain(errors_queue.get_nowait())
 
                 await anyio.sleep(0.01)
 
-        async with anyio.create_task_group() as tg:
-            try:
+        try:
+            async with anyio.create_task_group() as tg:
                 tg.start_soon(watch_queue)
-            except BaseException:
-                cancel_signal.set()
-                raise
+        except BaseException:
+            cancel_signal.set()
+            raise
 
     def cleanup(self, timeout: float = 5.0) -> None:
         if self._worker is None:

apps/bfd-pipeline-idr/src/idr_pipeline/exception_utils.py

@@ -0,0 +1,84 @@
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Any, cast
+
+from tblib import Traceback  # type: ignore
+
+type SerializedExceptionChain = list[SerializedException | SerializedExceptionGroup]
+
+
+@dataclass(frozen=True, eq=True)
+class SerializedException:
+    ex: BaseException
+    tb_dict: dict[str, Any]
+
+
+@dataclass(frozen=True, eq=True)
+class SerializedExceptionGroup:
+    ex_group: BaseExceptionGroup[Exception] | ExceptionGroup[Exception]
+    tb_dict: dict[str, Any]
+    exceptions: list[SerializedException]
+
+
+def get_exception_tb_dict(ex: BaseException) -> dict[str, Any]:
+    return Traceback(ex.__traceback__).as_dict()  # type: ignore
+
+
+def get_traceback_from_dict(tb_dict: dict[str, Any]) -> TracebackType | None:
+    return cast(Traceback, Traceback.from_dict(tb_dict)).as_traceback()  # type: ignore
+
+
+def serialize_exception_chain(exc: BaseException) -> SerializedExceptionChain:
+    chain: list[SerializedException | SerializedExceptionGroup] = []
+    current: BaseException | None = exc
+
+    while current is not None:
+        if isinstance(current, BaseExceptionGroup | ExceptionGroup):
+            current = cast(BaseExceptionGroup[Exception], current)
+            serialized_inner: list[SerializedException] = [
+                SerializedException(ex=inner, tb_dict=get_exception_tb_dict(inner))
+                for inner in current.exceptions
+            ]
+            chain.append(
+                SerializedExceptionGroup(
+                    ex_group=current,
+                    tb_dict=get_exception_tb_dict(current),
+                    exceptions=serialized_inner,
+                )
+            )
+        else:
+            chain.append(SerializedException(ex=current, tb_dict=get_exception_tb_dict(current)))
+
+        current = current.__cause__
+
+    return chain
+
+
+def rebuild_exception_chain(chain: SerializedExceptionChain) -> BaseException:
+    if not chain:
+        raise ValueError("Chain must contain at least one exception")
+
+    resolved: list[tuple[BaseException, dict[str, Any]]] = []
+    for entry in chain:
+        if isinstance(entry, SerializedExceptionGroup):
+            # Restore tracebacks on every inner exception.
+            restored_inners: list[BaseException] = []
+            for serialized_inner in entry.exceptions:
+                inner_exc = serialized_inner.ex
+                inner_exc.__traceback__ = get_traceback_from_dict(serialized_inner.tb_dict)
+                restored_inners.append(inner_exc)
+
+            # Re-create the group with the restored inner exceptions so that
+            # the group's own .exceptions tuple reflects the restored state.
+            restored_group = entry.ex_group.derive(restored_inners)
+            resolved.append((restored_group, entry.tb_dict))
+        else:
+            resolved.append((entry.ex, entry.tb_dict))
+
+    # Re-create the "exception chain" (of __cause__s) and restore top-level (non-inner Exception)
+    # tracebacks
+    for i, (exc, tb_dict) in enumerate(resolved):
+        exc.__traceback__ = get_traceback_from_dict(tb_dict)
+        exc.__cause__ = resolved[i + 1][0] if i < len(resolved) - 1 else None
+
+    return resolved[0][0]

apps/bfd-pipeline-idr/src/idr_pipeline/parallel_executor.py

@@ -1,10 +1,12 @@
 import contextlib
 import os
 import signal
+import sys
 import threading
 from collections.abc import Callable, Generator
 from concurrent.futures import Future, ProcessPoolExecutor
 from multiprocessing import Manager
+from pathlib import Path
 from queue import Empty, Queue
 from threading import Event
 from types import FrameType
@@ -13,6 +15,12 @@
 import anyio
 from loguru import logger
 
+from .exception_utils import (
+    SerializedExceptionChain,
+    rebuild_exception_chain,
+    serialize_exception_chain,
+)
+
 type StageTask[T] = Callable[[], T]
 type Stage[T] = Generator[StageTask[T]]
 
@@ -30,40 +38,47 @@ def __init__(self, max_workers: int | None = None) -> None:
     def _wrap[T](
         task: StageTask[T],
         index: int,
-        errors_queue: Queue[BaseException],
+        errors_queue: Queue[SerializedExceptionChain],
         cancel_signal: Event,
     ) -> tuple[int, T | None]:
+        sys.unraisablehook = lambda _: None
+        sys.excepthook = lambda _, __, ___: None
+        with Path(os.devnull).open("w") as devnull:
+            sys.stderr = devnull
 
-        def _watch_for_parent_cancel(cancel_signal: Event) -> None:
-            cancel_signal.wait()
-            os.kill(os.getpid(), signal.SIGUSR1)
+            def _watch_for_parent_cancel(cancel_signal: Event) -> None:
+                cancel_signal.wait()
+                os.kill(os.getpid(), signal.SIGUSR1)
 
-        def sigusr1_handler(signum: int, frame: FrameType | None) -> Never:  # noqa: ARG001
-            raise ExternallyCanceled("Externally canceled, interrupting")
+            def sigusr1_handler(signum: int, frame: FrameType | None) -> Never:  # noqa: ARG001
+                raise ExternallyCanceled("Externally canceled, interrupting")
 
-        signal.signal(signal.SIGUSR1, sigusr1_handler)
-        threading.Thread(
-            target=lambda: _watch_for_parent_cancel(cancel_signal), daemon=True
-        ).start()
+            signal.signal(signal.SIGUSR1, sigusr1_handler)
+            threading.Thread(
+                target=lambda: _watch_for_parent_cancel(cancel_signal), daemon=True
+            ).start()
+
+            try:
+                return (index, task())
+            except ExternallyCanceled:
+                pass
+            except BaseException as ex:
+                errors_queue.put(serialize_exception_chain(ex))
 
-        try:
-            return (index, task())
-        except BaseException as e:
-            errors_queue.put(e)
             return (index, None)
 
     async def _run_stage[T](
         self,
         stage: Stage[T],
-        errors_queue: Queue[BaseException],
+        errors_queue: Queue[SerializedExceptionChain],
         pool: ProcessPoolExecutor,
         results: dict[int, T | None],
     ) -> None:
         done = anyio.Event()
 
         cancel_signal = self._manager.Event()
-        async with anyio.create_task_group() as tg:
-            try:
+        try:
+            async with anyio.create_task_group() as tg:
                 tg.start_soon(self._poll_errors, errors_queue, done)
 
                 async with anyio.create_task_group() as tg2:
@@ -77,16 +92,12 @@ async def _task(task: StageTask[T] = task, idx: int = idx) -> None:
                         tg2.start_soon(_task)
 
                 done.set()
-            except BaseException:
-                cancel_signal.set()
-                raise
-
-        # Final drain in case an error landed after tasks completed
-        if not errors_queue.empty():
-            raise errors_queue.get()
+        except BaseException:
+            cancel_signal.set()
+            raise
 
     async def execute[T](self, stages: list[Stage[T]]) -> list[list[T | None]]:
-        errors_queue: Queue[BaseException] = self._manager.Queue()
+        errors_queue: Queue[SerializedExceptionChain] = self._manager.Queue()
         all_results: list[list[T | None]] = []
 
         with ProcessPoolExecutor(
@@ -100,10 +111,13 @@ async def execute[T](self, stages: list[Stage[T]]) -> list[list[T | None]]:
         return all_results
 
     @staticmethod
-    async def _poll_errors(errors_queue: Queue[BaseException], done: anyio.Event) -> None:
+    async def _poll_errors(
+        errors_queue: Queue[SerializedExceptionChain], done: anyio.Event
+    ) -> None:
         while not done.is_set():
             with contextlib.suppress(Empty):
-                raise errors_queue.get_nowait()
+                raise rebuild_exception_chain(errors_queue.get_nowait())
+
             await anyio.sleep(0.01)
 
     @staticmethod

apps/bfd-pipeline-idr/src/idr_pipeline/pipeline_utils.py

@@ -34,6 +34,10 @@
 from .settings import BENEFICIARY_PART_D_PRUNE_BATCH_LIMIT, BENEFICIARY_PRUNE_BATCH_LIMIT
 
 
+class ModelExtractError(Exception):
+    pass
+
+
 def get_progress(
     load_mode: LoadMode,
     source: Source,
@@ -124,8 +128,7 @@ def extract_and_load(
                 raise ex
             time.sleep(1)
         except Exception as ex:
-            logger.opt(exception=True).error("error loading {}", cls.table())
-            raise ex
+            raise ModelExtractError(f"error loading {cls.table()}-{partition.name}") from ex
 
 
 def prune_phase_1_ss_claims(