https://prefect.io logo
n

Nelson Griffiths

07/18/2023, 8:57 PM
Any best practices for how to profile a prefect flow effectively? Profilers generally don't seem to like Prefect that much. Any easy way to run a flow as a function outside of the prefect engine and just a normal python function so it is easier to profile?
n

Nate

07/18/2023, 9:24 PM
using
memory-profiler
you could do stuff like
Copy code
from io import TextIOWrapper
from typing import Any, Callable

from prefect import flow, Flow
from memory_profiler import profile


class ProfiledFlow(Flow):
    """
    Prefect Flow decorated with memory_profiler's `profile`
    """
    def __init__(
            self,
            fn: Callable = None,
            stream: TextIOWrapper = None,
            **kwargs: Any
    ):
        # `fn` is the callable to be profiled, writing output to `stream`
        profiled_fn = profile(fn, stream=stream) if stream else profile(fn)

        # we ultimately want this to behave like a `Flow`
        super().__init__(profiled_fn, **kwargs)


def profiled_flow(
    fn: Callable = None, stream: TextIOWrapper = None, **flow_init_kwargs: Any
):
    """Prefect Flow decorated with memory_profiler's `profile`

    Args:
        fn (Callable): some function to be profiled and interpreted as a Prefect Flow
        stream (TextIOWrapper): some file stream to write profiling output to
        flow_init_kwargs (Any): Flow kwargs to pass to parent class
        
    """
    if fn is None:
        return lambda fn: ProfiledFlow(
            fn=fn,
            stream=stream,
            **flow_init_kwargs,
        )
    return ProfiledFlow(fn=fn, stream=stream, **flow_init_kwargs)

def my_func():
    a = [1] * (10 ** 6)
    b = [2] * (2 * 10 ** 7)
    del b
    return a

if __name__ == '__main__':
    # profile just the my_func function - you can get a fn from a flow like my_flow.fn
    profile(my_func)() 
    
    # profile the entire flow
    profiled_flow(my_func)()
Copy code
❯ python profiling.py
Filename: /Users/nate/src/play/testing/prefect-sandbox/testing-prefect/profiling.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    44    116.7 MiB    116.7 MiB           1   def my_func():
    45    124.3 MiB      7.6 MiB           1       a = [1] * (10 ** 6)
    46    276.9 MiB    152.6 MiB           1       b = [2] * (2 * 10 ** 7)
    47    124.3 MiB   -152.6 MiB           1       del b
    48    124.3 MiB      0.0 MiB           1       return a


16:23:57.724 | INFO    | prefect.engine - Created flow run 'ivory-axolotl' for flow 'my-func'
Filename: /Users/nate/src/play/testing/prefect-sandbox/testing-prefect/profiling.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
    44    165.1 MiB    165.1 MiB           1   def my_func():
    45    172.8 MiB      7.6 MiB           1       a = [1] * (10 ** 6)
    46    325.4 MiB    152.6 MiB           1       b = [2] * (2 * 10 ** 7)
    47    172.8 MiB   -152.6 MiB           1       del b
    48    172.8 MiB      0.0 MiB           1       return a


16:24:02.124 | INFO    | Flow run 'ivory-axolotl' - Finished in state Completed()
👏 1
🎉 1