Use memory_profiler to profile lines in a Flask app

Question

UPDATE: In memory_profiler version 0.53 and later one can @profile decorate as many routes as you want. Earlier versions only allowed decorating one route. The below question only applies to those versions of memory_profiler with version <= 0.52

Using the normal @profile decorator doesn't work on two or more Flask routes. How do I still get line by line memory usage profiling in two or more Flask routes?

I want to profile /route_one and /route_two:

from functools import wraps

from memory_profiler import profile

@app.route("/route_one", methods=["GET"])
@profile
def route_one():
    api_live = ping_api()
    if not api_live:
        return make_response('', 503)
    return make_response('', 200)

@app.route("/route_two", methods=["GET"])
@profile
def route_two():
    api_live = ping_api()
    if not api_live:
        return make_response('', 503)
    return make_response('', 200)

When I start up my flask app with the above decorated routes I get this error:

AssertionError: View function mapping is overwriting an existing endpoint function: wrapper

Answer 1

memory_profiler has been updated to include code to use @profile on Flask routes. memory_profiler version >= 0.53 will not have this problem. See this GitHub issue for more information.

---

The error is telling us that the same function wrapper is used on two routes that we are attempting to map/decorate over. The way to fix this is to use @wraps. As described here: What does functools.wraps do? @wraps copies the name and the docstring from the inner function on to the outer wrapped function. So if we use @wraps, we can avoid the above error.

But we need to use @wraps inside our decorator definition. Our profile decorator is defined in the memory_profiler library, so we need to re-write that function to include @wraps. The memory_profiler profiler function is here https://github.com/pythonprofilers/memory_profiler/blob/master/memory_profiler.py and we will use a modified version of it below which uses @wraps.

Use the below code in your flask app Decorate your route with @my_profiler

from functools import wraps

import memory_profiler
try:
    import tracemalloc
    has_tracemalloc = True
except ImportError:
    has_tracemalloc = False


def my_profiler(func=None, stream=None, precision=1, backend='psutil'):
    """
    Decorator that will run the function and print a line-by-line profile
    """
    backend = memory_profiler.choose_backend(backend)
    if backend == 'tracemalloc' and has_tracemalloc:
        if not tracemalloc.is_tracing():
            tracemalloc.start()
    if func is not None:
        @wraps(func)
        def wrapper(*args, **kwargs):
            prof = memory_profiler.LineProfiler(backend=backend)
            val = prof(func)(*args, **kwargs)
            memory_profiler.show_results(prof, stream=stream,
                                         precision=precision)
            return val

        return wrapper
    else:
        def inner_wrapper(f):
            return profile(f, stream=stream, precision=precision,
                           backend=backend)

        return inner_wrapper

We can now use our fixed profiler with

@app.route("/route_one", methods=["GET"])
@my_profiler
def route_one():
    api_live = ping_api()
    if not api_live:
        return make_response('', 503)
    return make_response('', 200)

@app.route("/route_two", methods=["GET"])
@my_profiler
def route_two():
    api_live = ping_api()
    if not api_live:
        return make_response('', 503)
    return make_response('', 200)

Answer 2

The profile decorator isn't using functools.wraps, so it doesn't preserve the function name. app.route sees both functions as named 'profile' by the time it executes. Since Flask doesn't allow the same name for multiple routes, it raises an error.

Until memory_profiler fixes this, you can fix this by temporarily specifying the endpoint name when you want to apply @profile decorator.

@app.route('/one', endpoint='one')
@profile
@login_required
def one():
    return 'one'

@profile should be above any decorators you want to include in your profiling, and @route should be above @profile so that it routes to the profiled version of the view.