优化

src/moldinsight/services/processing_service.py

@@ -3,14 +3,16 @@
 
 import asyncio
 import os
+import shutil
+import tempfile
 import time
-import traceback
 from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from pathlib import Path
-from typing import Optional, Dict, Any, List
+from typing import Optional, Dict, Any, List, Tuple
 
+from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from moldinsight.core.stp_parser import STPParser
@@ -19,11 +21,13 @@ from moldinsight.core.mesh_generator import MeshGenerator
 from moldinsight.core.multi_scheme_planner import MultiSchemeMoldPlanner
 from moldinsight.core.cad_exporter import CADExporter
 from moldinsight.services.storage_integration_rustfs import StorageIntegrationService
+from moldinsight.storage.rustfs_storage import rustfs_manager
 from shared.services.redis_task_manager import redis_task_manager
 from moldinsight.services.material_service import MaterialService
 from moldinsight.services.calculation_service import CalculationService
 from moldinsight.services.llm_service import llm_service
 from shared.models.schemas import ProcessingStatus
+from shared.models.database import STPFile
 from shared.database.database import db_manager
 from shared.utils.html_generator import HTMLGenerator
 from shared.utils.logger import get_logger
@@ -71,23 +75,73 @@ class ProcessingService:
         loop = asyncio.get_running_loop()
         return await loop.run_in_executor(self._occ_executor, fn, *args)
 
+    async def _materialize_source_file(self, stp_file: STPFile) -> Tuple[Path, Optional[Path]]:
+        """把待处理文件落到本地磁盘，返回 (本地路径, 临时目录或 None)。
+
+        RustFS 为主存储：处理方按 object_key 下载到任务专属临时目录
+        （文件名保留原始名——下游产物命名依赖 Path(file_path).name）。
+        RustFS 不可用或对象缺失时回退 STPFile.file_path 记录的节点本地路径
+        （依赖 compose 共享卷，属过渡方案）；两者皆不可用则抛错置任务失败。
+        """
+        original_name = Path(stp_file.original_filename or "model.stp").name or "model.stp"
+
+        if stp_file.object_key:
+            temp_dir = Path(tempfile.mkdtemp(prefix=f"moldinsight_{stp_file.id}_"))
+            try:
+                data = await rustfs_manager.download_file(
+                    file_type="stp_files", object_key=stp_file.object_key
+                )
+                target = temp_dir / original_name
+                target.write_bytes(data)
+                return target, temp_dir
+            except Exception as exc:
+                shutil.rmtree(temp_dir, ignore_errors=True)
+                logger.warning(
+                    f"RustFS 源文件下载失败 (object_key={stp_file.object_key})，"
+                    f"回退节点本地路径: {exc}"
+                )
+
+        local = Path(stp_file.file_path) if stp_file.file_path else None
+        if local and local.exists():
+            return local, None
+
+        raise RuntimeError(
+            f"源文件不可用：RustFS 对象 {stp_file.object_key!r} 下载失败，"
+            f"且节点本地路径不存在: {stp_file.file_path!r}"
+        )
+
     async def process_file_with_storage(
         self,
         task_id: str,
-        file_path: str,
         stp_file_id: int,
         process_params: Optional[Dict[str, Any]] = None,
     ):
-        """处理文件的后台任务 — 使用独立数据库会话"""
+        """处理文件的后台任务 — 使用独立数据库会话
+
+        分派入参只带 stp_file_id（D6）：源文件由本方法按 PG 元数据中的
+        object_key 从 RustFS 获取，不再依赖分派方传入节点本地路径
+        （API 与 Celery worker 容器文件系统不互通）。
+        """
 
         # 创建独立的数据库会话，避免请求范围会话关闭
         async with db_manager.session() as db_session:
+            temp_dir: Optional[Path] = None
             try:
+                result = await db_session.execute(
+                    select(STPFile).where(STPFile.id == stp_file_id)
+                )
+                stp_file = result.scalar_one_or_none()
+                if not stp_file:
+                    raise RuntimeError(f"STPFile 记录不存在: stp_file_id={stp_file_id}")
+
+                source_path, temp_dir = await self._materialize_source_file(stp_file)
+                file_path = str(source_path)
+
                 logger.info(f"开始处理文件并生成模具型腔: {file_path}")
 
                 from shared.config.settings import settings
 
-                file_size_bytes = Path(file_path).stat().st_size if Path(file_path).exists() else 0
+                file_size_bytes = source_path.stat().st_size
                 file_size_mb = max(file_size_bytes / (1024 * 1024), 1)
                 timeout_seconds = min(
                     max(settings.PROCESSING_TIMEOUT_BASE, int(file_size_mb * settings.PROCESSING_TIMEOUT_PER_MB)),
@@ -111,12 +165,16 @@ class ProcessingService:
             except Exception as e:
                 logger.error(f"模具型腔生成失败: {e}")
 
+                # D9：先丢弃未提交的数据本体，失败状态单独提交，
+                # 避免 failed 更新把半成品 flush 数据一起带上
+                await db_session.rollback()
+
                 await self.storage_service.update_stp_file_status(db_session, stp_file_id, "failed")
                 await self.storage_service.update_task_status(
                     db_session, task_id, "failed", error_message=str(e)
                 )
 
-                # 安全更新 Redis 任务状态
+                # 安全更新 Redis 任务状态（Redis 仅热缓存，写失败不影响 PG 事实）
                 task = await redis_task_manager.get_task(task_id)
                 if task:
                     await redis_task_manager.update_task(task_id, {
@@ -124,6 +182,10 @@ class ProcessingService:
                         "error": str(e),
                         "completed_at": str(datetime.now()),
                     })
+            finally:
+                # 任务专属临时目录必须清理，长期运行不允许残留下载副本
+                if temp_dir:
+                    shutil.rmtree(temp_dir, ignore_errors=True)
 
     async def process_file_core(
         self,
@@ -234,6 +296,10 @@ class ProcessingService:
                 geometry_data.get("analysis_method", "mold_cavity"),
             )
 
+            # 阶段 A 提交（D9）：几何 + 网格原子落库——解析后的确定成果，
+            # 后续型腔失败任务标 failed 时这些数据仍完整保留
+            await db_session.commit()
+
             # 7. 生成HTML可视化
             await self.storage_service.update_task_status(
                 db_session, task_id, "processing", 85, "生成可视化报告"
@@ -325,6 +391,11 @@ class ProcessingService:
                 ),
             )
 
+            # 9.65 阶段 B 提交（D9）：型腔 / HTML / 特征 / 指标 / 摘要 / 验证指标
+            # 作为完整结果包原子落库——置 completed 前必须全部就位，
+            # 期间任一步失败回滚后任务标 failed，不会出现"completed 但数据残缺"
+            await db_session.commit()
+
             # 9.7 FreeCAD 几何验证
             stage_started = time.perf_counter()
             verification_result = await self._step_verify(
@@ -347,11 +418,7 @@ class ProcessingService:
                 )
             stage_timings["generate_llm_report"] = round(time.perf_counter() - stage_started, 3)
 
-            # 10. 完成处理
-            await self.storage_service.update_stp_file_status(db_session, stp_file_id, "completed")
-            await self.storage_service.update_task_status(
-                db_session, task_id, "completed", 100, "模具型腔生成完成"
-            )
+            # 10. 完成处理——先 flush 任务参数，完成状态提交时一并原子落库（D9）
             await self.storage_service.update_task_parameters(
                 db_session,
                 task_id,
@@ -364,6 +431,10 @@ class ProcessingService:
                     **process_params,
                 },
             )
+            await self.storage_service.update_stp_file_status(db_session, stp_file_id, "completed")
+            await self.storage_service.update_task_status(
+                db_session, task_id, "completed", 100, "模具型腔生成完成"
+            )
 
             # 更新任务缓存状态（仅保留轻量摘要，完整数据由PG+RustFS持久化；
             # 完成态视图由 TaskQueryService 从 PG+RustFS 组装，Redis 不再存
@@ -390,6 +461,9 @@ class ProcessingService:
         except Exception as e:
             logger.error(f"模具型腔生成失败: {e}")
 
+            # D9：先丢弃未提交的数据本体再置失败（同外层说明）
+            await db_session.rollback()
+
             await self.storage_service.update_stp_file_status(db_session, stp_file_id, "failed")
             await self.storage_service.update_task_status(
                 db_session, task_id, "failed", error_message=str(e)
@@ -471,29 +545,29 @@ class ProcessingService:
 
     async def _step_generate_cavity(
         self, shape, selected_material: dict, is_foam_material: bool, process_params: Dict[str, Any],
-    ) -> Optional[Dict[str, Any]]:
-        """生成多方案分模结果"""
-        plan_result = None
-        try:
-            if shape:
-                loop = asyncio.get_running_loop()
-                plan_result = await loop.run_in_executor(
-                    self._occ_executor,
-                    lambda: self.multi_scheme_planner.generate_plan(
-                        shape=shape,
-                        material=selected_material,
-                        is_foam_material=is_foam_material,
-                        process_params=process_params,
-                    ),
-                )
-                logger.info(
-                    f"多方案分模完成: 生成 {len(plan_result.get('candidate_schemes', []))} 套方案"
-                )
-        except Exception as cavity_err:
-            logger.warning(f"多方案分模失败，使用简化数据: {cavity_err}")
-            traceback.print_exc()
-            plan_result = None
+    ) -> Dict[str, Any]:
+        """生成多方案分模结果。
 
+        D8：型腔是任务的核心产出，生成失败必须让任务 failed——
+        此前异常在此被吞掉置 plan_result=None 继续主流程，最终任务
+        completed，"完成"状态不可信。异常直接向编排层传播。
+        """
+        if not shape:
+            raise RuntimeError("无有效几何 shape，无法生成模具型腔")
+
+        loop = asyncio.get_running_loop()
+        plan_result = await loop.run_in_executor(
+            self._occ_executor,
+            lambda: self.multi_scheme_planner.generate_plan(
+                shape=shape,
+                material=selected_material,
+                is_foam_material=is_foam_material,
+                process_params=process_params,
+            ),
+        )
+        logger.info(
+            f"多方案分模完成: 生成 {len(plan_result.get('candidate_schemes', []))} 套方案"
+        )
         return plan_result
 
     def _cache_export_shapes(self, task_id: str, export_shapes: Dict[str, Dict[str, Any]]):
@@ -727,7 +801,8 @@ class ProcessingService:
         )
 
         session.add(metrics)
-        await session.commit()
+        # D9：flush 不 commit，随结果包（阶段 B）由编排层统一提交
+        await session.flush()
         logger.info(f"分析指标保存成功: {metrics.id}")
 
     async def _save_verification_metrics(self, session: AsyncSession, stp_file_id: int, verification_result: dict):
@@ -759,7 +834,8 @@ class ProcessingService:
             )
             session.add(metrics)
 
-        await session.commit()
+        # D9：flush 不 commit，随结果包（阶段 B）由编排层统一提交
+        await session.flush()
         logger.info(f"验证指标保存成功: stp_file_id={stp_file_id}")