vent_agent_2.0/extensions/main_fan/batch_service.py

"""网页及可选月度调度的本地批次队列；环境关闭时等待取消并保存进度。"""
import asyncio
from contextlib import asynccontextmanager
from datetime import datetime
import hashlib
import json
from pathlib import Path
from uuid import uuid4
from zoneinfo import ZoneInfo

import config
from extensions.main_fan.batch import BatchRunner
from extensions.main_fan.batch_models import BatchRequest
from extensions.main_fan.exploration import _write


class LocalBatchController:
    def __init__(self):
        output = Path(config.get_str('MAIN_FAN_BATCH_OUTPUT_DIR', 'outputs/main_fan_batches'))
        self.root = (output if output.is_absolute() else config.ROOT / output).resolve()
        source = Path(config.get_str('MAIN_FAN_MOCK_REPORT_PDF'))
        self.pdf = (source if source.is_absolute() else config.ROOT / source).resolve()
        if not self.pdf.is_file() or self.pdf.suffix.lower() != '.pdf':
            raise ValueError('启用本地批次需配置 MAIN_FAN_MOCK_REPORT_PDF 指向测试PDF')
        self.tasks = {}
        self.lock = asyncio.Semaphore(1)

    def submit(self, request, *, key=None):
        request = BatchRequest.model_validate(request)
        if request.all_producing_mines and not config.get_bool('MAIN_FAN_BATCH_ALLOW_ALL', False):
            raise ValueError('本地预览未开放全部煤矿执行，请显式选择一两个矿')
        if len(request.dept_ids) > config.get_int('MAIN_FAN_BATCH_MAX_MINES', 2):
            raise ValueError('本地批次超过配置的矿井数量上限')
        identifier = hashlib.sha256(key.encode()).hexdigest()[:32] if key else uuid4().hex
        folder = self.root / identifier
        job = folder / 'job.json'
        if job.is_file():
            original = json.loads(job.read_text())
            if original['request'] != request.model_dump(mode='json'):
                raise ValueError('同一批次标识的请求内容发生变化')
            state_file = folder / 'batch.json'
            if identifier in self.tasks or (state_file.is_file() and json.loads(state_file.read_text())['status'] == 'completed'):
                return {'batch_id': identifier, 'status': original['status'], 'output': str(folder)}
        _write(job, {'batch_id': identifier, 'status': 'queued', 'request': request.model_dump(mode='json'), 'publish': False})
        self.tasks[identifier] = asyncio.create_task(self.work(identifier, request), name='main-fan-batch-' + identifier)
        return {'batch_id': identifier, 'status': 'queued', 'output': str(folder)}

    async def work(self, identifier, request):
        folder = self.root / identifier
        job = json.loads((folder / 'job.json').read_text())
        try:
            async with self.lock:
                job['status'] = 'running'
                _write(folder / 'job.json', job)
                result = await BatchRunner(folder, self.pdf,
                    max_mines=config.get_int('MAIN_FAN_BATCH_MAX_MINES', 2)).run(request)
                job['status'] = result['status']
        except asyncio.CancelledError:
            job['status'] = 'interrupted'
            raise
        except Exception as exc:
            job.update(status='failed', error=type(exc).__name__ + ': ' + str(exc)[:500])
        finally:
            _write(folder / 'job.json', job)
            self.tasks.pop(identifier, None)

    def get(self, identifier):
        if len(identifier) != 32 or any(char not in '0123456789abcdef' for char in identifier):
            raise ValueError('无效批次标识')
        folder = self.root / identifier
        job = json.loads((folder / 'job.json').read_text())
        state = folder / 'batch.json'
        return {**job, 'output': str(folder), 'batch': json.loads(state.read_text()) if state.is_file() else None}

    async def close(self):
        pending = list(self.tasks.items())
        tasks = [task for _, task in pending]
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        # 尚未获得执行机会的Task在进入work前就可被取消，仍须保存明确状态。
        for identifier, _ in pending:
            path = self.root / identifier / 'job.json'
            job = json.loads(path.read_text())
            if job['status'] in {'queued', 'running'}:
                job['status'] = 'interrupted'
                _write(path, job)
        self.tasks.clear()


@asynccontextmanager
async def service(dependencies):
    controller = LocalBatchController()
    try:
        yield controller
    finally:
        await controller.close()


async def monthly_callback():
    from extensions import get_service
    depts = [value.strip() for value in config.get_str('MAIN_FAN_BATCH_DEPT_IDS').split(',') if value.strip()]
    request = BatchRequest(all_producing_mines=depts == ['all'], dept_ids=[] if depts == ['all'] else depts)
    month = datetime.now(ZoneInfo('Asia/Shanghai')).strftime('%Y-%m')
    key = 'main-fan-monthly:' + month + ':' + ','.join(depts)
    return get_service('main_fan_batches').submit(request, key=key)