geMoldInsight/src/moldinsight/core/occ_worker.py

# 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()}))