Python 非同期ジェネレータでないジェネレータを利用しての非同期処理

Question

実現したいこと

非同期ジェネレータでないジェネレータの戻り値を利用して非同期処理を行う
＊generate_date()は非同期ジェネレータに変更できないことを前提としています。

発生している問題・分からないこと

heavy_func()が呼び出されない

該当のソースコード

import asyncio
import datetime as dt
import time


def main():
   exec_time = dt.datetime.now()
   print(f'{exec_time}, start')
   
   data_list = generate_data()
   asyncio.run(calc_from_data(data_list))


def generate_data():    
   data_list = [
       {'key': 's1', 'val': 1.0},
       {'key': 's2', 'val': 2.0},
       {'key': 's3', 'val': 3.0},
   ]

   while True:
       yield data_list[0]
       yield data_list[1]
       yield data_list[2]
       time.sleep(10)


async def calc_from_data(data_list):
   tasks = []
   for data in data_list:
       key = data['key']
       val = data['val']
       exec_time = dt.datetime.now()
       print(f'{exec_time}, {key}: {val}')

       loop = asyncio.get_event_loop()
       task = loop.create_task(heavy_func(key, val))
       tasks.append(task)
   
   await asyncio.gather(*tasks)


async def heavy_func(key, val):
   exec_time = dt.datetime.now()
   print(f'{exec_time}, called, {key}: {val}')
   
   # 何か重い処理
   await asyncio.sleep(10)
   
   exec_time = dt.datetime.now()
   print(f'{exec_time}, end, {key}: {val}')
   

if __name__ == '__main__':
   main()

2024-12-11 03:10:33.140207, start
2024-12-11 03:10:33.140740, s1: 1.0
2024-12-11 03:10:33.140784, s2: 2.0
2024-12-11 03:10:33.140808, s3: 3.0
2024-12-11 03:10:43.150936, s1: 1.0
2024-12-11 03:10:43.151100, s2: 2.0
2024-12-11 03:10:43.151130, s3: 3.0

試したこと・調べたこと

• teratailやGoogle等で検索した
• ソースコードを自分なりに変更した
• 知人に聞いた
• その他

上記の詳細・結果

同様の問題が見つかりませんでした。

補足

特になし

Answer 1

問題の原因
asyncio はイベントループを持ち、
await のタイミングで処理を切り替え、他のタスクを進めます。

問題のコードでは、
for data in data_list: の data_list が無限ループの為、for 分を抜けられず
await asyncio.gather(*tasks) まで実行は到達しません。⇒ 他のタスクは実行されない

解決策A;
同期 <=> 非同期のやりとりとして、まず思いつくのは、
janus というライブラリの同期<=>非同期キューを用いると
例えば別スレッドで無限ループを処理し、非同期側で値を受け取ることが可能です。(追記: 標準ライブラリの asyncio.Queue でも工夫すれば可能)

実質、同期ジェネレーターで生成する値を、
コルーチン側で非同期ジェネレーターとして扱う事が可能になります。

別スレッドでの実行は to_thread, run_in_executor 等。

追記2: 要点が埋もれてしまったので、ブロッキング処理を含むジェネレータを別スレッドで実行し非同期読み出し。for文は使わずに

# ジェネレーターからの非同期読み出し （別スレッドでジェネレーター内部の処理は進む ⇒ asyncio のイベントループはブロックされない)
value = await asyncio.to_thread(next, gen)

---

解決策B
タスク生成スケジュールの見直し、
現状のコードではタスクの内容次第では、処理速度が追いつかないことがあるかもしれません。
（タスクを溢れさせないための sleep 10 だとは思いますが）
例えば 20個のタスクを生成し、10個のタスクが完了した時点でもう10個次のタスクを追加する等、

無限ループのコードを変更することなく、
next() でジェネレーターから値を順次呼び出すことは可能です。
（time.sleep(10) は不要になると思うので見直し候補）

追記: Semaphore 等で数量制限

---

# 何か重い処理
await asyncio.sleep(10)

何か重たい処理として仮置きしてるのかもしれませんが、
コルーチン内での重たい処理は別スレッドで行いましょう。
実際のコード次第では問題になる場合があります。

追記: 動作サンプルコード

import time
import asyncio
from concurrent.futures import ThreadPoolExecutor
import logging

logger = logging.getLogger(__name__)

TASK_LIMIT = 5
QUEUE_SIZE = 10
INTERVAL = 10
DATA_LIST = [
    {"key": "a", "value": 1},
    {"key": "b", "value": 2},
    {"key": "c", "value": 3},
]

def producer(data=DATA_LIST, interval=INTERVAL):
    # sync generator
    while True:
        yield from data
        logger.info(f"waiting {interval} sec")
        time.sleep(interval)

async def producer_task(queue, sync_gen):
    # 非同期: ジェネレーターから値を取り出す
    logger.info("producer start")
    while True:
        item = await asyncio.to_thread(next, sync_gen)
        logger.debug(f"P: {item=}")
        await queue.put(item) # キューの空を待つ (maxsize)

async def consumer(queue):
    # 非同期 generator: async for で扱う為の wrapper
    while True:
        item = await queue.get() # キューに値が入るのを待つ
        yield item

async def consumer_task(queue):
    logger.info("consumer start")
    semaphore = asyncio.Semaphore(TASK_LIMIT) # 同時実行数
    loop = asyncio.get_event_loop()

    def done(future):
        # 同期タスクの実行完了時
        result = future.result()
        logger.debug(f"task done {result=}")
        queue.task_done()
        semaphore.release()

    with ThreadPoolExecutor(thread_name_prefix="worker") as executor:
        async for item in consumer(queue):
            logger.debug(f"C: {item=}")
            await semaphore.acquire() # 複数実行中のタスク完了を待つ
            future = loop.run_in_executor(executor, heavy_task, item)
            future.add_done_callback(done)

def debug(func):
    # タスクの実行タイム計測
    from datetime import datetime
    from functools import wraps

    @wraps(func)
    def _func(*args, **kw):
        result = None
        start_time = datetime.now()
        logger.debug(f"start {func.__name__}")
        try:
            result = func(*args, **kw)
        except Exception as exc:
            logger.error(f"error ", exc_info=exc)
        end_time = datetime.now()
        logger.debug(f"{result=} (time: {end_time - start_time})")
        return result
    return _func

@debug
def heavy_task(arg):
    # 時間のかかるタスク。run_in_executor で実行する
    time.sleep(5)
    return arg

async def main():
    # producer->consumer へ値を渡す、非同期キュー
    queue = asyncio.Queue(maxsize=QUEUE_SIZE)

    async with asyncio.TaskGroup() as tg:
        task1 = tg.create_task(producer_task(queue, producer()))
        task2 = tg.create_task(consumer_task(queue))

if __name__ == '__main__':
    logging.basicConfig(
        format="[%(asctime)s][%(threadName)-10s][%(levelname)-8s] %(message)s",
        level=logging.DEBUG,
    )
    try:
        asyncio.run(main())
    except (KeyboardInterrupt, SystemExit):
        pass
    finally:
        logger.info("Done.")

Comments

[keruuuu]

説明とサンプルコードありがとうございます。