Simplifying traceback when raising DAGBuildError

[edublancas]

When users interact with Ploomber via the CLI and the pipeline fails, a nicely formatted traceback is formatted:

===================================================================== DAG build failed ======================================================================
------------------------- NotebookRunner: fit -> MetaProduct({'model': File('output/model.pickle'), 'nb': File('output/nb.html')}) --------------------------
----------------------------------------------------------- /Users/Edu/Desktop/ml/scripts/fit.py ------------------------------------------------------------
---------------------------------------------------------------------------
Exception encountered at "In [7]":
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/var/folders/3h/_lvh_w_x5g30rrjzb_xnn2j80000gq/T/ipykernel_37808/656624616.py in <cell line: 1>()
----> 1 raise ValueError
      2 y_pred = clf.predict(X_test)

ValueError:

ploomber.exceptions.TaskBuildError: Error when executing task 'fit'. Partially executed notebook available at /Users/Edu/Desktop/ml/output/nb.ipynb
ploomber.exceptions.TaskBuildError: Error building task "fit"
===================================================================== Summary (1 task) ======================================================================
NotebookRunner: fit -> MetaProduct({'model': File('output/model.pickle'), 'nb': File('output/nb.html')})
===================================================================== DAG build failed ======================================================================

Need help? https://ploomber.io/community

However, when they execute the pipeline using Python, a long traceback shows:

ploomber examples -n templates/ml-basic -o ml
cd ml
# then modify fit.py and raise an exception there
ploomber interact
# once the python session initializes...
dag.build()

This is the traceback:

---------------------------------------------------------------------------
DAGBuildError                             Traceback (most recent call last)
<ipython-input-1-e5f7b0299162> in <cell line: 1>()
----> 1 dag.build()

~/dev/ploomber/src/ploomber/dag/dag.py in build(self, force, show_progress, debug, close_clients)
    548         with dag_logger:
    549             try:
--> 550                 report = callable_()
    551             finally:
    552                 if close_clients:

~/dev/ploomber/src/ploomber/dag/dag.py in _build(self, force, show_progress)
    648
    649                 # on_failure hook executed, raise original exception
--> 650                 raise build_exception
    651
    652     def close_clients(self):

~/dev/ploomber/src/ploomber/dag/dag.py in _build(self, force, show_progress)
    580                 # vs as part of a dag execution
    581                 # FIXME: not passing force flag
--> 582                 task_reports = self._executor(dag=self,
    583                                               show_progress=show_progress)
    584

~/dev/ploomber/src/ploomber/executors/serial.py in __call__(self, dag, show_progress)
    161             else:
    162                 try:
--> 163                     raise DAGBuildError(str(exceptions_all))
    164                 except Exception as e:
    165                     e.__traceback__.tb_next = None

DAGBuildError:
===================================================================== DAG build failed ======================================================================
------------------------- NotebookRunner: fit -> MetaProduct({'model': File('output/model.pickle'), 'nb': File('output/nb.html')}) --------------------------
----------------------------------------------------------- /Users/Edu/Desktop/ml/scripts/fit.py ------------------------------------------------------------
---------------------------------------------------------------------------
Exception encountered at "In [7]":
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/var/folders/3h/_lvh_w_x5g30rrjzb_xnn2j80000gq/T/ipykernel_37981/656624616.py in <cell line: 1>()
----> 1 raise ValueError
      2 y_pred = clf.predict(X_test)

ValueError:

ploomber.exceptions.TaskBuildError: Error when executing task 'fit'. Partially executed notebook available at /Users/Edu/Desktop/ml/output/nb.ipynb
ploomber.exceptions.TaskBuildError: Error building task "fit"
===================================================================== Summary (1 task) ======================================================================
NotebookRunner: fit -> MetaProduct({'model': File('output/model.pickle'), 'nb': File('output/nb.html')})
===================================================================== DAG build failed ======================================================================

Need help? https://ploomber.io/community

The first half is full of Ploomber internals that is irrelevant to the user. We'd like to remove all of those things.

I did some work, but I haven't found a solution. Here's what I found.

Here's a solution proposed with a context manager, the problem is that the code has issues: it doesn't revert the changes, and we end up modifying the global sys.tracebacklimit.

I guess that we could modify the exception and reformat it ourselves, but I haven't found how:

  try:
      raise DAGBuildError(str(exceptions_all))
  except Exception as e:
      error = e

  error.__traceback__.tb_next = None
  error.__traceback__ = None
  raise error

tested it by modifying this part

but it didn't work.