How to specify batch-size of tasks that celery workers can accept?

user_4587151 0
Viewed 499

I'm trying to design a system where the caller sends data to celery workers which run on GPU machines, they do some processing and return the results to the caller. Since GPUs can process a batch-full of data in parallel (say N items in parallel), it makes sense for the celery task to accept N items of data and process all N items together.

On the caller side, I am able to use chunks to break apart my list of data items into sublists of some fixed BATCH_SIZE.

 job = tasks.predict.chunks(texts, BATCH_SIZE)
 result = job.apply_async()

However, my issue is that I would like to specify the BATCH_SIZE on the receiving end (worker) instead of the caller, because each worker may have a different limit on the number of data items it can process in one batch (because GPU memory size can vary across different workers).

What's the best way to go about doing this? I want the celery worker to be able to consume an arbitrary N items from the data workload during a single task.

celery parallel-processing gpu
2 Answers

I think it is probably the best that each Worker has BATCH_SIZE in its configuration. Your tasks can then easily grab this value from the config.

user_8764970

The best way to do it currently is through the not-so-friendly celery-batches package. The example of use is in the deprecated celery.contrib.batches package.

With model prediction, you currently should have something wrapped in a context class such as:

class Context(Task):
    def __init__(self):
        self._model = get_model()

@app.task(base=Context)
def predict(img):
    return predict._model(img)

With celery-batches, you rewrite your model as:

class Context(Batches):
    def __init__(self):
        super(Context, self).__init__()
        self._model = get_batch_model()

@app.task(base=Context,
          flush_every=BATCH_SIZE,
          flush_interval=MAX_IDLE_TIME)
def predict(requests):
    batch = []
    for request in requests:
        batch.append(request.args[0])
    reponses = predict_batched(predict._model, batch)
    for response, request in zip(reponses, requests):
        app.backend.mark_as_done(request.id, response)

def predict_batched(model, batch):
    return model(batch)

where BATCH_SIZE can be defined individually in each worker, and your worker will process either when it accumulates BATCH_SIZE tasks, or when MAX_IDLE_TIME has passed.

user_138063260
Related