Merge pull request #203 from TideDra/fix-issue-202-bufferedreader

Fix non-pickleable retriever worker failures

Author: TideDra

src/zotero_arxiv_daily/retriever/arxiv_retriever.py

@@ -76,7 +76,11 @@ def extract_text_from_pdf(paper: ArxivResult) -> str | None:
         if paper.pdf_url is None:
             logger.warning(f"No PDF URL available for {paper.title}")
             return None
-        urlretrieve(paper.pdf_url, path)
+        try:
+            urlretrieve(paper.pdf_url, path)
+        except Exception as e:
+            logger.warning(f"Failed to download pdf for {paper.title}: {type(e).__name__}: {e}")
+            return None
         try:
             full_text = extract_markdown_from_pdf(path)
         except Exception as e:
@@ -91,7 +95,11 @@ def extract_text_from_tar(paper: ArxivResult) -> str | None:
         if source_url is None:
             logger.warning(f"No source URL available for {paper.title}")
             return None
-        urlretrieve(source_url, path)
+        try:
+            urlretrieve(source_url, path)
+        except Exception as e:
+            logger.warning(f"Failed to download source for {paper.title}: {type(e).__name__}: {e}")
+            return None
         try:
             file_contents = extract_tex_code_from_tar(path, paper.entry_id)
             if "all" not in file_contents:
@@ -101,4 +109,4 @@ def extract_text_from_tar(paper: ArxivResult) -> str | None:
         except Exception as e:
             logger.warning(f"Failed to extract full text of {paper.title} from tar: {e}")
             full_text = None
-        return full_text
+        return full_text

src/zotero_arxiv_daily/retriever/base.py

@@ -5,6 +5,30 @@
 from tqdm import tqdm
 from typing import Type
 from loguru import logger
+
+
+def _describe_raw_paper(raw_paper: RawPaperItem) -> str:
+    title = getattr(raw_paper, "title", None)
+    if title:
+        return str(title)
+    if isinstance(raw_paper, dict):
+        for key in ("title", "entry_id", "id", "doi"):
+            value = raw_paper.get(key)
+            if value:
+                return str(value)
+    return repr(raw_paper)
+
+
+def _convert_to_paper_safe(retriever: "BaseRetriever", raw_paper: RawPaperItem) -> Paper | None:
+    try:
+        return retriever.convert_to_paper(raw_paper)
+    except Exception as exc:
+        logger.warning(
+            f"Skipping paper {_describe_raw_paper(raw_paper)}: {type(exc).__name__}: {exc}"
+        )
+        return None
+
+
 class BaseRetriever(ABC):
     name: str
     def __init__(self, config:DictConfig):
@@ -24,10 +48,17 @@ def retrieve_papers(self) -> list[Paper]:
         papers = []
         logger.info("Processing papers...")
         with ProcessPoolExecutor(max_workers=self.config.executor.max_workers) as exec_pool:
-            futures = {exec_pool.submit(self.convert_to_paper, rp): i for i, rp in enumerate(raw_papers)}
+            futures = {exec_pool.submit(_convert_to_paper_safe, self, rp): i for i, rp in enumerate(raw_papers)}
             papers = [None] * len(raw_papers)
             for future in tqdm(as_completed(futures), total=len(raw_papers), desc="Converting papers"):
-                papers[futures[future]] = future.result()
+                try:
+                    papers[futures[future]] = future.result()
+                except Exception as exc:
+                    raw_paper = raw_papers[futures[future]]
+                    logger.warning(
+                        f"Skipping paper {_describe_raw_paper(raw_paper)} after worker failure: "
+                        f"{type(exc).__name__}: {exc}"
+                    )
         return [p for p in papers if p is not None]
 
 registered_retrievers = {}
@@ -42,4 +73,4 @@ def decorator(cls):
 def get_retriever_cls(name:str) -> Type[BaseRetriever]:
     if name not in registered_retrievers:
         raise ValueError(f"Retriever {name} not found")
-    return registered_retrievers[name]
+    return registered_retrievers[name]

tests/retriever/test_arxiv_retriever.py

@@ -1,6 +1,10 @@
 from zotero_arxiv_daily.retriever.arxiv_retriever import ArxivRetriever
+from zotero_arxiv_daily.retriever.base import BaseRetriever, register_retriever
+from zotero_arxiv_daily.protocol import Paper
 import feedparser
-import pickle
+import io
+from omegaconf import open_dict
+from urllib.error import HTTPError
 
 def test_arxiv_retriever(config, monkeypatch):
 
@@ -18,4 +22,42 @@ def mock_feedparser_parse(url):
     assert len(papers) == len(parsed_results)
     paper_titles = [i.title for i in papers]
     parsed_titles = [i.title for i in parsed_results]
-    assert set(paper_titles) == set(parsed_titles)
+    assert set(paper_titles) == set(parsed_titles)
+
+
+@register_retriever("failing_test")
+class FailingTestRetriever(BaseRetriever):
+    def _retrieve_raw_papers(self) -> list[dict[str, str]]:
+        return [
+            {"title": "good paper", "mode": "ok"},
+            {"title": "bad paper", "mode": "fail"},
+        ]
+
+    def convert_to_paper(self, raw_paper: dict[str, str]) -> Paper | None:
+        if raw_paper["mode"] == "fail":
+            raise HTTPError(
+                url="https://example.com/paper.pdf",
+                code=404,
+                msg="not found",
+                hdrs=None,
+                fp=io.BufferedReader(io.BytesIO(b"missing")),
+            )
+        return Paper(
+            source=self.name,
+            title=raw_paper["title"],
+            authors=[],
+            abstract="",
+            url=f"https://example.com/{raw_paper['mode']}",
+        )
+
+
+def test_retrieve_papers_skips_non_pickleable_worker_errors(config):
+    with open_dict(config.source):
+        config.source.failing_test = {}
+    config.executor.max_workers = 2
+
+    retriever = FailingTestRetriever(config)
+
+    papers = retriever.retrieve_papers()
+
+    assert [paper.title for paper in papers] == ["good paper"]