diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index f76f959..6daa3d9 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -12,7 +12,7 @@ from typing import Optional, List, Tuple from pydantic import BaseModel from app.core.auth import verify_api_key, verify_app_key, get_admin_api_key from app.core.config import config, get_config -from app.core.batch_tasks import create_task, get_task, expire_task +from app.core.batch import create_task, get_task, expire_task from app.core.storage import get_storage, LocalStorage, RedisStorage, SQLStorage from app.core.exceptions import AppException from app.services.token.manager import get_token_manager diff --git a/app/core/batch_tasks.py b/app/core/batch.py similarity index 58% rename from app/core/batch_tasks.py rename to app/core/batch.py index ff564ff..4857250 100644 --- a/app/core/batch_tasks.py +++ b/app/core/batch.py @@ -1,13 +1,94 @@ """ -Batch task manager for admin batch operations (SSE progress). -""" +Batch utilities. -from __future__ import annotations +- run_batch: generic batch concurrency runner +- BatchTask: SSE task manager for admin batch operations +""" import asyncio import time import uuid -from typing import Any, Dict, List, Optional +from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar + +from app.core.logger import logger + +T = TypeVar("T") + + +async def run_batch( + items: List[str], + worker: Callable[[str], Awaitable[T]], + *, + max_concurrent: int = 10, + batch_size: int = 50, + task: Optional["BatchTask"] = None, + on_item: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None, + should_cancel: Optional[Callable[[], bool]] = None, +) -> Dict[str, Dict[str, Any]]: + """ + 分批并发执行，单项失败不影响整体 + + Args: + items: 待处理项列表 + worker: 异步处理函数 + max_concurrent: 最大并发数 + batch_size: 每批大小 + + Returns: + {item: {"ok": bool, "data": ..., "error": ...}} + """ + try: + max_concurrent = int(max_concurrent) + except Exception: + max_concurrent = 10 + try: + batch_size = int(batch_size) + except Exception: + batch_size = 50 + + max_concurrent = max(1, max_concurrent) + batch_size = max(1, batch_size) + + sem = asyncio.Semaphore(max_concurrent) + + async def _one(item: str) -> tuple[str, dict]: + if (should_cancel and should_cancel()) or (task and task.cancelled): + return item, {"ok": False, "error": "cancelled", "cancelled": True} + async with sem: + try: + data = await worker(item) + result = {"ok": True, "data": data} + if task: + task.record(True) + if on_item: + try: + await on_item(item, result) + except Exception: + pass + return item, result + except Exception as e: + logger.warning(f"Batch item failed: {item[:16]}... - {e}") + result = {"ok": False, "error": str(e)} + if task: + task.record(False, error=str(e)) + if on_item: + try: + await on_item(item, result) + except Exception: + pass + return item, result + + results: Dict[str, dict] = {} + + # 分批执行，避免一次性创建所有 task + for i in range(0, len(items), batch_size): + if (should_cancel and should_cancel()) or (task and task.cancelled): + break + chunk = items[i : i + batch_size] + pairs = await asyncio.gather(*(_one(x) for x in chunk)) + results.update(dict(pairs)) + + return results class BatchTask: @@ -150,3 +231,13 @@ def delete_task(task_id: str) -> None: async def expire_task(task_id: str, delay: int = 300) -> None: await asyncio.sleep(delay) delete_task(task_id) + + +__all__ = [ + "run_batch", + "BatchTask", + "create_task", + "get_task", + "delete_task", + "expire_task", +] diff --git a/app/services/grok/batch_services/assets.py b/app/services/grok/batch_services/assets.py index 5ba9ab3..ca87b3c 100644 --- a/app/services/grok/batch_services/assets.py +++ b/app/services/grok/batch_services/assets.py @@ -11,7 +11,7 @@ from app.core.config import get_config from app.core.logger import logger from app.services.reverse.assets_list import AssetsListReverse from app.services.reverse.assets_delete import AssetsDeleteReverse -from app.services.grok.utils.batch import run_in_batches +from app.core.batch import run_batch class BaseAssetsService: @@ -149,7 +149,7 @@ class ListService(BaseAssetsService): return {"detail": detail, "count": 0} try: - return await run_in_batches( + return await run_batch( tokens, _fetch_detail, max_concurrent=max_concurrent, @@ -219,7 +219,7 @@ class DeleteService(BaseAssetsService): return {"status": "error", "error": str(e)} try: - return await run_in_batches( + return await run_batch( tokens, _clear_one, max_concurrent=max_concurrent, diff --git a/app/services/grok/batch_services/nsfw.py b/app/services/grok/batch_services/nsfw.py index f53128b..facfbe5 100644 --- a/app/services/grok/batch_services/nsfw.py +++ b/app/services/grok/batch_services/nsfw.py @@ -12,7 +12,7 @@ from app.core.exceptions import UpstreamException from app.services.reverse.accept_tos import AcceptTosReverse from app.services.reverse.nsfw_mgmt import NsfwMgmtReverse from app.services.reverse.set_birth import SetBirthReverse -from app.services.grok.utils.batch import run_in_batches +from app.core.batch import run_batch class NSFWService: @@ -85,7 +85,7 @@ class NSFWService: logger.error(f"NSFW enable failed: {e}") return {"success": False, "http_status": 0, "error": str(e)[:100]} - return await run_in_batches( + return await run_batch( tokens, _enable, max_concurrent=max_concurrent, diff --git a/app/services/grok/batch_services/usage.py b/app/services/grok/batch_services/usage.py index ddefb74..9d10c98 100644 --- a/app/services/grok/batch_services/usage.py +++ b/app/services/grok/batch_services/usage.py @@ -10,7 +10,7 @@ from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config from app.services.reverse.rate_limits import RateLimitsReverse -from app.services.grok.utils.batch import run_in_batches +from app.core.batch import run_batch _USAGE_SEMAPHORE = asyncio.Semaphore(25) _USAGE_SEM_VALUE = 25 @@ -71,7 +71,7 @@ class UsageService: async def _refresh_one(t: str): return await mgr.sync_usage(t, consume_on_fail=False, is_usage=False) - return await run_in_batches( + return await run_batch( tokens, _refresh_one, max_concurrent=max_concurrent, diff --git a/app/services/grok/utils/batch.py b/app/services/grok/utils/batch.py deleted file mode 100644 index adb64ea..0000000 --- a/app/services/grok/utils/batch.py +++ /dev/null @@ -1,86 +0,0 @@ -""" -批量执行工具 - -提供分批并发、单项失败隔离的通用批量处理能力。 -""" - -import asyncio -from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar - -from app.core.logger import logger - -T = TypeVar("T") - - -async def run_in_batches( - items: List[str], - worker: Callable[[str], Awaitable[T]], - *, - max_concurrent: int = 10, - batch_size: int = 50, - on_item: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None, - should_cancel: Optional[Callable[[], bool]] = None, -) -> Dict[str, Dict[str, Any]]: - """ - 分批并发执行，单项失败不影响整体 - - Args: - items: 待处理项列表 - worker: 异步处理函数 - max_concurrent: 最大并发数 - batch_size: 每批大小 - - Returns: - {item: {"ok": bool, "data": ..., "error": ...}} - """ - try: - max_concurrent = int(max_concurrent) - except Exception: - max_concurrent = 10 - try: - batch_size = int(batch_size) - except Exception: - batch_size = 50 - - max_concurrent = max(1, max_concurrent) - batch_size = max(1, batch_size) - - sem = asyncio.Semaphore(max_concurrent) - - async def _one(item: str) -> tuple[str, dict]: - if should_cancel and should_cancel(): - return item, {"ok": False, "error": "cancelled", "cancelled": True} - async with sem: - try: - data = await worker(item) - result = {"ok": True, "data": data} - if on_item: - try: - await on_item(item, result) - except Exception: - pass - return item, result - except Exception as e: - logger.warning(f"Batch item failed: {item[:16]}... - {e}") - result = {"ok": False, "error": str(e)} - if on_item: - try: - await on_item(item, result) - except Exception: - pass - return item, result - - results: Dict[str, dict] = {} - - # 分批执行，避免一次性创建所有 task - for i in range(0, len(items), batch_size): - if should_cancel and should_cancel(): - break - chunk = items[i : i + batch_size] - pairs = await asyncio.gather(*(_one(x) for x in chunk)) - results.update(dict(pairs)) - - return results - - -__all__ = ["run_in_batches"]