Hybrid Logger: a logger that works both inside and outsite of a Prefect Context

Prefect 2.3 introduced a new context manager (PR#6575) to disable the logger (i.e., so the function won’t fail when running outside of a prefect context).

That solves half of the problem, because now, when testing outside of prefect you do not see the logs.

A better solution would be to use a logger that works both inside and outside of Prefect (if the logger is called inside: it acts like a Prefect logger and the logs are sent to Prefect Cloud, whereas if the logger is called outside: it acts like a regular logger and the logs can be seen in the terminal).

Here is a small proof-of-concept:

import inspect
import logging
import os

os.environ["PREFECT_LOGGING_EXTRA_LOGGERS"] = "my_hybrid_logger"

try:
    from prefect import context, flow, task
except ModuleNotFoundError:
    print("prefect is not installed")



class HybridLogger(logging.Logger):

    def __init__(self, level="info"):
        self._level = level.upper()
        self.caller_filename = inspect.stack()[1].filename.split("/")[-1]
        self.logger = self._get_regular_logger()

    def _get_regular_logger(self) -> logging.Logger:
        """Create, configure and return a regular logger instance."""

        logging.basicConfig()
        regular_logger = logging.getLogger("my_hybrid_logger")
        regular_logger.setLevel(logging.getLevelName(self._level))
        return regular_logger

    def _get_prefect_context(self):
        """Return the current Prefect context. Otherwise return an empty string."""

        # Check for existing contexts
        try:
            task_run_context = context.TaskRunContext.get()
            flow_run_context = context.FlowRunContext.get()
        except NameError:
            # prefect is not installed
            return ""

        # Determine if this is a task or flow run logger
        if task_run_context:
            return f"Task run '{task_run_context.task_run.name}' - "
        elif flow_run_context:
            return f"Flow run '{flow_run_context.flow_run.name}' - "
        else:
            return ""

    def _fmt_msg(self, msg, caller_function_name):
        """Format the message adding some context."""
        return f"{self._get_prefect_context()}{self.caller_filename}:{caller_function_name} - {msg}"

    def debug(self, msg, *args, **kwargs):
        msg = self._fmt_msg(msg, caller_function_name=inspect.stack()[1].function)
        self.logger.debug(msg, *args, **kwargs)

    def info(self, msg, *args, **kwargs):
        msg = self._fmt_msg(msg, caller_function_name=inspect.stack()[1].function)
        self.logger.info(msg, *args, **kwargs)

    def warning(self, msg, *args, **kwargs):
        msg = self._fmt_msg(msg, caller_function_name=inspect.stack()[1].function)
        self.logger.warning(msg, *args, **kwargs)

    def error(self, msg, *args, **kwargs):
        msg = self._fmt_msg(msg, caller_function_name=inspect.stack()[1].function)
        self.logger.error(msg, *args, **kwargs)

    def exception(self, msg, *args, **kwargs):
        msg = self._fmt_msg(msg, caller_function_name=inspect.stack()[1].function)
        self.logger.exception(msg, *args, **kwargs)

    def critical(self, msg, *args, **kwargs):
        msg = self._fmt_msg(msg, caller_function_name=inspect.stack()[1].function)
        self.logger.critical(msg, *args, **kwargs)


def demo_outside_prefect():
    logger = HybridLogger(level="info")
    logger.debug("not showing")
    logger.warning("warning logging outside of a Prefect context")
    logger = HybridLogger(level="debug")
    logger.debug("debug logging outside of a Prefect context")


@flow
def demo_inside_prefect():
    logger = HybridLogger(level="info")
    logger.debug("not showing")
    logger.warning("warning logging inside of a Prefect Flow")
    logger = HybridLogger(level="debug")
    logger.debug("debug logging inside of a Prefect Flow")
    demo_inside_prefect_task()

@task
def demo_inside_prefect_task():
    logger = HybridLogger(level="info")
    logger.debug("not showing")
    logger.warning("warning logging inside of a Prefect Task")
    logger = HybridLogger(level="debug")
    logger.debug("debug logging inside of a Prefect Task")


if __name__ == "__main__":
    demo_outside_prefect()
    demo_inside_prefect()

Here is why someone would want to use such a logger:

  1. Works both inside and outside of Prefect context (i.e. can be used for testing while still seeing the logs)
  2. You no longer have to write “logger = get_run_logger()” in every functions. Instead, you can simply instantiate the logger once in the top of the file (with “logger = HybridLogger()”)
  3. [extra] You can see the filename and the function_name that called the logger (very useful if the function that calls the logger is a regular function (i.e., not a flow or a task), because otherwise it’s hard to tell where the log is coming from)

Downsides:

  • You have to implement it yourself
  • You have to set PREFECT_LOGGING_EXTRA_LOGGERS
  • Not tested & not supported by Prefect

Questions:

  • Could Prefect make their logger hybrid (so it works like a regular logger when outside of a Prefect context)?
  • Could Prefect make their logger importable like that : “from prefect import logger” (so you don’t have to write “logger = get_run_logger(level)” in every functions)?
2 Likes

Thanks for contributing! To answer your questions:

  1. Yes, it’s technically possible for us to provide this functionality, but this way, we would have to inject fake task or flow run information or reroute it in a way that would be surprising given the name of the function call. We can solve it in the PR linked below
  2. Yes, that’s possible and we can add that - here is a feature request so that you can track the progress on that - this will solve both of your requests:
1 Like

For somebody who’ll have a problem, the order is important!

import inspect
import logging
import os

os.environ["PREFECT_LOGGING_EXTRA_LOGGERS"] = "wrapper_around_prefect_logger"
# Don't change the order here!
# It's important that extra logger will be before prefect context!
from prefect import context
2 Likes

thanks for sharing your solution