Python time and size batcher

user_758174 0
Viewed 64

I need a little util to batch messages by count or time duration, whichever comes first (application: sending messages to Kinesis, either one at a time if production is slow, or in batches if all of a sudden there are lots of messages to send).

There are many ways to skin a cat, but I came up with the following, which uses a deque and threading.Timer. The questions are:

  1. is it safe (this is used by the main thread)?
  2. is there a simpler or more pythonic way of doing this?
  3. profiling suggests that acquiring _thread.lock and _thread.start_new_thread take a while; is there a different way that would be faster? (Note: if Batcher(..., seconds=None) is used, there is no such cost).
import threading
import time
from collections import deque

class Batcher():

    def __init__(self, size=None, seconds=None, callback=None):
        self.batch = deque()
        self.size = size
        self.seconds = seconds
        self.callback = callback
        self.thread = None
    
    def flush(self):
        if self.thread:
            self.thread.cancel()
            self.thread = None
        if self.batch:
            a = list(self.batch)
            self.batch.clear()
            if self.callback:
                self.callback(a)

    def add(self, e):
        self.batch.append(e)
        if self.size is not None and len(self.batch) >= self.size:
            self.flush()
        elif self.seconds is not None and self.thread is None:
            self.thread = threading.Timer(self.seconds, self.flush)
            self.thread.start()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.flush()

Simple test:

origin = time.time()

def walltime(origin):
    dt = time.time() - origin
    return f'{dt:6.3f} s'

def foo(batch):
    print(f'now={walltime(origin)}, batch={batch}')

with Batcher(size=3, seconds=0.5, callback=foo) as b:
    for k in range(7):
        b.add(f'at {walltime(origin)}: {k}')
        time.sleep(0.3)

Out[ ]:
now= 0.501 s, batch=['at  0.000 s: 0', 'at  0.301 s: 1']
now= 1.101 s, batch=['at  0.601 s: 2', 'at  0.902 s: 3']
now= 1.702 s, batch=['at  1.202 s: 4', 'at  1.503 s: 5']
now= 2.103 s, batch=['at  1.803 s: 6']

Speed test:

In[ ]:
%%time
batch_stats = []

def proc(batch):
    batch_stats.append(len(batch))

with Batcher(size=100, seconds=5, callback=proc) as b:
    for k in range(120164):
        b.add(k)

Out[ ]:
CPU times: user 166 ms, sys: 74.7 ms, total: 240 ms
Wall time: 178 ms

In[ ]:
Counter(batch_stats)

Out[ ]:
Counter({100: 1201, 64: 1})
python
1 Answers

The reason the code spend so much time within acquire.lock and start_thread because you start a thread every time you need to start a timer for a delayed send.

Here's a solution with the thread kept constantly running in the background. Two condition are used two trigger and wait for the requested delay. It was twice as fast in my timing test but I could only test on one machine:

class Batcher:
    def __init__(self, size=None, seconds=None, callback=None):
        self._batch = []
        self._seconds = seconds
        self._size = size
        self._callback = callback
        self._cyclic_requested = False
        self._wait_for_start = False
        self._cancelled = False
        self._timer_started = False
        self._lock = threading.RLock()
        self._timer_condition = Condition(self._lock)
        self._finished = Condition(self._lock)
        self._finished_flag = False
        self._thread = threading.Thread(target=self._cycle_send)
        self._thread.start()

    def _cycle_send(self):
        while True:
            with self._lock:
                # Wait for the timer_condition to be set to start the timer
                self._wait_for_start = True
                # If a cyclic send was requested and the thread was not
                # not waiting we go directly to the wait time
                if not self._cyclic_requested:
                    self._timer_condition.wait()

                # If finished is set end the thread
                if self._finished_flag:
                    return

                # Reset the flags
                self._cyclic_requested = False
                self._wait_for_start = False
                self._cancelled = False
                self._timer_started = True

                # Wait for the finished timer to be set or the timeout
                self._finished.wait(self._seconds)
                # If finished is set end the thread
                if self._finished_flag:
                    return

                self._timer_started = False
                # If the time_condition has been clear no sending
                # is needed anymore, go back to waiting
                if self._cancelled:
                    continue
                self._timer_condition_flag = False


                batch = self._batch
                self._batch = []

            self._send_batch(batch)

    def _send_batch(self, batch):
        if self._callback:
            self._callback(batch)

    def add(self, e):
        batch = None
        with self._lock:
            # Unconditionally append to the batch
            self._batch.append(e)

            if  self._size is not None and len(self._batch) >= self._size:
                # If immediate send required, copy the batch and reset the shared variable
                # also cancel the cycle_send by clearing
                # the timer_condition and setting the finished event
                batch = self._batch
                self._batch = []
                self._cancelled = True
                self._cyclic_requested = False
                self._finished.notify_all()

            # If the batch is not full, set the timer send condition
            elif not self._timer_started:
                if self._wait_for_start:
                    self._timer_condition.notify_all()
                else:
                    self._cyclic_requested = True

        # the sending is done outside the lock to avoid keeping the lock for too long
        if batch is not None:
            self._send_batch(batch)


    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        with self._lock:
            # Set the finish and timer condition to let the thread terminate
            self._finished_flag = True
            self._timer_condition.notify_all()
            self._finished.notify_all()

        self._thread.join()
        # Send what is left
        self._send_batch(self._batch)
user_135379410
Related