"""启动批次、网页及月度调度共用串行队列；退出时取消并保存进度。""" 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()