cortexlabs/cortex

@manneshiva server-side batching is indeed only supported for the TensorFlow Predictor out of the box.

For the Python Predictor, there is a workaround, with the recommendation to set the processes_per_replica field to 1. You could still set it to a value higher than 1, but the likelihood of a request landing on that specific process on an API replica drops the higher the processes_per_replica value is. This is due to the fact that requests are randomly distributed to all processes on an API replica.

Here's a template of how you could implement server-side batching on the Python Predictor. This hasn't been tested, so let me know if you encounter any issue down the road:

import threading as td

class PythonPredictor:
    def __init__(self, config):
        self.model = None # initialize the model here

        self.waiter = td.Event()
        self.waiter.set()

        self.batch_max_size = config["batch_max_size"]
        self.batch_interval = config["batch_interval"] # measured in seconds
        self.barrier = td.Barrier(self.batch_max_size + 1)

        self.samples = {}
        self.predictions = {}
        td.Thread(target=self.batch_inference).start()

    def batch_inference(self):
        while True:
            self.barrier.wait(self.batch_interval)
            self.waiter.clear()
            self.predictions = {}

            # batch process self.samples
            # store results in self.predictions

            self.samples = {}
            self.waiter.set()

    def predict(self, payload):
        tID = td.get_ident()

        self.waiter.wait()
        self.samples[tID] = payload
        self.barrier.wait()

        prediction = self.predictions[tID]
        return prediction