vent_agent_2.0/extensions/main_fan/monthly_schedule.py

"""冻结月度清单，均分到多天；每天最多投递一份，重启仅续跑已投递的份额。"""
import asyncio
from datetime import datetime, timedelta
import hashlib
import json
from zoneinfo import ZoneInfo

import config
from extensions.main_fan import mine_catalog
from extensions.main_fan.batch_models import BatchRequest
from extensions.main_fan.exploration import _write
from extensions.main_fan.workflow import VERSION as WORKFLOW_VERSION
from obs import get_logger

logger = get_logger(__name__)
TIMEZONE = ZoneInfo('Asia/Shanghai')
FINISHED = {'completed', 'partial', 'failed'}


def configured_spread_days():
    days = config.get_int('MAIN_FAN_BATCH_SPREAD_DAYS', 1)
    if not 1 <= days <= 28:
        raise ValueError('MAIN_FAN_BATCH_SPREAD_DAYS 必须为1～28；1表示整批执行')
    return days


def _identifier(key):
    return hashlib.sha256(key.encode()).hexdigest()[:32]


class MonthlyBatchPlanner:
    def __init__(self, controller):
        self.controller = controller
        self.days = configured_spread_days()
        self.root = controller.root / 'monthly_plans'
        self.lock = asyncio.Lock()

    async def start(self, departments, *, now=None):
        """仅月度触发器创建计划；重复触发复用清单、分组和每日批次标识。"""
        now = (now or datetime.now(TIMEZONE)).astimezone(TIMEZONE)
        month = now.strftime('%Y-%m')
        key = 'main-fan-monthly:' + WORKFLOW_VERSION + ':' + month + ':' + ','.join(departments)
        plan_id = _identifier(key)
        path = self.root / (plan_id + '.json')
        async with self.lock:
            if self.controller.closed:
                raise RuntimeError('主通风机批次队列正在关闭')
            if path.is_file():
                plan = json.loads(path.read_text(encoding='utf-8'))
            else:
                self._check_all_allowed(departments)
                snapshot = await mine_catalog.refresh()
                available = {mine['dept_id']: mine for mine in snapshot['mines']}
                selected = list(available) if departments == ['all'] else list(departments)
                if len(set(selected)) != len(selected) or any(dept not in available for dept in selected):
                    raise ValueError('月度批次的部门编号重复或不在正常生产煤矿清单中')
                quotient, remainder = divmod(len(selected), self.days)
                scheduled_start = now.replace(second=0, microsecond=0)
                rows, offset = [], 0
                for index in range(self.days):
                    count = quotient + (index < remainder)
                    depts = selected[offset:offset + count]
                    offset += count
                    if self.controller.max_mines and count > self.controller.max_mines:
                        raise ValueError('月度计划的每日批次超过配置的矿井数量上限')
                    request = (BatchRequest(dept_ids=depts, snapshot_id=snapshot['snapshot_id'])
                               .model_dump(mode='json')) if depts else None
                    daily_key = 'main-fan-monthly-day:' + plan_id + ':' + str(index + 1)
                    rows.append({'day': index + 1, 'scheduled_at': (scheduled_start + timedelta(days=index)).isoformat(),
                        'dept_ids': depts, 'request': request, 'key': daily_key,
                        'batch_id': _identifier(daily_key), 'last_dispatched_on': None})
                plan = {'schema_version': 1, 'plan_id': plan_id, 'execution_mode': WORKFLOW_VERSION,
                    'month': month, 'created_at': now.isoformat(), 'departments': departments,
                    'snapshot_id': snapshot['snapshot_id'], 'spread_days': self.days,
                    'mine_count': len(selected), 'days': rows}
                _write(path, plan)
                logger.info('月度研判清单已冻结并分日', extra={'kv': {'event': 'batch.monthly.plan.created',
                    'plan_id': plan_id, 'month': month, 'mine_count': len(selected),
                    'spread_days': self.days, 'daily_mine_counts': [len(row['dept_ids']) for row in rows],
                    'plan_path': str(path)}})
            receipt = self._dispatch(now)
            return {'plan_id': plan_id, 'month': month, 'spread_days': plan['spread_days'],
                'mine_count': plan['mine_count'], 'plan_path': str(path), 'daily_batch': receipt}

    async def dispatch(self, *, now=None, resume_only=False):
        now = (now or datetime.now(TIMEZONE)).astimezone(TIMEZONE)
        async with self.lock:
            return self._dispatch(now, resume_only=resume_only)

    @staticmethod
    def _check_all_allowed(departments):
        if departments == ['all'] and not config.get_bool('MAIN_FAN_BATCH_ALLOW_ALL', False):
            raise ValueError('全矿研判未启用，请设置 MAIN_FAN_BATCH_ALLOW_ALL=True 或显式选择煤矿')

    def _receipt(self, row):
        try:
            job = self.controller.get(row['batch_id'])
        except FileNotFoundError:
            return None
        # batch.json 可能先于 job.json 完成落盘，避免此时重启被当作中断重跑。
        status = (job.get('batch') or {}).get('status')
        return {'batch_id': row['batch_id'], 'status': status if status in FINISHED else job['status'],
                'output': job['output']}

    def _dispatch(self, now, *, resume_only=False):
        if self.controller.closed:
            return None
        today = now.date().isoformat()
        rows = []
        for path in sorted(self.root.glob('*.json')):
            plan = json.loads(path.read_text(encoding='utf-8'))
            if plan['execution_mode'] == WORKFLOW_VERSION:
                rows.extend((path, plan, row) for row in plan['days'] if row['dept_ids'])
        rows.sort(key=lambda item: (item[1]['created_at'], item[2]['day']))

        # 不累积等待中的多日批次；先恢复已分配但未结束的那一份。
        for path, plan, row in rows:
            if row['last_dispatched_on']:
                receipt = self._receipt(row)
                if receipt is None or receipt['status'] not in FINISHED:
                    return self._submit(path, plan, row, today)
        # 同一天重复触发或重启，即使当日已完成，也不提前领取下一份。
        for _, _, row in rows:
            if row['last_dispatched_on'] == today:
                return self._receipt(row)
        if resume_only or self.controller.tasks:
            return None
        for path, plan, row in rows:
            if not row['last_dispatched_on'] and datetime.fromisoformat(row['scheduled_at']) <= now:
                return self._submit(path, plan, row, today)
        return None

    def _submit(self, path, plan, row, today):
        if self.controller.tasks and row['batch_id'] not in self.controller.tasks:
            return None
        self._check_all_allowed(plan['departments'])
        if self.controller.max_mines and len(row['dept_ids']) > self.controller.max_mines:
            raise ValueError('月度计划的每日批次超过配置的矿井数量上限')
        # 先保存日期，突然退出后仍通过相同 key 恢复，不会重复领取下一份。
        row['last_dispatched_on'] = today
        _write(path, plan)
        receipt = self.controller.submit(row['request'], key=row['key'], trigger='monthly_daily')
        logger.info('月度分日研判批次已入队', extra={'kv': {'event': 'batch.monthly.day.queued',
            'plan_id': plan['plan_id'], 'month': plan['month'], 'day': row['day'],
            'scheduled_at': row['scheduled_at'], 'dispatched_on': today,
            'mine_count': len(row['dept_ids']), 'batch_id': receipt['batch_id'], 'status': receipt['status']}})
        return receipt