Async aligner work tracking issue fixed

This commit is contained in:
Harrison Deng 2025-02-05 21:47:51 +00:00
parent 1c6e1cfb35
commit 49f31b7943

View File

@ -26,6 +26,7 @@ class AsyncPairwiseAlignmentEngine(AbstractContextManager):
self._work_left.add(work) self._work_left.add(work)
def _on_complete(self, future: Future): def _on_complete(self, future: Future):
self._work_left.remove(future)
self._work_complete.put(future) self._work_complete.put(future)
def work(self, reference, query, **associated_data): def work(self, reference, query, **associated_data):
@ -52,9 +53,7 @@ class AsyncPairwiseAlignmentEngine(AbstractContextManager):
if self._work_complete.empty() and len(self._work_left): if self._work_complete.empty() and len(self._work_left):
return None return None
future_now = await asyncio.wrap_future(self._work_complete.get()) future_now = await asyncio.wrap_future(self._work_complete.get())
completed: tuple[PairwiseAlignment, dict[str, Any]] = future_now return future_now
self._work_left.remove(future_now)
return completed
def __exit__(self, exc_type, exc_value, traceback): def __exit__(self, exc_type, exc_value, traceback):
self.shutdown() self.shutdown()