Refa:replace trio with asyncio (#11831)

### What problem does this PR solve?

change:
replace trio with asyncio

### Type of change
- [x] Refactoring
This commit is contained in:
buua436
2025-12-09 19:23:14 +08:00
committed by GitHub
parent ca2d6f3301
commit 65a5a56d95
31 changed files with 821 additions and 429 deletions

View File

@ -14,6 +14,8 @@
# limitations under the License.
#
import asyncio
import logging
import os
import sys
sys.path.insert(
@ -28,7 +30,6 @@ from deepdoc.vision.seeit import draw_box
from deepdoc.vision import OCR, init_in_out
import argparse
import numpy as np
import trio
# os.environ['CUDA_VISIBLE_DEVICES'] = '0,2' #2 gpus, uncontinuous
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #1 gpu
@ -39,7 +40,7 @@ def main(args):
import torch.cuda
cuda_devices = torch.cuda.device_count()
limiter = [trio.CapacityLimiter(1) for _ in range(cuda_devices)] if cuda_devices > 1 else None
limiter = [asyncio.Semaphore(1) for _ in range(cuda_devices)] if cuda_devices > 1 else None
ocr = OCR()
images, outputs = init_in_out(args)
@ -62,22 +63,29 @@ def main(args):
async def __ocr_thread(i, id, img, limiter = None):
if limiter:
async with limiter:
print("Task {} use device {}".format(i, id))
await trio.to_thread.run_sync(lambda: __ocr(i, id, img))
print(f"Task {i} use device {id}")
await asyncio.to_thread(__ocr, i, id, img)
else:
__ocr(i, id, img)
await asyncio.to_thread(__ocr, i, id, img)
async def __ocr_launcher():
if cuda_devices > 1:
async with trio.open_nursery() as nursery:
for i, img in enumerate(images):
nursery.start_soon(__ocr_thread, i, i % cuda_devices, img, limiter[i % cuda_devices])
await trio.sleep(0.1)
else:
for i, img in enumerate(images):
await __ocr_thread(i, 0, img)
tasks = []
for i, img in enumerate(images):
dev_id = i % cuda_devices if cuda_devices > 1 else 0
semaphore = limiter[dev_id] if limiter else None
tasks.append(asyncio.create_task(__ocr_thread(i, dev_id, img, semaphore)))
trio.run(__ocr_launcher)
try:
await asyncio.gather(*tasks, return_exceptions=False)
except Exception as e:
logging.error("OCR tasks failed: {}".format(e))
for t in tasks:
t.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
raise
asyncio.run(__ocr_launcher())
print("OCR tasks are all done")