lllin000_PaperForge/docs/plans/pr9b-embed-parallel-encode.md

PR9B: Embed Parallel Encode

Depends on: PR9A (resume/rebuild correctness)
Scope: 并行 embedding API encode + 串行 ChromaDB write
不改变 resume/selection 逻辑，只加速

---

瓶颈分析

当前时序（--force, 729 papers）

Total: ~30 min
  Per paper: ~2.5s
    API encode:              ~1.5s    ← 纯网络 I/O，可并行
    ChromaDB write:          ~0.8s    ← HNSW 构建，必须串行
    Metadata + state:        ~0.2s

目标

4 worker 并行 → 729 ÷ 4 × (1.5s + 0.8s) ≈ 10 min（~3× 加速）

---

设计原则

1. paper 级原子性：同一篇论文的 body/object/legacy 全部 encode 成功后才 replace
2. prepare 在主线程（读 DB、chunk fulltext），不参与并行
3. encode 在线程池（纯网络 I/O），每个 worker 创建自己的 provider
4. write 在主线程串行（ChromaDB HNSW 构建不是线程安全）
5. delete old vectors 在 encode 成功后、write 前执行

核心 contract

A paper is deleted and rewritten only after ALL payloads for that paper
have been successfully encoded.

---

数据结构

@dataclass
class EmbeddingPayload:
    """准备阶段产出的载荷：texts 和 metadatas，未 encode。"""
    collection_name: str
    texts: list[str]
    ids: list[str]
    metadatas: list[dict]

@dataclass
class EncodedPayload:
    """encode 后的载荷：包含 embeddings。必须带 texts（不可从 metadata 反推）。"""
    collection_name: str
    texts: list[str]
    ids: list[str]
    metadatas: list[dict]
    embeddings: list[list[float]]

@dataclass
class PaperEmbeddingJob:
    """一篇论文的所有载荷，打包成一个 job。"""
    paper_id: str
    payloads: list[EmbeddingPayload]

@dataclass
class PaperEncodedBundle:
    """一篇论文 encode 完成后的结果。"""
    paper_id: str
    payloads: list[EncodedPayload]
    chunk_count: int

---

三段式 pipeline

Phase 1: PREPARE（主线程串行）

为每篇论文准备 1–3 个 EmbeddingPayload（legacy / body / object），分别对应三个 collection。

def prepare_payloads_for_entry(
    vault: Path, entry: dict
) -> list[EmbeddingPayload] | None:
    """为一篇论文准备所有 payload。不写 ChromaDB，不改状态，不删 vectors。"""
    key = entry.get("zotero_key")
    payloads = []

    has_body = _has_body_units_in_db(vault, key)
    has_object = _has_object_units_in_db(vault, key)

    if has_body or has_object:
        if has_body:
            body_units = get_body_units_for_embedding(vault, key)
            payloads.append(prepare_body_payload(key, body_units))
        if has_object:
            object_units = get_object_units_for_embedding(vault, key)
            payloads.append(prepare_object_payload(key, object_units))
    else:
        # Legacy path
        fulltext_rel = entry.get("fulltext_path", "")
        if not fulltext_rel:
            return None
        fulltext_path = vault / fulltext_rel
        chunks = chunk_fulltext(fulltext_path)          # chunk 在主线程做
        if chunks:
            payloads.append(prepare_legacy_payload(key, chunks))

    return payloads if payloads else None

prepare_legacy_payload()

def prepare_legacy_payload(zotero_key: str, chunks: list[dict]) -> EmbeddingPayload:
    texts = [c["text"] for c in chunks]
    ids = [f"{zotero_key}_{c['chunk_index']}" for c in chunks]
    metadatas = [
        {
            "paper_id": zotero_key,
            "section": c.get("section", ""),
            "page_number": c.get("page_number", 0),
            "chunk_index": c.get("chunk_index", i),
            "token_estimate": c.get("token_estimate", 0),
        }
        for i, c in enumerate(chunks)
    ]
    return EmbeddingPayload(
        collection_name="paperforge_fulltext",
        texts=texts, ids=ids, metadatas=metadatas,
    )

prepare_body_payload()

def prepare_body_payload(zotero_key: str, body_units: list[dict]) -> EmbeddingPayload:
    current_hash = compute_body_units_hash(body_units)
    texts = [u["unit_text"] for u in body_units]
    ids = [u["unit_id"] for u in body_units]
    metadatas = [
        {
            "paper_id": zotero_key,
            "section_path": u.get("section_path", ""),
            "unit_id": u["unit_id"],
            "unit_kind": "body",
            "body_units_hash": current_hash,
            "retrieval_policy_version": RETRIEVAL_POLICY_VERSION,
            "token_estimate": u.get("token_estimate", 0),
        }
        for u in body_units
    ]
    return EmbeddingPayload(collection_name="paperforge_body", texts=texts, ids=ids, metadatas=metadatas)

prepare_object_payload() 同理

collection_name="paperforge_objects"，metadata 包含 unit_kind="object", object_kind, object_label, object_units_hash。

向下兼容

def embed_body_units(vault, key, body_units):
    payload = prepare_body_payload(key, body_units)
    encoded = _encode_payload(vault, payload)
    _write_encoded_payload(vault, encoded)
    return len(body_units)

def embed_paper(vault, key, chunks):
    payload = prepare_legacy_payload(key, chunks)
    encoded = _encode_payload(vault, payload)
    _write_encoded_payload(vault, encoded)
    return len(chunks)

旧函数保留为兼容 wrapper，不破坏已有调用。

Phase 2: ENCODE（线程池，并行）

def _encode_payload(vault: Path, payload: EmbeddingPayload) -> EncodedPayload:
    """纯函数：encode texts → embeddings。每个 worker 内创建自己的 provider。"""
    from paperforge.embedding.providers.openai_compatible import OpenAICompatibleProvider
    provider = OpenAICompatibleProvider(vault)
    embeddings = provider.encode(payload.texts)
    return EncodedPayload(
        collection_name=payload.collection_name,
        texts=payload.texts,
        ids=payload.ids,
        metadatas=payload.metadatas,
        embeddings=embeddings,
    )

def _encode_paper_job(vault: Path, job: PaperEmbeddingJob) -> PaperEncodedBundle:
    """将一篇论文的所有 payload 顺序 encode（worker 线程内）。"""
    encoded_payloads = []
    total_chunks = 0
    for payload in job.payloads:
        encoded = _encode_payload(vault, payload)
        encoded_payloads.append(encoded)
        total_chunks += len(payload.ids)
    return PaperEncodedBundle(
        paper_id=job.paper_id,
        payloads=encoded_payloads,
        chunk_count=total_chunks,
    )

注意：

• OpenAICompatibleProvider 在线程内创建，不共享 client
• 同一篇论文的多个 payload（body + object）在同一个 worker 线程内顺序 encode
• 一篇论文的所有 payload 全部 encode 成功后才返回 → paper 级原子

Phase 3: WRITE（主线程串行）

def _write_encoded_payload(vault: Path, encoded: EncodedPayload):
    """写入 ChromaDB。主线程串行，HNSW 构建不是线程安全。"""
    col = get_collection(vault, name=encoded.collection_name)
    col.add(
        ids=encoded.ids,
        embeddings=encoded.embeddings,
        documents=encoded.texts,
        metadatas=encoded.metadatas,
    )

---

主流程

from concurrent.futures import ThreadPoolExecutor, as_completed
from itertools import batched

BATCH_SIZE = 8      # 每批最多 8 篇 in-flight
MAX_WORKERS = 4     # 并行 encode 线程数

papers_embedded = 0
chunks_embedded = 0

# Phase 1: PREPARE（主线程）
# resume/hash skip 已在主线程完成
jobs: list[PaperEmbeddingJob] = []
for entry in papers_iter:
    payloads = prepare_payloads_for_entry(vault, entry)
    if payloads:
        jobs.append(PaperEmbeddingJob(
            paper_id=entry.get("zotero_key"),
            payloads=payloads,
        ))

# Phase 2+3: ENCODE（线程池）+ WRITE（主线程）
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
    for batch in batched(jobs, BATCH_SIZE):
        futures = [
            pool.submit(_encode_paper_job, vault, job)
            for job in batch
        ]

        for future in as_completed(futures):
            try:
                bundle = future.result()    # 如果抛异常，不删旧 vectors
            except Exception as e:
                # 不处理这篇，继续下一篇（旧 vectors 还在）
                logger.error(f"Encode failed for {bundle.paper_id}: {e}")
                continue

            # encode 全部成功后才 replace old vectors
            delete_paper_vectors(vault, bundle.paper_id)

            for payload in bundle.payloads:
                try:
                    _write_encoded_payload(vault, payload)
                except Exception as e:
                    # write 失败 → 已进入 replace 状态
                    # 写 failed state，下一次 resume 用 hash 修复
                    logger.error(f"Write failed for {bundle.paper_id}: {e}")
                    mark_vector_build_state(vault, status="failed", ...)
                    raise

            chunks_embedded += bundle.chunk_count
            papers_embedded += 1
            mark_vector_build_state(vault, current=papers_embedded, ...)

---

边界处理

情况	行为
API key invalid	provider.encode() 抛异常 → future.result() 抛 → 不删旧 vectors，跳过这篇
某篇 body encode 成功、object encode 失败	worker 线程抛异常 → bundle 不返回 → 旧 vectors 完整保留
_write_encoded_payload 失败	已执行 delete → 写 failed state
线程池满载	batched(jobs, BATCH_SIZE) 限制 in-flight 数量
provider 被多个 worker 共享？	不会，每个 worker 内部创建自己的 provider
ChromaDB add() 异常	主线程 catch → 写 failed state → 当前批次的 futures 取消

---

文件清单

File	Change
paperforge/embedding/builder.py	新增 prepare_legacy_payload(), prepare_body_payload(), prepare_object_payload()；embed_body_units/embed_paper 保留为兼容 wrapper
paperforge/embedding/__init__.py	导出新函数（如需要）
paperforge/commands/embed.py	主循环改为三段式 pipeline + ThreadPoolExecutor

---

风险

风险	缓解
httpx.Client 跨线程安全？	每个 worker 创建独立 provider，不共享 client
ChromaDB add() 线程安全？	只在主线程串行调用
API rate limit	MAX_WORKERS=4 参数化，可调小
batched 需要 Python 3.12	如果 3.12 以下手动实现 batched()
legacy chunk_fulltext 本身慢	prepare 在主线程做，如果成为瓶颈再考虑并行 prepare