mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-20 12:56:55 +08:00
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:
@ -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")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user