MCPcopy Create free account
hub / github.com/Project-MONAI/MONAI / ThreadBuffer

Class ThreadBuffer

monai/data/thread_buffer.py:23–80  ·  view source on GitHub ↗

Iterates over values from self.src in a separate thread but yielding them in the current thread. This allows values to be queued up asynchronously. The internal thread will continue running so long as the source has values or until the stop() method is called. One issue raised by u

Source from the content-addressed store, hash-verified

21
22
23class ThreadBuffer:
24 """
25 Iterates over values from self.src in a separate thread but yielding them in the current thread. This allows values
26 to be queued up asynchronously. The internal thread will continue running so long as the source has values or until
27 the stop() method is called.
28
29 One issue raised by using a thread in this way is that during the lifetime of the thread the source object is being
30 iterated over, so if the thread hasn't finished another attempt to iterate over it will raise an exception or yield
31 unexpected results. To ensure the thread releases the iteration and proper cleanup is done the stop() method must
32 be called which will join with the thread.
33
34 Args:
35 src: Source data iterable
36 buffer_size: Number of items to buffer from the source
37 timeout: Time to wait for an item from the buffer, or to wait while the buffer is full when adding items
38 """
39
40 def __init__(self, src, buffer_size: int = 1, timeout: float = 0.01):
41 self.src = src
42 self.buffer_size = buffer_size
43 self.timeout = timeout
44 self.buffer: Queue = Queue(self.buffer_size)
45 self.gen_thread: Thread | None = None
46 self.is_running = False
47
48 def enqueue_values(self):
49 for src_val in self.src:
50 while self.is_running:
51 try:
52 self.buffer.put(src_val, timeout=self.timeout)
53 except Full:
54 pass # try to add the item again
55 else:
56 break # successfully added the item, quit trying
57 else: # quit the thread cleanly when requested to stop
58 break
59
60 def stop(self):
61 self.is_running = False # signal the thread to exit
62
63 if self.gen_thread is not None:
64 self.gen_thread.join()
65
66 self.gen_thread = None
67
68 def __iter__(self):
69 self.is_running = True
70 self.gen_thread = Thread(target=self.enqueue_values, daemon=True)
71 self.gen_thread.start()
72
73 try:
74 while self.is_running and (self.gen_thread.is_alive() or not self.buffer.empty()):
75 try:
76 yield self.buffer.get(timeout=self.timeout)
77 except Empty:
78 pass # queue was empty this time, try again
79 finally:
80 self.stop() # ensure thread completion

Callers 4

_batch_samplerMethod · 0.90
test_valuesMethod · 0.90
test_timeMethod · 0.90
buffer_iteratorFunction · 0.85

Calls

no outgoing calls

Tested by 2

test_valuesMethod · 0.72
test_timeMethod · 0.72

Used in the wild real call sites across dependent graphs

searching dependent graphs…