Profile memory in Metaflow

profileflow.py

from metaflow import FlowSpec, step
from functools import wraps

def profile_memory(f):
    @wraps(f)
    def func(self):
        from memory_profiler import memory_usage
        self.mem_usage = memory_usage((f, (self,), {}), timeout=6000000, interval=1)
    return func

class ProfileFlow(FlowSpec):

    @profile_memory
    @step
    def start(self):
        import time
        arr = []
        for i in range(5):
            arr.append('a' * int((100 * 1024**2)))
            time.sleep(1)
        self.next(self.end)

    @step
    def end(self):
        print('memory usage over time (in MB):', self.mem_usage)

if __name__ == '__main__':
    ProfileFlow()