Spaces:
Running
Running
| """后台任务的所属启动、停止信号与取消安全等待（保留周期任务导出名）。""" | |
| import asyncio | |
| import logging | |
| logger = logging.getLogger("ComfyUI-Ranking.PeriodicTasks") | |
| def start_periodic_task(state, name, factory): | |
| current = getattr(state, name + "_worker", None) | |
| if current is not None and not current.done(): | |
| return current | |
| stop = asyncio.Event() | |
| coroutine = factory(stop) | |
| try: | |
| task = asyncio.create_task(coroutine) | |
| except BaseException: | |
| coroutine.close() | |
| raise | |
| setattr(state, name + "_stop", stop) | |
| setattr(state, name + "_worker", task) | |
| def consume(completed): | |
| if not completed.cancelled(): | |
| error = completed.exception() | |
| if error is not None: | |
| logger.warning("后台任务结束异常: task=%s | error_type=%s", name, type(error).__name__) | |
| task.add_done_callback(consume) | |
| return task | |
| def request_periodic_stop(state, names): | |
| tasks = [] | |
| for name in names: | |
| stop = getattr(state, name + "_stop", None) | |
| if stop is not None: | |
| stop.set() | |
| task = getattr(state, name + "_worker", None) | |
| if task is not None: | |
| tasks.append(task) | |
| return tasks | |
| async def wait_periodic_tasks(tasks): | |
| cancelled = None | |
| for task in tasks: | |
| while not task.done(): | |
| try: | |
| await asyncio.shield(task) | |
| except asyncio.CancelledError as error: | |
| # worker自身取消与关闭等待者取消分别处理。 | |
| if asyncio.current_task().cancelling() or not task.cancelled(): | |
| cancelled = error | |
| except Exception: | |
| break | |
| if not task.cancelled(): | |
| # 业务异常已有所属回调诊断，不阻止其它任务收尾和最终同步。 | |
| task.exception() | |
| if cancelled is not None: | |
| raise cancelled | |