An multithreaded online async predictor which runs a list of OnlinePredictor. It would do an extra batching internally.
| 129 | |
| 130 | |
| 131 | class MultiThreadAsyncPredictor(AsyncPredictorBase): |
| 132 | """ |
| 133 | An multithreaded online async predictor which runs a list of OnlinePredictor. |
| 134 | It would do an extra batching internally. |
| 135 | """ |
| 136 | |
| 137 | def __init__(self, predictors, batch_size=5): |
| 138 | """ |
| 139 | Args: |
| 140 | predictors (list): a list of OnlinePredictor available to use. |
| 141 | batch_size (int): the maximum of an internal batch. |
| 142 | """ |
| 143 | assert len(predictors) |
| 144 | self._need_default_sess = False |
| 145 | for k in predictors: |
| 146 | assert isinstance(k, OnlinePredictor), type(k) |
| 147 | if k.sess is None: |
| 148 | self._need_default_sess = True |
| 149 | # TODO support predictors.return_input here |
| 150 | assert not k.return_input |
| 151 | self.input_queue = queue.Queue(maxsize=len(predictors) * 100) |
| 152 | self.threads = [ |
| 153 | PredictorWorkerThread( |
| 154 | self.input_queue, f, id, batch_size=batch_size) |
| 155 | for id, f in enumerate(predictors)] |
| 156 | |
| 157 | def start(self): |
| 158 | if self._need_default_sess: |
| 159 | assert tfv1.get_default_session() is not None, \ |
| 160 | "Not session is bind to predictors, " \ |
| 161 | "MultiThreadAsyncPredictor.start() has to be called under a default session!" |
| 162 | for t in self.threads: |
| 163 | t.start() |
| 164 | |
| 165 | def put_task(self, dp, callback=None): |
| 166 | """ |
| 167 | Args: |
| 168 | dp (list): A datapoint as inputs. It could be either batched or not |
| 169 | batched depending on the predictor implementation). |
| 170 | callback: a thread-safe callback. When the results are ready, it will be called |
| 171 | with the "future" object. |
| 172 | Returns: |
| 173 | concurrent.futures.Future: a Future of results. |
| 174 | """ |
| 175 | f = Future() |
| 176 | if callback is not None: |
| 177 | f.add_done_callback(callback) |
| 178 | self.input_queue.put((dp, f)) |
| 179 | return f |
no outgoing calls
no test coverage detected
searching dependent graphs…