obsidian-plugin/lllin000_PaperForge
1
0
Fork 0
mirror of https://github.com/lllin000/PaperForge.git synced 2026-07-22 06:50:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

PR9A: Resume & Rebuild Correctness

Browse source 
- OCR rebuild --all: version/artifact-based selection via _needs_derived_rebuild()
- --status and explicit keys: manual override, no version filter
- Removed .done.{key} checkpoint markers from selection logic
- _apply_post_rebuild_version_flags now writes derived_version
- Embed resume entry: three-gate protection (stale state / missing DB / corrupt DB)
- _pid_alive() cross-platform PID health check
- _assert_collections_healthy() explicit ChromaDB probe
- 14 regression tests for selection logic, version flags, and resume guards
This commit is contained in:
LLLin000 2026-07-06 22:33:08 +08:00
parent 6afcdc1830
commit 89f5a3fce6
4 changed files with 462 additions and 56 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

75
paperforge/commands/embed.py
View file

@ -56,6 +56,37 @@ def _has_object_units_in_db(vault: Path, key: str) -> bool:
finally:
conn.close()
def _pid_alive(pid: int) -> bool:
if pid <= 0:
return False
if os.name == "nt":
try:
import subprocess
r = subprocess.run(
["tasklist", "/FI", f"PID eq {pid}"],
capture_output=True, text=True, timeout=5
)
return str(pid) in r.stdout
except:
return False
else:
try:
os.kill(pid, 0)
return True
except OSError:
return False
def _assert_collections_healthy(vault: Path) -> tuple[bool, str]:
"""Probe three collections. Doesn't depend on get_embed_status."""
for name in ("paperforge_fulltext", "paperforge_body", "paperforge_objects"):
try:
col = get_collection(vault, name=name)
col.count()
except Exception as exc:
return False, f"{name}: {exc}"
return True, ""
def run(args: argparse.Namespace) -> int:
vault = args.vault_path
sub = getattr(args, "embed_subcommand", "build")
@ -191,12 +222,48 @@ def run(args: argparse.Namespace) -> int:
if resume:
build_state = read_vector_build_state(vault)
stored_model = build_state.get("model", "")
if stored_model and _current_model and stored_model != _current_model:
msg = f"Model changed: {stored_model} -> {_current_model}. Re-embedding all papers."
if not getattr(args, "json", False):
# 门一：stale running state 检测
if build_state.get("status") == "running":
stale = False
pid = build_state.get("pid", 0)
if not pid:
stale = True
elif not _pid_alive(pid):
stale = True
else:
started = build_state.get("started_at", "")
if started:
try:
dt = __import__('datetime').datetime.fromisoformat(started)
if (__import__('datetime').datetime.now(__import__('datetime').timezone.utc) - dt).total_seconds() > 43200:
stale = True
except:
pass
if stale:
msg = "Previous build appears stale (crashed?). Use --force to rebuild."
print(msg)
return 1
# 门二：missing DB → fresh build（不是 error）
db_path = get_vector_db_path(vault)
if not db_path.exists():
resume = False
else:
# 门三：corrupted DB
ok, err = _assert_collections_healthy(vault)
if not ok:
msg = f"Vector DB corrupted ({err}). Use --force to rebuild."
print(msg)
return 1
# 过三道门后，正常 model check
stored_model = build_state.get("model", "")
if stored_model and _current_model and stored_model != _current_model:
msg = f"Model changed: {stored_model} -> {_current_model}. Re-embedding all papers."
if not getattr(args, "json", False):
print(msg)
resume = False
_force_rebuild = args.force or (resume is False and getattr(args, "resume", False))
if _force_rebuild:

119
paperforge/commands/ocr.py
View file

@ -9,6 +9,11 @@ from paperforge import __version__
from paperforge.core.errors import ErrorCode
from paperforge.core.result import PFError, PFResult
from paperforge.worker.ocr_artifacts import artifact_paths_for_root
from paperforge.worker.ocr_versions import classify_version_state, expected_derived_payload
from paperforge.worker.ocr_maintenance import _can_rebuild
from paperforge.core.io import read_json
logger = logging.getLogger(__name__)
@ -316,6 +321,82 @@ def _run_ocr_list(vault: Path, json_output: bool = False, output_file: str | Non
return 0
def _needs_derived_rebuild(vault: Path, key: str) -> tuple[bool, str]:
"""检测一篇论文是否需要重建。返回 (need, reason)。"""
from paperforge.worker._utils import pipeline_paths
ocr_root = Path(pipeline_paths(vault)["ocr"])
artifacts = artifact_paths_for_root(ocr_root, key)
paper_dir = artifacts.paper_root
if not artifacts.meta_json.exists():
return False, "no_meta"
meta = read_json(artifacts.meta_json)
has_raw = artifacts.blocks_raw.exists()
has_source_meta = artifacts.source_metadata.exists()
if not _can_rebuild(meta, has_raw, has_source_meta):
return False, "cannot_rebuild"
# 版本检测（运行时比较，不依赖 meta.derived_stale）
state = classify_version_state(
meta,
expected_raw={},
expected_derived=expected_derived_payload(),
)
if state["derived_stale"]:
return True, "version_mismatch"
# 产物完整性检测
required = [
"structure/blocks.structured.jsonl",
"render/render-map.json",
"index/structure-tree.json",
"index/role-index.json",
"fulltext.md",
"health/ocr_health.json",
]
for rel in required:
if not (paper_dir / rel).exists():
return True, f"missing:{rel.split('/')[-1]}"
return False, "current"
def _select_rebuild_keys(vault, rows, all_papers, status_filter, keys):
"""确定需要重建的论文列表。
--all: 只选 _needs_derived_rebuild()=True 的论文
--status: 按用户指定状态，不过滤版本
explicit keys: manual override，不过滤版本
Returns (selected_keys: list[str], reasons: dict[str, str])
"""
by_key = {r.key: r for r in rows}
if all_papers:
selected = []
reasons = {}
for r in rows:
if not r.can_rebuild:
continue
need, reason = _needs_derived_rebuild(vault, r.key)
if need:
selected.append(r.key)
reasons[r.key] = reason
return selected, reasons
if status_filter:
selected = [r.key for r in rows if r.status == status_filter and r.can_rebuild]
return selected, {}
if keys:
selected = [k for k in keys if k in by_key and by_key[k].can_rebuild]
return selected, {}
return [], {}
def _run_ocr_rebuild(
vault: Path,
keys: list[str] | None = None,
@ -330,46 +411,26 @@ def _run_ocr_rebuild(
from paperforge.worker.ocr_rebuild import run_derived_rebuild_for_keys
rows = collect_maintenance_rows(vault)
selected, reasons = _select_rebuild_keys(vault, rows, all_papers, status_filter, keys)
if all_papers:
keys = [r.key for r in rows if r.can_rebuild]
elif status_filter:
keys = [r.key for r in rows if r.status == status_filter and r.can_rebuild]
elif keys:
valid = {r.key for r in rows}
keys = [k for k in keys if k in valid]
else:
print("Specify paper keys, --all, or --status")
return 1
if not keys:
if not selected:
print("No papers matched for rebuild.")
return 0
# Resume: skip keys already in checkpoint (.done.* markers)
cp_dir = vault / "System" / "PaperForge" / ".ocr_rebuild_checkpoint"
if resume and cp_dir.exists():
done = {p.name.removeprefix(".done.") for p in cp_dir.glob(".done.*")}
skipped = [k for k in keys if k in done]
keys = [k for k in keys if k not in done]
if skipped:
print(f"Skipped {len(skipped)} paper(s) already in checkpoint.")
if not keys:
print("No papers to rebuild (all done).")
return 0
if resume:
print("Note: OCR rebuild resume is now version/artifact based; .done markers are ignored.")
if dry_run:
print(f"Would rebuild {len(keys)} paper(s):")
for k in keys:
print(f" - {k}")
print(f"Would rebuild {len(selected)} paper(s):")
for k in selected:
reason = reasons.get(k, "manual_override")
print(f" - {k}: {reason}")
return 0
from paperforge.worker._progress import progress_bar
result = run_derived_rebuild_for_keys(
vault, keys,
vault, selected,
progress_bar=progress_bar,
checkpoint_dir=cp_dir if resume else None,
parallel=parallel_workers,
)
count = result.get("rebuild_count", 0)

30
paperforge/worker/ocr_rebuild.py
View file

@ -101,7 +101,9 @@ def _apply_post_rebuild_version_flags(meta: dict) -> dict:
- raw_upgradable is preserved (only raw OCR can clear this)
- version_state_updated_at is refreshed
"""
from paperforge.worker.ocr_versions import expected_derived_payload
updated = dict(meta)
updated["derived_version"] = expected_derived_payload()
updated["derived_stale"] = False
updated["version_state_updated_at"] = datetime.datetime.now().isoformat()
return updated
@ -121,22 +123,8 @@ def select_papers_for_derived_rebuild(papers: list[dict]) -> list[str]:
def _filter_completed_keys(checkpoint_dir: Path | None, keys: list[str]) -> list[str]:
"""Return keys that do not have a .done.<key> marker in checkpoint_dir."""
if not checkpoint_dir:
return keys
cp = Path(checkpoint_dir)
if not cp.exists():
return keys
done = {p.name.removeprefix(".done.") for p in cp.glob(".done.*")}
return [k for k in keys if k not in done]
def _write_done_marker(checkpoint_dir: Path | None, key: str) -> None:
"""Write a completion marker for a successfully rebuilt paper."""
if not checkpoint_dir:
return
(Path(checkpoint_dir) / f".done.{key}").touch()
def _rebuild_one_paper(vault: Path, key: str) -> dict:
@ -578,7 +566,7 @@ def _rebuild_one_paper(vault: Path, key: str) -> dict:
_phase5_finalize(resolved, structured, rendered, span_meta_patch, health_overall=health_overall)
return {"key": key, "status": "ok"}
def _run_parallel_rebuild(vault: Path, keys: list[str], workers: int, checkpoint_dir: Path | None) -> list[dict]:
def _run_parallel_rebuild(vault: Path, keys: list[str], workers: int) -> list[dict]:
"""Run rebuild in parallel using a process pool."""
from concurrent.futures import ProcessPoolExecutor, as_completed
@ -589,8 +577,6 @@ def _run_parallel_rebuild(vault: Path, keys: list[str], workers: int, checkpoint
key = futures[future]
try:
result = future.result()
if result.get("status") == "ok":
_write_done_marker(checkpoint_dir, key)
results.append(result)
except Exception as e:
results.append({"key": key, "status": "failed", "error": str(e)})
@ -609,24 +595,23 @@ def run_derived_rebuild_for_keys(
Rebuilds: structured blocks, metadata, figure/table inventories, objects,
render outputs, and health — from stored raw blocks only.
If checkpoint_dir is provided, .done.<key> marker files track progress so
interrupted runs can skip completed work via --resume.
Note: checkpoint_dir / .done.<key> markers are no longer used; kept for
backward compatibility.
Args:
vault: Vault root path.
keys: Paper keys to rebuild.
progress_bar: Optional progress bar wrapper (tqdm-style).
checkpoint_dir: Directory for .done.<key> completion markers.
checkpoint_dir: Deprecated, kept for backward compatibility.
parallel: Number of parallel workers (0 = serial). Default 4.
"""
keys = _filter_completed_keys(checkpoint_dir, keys)
if not keys:
return {"rebuild_count": 0}
workers = int(parallel) if parallel else 0
if workers > 0 and len(keys) > 1:
results = _run_parallel_rebuild(vault, keys, workers, checkpoint_dir)
results = _run_parallel_rebuild(vault, keys, workers)
return {"rebuild_count": sum(1 for r in results if r.get("status") == "ok")}
rebuilt_count = 0
@ -635,7 +620,6 @@ def run_derived_rebuild_for_keys(
result = _rebuild_one_paper(vault, key)
if result.get("status") == "ok":
rebuilt_count += 1
_write_done_marker(checkpoint_dir, key)
return {"rebuild_count": rebuilt_count}

294
tests/test_pr9a_resume_rebuild.py Normal file
View file

@ -0,0 +1,294 @@
"""Tests for PR9A: Resume & Rebuild Correctness."""
import json
from pathlib import Path
from unittest.mock import patch, MagicMock, PropertyMock
import pytest
# ---------------------------------------------------------------------------
# Import shim: paperforge.commands.ocr lazily imports pipeline_paths from
# paperforge.worker._utils; provide the re-export before the ocr module loads.
# ---------------------------------------------------------------------------
import paperforge.config
from paperforge.worker._utils import pipeline_paths as _pipeline_paths_impl
paperforge.config.pipeline_paths = _pipeline_paths_impl
from paperforge.commands.ocr import _needs_derived_rebuild, _select_rebuild_keys
from paperforge.commands.embed import _pid_alive, _assert_collections_healthy
from paperforge.worker.ocr_rebuild import _apply_post_rebuild_version_flags
# ---------------------------------------------------------------------------
# _needs_derived_rebuild
# ---------------------------------------------------------------------------
# These tests patch paperforge.worker._utils.pipeline_paths because
# _needs_derived_rebuild imports it lazily at runtime.
# They also patch artifact_paths_for_root to avoid the real filesystem.
def _make_artifact_paths_mock(tmp_path, key, meta_content="{}",
blocks_raw_exists=True, source_meta_exists=True):
"""Build a mock for artifact_paths_for_root returning structured paths."""
paper_dir = tmp_path / key
paper_dir.mkdir(parents=True, exist_ok=True)
meta_file = paper_dir / "meta.json"
meta_file.write_text(meta_content)
blocks_raw = paper_dir / "blocks_raw"
if blocks_raw_exists:
blocks_raw.mkdir(parents=True, exist_ok=True)
source_meta = paper_dir / "source_metadata.json"
if source_meta_exists:
source_meta.write_text("{}")
mock_artifacts = MagicMock()
mock_artifacts.meta_json = meta_file
mock_artifacts.paper_root = paper_dir
mock_artifacts.blocks_raw = blocks_raw
mock_artifacts.source_metadata = source_meta
return mock_artifacts, paper_dir
class TestNeedsDerivedRebuild:
"""Tests for _needs_derived_rebuild."""
PATCH_PP = "paperforge.worker._utils.pipeline_paths"
PATCH_ART = "paperforge.commands.ocr.artifact_paths_for_root"
@patch(PATCH_PP)
@patch(PATCH_ART)
@patch("paperforge.commands.ocr._can_rebuild", return_value=True)
@patch("paperforge.commands.ocr.classify_version_state")
def test_version_mismatch(self, mock_cvs, mock_cr, mock_art, mock_pp):
"""Stale derived_version triggers (True, 'version_mismatch')."""
tmp_path = Path("/tmp/test_vm")
mock_pp.return_value = {"ocr": tmp_path}
mock_artifacts, _ = _make_artifact_paths_mock(tmp_path, "key")
mock_art.return_value = mock_artifacts
mock_cvs.return_value = {"derived_stale": True}
assert _needs_derived_rebuild(Path("/vault"), "key") == (True, "version_mismatch")
@patch(PATCH_PP)
@patch(PATCH_ART)
@patch("paperforge.commands.ocr._can_rebuild", return_value=True)
@patch("paperforge.commands.ocr.classify_version_state")
def test_missing_artifact(self, mock_cvs, mock_cr, mock_art, mock_pp):
"""Missing required artifact returns (True, 'missing:<name>')."""
tmp_path = Path("/tmp/test_ma")
mock_pp.return_value = {"ocr": tmp_path}
mock_artifacts, paper_dir = _make_artifact_paths_mock(tmp_path, "key")
mock_art.return_value = mock_artifacts
mock_cvs.return_value = {"derived_stale": False}
# Create all required files except index/role-index.json
present = [
"structure/blocks.structured.jsonl",
"render/render-map.json",
"index/structure-tree.json",
"fulltext.md",
"health/ocr_health.json",
]
for rel in present:
(paper_dir / rel).parent.mkdir(parents=True, exist_ok=True)
(paper_dir / rel).write_text("")
ok, reason = _needs_derived_rebuild(Path("/vault"), "key")
assert ok is True
assert reason == "missing:role-index.json"
@patch(PATCH_PP)
@patch(PATCH_ART)
@patch("paperforge.commands.ocr._can_rebuild", return_value=True)
@patch("paperforge.commands.ocr.classify_version_state")
def test_current(self, mock_cvs, mock_cr, mock_art, mock_pp):
"""All checks pass -> (False, 'current')."""
tmp_path = Path("/tmp/test_cu")
mock_pp.return_value = {"ocr": tmp_path}
mock_artifacts, paper_dir = _make_artifact_paths_mock(tmp_path, "key")
mock_art.return_value = mock_artifacts
mock_cvs.return_value = {"derived_stale": False}
required = [
"structure/blocks.structured.jsonl",
"render/render-map.json",
"index/structure-tree.json",
"index/role-index.json",
"fulltext.md",
"health/ocr_health.json",
]
for rel in required:
(paper_dir / rel).parent.mkdir(parents=True, exist_ok=True)
(paper_dir / rel).write_text("")
assert _needs_derived_rebuild(Path("/vault"), "key") == (False, "current")
@patch(PATCH_PP)
@patch(PATCH_ART)
@patch("paperforge.commands.ocr._can_rebuild", return_value=False)
def test_cannot_rebuild(self, mock_cr, mock_art, mock_pp):
"""_can_rebuild=False -> (False, 'cannot_rebuild')."""
tmp_path = Path("/tmp/test_cnr")
mock_pp.return_value = {"ocr": tmp_path}
mock_artifacts, _ = _make_artifact_paths_mock(tmp_path, "key")
mock_art.return_value = mock_artifacts
assert _needs_derived_rebuild(Path("/vault"), "key") == (False, "cannot_rebuild")
# ---------------------------------------------------------------------------
# _select_rebuild_keys
# ---------------------------------------------------------------------------
class FakeRow:
"""Minimal replacement for OCRMaintenanceRow."""
def __init__(self, key, status="done", can_rebuild=True):
self.key = key
self.status = status
self.can_rebuild = can_rebuild
class TestSelectRebuildKeys:
"""Tests for _select_rebuild_keys."""
@patch("paperforge.commands.ocr._needs_derived_rebuild")
def test_all_papers_selects_stale_only(self, mock_need):
"""all_papers=True selects only keys where _needs_derived_rebuild is True."""
rows = [
FakeRow("stale1"),
FakeRow("stale2"),
FakeRow("current1"),
]
def side_effect(vault, key):
return (True, "version_mismatch") if key.startswith("stale") else (False, "current")
mock_need.side_effect = side_effect
selected, reasons = _select_rebuild_keys(
Path("/vault"), rows, all_papers=True, status_filter=None, keys=None
)
assert selected == ["stale1", "stale2"]
assert reasons == {"stale1": "version_mismatch", "stale2": "version_mismatch"}
def test_explicit_keys_skip_version_check(self):
"""Explicit keys are selected regardless of version state."""
rows = [
FakeRow("current_key"),
FakeRow("other_key"),
]
selected, reasons = _select_rebuild_keys(
Path("/vault"), rows, all_papers=False, status_filter=None, keys=["current_key"]
)
assert selected == ["current_key"]
assert reasons == {}
def test_status_filter_skips_version_check(self):
"""status_filter selects matching rows regardless of version state."""
rows = [
FakeRow("k1", status="done_degraded"),
FakeRow("k2", status="done"),
]
selected, reasons = _select_rebuild_keys(
Path("/vault"), rows, all_papers=False, status_filter="done_degraded", keys=None
)
assert selected == ["k1"]
assert reasons == {}
def test_cannot_rebuild_excluded(self):
"""Keys with can_rebuild=False are excluded from explicit keys."""
rows = [
FakeRow("good", can_rebuild=True),
FakeRow("bad", can_rebuild=False),
]
selected, reasons = _select_rebuild_keys(
Path("/vault"), rows, all_papers=False, status_filter=None, keys=["good", "bad"]
)
assert selected == ["good"]
def test_empty_all_returns_empty(self):
"""all_papers=True with all current returns empty list."""
rows = [FakeRow("p1"), FakeRow("p2")]
with patch("paperforge.commands.ocr._needs_derived_rebuild", return_value=(False, "current")):
selected, reasons = _select_rebuild_keys(
Path("/vault"), rows, all_papers=True, status_filter=None, keys=None
)
assert selected == []
assert reasons == {}
# ---------------------------------------------------------------------------
# _apply_post_rebuild_version_flags
# ---------------------------------------------------------------------------
class TestPostRebuildVersionFlags:
"""Tests for _apply_post_rebuild_version_flags."""
def test_writes_derived_version(self):
"""Rebuild completion writes derived_version dict with structure_version."""
meta = {"derived_stale": True, "ocr_status": "done", "version_state_updated_at": ""}
result = _apply_post_rebuild_version_flags(meta)
assert "derived_version" in result
assert isinstance(result["derived_version"], dict)
assert "structure_version" in result["derived_version"]
assert not result["derived_stale"]
assert result["version_state_updated_at"] != ""
def test_preserves_existing_keys(self):
"""Function preserves other meta keys."""
meta = {"zotero_key": "ABC123", "ocr_status": "done", "derived_stale": True}
result = _apply_post_rebuild_version_flags(meta)
assert result["zotero_key"] == "ABC123"
# ---------------------------------------------------------------------------
# _pid_alive
# ---------------------------------------------------------------------------
class TestPidAlive:
"""Tests for _pid_alive."""
def test_dead_pid(self):
"""PID 0 or negative is always dead."""
assert not _pid_alive(0)
assert not _pid_alive(-1)
# ---------------------------------------------------------------------------
# _assert_collections_healthy
# ---------------------------------------------------------------------------
class TestAssertCollectionsHealthy:
"""Tests for _assert_collections_healthy."""
@patch("paperforge.commands.embed.get_collection")
def test_healthy_collections(self, mock_get_collection):
"""Three healthy collections -> (True, '')."""
mock_col = MagicMock()
mock_col.count.return_value = 42
mock_get_collection.return_value = mock_col
ok, msg = _assert_collections_healthy(Path("/vault"))
assert ok is True
assert msg == ""
@patch("paperforge.commands.embed.get_collection")
def test_corrupted_collection(self, mock_get_collection):
"""One corrupted collection -> (False, 'name: error')."""
def mock_get(vault_arg, *, name):
if name == "paperforge_body":
raise RuntimeError("connection refused")
mock_col = MagicMock()
mock_col.count.return_value = 42
return mock_col
mock_get_collection.side_effect = mock_get
ok, msg = _assert_collections_healthy(Path("/vault"))
assert ok is False
assert "paperforge_body" in msg
assert "connection refused" in msg