252 lines
9.2 KiB
Python
252 lines
9.2 KiB
Python
|
# 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()}))
|