# core/occ_worker.py """OCC 常驻工作进程入口（方案 B，见 docs/topics/performance/OCC_THROUGHPUT.md）。 子进程经 multiprocessing spawn 拉起后运行 worker_main 消息循环：父进程经管道 下发 (op_name, payload)，本进程从 _OPS 注册表取 handler 执行并回传结果。 所有操作输入输出均为文件路径 + 可 pickle 的普通字典，**杜绝 pickle OCC 对象** （TopoDS_Shape 为 C++ 原生内存对象，跨进程传输的唯一干净方式是经文件中转）。 OCC 模块在 handler 内惰性导入：本模块在无 OCC 的父进程（pip 环境）导入不报错， 子进程（conda 环境）首次执行某操作时才加载对应模块——保持 processing_service 可无 OCC 导入（OCC 契约测试在无 OCC 环境 skip）。 """ import os import traceback from datetime import datetime from pathlib import Path from typing import Any, Callable, Dict # ─── 操作注册表 ───────────────────────────────────────────────── # 新增 OCC 操作必须在此注册：op_name -> 单 payload 字典的 handler。 # handler 内惰性 import OCC 相关核心模块。 _OPS: Dict[str, Callable[[Dict[str, Any]], Any]] = {} def _op(name: str): """注册操作到 _OPS。""" def deco(fn): _OPS[name] = fn return fn return deco # 进程内惰性单例缓存（OCC 实例化成本高，子进程常驻期间复用） _OCC_CACHE: Dict[str, Any] = {} def _cached(key: str, factory: Callable[[], Any]) -> Any: if key not in _OCC_CACHE: _OCC_CACHE[key] = factory() return _OCC_CACHE[key] def _get_parser(): from moldinsight.core.stp_parser import STPParser return STPParser() def _get_planner(): from moldinsight.core.multi_scheme_planner import MultiSchemeMoldPlanner return MultiSchemeMoldPlanner() def _get_geometry_analyzer(): from moldinsight.core.geometry_analyzer import GeometryAnalyzer return GeometryAnalyzer() def _get_mesh_generator(): from moldinsight.core.mesh_generator import MeshGenerator return MeshGenerator(quality="medium") def _get_side_action_designer(): from moldinsight.core.side_action_designer import SideActionDesigner return SideActionDesigner() def _get_cad_exporter(output_dir: str): from moldinsight.core.cad_exporter import CADExporter return CADExporter(output_dir=output_dir) # ─── 操作实现 ─────────────────────────────────────────────────── @_op("ping") def _op_ping(payload): return {"pong": True} @_op("sleep") def _op_sleep(payload): """诊断/测试用：在子进程内挂起指定秒数。""" import time time.sleep(float(payload.get("seconds", 0))) return {"slept": True} @_op("warmup") def _op_warmup(payload): """预热：子进程启动后触发，把 OCC 模块加载成本放到池启动而非首个分析任务。""" _cached("parser", _get_parser) _cached("planner", _get_planner) _cached("geometry_analyzer", _get_geometry_analyzer) _cached("mesh_generator", _get_mesh_generator) return {"warmed": True} @_op("parse_stp") def _op_parse_stp(payload): """STP 解析 + 几何分析（原主进程 load_step_file → analyze_geometry 两步合一）。 形状在子进程内创建并即刻消费，不跨进程传输。 """ parser = _cached("parser", _get_parser) shape = parser.load_step_file(Path(payload["stp_path"])) return parser.analyze_geometry(shape) @_op("generate_mesh") def _op_generate_mesh(payload): parser = _cached("parser", _get_parser) mesh_gen = _cached("mesh_generator", _get_mesh_generator) shape = parser.load_step_file(Path(payload["stp_path"])) return mesh_gen.generate_multi_lod_mesh(shape) @_op("generate_cavity") def _op_generate_cavity(payload): """多方案分模 + 方案形状 STEP 导出，全部在子进程内完成。 plan_result 里携带的 _export_shapes（TopoDS 对象）无法跨进程，子进程直接 经 CADExporter 落盘为持久化 STEP，返回文件 manifest——与旧 _persist_step_exports 产物结构一致，主进程原样存入 export_artifacts。 """ parser = _cached("parser", _get_parser) planner = _cached("planner", _get_planner) shape = parser.load_step_file(Path(payload["stp_path"])) plan_result = planner.generate_plan( shape=shape, material=payload["material"], is_foam_material=payload.get("is_foam_material", False), process_params=payload.get("process_params"), ) export_shapes = plan_result.pop("_export_shapes", {}) or {} export_manifest = _persist_export_shapes(payload, export_shapes) return {"plan_result": plan_result, "export_manifest": export_manifest} def _persist_export_shapes(payload: Dict[str, Any], export_shapes: Dict[str, Any]) -> Any: """在子进程内把各方案的 TopoDS 形状导出为持久化 STEP，返回 manifest。""" if not export_shapes: return None task_id = payload["task_id"] stp_path = payload["stp_path"] export_out_dir = payload.get("export_out_dir") if not export_out_dir: raise ValueError("generate_cavity 缺少 export_out_dir") # 每次按任务目录新建 exporter（output_dir 依任务固定；模块 import 由 importlib 缓存） exporter = _get_cad_exporter(export_out_dir) base_filename = Path(stp_path).stem or f"mold_{task_id}" components = ["cavity", "core", "parting_surface", "product", "a_plate", "b_plate"] manifest = { "version": 1, "task_id": task_id, "generated_at": datetime.now().isoformat(), "schemes": {}, } for scheme_id, cavity_data in export_shapes.items(): try: result = exporter.export_persisted_steps( cavity_data=cavity_data, base_filename=base_filename, components=components, task_id=task_id, scheme_id=scheme_id, ) manifest["schemes"][scheme_id] = { "base_filename": result.get("base_filename"), "generated_at": datetime.now().isoformat(), "files": result.get("files", []), "errors": result.get("errors", []), "total_files": result.get("total_files", 0), "total_errors": result.get("total_errors", 0), } except Exception as exc: logger = _get_logger() logger.warning("持久化 STEP 导出失败: task=%s scheme=%s error=%s", task_id, scheme_id, exc) manifest["schemes"][scheme_id] = { "base_filename": base_filename, "generated_at": datetime.now().isoformat(), "files": [], "errors": [str(exc)], "total_files": 0, "total_errors": 1, } return manifest def _get_logger(): from shared.utils.logger import get_logger return get_logger(__name__) @_op("analyze_mold_design") def _op_analyze_mold_design(payload): parser = _cached("parser", _get_parser) analyzer = _cached("geometry_analyzer", _get_geometry_analyzer) shape = parser.load_step_file(Path(payload["stp_path"])) return analyzer.analyze_mold_design( payload["geometry_data"], product_material=payload.get("product_material", "ABS"), shape=shape, ) @_op("detect_undercuts") def _op_detect_undercuts(payload): parser = _cached("parser", _get_parser) designer = _cached("side_action_designer", _get_side_action_designer) shape = parser.load_step_file(Path(payload["stp_path"])) return designer.analyze_and_design( shape, payload["parting_direction"], payload["mold_size"], ) @_op("convert_component_step") def _op_convert_component_step(payload): exporter = _get_cad_exporter(os.path.dirname(payload["step_path"])) return exporter.convert_component_step( payload["step_path"], payload["out_path"], payload["fmt"] ) # ─── 子进程主循环 ─────────────────────────────────────────────── def worker_main(conn) -> None: """OCC 子进程消息循环。conn 为 multiprocessing.Connection（双工）。 收到 None 或管道断开即退出；单条操作异常回传 error 不杀循环 （进程级隔离意味着一个操作 segfault 才会杀死本进程，由父进程补位）。 """ while True: try: msg = conn.recv() except (EOFError, KeyboardInterrupt, ConnectionResetError, OSError): break if msg is None: break op_name, payload = msg handler = _OPS.get(op_name) if handler is None: conn.send(("error", {"error": f"未知 OCC 操作: {op_name}"})) continue try: result = handler(payload) conn.send(("ok", result)) except BaseException as exc: # noqa: BLE001 子进程内兜底，保证循环存活 conn.send(("error", {"error": f"{type(exc).__name__}: {exc}", "traceback": traceback.format_exc()}))