Nelson Griffiths
07/18/2023, 8:57 PMNate
07/18/2023, 9:24 PMmemory-profiler
you could do stuff like
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)()
❯ 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()