| 4 | |
| 5 | |
| 6 | class QueueStorage(object): |
| 7 | def __init__(self, threshold=15, size=20): |
| 8 | """Queue storage |
| 9 | Parameters |
| 10 | ---------- |
| 11 | threshold: int |
| 12 | if the current size if larger than threshold, the data won't be collected |
| 13 | size: int |
| 14 | the size of the queue |
| 15 | """ |
| 16 | self.threshold = threshold |
| 17 | self.queue = Queue(maxsize=size) |
| 18 | |
| 19 | def push(self, batch): |
| 20 | if self.queue.qsize() <= self.threshold: |
| 21 | self.queue.put(batch) |
| 22 | |
| 23 | def pop(self): |
| 24 | if self.queue.qsize() > 0: |
| 25 | return self.queue.get() |
| 26 | else: |
| 27 | return None |
| 28 | |
| 29 | def get_len(self): |
| 30 | return self.queue.qsize() |
| 31 | |
| 32 | |
| 33 | @ray.remote |