diff --git a/.gitignore b/.gitignore index 296bf1d..42fbe7d 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,7 @@ logs/*.log.* data/raw/ data/processed/ data/cache/ +data/pipeline/ *.sqlite *.sqlite3 *.db diff --git a/a_share_cli/main.py b/a_share_cli/main.py index a581eab..c00c98b 100644 --- a/a_share_cli/main.py +++ b/a_share_cli/main.py @@ -179,6 +179,8 @@ def cmd_pipeline(args: argparse.Namespace) -> int: return _run_module("scripts.run_scheduler", extra, timeout=pipeline_timeout) if args.once: extra = ["--once", "--date", args.date or _today()] + if args.resume: + extra.append("--resume") steps = args.steps if args.report: steps = (steps + ",report") if steps else "report" @@ -878,6 +880,8 @@ def main() -> int: p.add_argument("--once", action="store_true", help="立即执行一次") p.add_argument("--date", default=None) p.add_argument("--steps", default=None, help="指定步骤(crawler,extractor,...,report)") + p.add_argument("--resume", action="store_true", + help="断点续跑(仅 --once):跳过连续成功步骤,从上次失败/未执行步骤继续") p.add_argument("--report", action="store_true", help="全链路末尾生成日报") p.add_argument("--cninfo-once", action="store_true", help="cninfo watchlist 全链路(公告+调研+IRM)") p.set_defaults(func=cmd_pipeline) diff --git a/scheduler/pipeline.py b/scheduler/pipeline.py index c37251d..a16b1d8 100644 --- a/scheduler/pipeline.py +++ b/scheduler/pipeline.py @@ -2,17 +2,27 @@ 编排 M1→M6 全链路,每一步调用已有脚本。 单步失败记录日志但不阻断后续(后续步骤可能使用旧缓存数据,降级继续)。 + +断点恢复: + 每次运行把各步骤结果写入 data/pipeline/state.json(按日期隔离); + run_pipeline(resume=True) 时跳过连续成功的步骤,从第一个失败/未执行 + 的步骤继续,实现 `pipeline --once --resume` 断点续跑。 """ from __future__ import annotations +import json import subprocess import time from dataclasses import dataclass, field from datetime import date, datetime +from pathlib import Path from loguru import logger +# 断点状态文件(按日期隔离,记录每步骤结果) +DEFAULT_STATE_PATH = Path("data/pipeline/state.json") + # 步骤超时(秒) STEP_TIMEOUTS: dict[str, int] = { "crawler": 900, # M1 抓取(含 Playwright 浏览器,13 源约 8-12 min) @@ -64,6 +74,54 @@ class PipelineResult: return all(s.success for s in self.steps) +def _load_pipeline_state(path: Path = DEFAULT_STATE_PATH) -> dict: + """读取断点状态文件;不存在或损坏时返回空 dict。""" + if not path.is_file(): + return {} + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError) as e: + logger.warning("pipeline 状态文件损坏,忽略: {} ({})", path, e) + return {} + return data if isinstance(data, dict) else {} + + +def _save_pipeline_state(state: dict, path: Path = DEFAULT_STATE_PATH) -> None: + """原子写状态文件(tmp + rename,避免中断写坏)。""" + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_suffix(".json.tmp") + tmp.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8") + tmp.replace(path) + + +def _update_step_state(state: dict, date_str: str, sr: StepResult) -> None: + """把单步结果写入状态(ok/failed,含退出码与耗时)。""" + day_state = state.setdefault(date_str, {}) + day_state[sr.name] = { + "status": "ok" if sr.success else "failed", + "exit_code": sr.exit_code, + "started_at": sr.started_at.isoformat() if sr.started_at else None, + "elapsed_sec": round(sr.elapsed_sec, 1), + } + + +def _resume_start_index( + names: list[str], + date_str: str, + state: dict, +) -> int: + """计算断点续跑起始下标:跳过连续 ok 前缀,从首个失败/未记录步骤开始。 + + 返回 0..len(names)-1;全部成功时返回 len(names)(表示无需续跑)。 + """ + day_state = state.get(date_str, {}) + for i, name in enumerate(names): + rec = day_state.get(name) + if rec is None or rec.get("status") != "ok": + return i + return len(names) + + def run_step(name: str, date_str: str) -> StepResult: """执行单个 pipeline 步骤。 @@ -158,19 +216,45 @@ def run_step(name: str, date_str: str) -> StepResult: tail_msg=str(e)[:200], started_at=started) -def run_pipeline(date_str: str, *, steps: list[str] | None = None) -> PipelineResult: +def run_pipeline( + date_str: str, + *, + steps: list[str] | None = None, + resume: bool = False, + state_path: Path = DEFAULT_STATE_PATH, +) -> PipelineResult: """串联执行全链路(M1→M6)。 参数: date_str: YYYYMMDD。 steps: 可选步骤列表,默认全部 6 步。 + resume: True 时断点续跑——读取 data/pipeline/state.json 中该日期的 + 记录,跳过连续成功的步骤,从第一个失败/未执行步骤继续。 + state_path: 断点状态文件路径(测试可注入)。 """ names = steps or [k for k in STEP_COMMANDS if k not in ("report", "cninfo_crawl", "cninfo_extract", "cninfo_pdf")] result = PipelineResult(started_at=datetime.now()) - for name in names: + state = _load_pipeline_state(state_path) + start_idx = 0 + if resume: + start_idx = _resume_start_index(names, date_str, state) + if start_idx >= len(names): + logger.info("resume: {} 的所有步骤均已完成,无需续跑", date_str) + result.finished_at = datetime.now() + return result + logger.info( + "resume: 从步骤 {} 继续{}", + names[start_idx], + f" (跳过已成功 {names[:start_idx]})" if start_idx > 0 else "", + ) + + for name in names[start_idx:]: sr = run_step(name, date_str) result.steps.append(sr) + # 记录断点状态(无论成败,便于下次 resume) + _update_step_state(state, date_str, sr) + _save_pipeline_state(state, state_path) if not sr.success: logger.warning("步骤 {} 失败,后续步骤继续(可能降级)", name) # 步间留一点缓冲 diff --git a/scripts/run_embedding.py b/scripts/run_embedding.py index 7f7c00d..5e5e088 100644 --- a/scripts/run_embedding.py +++ b/scripts/run_embedding.py @@ -13,6 +13,9 @@ data/embeddings/{day}/{url_hash}.json (含 vector 完整内容) data/embeddings/{day}/index.jsonl (扁平摘要,不含向量,便于检索/调试) data/embeddings/{day}/failed.jsonl (失败列表) + +增量: 默认跳过已嵌入的文章(输出目录已有 {url_hash}.json 视为已处理), + 断点续跑/失败重试不会重复调用 embed API;--force 强制全量重嵌入。 """ from __future__ import annotations @@ -142,6 +145,26 @@ def _collect_inputs(args: argparse.Namespace) -> list[tuple[Path, str]]: return files +def _filter_existing( + files: list[tuple[Path, str]], out_dir: Path +) -> tuple[list[tuple[Path, str]], int]: + """过滤掉已有产物(输出目录存在同名 {url_hash}.json)的输入。 + + 输入与输出文件名均为 {url_hash}.json,直接比对 stem。 + 返回 (待处理, 跳过数);断点续跑/失败重试借此避免重复调用 embed API。 + """ + pending: list[tuple[Path, str]] = [] + skipped = 0 + for fp, kind in files: + if (out_dir / f"{fp.stem}.json").exists(): + skipped += 1 + else: + pending.append((fp, kind)) + if skipped: + logger.info("跳过已嵌入 {} 篇(产物已存在),待处理 {}", skipped, len(pending)) + return pending, skipped + + # --------------------------------------------------------------------------- # # 主流程 # --------------------------------------------------------------------------- # @@ -150,12 +173,24 @@ async def _run(args: argparse.Namespace) -> int: load_dotenv() files = _collect_inputs(args) - if args.limit: - files = files[: args.limit] if not files: logger.error("未发现任何输入文件: {} ({})", args.input, args.date) return 2 - logger.info("待嵌入文章数: {} (input={})", len(files), args.input) + + out_dir = Path(args.out_root) / args.date + out_dir.mkdir(parents=True, exist_ok=True) + # 增量:跳过已有产物(断点续跑/失败重试不重复调用 embed API),--force 全量 + skipped = 0 + if not args.force: + files, skipped = _filter_existing(files, out_dir) + if args.limit: + files = files[: args.limit] + if not files: + logger.info( + "无待嵌入文章(全部已处理,跳过 {} 篇),如需重新嵌入请加 --force", skipped + ) + return 0 + logger.info("待嵌入文章数: {} (跳过已处理 {}; input={})", len(files), skipped, args.input) # 准备每篇文本 prepared: list[tuple[str, Article, str | None]] = [] @@ -170,13 +205,17 @@ async def _run(args: argparse.Namespace) -> int: logger.error("所有输入文件均无法解析") return 2 - out_dir = Path(args.out_root) / args.date - out_dir.mkdir(parents=True, exist_ok=True) index_path = out_dir / "index.jsonl" failed_path = out_dir / "failed.jsonl" - for p in (index_path, failed_path): - if p.exists(): - p.unlink() + if args.force: + # 全量模式:重建 index / failed + for p in (index_path, failed_path): + if p.exists(): + p.unlink() + else: + # 增量模式:index 累积追加;failed 只保留本次运行失败的 + if failed_path.exists(): + failed_path.unlink() started = time.time() succ_cnt = 0 @@ -282,6 +321,8 @@ def main() -> int: help="每批送 embed 的条数(DashScope 上限 10)") parser.add_argument("--limit", type=int, default=0, help="最多处理 N 篇,0=不限") + parser.add_argument("--force", action="store_true", + help="强制全量重嵌入(默认跳过已嵌入文章)") parser.add_argument("--log-level", default="INFO") args = parser.parse_args() diff --git a/scripts/run_event_extraction.py b/scripts/run_event_extraction.py index c7de976..5b717ac 100644 --- a/scripts/run_event_extraction.py +++ b/scripts/run_event_extraction.py @@ -6,12 +6,16 @@ data/events/{YYYYMMDD}/index.jsonl (扁平摘要) data/events/{YYYYMMDD}/failed.jsonl (失败列表) +增量: 默认跳过已抽取的文章(输出目录已有 {url_hash}.json 视为已处理), + 断点续跑/失败重试不会重复调用 LLM API;--force 强制全量重抽。 + 用法: uv run python -m scripts.run_event_extraction uv run python -m scripts.run_event_extraction --date 20260616 uv run python -m scripts.run_event_extraction --provider qwen --model qwen-plus uv run python -m scripts.run_event_extraction --concurrency 5 --limit 10 uv run python -m scripts.run_event_extraction --input-root data/processed --no-deduped + uv run python -m scripts.run_event_extraction --force # 全量重抽 """ from __future__ import annotations @@ -79,6 +83,24 @@ def _collect_inputs( return files +def _filter_existing(files: list[Path], out_dir: Path) -> tuple[list[Path], int]: + """过滤掉已有产物(输出目录存在同名 {url_hash}.json)的输入。 + + 输入文件名即 url_hash(如 {url_hash}.json),与 M4 产物命名一致。 + 返回 (待处理文件, 跳过数);断点续跑/失败重试借此避免重复调用 LLM API。 + """ + pending: list[Path] = [] + skipped = 0 + for fp in files: + if (out_dir / f"{fp.stem}.json").exists(): + skipped += 1 + else: + pending.append(fp) + if skipped: + logger.info("跳过已处理 {} 篇(产物已存在),待处理 {}", skipped, len(pending)) + return pending, skipped + + def _load_article(p: Path) -> Article | None: try: return Article.model_validate(json.loads(p.read_text(encoding="utf-8"))) @@ -101,24 +123,40 @@ async def _run(args: argparse.Namespace) -> int: input_root = Path(args.input_root) use_deduped = not args.no_deduped files = _collect_inputs(input_root, args.date, use_deduped, args.source) - if args.limit: - files = files[: args.limit] if not files: logger.error( "{} 下未发现 {} 的文章(use_deduped={})", input_root, args.date, use_deduped, ) return 2 - logger.info("待处理文章数: {}", len(files)) out_dir = Path(args.out_root) / args.date out_dir.mkdir(parents=True, exist_ok=True) + # 增量:跳过已有产物(断点续跑/失败重试不重复调用 LLM API),--force 全量 + skipped = 0 + if not args.force: + files, skipped = _filter_existing(files, out_dir) + if args.limit: + files = files[: args.limit] + if not files: + logger.info( + "无待处理文章(全部已抽取,跳过 {} 篇),如需重抽请加 --force", skipped + ) + return 0 + + logger.info("待处理文章数: {} (跳过已处理 {})", len(files), skipped) + index_path = out_dir / "index.jsonl" failed_path = out_dir / "failed.jsonl" - # 重跑时清掉旧的 jsonl,避免重复追加 - for p in (index_path, failed_path): - if p.exists(): - p.unlink() + if args.force: + # 全量模式:重建 index / failed + for p in (index_path, failed_path): + if p.exists(): + p.unlink() + else: + # 增量模式:index 累积追加;failed 只保留本次运行失败的 + if failed_path.exists(): + failed_path.unlink() template = PromptTemplate(args.prompt) semaphore = asyncio.Semaphore(args.concurrency) @@ -209,6 +247,8 @@ def main() -> int: help="LLM 异步并发上限") parser.add_argument("--max-attempts", type=int, default=3, help="单篇文章最大重试次数") + parser.add_argument("--force", action="store_true", + help="强制全量重抽(默认跳过已抽取文章)") parser.add_argument("--limit", type=int, default=0, help="最多处理 N 篇,0=不限制(用于联调)") parser.add_argument("--prompt", default="prompts/event_extraction.md", diff --git a/scripts/run_extractor.py b/scripts/run_extractor.py index 245f0fd..135b21c 100644 --- a/scripts/run_extractor.py +++ b/scripts/run_extractor.py @@ -14,12 +14,14 @@ from __future__ import annotations import argparse +import contextlib import json import sys from datetime import date, datetime from pathlib import Path from loguru import logger +from pydantic import ValidationError from extractor import Article, ExtractError, extract_article from extractor.parser import _url_hash @@ -238,16 +240,21 @@ def _process_cninfo_v2(rec: dict, json_path: Path, out_dir: Path) -> Article | N return _save_article(article, out_dir) +def _append_index(article: Article, out_dir: Path) -> None: + """把 Article 的扁平摘要追加到 index.jsonl。""" + flat = article.model_dump(exclude={"content", "images"}, mode="json") + flat["article_file"] = f"{article.url_hash}.json" + flat["content_preview"] = article.content[:80] + with (out_dir / "index.jsonl").open("a", encoding="utf-8") as f: + f.write(json.dumps(flat, ensure_ascii=False) + "\n") + + def _save_article(article: Article, out_dir: Path) -> Article: """保存 Article JSON 并追加 index。""" out_dir.mkdir(parents=True, exist_ok=True) article_path = out_dir / f"{article.url_hash}.json" article_path.write_text(article.model_dump_json(indent=2), encoding="utf-8") - flat = article.model_dump(exclude={"content", "images"}, mode="json") - flat["article_file"] = article_path.name - flat["content_preview"] = article.content[:80] - with (out_dir / "index.jsonl").open("a", encoding="utf-8") as f: - f.write(json.dumps(flat, ensure_ascii=False) + "\n") + _append_index(article, out_dir) return article @@ -257,37 +264,57 @@ def _process_source_day( raw_root: Path, out_root: Path, body_xpath_map: dict[str, str] | None = None, -) -> tuple[int, int]: - """处理单个源单日。返回 (成功数, 总数)。""" + *, + force: bool = False, +) -> tuple[int, int, int]: + """处理单个源单日。返回 (成功数, 总数, 跳过数)。 + + 默认增量:已提取的文章(输出目录已有 {url_hash}.json)跳过提取,仅回补 index 行; + force=True 时全量重提取并重建 index。 + """ raw_dir = raw_root / source_id / day out_dir = out_root / source_id / day records = _iter_article_records(raw_dir) if not records: logger.info("源 {} 日期 {} 无可处理记录", source_id, day) - return 0, 0 + return 0, 0, 0 - # 清理同日旧的 index.jsonl,避免重复追加 + # 全量模式:重建 index;增量模式:保留旧 index 追加新条目 old_index = out_dir / "index.jsonl" - if old_index.exists(): + if force and old_index.exists(): old_index.unlink() succ = 0 + skipped = 0 for rec in records: + url_hash = rec.get("url_hash") or _url_hash(rec.get("url") or "") + existing = out_dir / f"{url_hash}.json" + if not force and existing.is_file(): + # 增量:跳过已提取,回补 index 行保持摘要完整 + skipped += 1 + with contextlib.suppress(json.JSONDecodeError, ValidationError, OSError): + _append_index( + Article.model_validate(json.loads(existing.read_text(encoding="utf-8"))), + out_dir, + ) + continue article = _process_one(rec, raw_dir, out_dir, body_xpath_map) if article is not None: succ += 1 total = len(records) - rate = succ / max(total, 1) + # 增量模式下「跳过已提取」视为已成功处理,避免全跳过时误报成功率 0% + rate = (succ + skipped) / max(total, 1) logger.info( - "源 {} 日期 {} 提取完成: {}/{} 成功率 {:.0%}", + "源 {} 日期 {} 提取完成: {}/{} 成功率 {:.0%} (跳过已提取 {})", source_id, day, succ, total, rate, + skipped, ) - return succ, total + return succ, total, skipped def _list_source_dirs(raw_root: Path) -> list[str]: @@ -307,6 +334,8 @@ def main() -> int: default=date.today().strftime("%Y%m%d"), help="处理日期 YYYYMMDD,默认今日", ) + parser.add_argument("--force", action="store_true", + help="强制全量重提取(默认跳过已提取文章)") parser.add_argument("--log-level", default="INFO") args = parser.parse_args() @@ -335,18 +364,24 @@ def main() -> int: started = datetime.now() total_succ = 0 total_all = 0 + total_skipped = 0 for src in sources: - succ, total = _process_source_day(src, args.date, raw_root, out_root, body_xpath_map) + succ, total, skipped = _process_source_day( + src, args.date, raw_root, out_root, body_xpath_map, force=args.force + ) total_succ += succ total_all += total + total_skipped += skipped elapsed = (datetime.now() - started).total_seconds() - rate = total_succ / max(total_all, 1) + # 增量模式下跳过已提取视为成功 + rate = (total_succ + total_skipped) / max(total_all, 1) logger.info( - "全部完成: {}/{} 成功率 {:.0%} 用时 {:.1f}s", + "全部完成: {}/{} 成功率 {:.0%} (跳过已提取 {}) 用时 {:.1f}s", total_succ, total_all, rate, + total_skipped, elapsed, ) return 0 if rate >= 0.9 or total_all == 0 else 1 diff --git a/scripts/run_scheduler.py b/scripts/run_scheduler.py index 992244a..5d11f4f 100644 --- a/scripts/run_scheduler.py +++ b/scripts/run_scheduler.py @@ -59,11 +59,17 @@ def _parse_schedule_times(raw: str) -> list[tuple[int, int]]: def _once(args: argparse.Namespace) -> int: - """单次执行模式。""" + """单次执行模式。 + + 默认全量执行;--resume 时断点续跑(跳过连续成功步骤,从失败/未执行步骤继续)。 + """ steps = None if args.steps: steps = [s.strip() for s in args.steps.split(",")] - run_pipeline(args.date, steps=steps) + if args.resume and args.steps: + logger.error("--resume 与 --steps 不能同时使用(断点续跑针对全链路)") + return 2 + run_pipeline(args.date, steps=steps, resume=args.resume) return 0 @@ -185,6 +191,10 @@ def main() -> int: ) parser.add_argument("--steps", default=None, help="仅执行指定步骤,逗号分隔 (如 crawler,extractor)") + parser.add_argument( + "--resume", action="store_true", + help="断点续跑(仅 --once):跳过连续成功步骤,从上次失败/未执行步骤继续", + ) parser.add_argument("--log-level", default="INFO") args = parser.parse_args() diff --git a/tests/test_incremental.py b/tests/test_incremental.py new file mode 100644 index 0000000..38833cd --- /dev/null +++ b/tests/test_incremental.py @@ -0,0 +1,268 @@ +"""增量处理与 pipeline 断点续跑测试。 + +覆盖: + - M2/M4/M5 脚本的「产物存在即跳过」过滤逻辑 + - scheduler.pipeline 断点状态记录与 --resume 续跑逻辑 +""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +# --------------------------------------------------------------------------- # +# M4 / M5: 产物存在即跳过 +# --------------------------------------------------------------------------- # + +def test_llm_filter_existing_skips_done(tmp_path: Path) -> None: + """M4:输出目录已有 {url_hash}.json 的输入被过滤,不重复调用 LLM API。""" + from scripts.run_event_extraction import _filter_existing + + out_dir = tmp_path / "out" + out_dir.mkdir() + # 已处理 + (out_dir / "aaa.json").write_text("{}", encoding="utf-8") + (out_dir / "ccc.json").write_text("{}", encoding="utf-8") + files = [ + tmp_path / "in" / "aaa.json", # 已处理 → 跳过 + tmp_path / "in" / "bbb.json", # 未处理 → 待处理 + tmp_path / "in" / "ccc.json", # 已处理 → 跳过 + ] + pending, skipped = _filter_existing(files, out_dir) + assert skipped == 2 + assert [p.stem for p in pending] == ["bbb"] + + +def test_llm_filter_existing_force_keeps_all(tmp_path: Path) -> None: + """M4:--force 时不做过滤(全量重抽由调用方控制)。""" + from scripts.run_event_extraction import _filter_existing + + out_dir = tmp_path / "out" + out_dir.mkdir() + (out_dir / "aaa.json").write_text("{}", encoding="utf-8") + files = [tmp_path / "in" / "aaa.json"] + # _filter_existing 本身不含 force 逻辑,验证在 force 下不会被调用: + # 直接验证「已存在也被返回」需由上层跳过调用,这里仅确认过滤函数行为。 + pending, skipped = _filter_existing(files, out_dir) + assert skipped == 1 + assert pending == [] + + +def test_embedding_filter_existing_skips_done(tmp_path: Path) -> None: + """M5:输出目录已有 {url_hash}.json 的输入被过滤,不重复调用 embed API。""" + from scripts.run_embedding import _filter_existing + + out_dir = tmp_path / "emb" + out_dir.mkdir() + (out_dir / "h1.json").write_text("{}", encoding="utf-8") + files = [ + (tmp_path / "in" / "h1.json", "event"), # 已处理 → 跳过 + (tmp_path / "in" / "h2.json", "event"), # 未处理 → 待处理 + ] + pending, skipped = _filter_existing(files, out_dir) + assert skipped == 1 + assert [p.stem for p, _ in pending] == ["h2"] + + +# --------------------------------------------------------------------------- # +# M2: 已提取文章跳过 +# --------------------------------------------------------------------------- # + +def test_extractor_process_source_day_skips_existing(tmp_path: Path) -> None: + """M2:输出目录已有产物的记录被跳过提取,且 index 回补完整。""" + from scripts.run_extractor import _process_source_day + + raw_dir = tmp_path / "raw" / "cls" / "20260616" + raw_dir.mkdir(parents=True) + # 两条 raw 记录(url_hash 与产物文件名一致) + rec1 = {"source_id": "cls", "url": "https://a/1", "url_hash": "aaa1111111111111", + "stage": "article", "success": True, "html_file": "aaa1111111111111.html"} + rec2 = {"source_id": "cls", "url": "https://b/2", "url_hash": "bbb2222222222222", + "stage": "article", "success": True, "html_file": "bbb2222222222222.html"} + with (raw_dir / "index.jsonl").open("a", encoding="utf-8") as f: + f.write(json.dumps(rec1, ensure_ascii=False) + "\n") + f.write(json.dumps(rec2, ensure_ascii=False) + "\n") + + out_dir = tmp_path / "proc" / "cls" / "20260616" + out_dir.mkdir(parents=True) + # 预置一条已有产物(视为已提取) + article = { + "source_id": "cls", "url": "https://a/1", "url_hash": "aaa1111111111111", + "title": "已有", "content": "内容", "word_count": 2, + } + (out_dir / "aaa1111111111111.json").write_text( + json.dumps(article, ensure_ascii=False), encoding="utf-8" + ) + + # 另一条无 html 文件 → 提取失败(但不影响跳过逻辑断言) + succ, total, skipped = _process_source_day( + "cls", "20260616", tmp_path / "raw", tmp_path / "proc" + ) + assert total == 2 + assert skipped == 1 # 已有产物被跳过 + assert succ == 0 # 另一条因 html 缺失提取失败 + # index 回补了被跳过条目的行 + idx = (out_dir / "index.jsonl").read_text(encoding="utf-8").strip() + assert "aaa1111111111111" in idx + + +def test_extractor_process_source_day_force_rebuilds(tmp_path: Path) -> None: + """M2:--force 时不做跳过,并重建 index。""" + from scripts.run_extractor import _process_source_day + + raw_dir = tmp_path / "raw" / "cls" / "20260616" + raw_dir.mkdir(parents=True) + rec = {"source_id": "cls", "url": "https://a/1", "url_hash": "aaa1111111111111", + "stage": "article", "success": True, "html_file": "aaa1111111111111.html"} + with (raw_dir / "index.jsonl").open("a", encoding="utf-8") as f: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + out_dir = tmp_path / "proc" / "cls" / "20260616" + out_dir.mkdir(parents=True) + (out_dir / "aaa1111111111111.json").write_text("{}", encoding="utf-8") + (out_dir / "index.jsonl").write_text("旧内容", encoding="utf-8") + + succ, total, skipped = _process_source_day( + "cls", "20260616", tmp_path / "raw", tmp_path / "proc", force=True + ) + assert skipped == 0 + # force 模式重建 index(旧内容被清掉;此处无 html 提取失败,index 为空或不存在) + idx = out_dir / "index.jsonl" + assert not idx.exists() or idx.read_text(encoding="utf-8") == "" + + +# --------------------------------------------------------------------------- # +# pipeline 断点状态与 --resume +# --------------------------------------------------------------------------- # + +@pytest.fixture +def fake_subprocess(monkeypatch: pytest.MonkeyPatch): + """mock subprocess.run,按步骤名返回 returncode,并记录调用顺序。""" + from scheduler import pipeline + + calls: list[str] = [] + + def _step_name(cmd: list[str]) -> str: + """从命令中提取脚本名,如 scripts.run_extractor → run_extractor。""" + return next(c.split(".")[-1] for c in cmd if "scripts.run_" in c) + + def _fake_run(cmd, timeout=None): # noqa: ARG001 + calls.append(_step_name(cmd)) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(pipeline.subprocess, "run", _fake_run) + return calls + + +def _run_with_steps(monkeypatch: pytest.MonkeyPatch, failures: set[str]): + """构造 run_step:指定步骤(如 'extractor')返回失败。""" + from scheduler import pipeline + + def _fake_run(cmd, timeout=None): # noqa: ARG001 + name = next(c.split(".")[-1] for c in cmd if "scripts.run_" in c) + name = name.replace("run_", "") # run_extractor → extractor + return SimpleNamespace(returncode=1 if name in failures else 0) + + monkeypatch.setattr(pipeline.subprocess, "run", _fake_run) + + +def test_pipeline_records_state(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """全量运行后状态文件按日期记录每个步骤的 ok/failed。""" + from scheduler import pipeline + + _run_with_steps(monkeypatch, failures={"extractor"}) + state_path = tmp_path / "state.json" + steps = ["extractor", "dedup", "llm"] + pipeline.run_pipeline("20260616", steps=steps, state_path=state_path) + + state = pipeline._load_pipeline_state(state_path) + day = state["20260616"] + assert day["extractor"]["status"] == "failed" + assert day["dedup"]["status"] == "ok" + assert day["llm"]["status"] == "ok" + # dedup 返回 1 被特判为成功,故用 extractor 制造失败 + assert day["extractor"]["exit_code"] == 1 + + +def test_pipeline_resume_skips_success_prefix( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + """resume 跳过连续成功步骤,从失败步骤继续。""" + from scheduler import pipeline + + state_path = tmp_path / "state.json" + # 预置状态:extractor 失败,dedup/llm 成功(模拟上次运行) + state = {"20260616": { + "extractor": {"status": "failed", "exit_code": 1}, + "dedup": {"status": "ok", "exit_code": 0}, + "llm": {"status": "ok", "exit_code": 0}, + }} + pipeline._save_pipeline_state(state, state_path) + + calls: list[str] = [] + + def _fake_run(cmd, timeout=None): # noqa: ARG001 + name = next(c.split(".")[-1] for c in cmd if "scripts.run_" in c) + calls.append(name) + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(pipeline.subprocess, "run", _fake_run) + steps = ["extractor", "dedup", "llm"] + result = pipeline.run_pipeline( + "20260616", steps=steps, resume=True, state_path=state_path + ) + # 从 extractor 开始重跑全部(extractor 之后的 dedup/llm 需重跑以覆盖降级数据) + assert calls == ["run_extractor", "run_dedup", "run_event_extraction"] + assert all(s.success for s in result.steps) + + +def test_pipeline_resume_all_done_noop(tmp_path: Path) -> None: + """resume 且所有步骤均已成功时,不执行任何步骤。""" + from scheduler import pipeline + + state_path = tmp_path / "state.json" + state = {"20260616": { + "extractor": {"status": "ok", "exit_code": 0}, + "dedup": {"status": "ok", "exit_code": 0}, + "llm": {"status": "ok", "exit_code": 0}, + }} + pipeline._save_pipeline_state(state, state_path) + + steps = ["extractor", "dedup", "llm"] + result = pipeline.run_pipeline( + "20260616", steps=steps, resume=True, state_path=state_path + ) + assert result.steps == [] + assert result.all_success # 空步骤视为成功 + + +def test_pipeline_resume_missing_step_starts_from_first_missing( + tmp_path: Path, +) -> None: + """resume:部分步骤无历史记录时,从首个缺失步骤开始。""" + from scheduler import pipeline + + state_path = tmp_path / "state.json" + state = {"20260616": { + "extractor": {"status": "ok", "exit_code": 0}, + }} + pipeline._save_pipeline_state(state, state_path) + + idx = pipeline._resume_start_index( + ["extractor", "dedup", "llm"], "20260616", + pipeline._load_pipeline_state(state_path), + ) + assert idx == 1 # dedup 缺失 → 从它开始 + + +def test_once_rejects_resume_with_steps(monkeypatch: pytest.MonkeyPatch) -> None: + """--resume 与 --steps 同时使用时报错。""" + import scripts.run_scheduler as rs + + args = SimpleNamespace(steps="crawler,extractor", resume=True, + date="20260616") + rc = rs._once(args) + assert rc == 2