"""冻结月度清单，均分到多天；每天最多投递一份，重启仅续跑已投递的份额。""" 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