Logging Full Traceback with Custom Task Decorator

This is an example code snippet that will log the full traceback of an error to Prefect Cloud. It is also a good guide how to make your own task decorator.

import traceback
from prefect import task, Flow
from functools import partial, wraps
from prefect.executors import LocalDaskExecutor

def custom_task(func=None, **task_init_kwargs):
    if func is None:
        return partial(custom_task, **task_init_kwargs)

    @wraps(func)
    def safe_func(**kwargs):
        try:
            return func(**kwargs)
        except Exception as e:
            print(f"Full Traceback: {traceback.format_exc()}")
            raise RuntimeError(type(e)) from None  # from None is necessary to not log the stacktrace

    safe_func.__name__ = func.__name__
    return task(safe_func, **task_init_kwargs)

@custom_task
def abc(x):
    return x

with Flow("custom-decorator-test") as flow:
    abc.map([1,2,3,4,5])

flow.run()
1 Like