Clerk2.5/pdf_converter_v2/utils/mineru_url_selector.py

# Copyright (c) Opendatalab. All rights reserved.

"""
MinerU API 多实例 URL 轮询
当 API_URL 配置为逗号分隔的多个地址时（多卡多实例），按请求轮询使用。
"""

import os
import threading
from typing import List

_DEFAULT_SINGLE = "http://127.0.0.1:5282"
_URL_LIST: List[str] = []
_URL_INDEX: int = 0
_URL_LOCK = threading.Lock()


def _parse_api_url_list() -> List[str]:
    """从环境变量 API_URL 解析 URL 列表（逗号分隔，去空格）。"""
    raw = os.getenv("API_URL", _DEFAULT_SINGLE).strip()
    if not raw:
        return [_DEFAULT_SINGLE]
    return [u.strip() for u in raw.split(",") if u.strip()] or [_DEFAULT_SINGLE]


def _get_url_list() -> List[str]:
    """解析并缓存 API_URL 列表（逗号分隔）。"""
    global _URL_LIST
    with _URL_LOCK:
        if not _URL_LIST:
            _URL_LIST[:] = _parse_api_url_list()
        return list(_URL_LIST)


def get_mineru_api_url_list() -> List[str]:
    """
    返回 MinerU API 地址列表（用于单任务多卡：按页拆分后并行发到各实例）。
    当 API_URL 为逗号分隔多地址时返回多元素；单地址时返回单元素列表。
    """
    return _get_url_list()


def get_next_mineru_api_url() -> str:
    """
    获取下一次应使用的 MinerU API 地址（线程安全轮询）。
    当 API_URL 为多个地址（逗号分隔）时轮流返回；单地址则始终返回该地址。
    """
    global _URL_INDEX
    urls = _get_url_list()
    with _URL_LOCK:
        if len(urls) == 1:
            return urls[0]
        idx = _URL_INDEX % len(urls)
        _URL_INDEX += 1
        return urls[idx]