From 94fed47a754d64614e57eaeec3377fa256361700 Mon Sep 17 00:00:00 2001 From: lensen Date: Mon, 29 Jun 2026 21:36:10 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=20rerank=20=E5=93=8D?= =?UTF-8?q?=E5=BA=94=E8=A7=A3=E6=9E=90=EF=BC=8C=E6=94=AF=E6=8C=81=E5=A4=9A?= =?UTF-8?q?=E7=A7=8D=E6=95=B0=E6=8D=AE=E7=B1=BB=E5=9E=8B=E5=B9=B6=E6=B7=BB?= =?UTF-8?q?=E5=8A=A0=E8=AF=B7=E6=B1=82=E6=97=B6=E9=97=B4=E7=BB=9F=E8=AE=A1?= =?UTF-8?q?=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pipeline/modules/search/atomic.py | 15 +- pipeline/modules/search/multi.py | 16 +- pipeline/modules/search/multi_vector.py | 204 +++++++++++++++++++++--- scripts/run_search.py | 25 +++ scripts/run_search_benchmark.py | 122 +++++++++++++- 5 files changed, 345 insertions(+), 37 deletions(-) diff --git a/pipeline/modules/search/atomic.py b/pipeline/modules/search/atomic.py index b9838b3..688710a 100644 --- a/pipeline/modules/search/atomic.py +++ b/pipeline/modules/search/atomic.py @@ -660,14 +660,18 @@ def _correct_rerank_line( def _parse_rerank_response( self, - useful_relations: List[str], + useful_relations: List[Any], valid_ids: set, relation_ids: List[str], relation_texts: List[str], ) -> List[str]: - """解析 LLM 返回的 useful_relations，提取 [id] 并纠错""" + """解析 LLM 返回的 useful_relations，提取纯数字或 [id] 并纠错""" selected: List[str] = [] for line in useful_relations: + line = str(line).strip() + if line in valid_ids and line not in selected: + selected.append(line) + continue if "[" not in line or "]" not in line: continue rel_id = line[line.find("[") + 1: line.find("]")].strip() @@ -749,7 +753,12 @@ async def step7_llm_rerank( "thought_process": {"type": "string"}, "useful_relations": { "type": "array", - "items": {"type": "string"}, + "items": { + "anyOf": [ + {"type": "string"}, + {"type": "integer"}, + ], + }, }, }, "required": ["thought_process", "useful_relations"], diff --git a/pipeline/modules/search/multi.py b/pipeline/modules/search/multi.py index 2ff46ba..f6afa81 100644 --- a/pipeline/modules/search/multi.py +++ b/pipeline/modules/search/multi.py @@ -592,15 +592,18 @@ def _correct_rerank_line( def _parse_rerank_response( self, - useful_relations: List[str], + useful_relations: List[Any], valid_ids: set, relation_ids: List[str], relation_texts: List[str], ) -> List[str]: - """解析 LLM 返回的 useful_relations，提取 [id] 并纠错""" + """解析 LLM 返回的 useful_relations，提取纯数字或 [id] 并纠错""" selected: List[str] = [] for line in useful_relations: - line = str(line) + line = str(line).strip() + if line in valid_ids and line not in selected: + selected.append(line) + continue if "[" not in line or "]" not in line: continue rel_id = line[line.find("[") + 1: line.find("]")].strip() @@ -672,7 +675,12 @@ async def step7_llm_rerank( "thought_process": {"type": "string"}, "useful_relations": { "type": "array", - "items": {"type": "string"}, + "items": { + "anyOf": [ + {"type": "string"}, + {"type": "integer"}, + ], + }, }, }, "required": ["thought_process", "useful_relations"], diff --git a/pipeline/modules/search/multi_vector.py b/pipeline/modules/search/multi_vector.py index 02bb18b..62617a2 100644 --- a/pipeline/modules/search/multi_vector.py +++ b/pipeline/modules/search/multi_vector.py @@ -45,6 +45,36 @@ PRECISE_ENTITY_EVENT_TOP_K = 40 +FAST_STAGE_TIMING_KEYS = [ + ("step1_entity_recall", "S1 entity", ("step1_entity_bm25",)), + ("step2_query_embedding", "S2 embed", ("step2_embedding",)), + ("step3_seed_recall", "S3 seed", ("step3_fast_recall",)), + ( + "step4_expand_and_rank", + "S4 expand+rank", + ( + "step4_fast_event_entities", + "step5_fast_expand", + "step6_fast_expand_rank", + "step6_fast_final_rank", + ), + ), + ("step5_chunk_fetch", "S5 chunk", ("step8_chunks",)), +] + +PRECISE_STAGE_TIMING_KEYS = [ + ("step1_entity_recall", "S1 entity", ("step1_entity_bm25",)), + ("step2_query_embedding", "S2 embed", ("step2_embedding",)), + ("step3_dual_recall", "S3 recall", ("step3_dual_recall",)), + ( + "step4_expand_and_coarse_rank", + "S4 expand+rank", + ("step4_event_entities", "step5_expand", "step6_coarse_rank"), + ), + ("step5_llm_filter", "S5 llm", ("step7_llm_filter",)), + ("step6_chunk_fetch", "S6 chunk", ("step8_chunks", "native_chunk_total")), +] + # ── LLM 过滤提示词（本地版本：去掉 thought_process，只要求返回 ID 列表）── _RERANK_SYSTEM_PROMPT_LOCAL = """I will provide you with a set of relationship descriptions from a knowledge graph. \ @@ -197,6 +227,74 @@ async def warmup(self, config: Optional[MultiConfig] = None) -> None: await self._get_llm_client() logger.info(f"[multi_vector] mode={mode}, entity_recall=bm25, ranking={ranking_strategy}") + def _build_stage_timings( + self, + mode: str, + timings: Dict[str, float], + ) -> Dict[str, float]: + stage_definitions = ( + FAST_STAGE_TIMING_KEYS if mode == "fast" else PRECISE_STAGE_TIMING_KEYS + ) + stage_timings: Dict[str, float] = {} + for key, _, timing_keys in stage_definitions: + stage_timings[key] = sum( + float(timings.get(timing_key, 0.0) or 0.0) + for timing_key in timing_keys + ) + return stage_timings + + def _format_stage_timings( + self, + mode: str, + stage_timings: Dict[str, float], + ) -> str: + stage_definitions = ( + FAST_STAGE_TIMING_KEYS if mode == "fast" else PRECISE_STAGE_TIMING_KEYS + ) + return ", ".join( + f"{label}={stage_timings.get(key, 0.0):.3f}s" + for key, label, _ in stage_definitions + ) + + def _format_timing_counts(self, stats: Dict[str, Any]) -> str: + count_labels = { + "entities": "entities", + "initial_events": "initial", + "expanded": "expanded", + "candidates": "candidates", + "seed": "seed", + "final_events": "events", + "selected": "selected", + "chunks": "chunks", + "final_chunks": "chunks", + "native_added": "native", + } + parts = [] + for key, label in count_labels.items(): + value = stats.get(key) + if value is not None: + parts.append(f"{label}={value}") + return " ".join(parts) + + def _log_search_summary( + self, + *, + mode: str, + stats: Dict[str, Any], + stage_timings: Dict[str, float], + total_time: float, + ) -> None: + timing_text = self._format_stage_timings(mode, stage_timings) + counts_text = self._format_timing_counts(stats) + sections = stats.get("sections") + result_text = f" sections={sections}" if sections is not None else "" + counts_line = f"\n counts: {counts_text}" if counts_text else "" + logger.info( + f"[multi_es.timing] mode={mode} total={total_time:.3f}s{result_text}" + f"{counts_line}\n" + f" steps: {timing_text}" + ) + # ── Step1: BM25 实体召回 ──────────────────────────────────── async def step1_retrieve_entities_bm25( @@ -243,11 +341,11 @@ async def step1_retrieve_entities_bm25( f"{rank}:{eid} name={name!r} score={score:.4f} kept={kept}" ) - logger.info( + logger.debug( f"[entity.bm25] input=1, candidates={len(hits)}, output={len(entity_ids)}, " f"top_k={entity_top_k}" ) - logger.info( + logger.debug( f"[entity.bm25.entities] details={'; '.join(details_for_log)}" ) return entity_ids @@ -353,7 +451,7 @@ async def step3_retrieve_events( items = [{"event_id": eid, "score": score} for eid, score in merged.items()] - logger.info( + logger.debug( f"[event.recall] entity_to_event={db_count}/{entity_event_top_k}, " f"query_to_event={es_new_count}/{multi_top_k}, " f"merged={len(items)}" @@ -571,7 +669,7 @@ async def step3_fast_recall( seed_items = scored[: config.fast_answer_k] timings["step3_seed_score"] = time.perf_counter() - t0 - logger.info( + logger.debug( f"[event.seed] entities={len(seed_entity_ids)}/{len(entity_ids)}, " f"entity_candidates={len(candidate_event_ids)}, " f"event1={len(event1)}, event2={len(event2)}, " @@ -638,7 +736,7 @@ async def step4_fetch_events( if "entity_ids" in includes: extra = f", event_entity_relations={event_entity_relations}" - logger.info( + logger.debug( f"[event.fetch] label={log_label}, input_event_ids={len(event_ids)}, " f"found_events={len(event_map)}{extra}" ) @@ -673,7 +771,7 @@ def get_new_entity_ids( total += 1 if eid not in state.entity_ids: new_ids.append(eid) - logger.info( + logger.debug( f"[entity.dedupe] total={total}, " f"already_tracked={total - len(new_ids)}, " f"new={len(new_ids)}" @@ -740,7 +838,7 @@ async def step5_expand( ) if not new_entity_ids: - logger.info( + logger.debug( f"[event.expand] hop={hop+1}/{max_hops} " f"no_new_entities tracked_entities={len(state.entity_ids)}" ) @@ -755,7 +853,7 @@ async def step5_expand( + time.perf_counter() - t_step ) - logger.info( + logger.debug( f"[event.expand] hop={hop+1}/{max_hops} " f"entities: {pre_entities} -> +{len(new_entity_ids)} new, " f"total={len(state.entity_ids)}" @@ -779,7 +877,7 @@ async def step5_expand( new_event_ids = all_new_event_ids if not new_event_ids: - logger.info( + logger.debug( f"[event.expand] hop={hop+1}/{max_hops} " f"no_new_events tracked_events={len(state.relation_ids)}" ) @@ -798,7 +896,7 @@ async def step5_expand( + time.perf_counter() - t_step ) - logger.info( + logger.debug( f"[event.expand] hop={hop+1}/{max_hops} done: " f"events {pre_events} -> {len(state.relation_ids)} (+{len(new_event_ids)}), " f"entities {pre_entities} -> {len(state.entity_ids)}, " @@ -837,7 +935,7 @@ async def step5_expand( + time.perf_counter() - t_step ) - logger.info( + logger.debug( f"[event.expand] hop={hop+1}/{max_hops} done: " f"events {pre_events} -> {len(state.relation_ids)} (+{len(new_event_ids)}), " f"entities {pre_entities} -> {len(state.entity_ids)}, " @@ -891,7 +989,7 @@ async def step6_coarse_rank( scored.append({"event_id": eid, "score": score}) top_score_str = f"{scored[0]['score']:.4f}" if scored else "0" - logger.info( + logger.debug( f"[event.rank.coarse] input={len(event_ids)}, " f"returned={len(scored)}, " f"top_score={top_score_str}" @@ -903,7 +1001,7 @@ async def step6_coarse_rank( def _parse_llm_filter_response( self, - useful_relations: List[str], + useful_relations: List[Any], valid_ids: set, ) -> List[str]: """解析 LLM 返回的 useful_relations，去重 + 校验。 @@ -992,7 +1090,12 @@ async def step7_llm_filter( "properties": { "useful_relations": { "type": "array", - "items": {"type": "string"}, + "items": { + "anyOf": [ + {"type": "string"}, + {"type": "integer"}, + ], + }, }, }, "required": ["useful_relations"], @@ -1001,7 +1104,7 @@ async def step7_llm_filter( ) # 打印完整输出 - logger.info(f"[event.filter.llm.raw] raw_response={response}") + logger.debug(f"[event.filter.llm.raw] raw_response={response}") # 4. 解析 + 去重校验 useful_relations = response.get("useful_relations", []) @@ -1081,7 +1184,7 @@ async def step8_fetch_chunks( if chunk_id in chunk_map: result_map[eid] = chunk_map[chunk_id] - logger.info( + logger.debug( f"[chunk.fetch] events={len(event_ids)} -> " f"chunk_ids={len(chunk_ids)}, matched={len(result_map)}" ) @@ -1125,6 +1228,15 @@ async def search_fast( "items": [], "_timings": timings, "_query_vector": query_vector, + "stage_timings_seconds": self._build_stage_timings("fast", timings), + "_stats": { + "mode": "fast", + "entities": len(entity_ids), + "seed": 0, + "expanded": 0, + "final_events": 0, + "final_chunks": 0, + }, } t0 = time.perf_counter() @@ -1178,7 +1290,7 @@ async def search_fast( ] min_output_score = min(output_scores) if output_scores else 0.0 max_output_score = max(output_scores) if output_scores else 0.0 - logger.info( + logger.debug( f"[fast.rank] input={rank_input_count}, output={len(final_items)}, " f"score_min={min_output_score:.4f}, score_max={max_output_score:.4f}" ) @@ -1204,7 +1316,7 @@ async def search_fast( deduped.append(item) timings["total"] = time.perf_counter() - t_total - logger.info( + logger.debug( f"[fast.done] seed={len(seed_items)}, expanded_selected={len(expanded_items)}, " f"final_events={len(final_items)}, final_chunks={len(deduped)}" ) @@ -1212,6 +1324,15 @@ async def search_fast( "items": deduped, "_timings": timings, "_query_vector": query_vector, + "stage_timings_seconds": self._build_stage_timings("fast", timings), + "_stats": { + "mode": "fast", + "entities": len(entity_ids), + "seed": len(seed_items), + "expanded": len(expanded_items), + "final_events": len(final_items), + "final_chunks": len(deduped), + }, } # ── 主搜索接口 ────────────────────────────────────────────── @@ -1254,7 +1375,7 @@ async def search( else None ) - logger.info( + logger.debug( f"[multi_es.start] mode={config.mode}, ranking={ranking_strategy}, " f"entity_top_k={config.entity_top_k}, multi_top_k={config.multi_top_k}, " f"entity_event_top_k={precise_entity_event_top_k or config.fast_entity_event_candidate_k}" @@ -1310,6 +1431,16 @@ async def search( "items": [], "_timings": timings, "_query_vector": query_vector, + "stage_timings_seconds": self._build_stage_timings("precise", timings), + "_stats": { + "mode": "precise", + "entities": len(entity_ids), + "initial_events": 0, + "expanded": 0, + "candidates": 0, + "selected": 0, + "chunks": 0, + }, } t0 = time.perf_counter() @@ -1357,7 +1488,7 @@ async def search( ] min_candidate_score = min(candidate_scores) if candidate_scores else 0.0 max_candidate_score = max(candidate_scores) if candidate_scores else 0.0 - logger.info( + logger.debug( f"[event.candidates] initial={len(event_ids)}, expanded={len(expanded_event_ids)}, " f"input={len(all_event_ids)}, output={len(candidate_items)}, " f"score_min={min_candidate_score:.4f}, score_max={max_candidate_score:.4f}" @@ -1415,7 +1546,7 @@ async def search( deduped.append(item) items = deduped - logger.info( + logger.debug( f"[precise.done] events={len(candidate_items)}, selected={len(items)}, " f"chunks={sum(1 for item in items if item.get('chunk'))}" ) @@ -1425,6 +1556,16 @@ async def search( "items": items, "_timings": timings, "_query_vector": query_vector, + "stage_timings_seconds": self._build_stage_timings("precise", timings), + "_stats": { + "mode": "precise", + "entities": len(entity_ids), + "initial_events": len(event_ids), + "expanded": len(expanded_event_ids), + "candidates": len(candidate_items), + "selected": len(items), + "chunks": sum(1 for item in items if item.get("chunk")), + }, } # ── 段落返回兼容接口 ─────────────────────────────────────── @@ -1457,6 +1598,7 @@ async def search_for_sections( timings = result.get("_timings", {}).copy() if "total" in timings: timings.pop("total") + stats = dict(result.get("_stats", {})) seen_chunk_ids: set = set() sections = [] @@ -1502,16 +1644,28 @@ async def search_for_sections( native_added += 1 if len(sections) >= target: break - logger.info( + logger.debug( f"[native.fill] multi={multi_count}, native=+{native_added}, " f"total={len(sections)}" ) - - timings["total"] = time.perf_counter() - t_total + stats["native_added"] = native_added + + total_time = time.perf_counter() - t_total + timings["total"] = total_time + mode = self._resolve_search_mode(multi_config) + stage_timings = self._build_stage_timings(mode, timings) + stats["sections"] = min(len(sections), target) + self._log_search_summary( + mode=mode, + stats=stats, + stage_timings=stage_timings, + total_time=total_time, + ) return { "sections": sections[:target], "_timings": timings, + "stage_timings_seconds": stage_timings, } # 兼容旧 pipelineEngine 接口名；内部不执行模型 rerank。 @@ -1569,7 +1723,7 @@ async def search_chunks( ] total_time = time.perf_counter() - start_time - logger.info( + logger.debug( f"[query.chunk] returned={len(sections)}, total_time={total_time:.3f}s" ) diff --git a/scripts/run_search.py b/scripts/run_search.py index 0bf2033..86ab011 100755 --- a/scripts/run_search.py +++ b/scripts/run_search.py @@ -1,6 +1,7 @@ import argparse import asyncio import json +import logging import sys import time from datetime import datetime @@ -24,6 +25,29 @@ logger = get_logger("scripts.run_search_only") +def configure_logging() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + logging.getLogger("pipeline").setLevel(logging.WARNING) + logging.getLogger("pipeline.search.multi_es").setLevel(logging.INFO) + logging.getLogger("pipeline.ai.llm").setLevel(logging.INFO) + for noisy_logger_name in ( + "elastic_transport", + "elastic_transport.transport", + ): + logging.getLogger(noisy_logger_name).setLevel(logging.ERROR) + for noisy_logger_name in ( + "elasticsearch", + "httpx", + "httpcore", + "urllib3", + ): + logging.getLogger(noisy_logger_name).setLevel(logging.WARNING) + + def get_dataset_path_for_name(dataset_name: str) -> str: """根据数据集名称获取默认数据集路径""" return str(project_root / "pipeline" / "evaluation" / "dataset" / f"{dataset_name}.json") @@ -225,5 +249,6 @@ async def search_one(idx: int, item: Dict) -> Dict: if __name__ == "__main__": + configure_logging() exit_code = asyncio.run(main()) sys.exit(exit_code) diff --git a/scripts/run_search_benchmark.py b/scripts/run_search_benchmark.py index 82e9720..617545e 100644 --- a/scripts/run_search_benchmark.py +++ b/scripts/run_search_benchmark.py @@ -45,6 +45,20 @@ logging.getLogger("pipeline.search.atomic").setLevel(logging.INFO) # 放开 LLM 重试日志，方便观察重试次数和等待时间 logging.getLogger("pipeline.ai.llm").setLevel(logging.INFO) +# 压制 ES 传输层请求日志，避免淹没 multi_es 阶段耗时摘要 +for noisy_logger_name in ( + "elastic_transport", + "elastic_transport.transport", +): + logging.getLogger(noisy_logger_name).setLevel(logging.ERROR) +# 其他 HTTP 客户端保留 WARNING 以上，便于保留真实异常信号。 +for noisy_logger_name in ( + "elasticsearch", + "httpx", + "httpcore", + "urllib3", +): + logging.getLogger(noisy_logger_name).setLevel(logging.WARNING) def calculate_precision_f1_at_k( @@ -89,6 +103,72 @@ def calculate_precision_f1_at_k( return pooled_results +MULTI_ES_STAGE_KEYS = { + "fast": [ + "step1_entity_recall", + "step2_query_embedding", + "step3_seed_recall", + "step4_expand_and_rank", + "step5_chunk_fetch", + ], + "precise": [ + "step1_entity_recall", + "step2_query_embedding", + "step3_dual_recall", + "step4_expand_and_coarse_rank", + "step5_llm_filter", + "step6_chunk_fetch", + ], +} + + +def normalize_stage_timings(stage_timings: Any) -> Dict[str, float]: + if not isinstance(stage_timings, dict): + return {} + + normalized: Dict[str, float] = {} + for key, value in stage_timings.items(): + try: + normalized[str(key)] = float(value) + except (TypeError, ValueError): + continue + return normalized + + +def calculate_multi_es_stage_averages( + results: List[Dict], + mode: str, +) -> Dict[str, float]: + stage_keys = MULTI_ES_STAGE_KEYS.get(mode, []) + if not stage_keys: + return {} + + result_stage_timings = [ + normalize_stage_timings(result.get("stage_timings_seconds")) + for result in results + ] + result_stage_timings = [timings for timings in result_stage_timings if timings] + + stage_averages: Dict[str, float] = {} + for stage_key in stage_keys: + values = [ + timings[stage_key] + for timings in result_stage_timings + if stage_key in timings + ] + if values: + stage_averages[stage_key] = sum(values) / len(values) + + return stage_averages + + +def build_multi_es_stage_metrics(mode: str, stage_averages: Dict[str, float]) -> Dict[str, float]: + return { + f"multi_es_{mode}_{stage_key}_avg_seconds": value + for stage_key, value in stage_averages.items() + } + + # ───────────────────────────────────────────────────────────────────────────── # 批次回调：与 benchmark.py 的 _handle_bench_callback 完全一致 # ───────────────────────────────────────────────────────────────────────────── @@ -99,6 +179,8 @@ async def handle_bench_callback( results_so_far: List[Dict], # 已完成的 search_result 列表 gold_docs_for_recall: List[List[str]], bench_size: int, + strategy: str, + mode: str, mlflow_tracker: Optional[Any], bench_logger, ) -> None: @@ -183,6 +265,11 @@ async def handle_bench_callback( avg_request_time, step=batch_index, ) + if strategy == "multi_es": + stage_averages = calculate_multi_es_stage_averages(results_so_far, mode) + stage_metrics = build_multi_es_stage_metrics(mode, stage_averages) + if stage_metrics: + mlflow_tracker.log_batch_metrics(stage_metrics, step=batch_index) # ───────────────────────────────────────────────────────────────────────────── @@ -275,12 +362,18 @@ async def search_one(idx: int, question: str) -> Dict: raw_sections = raw.get("sections", []) sections = [normalize_section(s) for s in raw_sections] sections = sections[:top_k] - return { + result = { "question_index": idx + 1, "question": question, "sections": sections, "request_time_seconds": time.perf_counter() - request_start, } + stage_timings = normalize_stage_timings( + raw.get("stage_timings_seconds") + ) + if stage_timings: + result["stage_timings_seconds"] = stage_timings + return result # 通过引擎执行搜索，不直接访问底层 searcher engine = pipelineEngine( @@ -351,6 +444,8 @@ async def search_one(idx: int, question: str) -> Dict: results_so_far=results_so_far, gold_docs_for_recall=gold_docs_for_recall, bench_size=effective_bench_size, + strategy=strategy, + mode=mode, mlflow_tracker=mlflow_tracker, bench_logger=bench_logger, ) @@ -793,15 +888,21 @@ async def main(): # ── 保存搜索原始结果 ────────────────────────────────────── search_output = output_dir / "search_results.json" - serializable = [ - { + serializable = [] + for r in search_results: + item = { "question_index": r["question_index"], "question": r["question"], "request_time_seconds": round(r.get("request_time_seconds", 0.0), 3), "retrieved_docs": r["sections"], } - for r in search_results - ] + stage_timings = normalize_stage_timings(r.get("stage_timings_seconds")) + if stage_timings: + item["stage_timings_seconds"] = { + key: round(value, 3) + for key, value in stage_timings.items() + } + serializable.append(item) with open(search_output, "w", encoding="utf-8") as f: json.dump(serializable, f, ensure_ascii=False, indent=2) logger.info(f"\n💾 搜索结果已保存: {search_output}") @@ -869,6 +970,17 @@ async def main(): avg_request_time, step=final_batch_index, ) + if args.strategy == "multi_es": + stage_averages = calculate_multi_es_stage_averages( + search_results, + args.mode, + ) + stage_metrics = build_multi_es_stage_metrics(args.mode, stage_averages) + if stage_metrics: + mlflow_tracker.log_batch_metrics( + stage_metrics, + step=final_batch_index, + ) mlflow.log_param("output_dir", str(output_dir)) logger.info("✓ MLflow metrics/params 记录完成") except Exception as e: