Celery throttling — setting rate limit for queues
By: Magomed Aliev
Tags: Distributed Systems, Python, Software Development, Celery, Rabbitmq

Writing some code

Well, let’s write some code. Create the main.py file and set the basic settings:

from celery import Celery

from kombu import Queue

​

app = Celery('Test app', broker='amqp://guest@localhost//')

​

# 1 queue for tasks and 1 queue for tokens

app.conf.task_queues = [

Queue('github'),

# I limited the queue length to 2, so that tokens do not accumulate

# otherwise this could lead to a breakdown of our rate limit

Queue('github_tokens', max_length=2)

]

​

# this task will play the role of our token

# it will never be executed, we will just pull it as a message from the queue

@app.task

def token():

return 1

​

# setting up a constant issue of our token

@app.on_after_configure.connect

def setup_periodic_tasks(sender, **kwargs):

# we will issue 1 token per second

# that means rate limit for github queue is 60 tasks per minute

sender.add_periodic_task(1.0, token.signature(queue='github_tokens'))

Do not forget to launch Rabbit, I prefer to do this with docker:

docker run -d --rm --name rabbit -p 15672:15672 -p 5672:5672 rabbitmq:3-management

Now let’s run celery beat - special celery worker, that is always launched and responsible for running periodic tasks.

celery -A main beat --loglevel=info

After that, messages will appear in the console once a second:

[2020-03-22 22:49:00,992: INFO/MainProcess] Scheduler: Sending due task main.token() (main.token)

Well, we have set up the issue of tokens for our ‘bucket’. Now all we have to do is to learn how to pull tokens. Let’s try to optimize the code that we wrote earlier for requests to github. Add these lines to main.py :

# function for pulling tokens from queue

def rate_limit(task, task_group):

# acquiring broker connection from pool

with task.app.connection_for_read() as conn:

# getting token

msg = conn.default_channel.basic_get(task_group+'_tokens', no_ack=True)

# received None - queue is empty, no tokens

if msg is None:

# repeat task after 1 second

task.retry(countdown=1)

​

# Added some prints for logging

# I set max_retries=None, so that tasks will repeat until complete

@app.task(bind=True)

def get_github_api1(self, max_retries=None):

rate_limit(self, 'github')

print ('Called Api 1')

​

​

@app.task(bind=True)

def get_github_api2(self, max_retries=None):

rate_limit(self, 'github')

print ('Called Api 2')

Now lets check how it works. In addition to the beat process, add 8 workers:

celery -A main worker -c 8 -Q github

And create a separate little script to run these tasks, call it producer.py :

from main import get_github_api1, get_github_api2

​

tasks = [get_github_api1, get_github_api2]

​

for i in range(100):

# launching tasks one by one

tasks[i % 2].apply_async(queue='github')

Start it with python producer.py , and look at logs of workers:

[2020-03-23 13:04:15,017: WARNING/ForkPoolWorker-3] Called Api 2

[2020-03-23 13:04:16,053: WARNING/ForkPoolWorker-8] Called Api 2

[2020-03-23 13:04:17,112: WARNING/ForkPoolWorker-1] Called Api 2

[2020-03-23 13:04:18,187: WARNING/ForkPoolWorker-1] Called Api 1

... (96 more lines)

Despite the fact that we have 8 workers, tasks are executed approximately once per second. If there was no token at the time task reached the worker, task is rescheduled. Also, I think you have already noticed, that in fact we throttle not queue, but some logical group of tasks, that can actually be located in different queues. Thus, our control becomes even more detailed and granular.

Putting it all together

Of course, the number of such task groups is not limited (only by capabilities of the broker). Putting the whole code together, expanding and ‘beautifying’ it:

from celery import Celery

from kombu import Queue

from queue import Empty

from functools import wraps

​

app = Celery('hello', broker='amqp://guest@localhost//')

​

task_queues = [

Queue('github'),

Queue('google')

]

​

# per minute rate

rate_limits = {

'github': 60,

'google': 100

}

​

# generating queues for all groups with limits, that we defined in dict above

task_queues += [Queue(name+'_tokens', max_length=2) for name, limit in rate_limits.items()]

​

app.conf.task_queues = task_queues

​

@app.task

def token():

return 1

​

@app.on_after_configure.connect

def setup_periodic_tasks(sender, **kwargs):

# generating auto issuing of tokens for all lmited groups

for name, limit in rate_limits.items():

sender.add_periodic_task(60 / limit, token.signature(queue=name+'_tokens'))

​

# I really like decorators ;)

def rate_limit(task_group):

def decorator_func(func):

@wraps(func)

def function(self, *args, **kwargs):

with self.app.connection_for_read() as conn:

# Here I used another higher level method

# We are getting complete queue interface

# but in return losing some perfomance because

# under the hood there is additional work done

with conn.SimpleQueue(task_group+'_tokens', no_ack=True, queue_opts={'max_length':2}) as queue:

try:

# Another advantage is that we can use blocking call

# It can be more convenient than calling retry() all the time

# However, it depends on the specific case

queue.get(block=True, timeout=5)

return func(self, *args, **kwargs)

except Empty:

self.retry(countdown=1)

return function

return decorator_func

​

# much more beautiful and readable with decorators, agree?

@app.task(bind=True, max_retries=None)

@rate_limit('github')

def get_github_api1(self):

print ('Called github Api 1')

​

@app.task(bind=True, max_retries=None)

@rate_limit('github')

def get_github_api2(self):

print ('Called github Api 2')

​

@app.task(bind=True, max_retries=None)

@rate_limit('google')

def query_google_api1(self):

print ('Called Google Api 1')

​

@app.task(bind=True, max_retries=None)

@rate_limit('google')

def query_google_api1(self):

print ('Called Google Api 2')

Thus, the total task calls of the google group will not exceed 100/min, and the github group — 60/min. Note that in order to set up such throttling, it took less than 50 lines of code. Is it possible to make it even simpler?