Search code examples
pythonloggingfastapibackground-task

How to log the return value of a POST method after returning the response?

I'm working on my first ever REST API, so apologies in advance if I've missed something basic. I have a function that takes a JSON request from another server, processes it (makes a prediction based on the data), and returns another JSON with the results. I'd like to keep a log on the server's local disk of all requests to this endpoint along with their results, for evaluation purposes and for retraining the model. However, for the purposes of minimising the latency of returning the result to the user, I'd like to return the response data first, and then write it to the local disk. It's not obvious to me how to do this properly, as the FastAPI paradigm necessitates that the result of a POST method is the return value of the decorated function, so anything I want to do with the data has to be done before it is returned.

Below is a minimal working example of what I think is my closest attempt at getting it right so far, using a custom object with a log decorator - my idea was just to assign the result to the log object as a class attribute, then use another method to write it to disk, but I can't figure out how to make sure that that function gets called after get_data every time.

import json
import uvicorn
from fastapi import FastAPI, Request
from functools import wraps
from pydantic import BaseModel

class Blob(BaseModel):
    id: int
    x: float

def crunch_numbers(data: Blob) -> dict:
    # does some stuff
    return {'foo': 'bar'}

class PostResponseLogger:

    def __init__(self) -> None:
        self.post_result = None

    def log(self, func, *args, **kwargs):
        @wraps(func)
        def func_to_log(*args, **kwargs):
            post_result = func(*args, **kwargs)
            self.post_result = post_result

            # how can this be done outside of this function ???
            self.write_data()

            return post_result
        return func_to_log

    def write_data(self):
        if self.post_result:
            with open('output.json', 'w') as f:
                    json.dump(self.post_result, f)

def main():
    app = FastAPI()
    logger = PostResponseLogger()

    @app.post('/get_data/')
    @logger.log
    def get_data(input_json: dict, request: Request):
        result = crunch_numbers(input_json)
        return result

    uvicorn.run(app=app)

if __name__ == '__main__':
    main()

Basically, my question boils down to: "is there a way, in the PostResponseLogger class, to automatically call self.write_data after every call to self.log?", but if I'm using the wrong approach altogether, any other suggestions are also welcome.

Solution

  • You could have a Background Task for that purpose. A background task "will run only once the response has been sent" (see Starlette documentation). "This is useful for operations that need to happen after a request, but that the client doesn't really have to be waiting for the operation to complete before receiving the response" (see FastAPI documentation).

    You can define a task function to run in the background for writing the log data, as shown below:

    def write_log_data():
    logger.write_data()

    Then, import BackgroundTasks and define a parameter in your endpoint with a type declaration of BackgroundTasks. Inside of your endpoint, pass your task function (i.e., write_log_data, as defined above) to the background_tasks object with the method .add_task(), as shown below. More details on how to use BackgroundTasks can be found in this answer.

    from fastapi import BackgroundTasks

@app.post('/get_data/')
def get_data(input_json: dict, request: Request, background_tasks: BackgroundTasks):
    result = crunch_numbers(input_json)
    background_tasks.add_task(write_log_data)
    return result

    The same principle could be applied if a middleware was used to capture and log the response data, as described in this answer, or a custom APIRoute class, as demonstrated in this answer.

    For future reference, if you (or anyone) ever need to use async/await syntax, and run into concurrency issues (such as the event loop getting blocked) while performing some heavy background computation, please have a look at this answer, which explains the difference between defining an endpoint or a background task function with async def and def (briefly, both async def endpoints and background task functions will run directly in the event loop, whereas normal def functions will run in a separate thread from an external threadpool and then be awaited), as well as provides solutions when it comes to running blocking I/O-bound or CPU-bound operations in such functions.