Uvicorn/FastAPI - Execute code on a single worker only

user_12091631 0
Viewed 584

Is it possible to run code on only a specific worker in Uvicorn?

I have the following functions currently being executed on application startup.

@app.on_event('startup')
async def app_startup():
    asyncio.create_task(check_idle())
    asyncio.create_task(calculate_eta())

However, I'd only like a single worker (e.g. worker 0) to run those functions to avoid straining the database.

Is this possible? Thanks in advance.

uvicorn fastapi python
1 Answers

So I wrote something for this:

The idea is to use the pid of a uvicorn worker as a "uniquifier".

I have a UniqueWorker class, which basically creates in every process a worker, tho only one gets randomly assigned (probably the last one who writes to the pid file)

It's not very cool that the function still gets called everytime, but at least the part, which you don't want to be executed often isn't called. I think it's a tolerable overhead.

from pathlib import Path
import os
from functools import cached_property


class UniqueWorker:
    def __init__(self, path):
        self.path = Path(path)
        self.id = str(os.getpid())
        self._set_id()

    def _set_id(self):
        """create the path to the pid file if it doesn't exist, 
        write the pid to the file
        """
        if self.path.exists():
            self.path.parents[0].mkdir(parents=True, exist_ok=True)

        with open(self.path, "w") as f:
            f.write(self.id)

    @cached_property
    def is_assiged(self):
        """check if the worker in this process has been assigned 
        as the unique worker
        """
        with open(self.path, "r") as f:
            assigned_worker = f.read()

        return assigned_worker == self.id

# bake this badboy into a decorator
def by_one_worker(worker_pid_path):
    unique_worker = UniqueWorker(worker_pid_path)

    def deco(f):
        def wrapped(*args, **kwargs):
            if not unique_worker.is_assiged:
                return

            return f(*args, **kwargs)

        return wrapped

    return deco

# then with fastapi_utils we can:
from fastapi_utils.tasks import repeat_every

@app.on_event("startup")
@repeat_every(seconds=2)
@by_one_worker("/dev/shm/path/to/you/pid/sync/file")
def do_something_repeatable():
    print("blabla I love decorator n stuff")

Also available in my package:

pip3 install --user dimfred
from dimfred.fastapi_utils import by_one_worker
user_72001390
Related