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
from extensions.main_fan import mine_catalog
from extensions.main_fan.monthly_schedule import MonthlyBatchPlanner
from extensions.main_fan.workflow import VERSION as WORKFLOW_VERSION, publish_enabled
from obs import get_logger

logger = get_logger(__name__)


def configured_departments():
    depts = [value.strip() for value in config.get_str('MAIN_FAN_BATCH_DEPT_IDS', 'all').split(',') if value.strip()]
    if not depts or ('all' in depts and depts != ['all']):
        raise ValueError('MAIN_FAN_BATCH_DEPT_IDS 必须为 all 或逗号分隔的部门编号')
    return depts


def configured_startup_exclusions():
    name = 'MAIN_FAN_BATCH_STARTUP_EXCLUDE_DEPT_IDS'
    try:
        values = json.loads(config.get_str(name, '[]').strip() or '[]')
    except json.JSONDecodeError as exc:
        raise ValueError(name + ' 必须为 JSON 数组，例如 [1045,1106]') from exc
    if not isinstance(values, list):
        raise ValueError(name + ' 必须为 JSON 数组，例如 [1045,1106]')
    excluded = set()
    for value in values:
        if type(value) not in (int, str):
            raise ValueError(name + ' 的每项必须是数字部门编号或其字符串')
        dept = str(value).strip()
        if not dept.isascii() or not dept.isdigit() or len(dept) > 64 or int(dept) == 0:
            raise ValueError(name + ' 的每项必须是非零、最多64位的数字部门编号')
        excluded.add(dept)
    return excluded


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()
        self.pdf = None
        if config.get_str('MAIN_FAN_REPORT_MODE', 'crawler').lower() == 'mock':
            source = Path(config.get_str('MAIN_FAN_MOCK_REPORT_PDF'))
            self.pdf = (source if source.is_absolute() else config.ROOT / source).resolve()
        from extensions.main_fan.crawler_reports import report_source
        report_source(self.pdf)  # 启动时校验配置，不发请求。
        self.max_mines = config.get_int('MAIN_FAN_BATCH_MAX_MINES', 2)
        if self.max_mines < 0:
            raise ValueError('MAIN_FAN_BATCH_MAX_MINES 不能小于0；0表示不限制批次矿数')
        self.startup_enabled = config.get_bool('MAIN_FAN_BATCH_RUN_ON_STARTUP', True)
        self.startup_max_mines = config.get_int('MAIN_FAN_BATCH_STARTUP_MAX_MINES', 2)
        if self.startup_enabled and self.startup_max_mines <= 0:
            raise ValueError('MAIN_FAN_BATCH_STARTUP_MAX_MINES 必须大于0')
        self.tasks = {}
        self.lock = asyncio.Semaphore(1)
        self.startup_task = None
        self.monthly_task = None
        self.closed = False
        self.monthly = MonthlyBatchPlanner(self)

    def submit(self, request, *, key=None, trigger='manual'):
        if self.closed:
            raise RuntimeError('主通风机批次队列正在关闭')
        request = BatchRequest.model_validate(request)
        if request.all_producing_mines and not config.get_bool('MAIN_FAN_BATCH_ALLOW_ALL', False):
            raise ValueError('全矿研判未启用，请设置 MAIN_FAN_BATCH_ALLOW_ALL=True 或显式选择煤矿')
        if self.max_mines and len(request.dept_ids) > self.max_mines:
            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': publish_enabled(), 'trigger': trigger})
        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=self.max_mines or None).run(request)
                job['status'] = result['status']
        except asyncio.CancelledError:
            job['status'] = 'interrupted'
            raise
        except Exception as exc:
            logger.exception('批次队列执行失败', extra={'kv': {'event': 'batch.queue.failed', 'batch_id': identifier}})
            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}

    def start_on_startup(self):
        """服务和任务执行器就绪后调用；每个进程生命周期只启动一批，不阻塞就绪。"""
        if (not self.closed and self.monthly.days > 1 and config.get_str('MAIN_FAN_BATCH_CRON')
                and self.monthly_task is None):
            self.monthly_task = asyncio.create_task(self._resume_monthly(), name='main-fan-monthly-resume')
        if not self.closed and self.startup_enabled and self.startup_task is None:
            self.startup_task = asyncio.create_task(self._startup_once(), name='main-fan-startup-selection')

    async def _resume_monthly(self):
        try:
            await self.monthly.dispatch(resume_only=True)
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception('恢复月度分日批次失败，日调度仍保留',
                extra={'kv': {'event': 'batch.monthly.resume.failed'}})

    async def _startup_once(self):
        try:
            depts = configured_departments()
            excluded = configured_startup_exclusions()
            snapshot = await mine_catalog.refresh()
            available = {mine['dept_id']: mine for mine in snapshot['mines']}
            if depts == ['all']:
                eligible = list(available)
            else:
                if len(set(depts)) != len(depts) or any(dept not in available for dept in depts):
                    raise ValueError('启动批次的部门编号重复或不在正常生产煤矿清单中')
                eligible = depts
            skipped = [dept for dept in eligible if dept in excluded]
            eligible = [dept for dept in eligible if dept not in excluded]
            if excluded:
                logger.info('启动批次已应用部门排除列表', extra={'kv': {
                    'event': 'batch.startup.filtered', 'configured_excluded_dept_ids': sorted(excluded),
                    'excluded_dept_ids': skipped, 'excluded_count': len(skipped), 'eligible_count': len(eligible)}})
            limit = min(self.startup_max_mines, self.max_mines) if self.max_mines else self.startup_max_mines
            selected = eligible[:limit]
            if not selected:
                logger.info('没有可研判的正常生产煤矿，跳过启动批次',
                    extra={'kv': {'event': 'batch.startup.empty'}})
                return
            request = BatchRequest(dept_ids=selected, snapshot_id=snapshot['snapshot_id'])
            receipt = self.submit(request, trigger='startup')
            logger.info('启动研判批次已入队', extra={'kv': {'event': 'batch.startup.queued',
                'batch_id': receipt['batch_id'], 'selected_dept_ids': selected, 'mine_count': len(selected),
                'excluded_dept_ids': skipped,
                'eligible_count': len(eligible), 'startup_max_mines': self.startup_max_mines,
                'effective_limit': limit}})
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception('启动研判批次创建失败，月度定时仍保留',
                extra={'kv': {'event': 'batch.startup.failed'}})

    async def close(self):
        self.closed = True
        selections = [task for task in (self.startup_task, self.monthly_task) if task is not None]
        for task in selections:
            task.cancel()
        if selections:
            await asyncio.gather(*selections, return_exceptions=True)
        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 = configured_departments()
    controller = get_service('main_fan_batches')
    if controller.monthly.days > 1:
        return await controller.monthly.start(depts)
    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:' + WORKFLOW_VERSION + ':' + month + ':' + ','.join(depts)
    receipt = controller.submit(request, key=key, trigger='monthly')
    logger.info('月度研判批次已入队', extra={'kv': {'event': 'batch.monthly.queued',
        'batch_id': receipt['batch_id'], 'month': month, 'departments': depts}})
    return receipt


async def daily_callback():
    from extensions import get_service
    return await get_service('main_fan_batches').monthly.dispatch()