add greptile feedback

Signed-off-by: Sarah Yurick <sarahyurick@gmail.com>

Author: sarahyurick

nemo_curator/stages/text/download/wikipedia/url_generation.py

@@ -97,7 +97,20 @@ def _get_wikipedia_urls(self) -> list[str]:
             dump_data = None
             for candidate_date in sorted(dump_dates, reverse=True):
                 candidate_dump_date = f"{candidate_date}/"
-                candidate_dump_data = self._get_data_for_dump(candidate_dump_date, wiki_index_url)
+                try:
+                    candidate_dump_data = self._get_data_for_dump(candidate_dump_date, wiki_index_url)
+                except requests.HTTPError as e:
+                    status_code = e.response.status_code if e.response is not None else None
+                    if status_code is not None and (
+                        status_code == requests.codes.too_many_requests
+                        or status_code >= requests.codes.internal_server_error
+                    ):
+                        logger.warning(
+                            f"Unable to load dump data for {candidate_date} due to HTTP {status_code}; "
+                            "trying next dump"
+                        )
+                        continue
+                    raise
                 if candidate_dump_data is None:
                     logger.warning(f"Cannot load dump data for {candidate_date}")
                     continue

tests/stages/text/download/wikipedia/test_url_generation.py

@@ -142,6 +142,48 @@ def mock_get_side_effect(url: str, **kwargs) -> Mock:  # noqa: ARG001
             "enwiki-20230501-pages-articles-multistream1.xml.bz2"
         ]
 
+    @pytest.mark.parametrize("transient_status", [429, 500, 503])
+    @patch("requests.get")
+    def test_get_latest_dump_date_skips_transient_status_error(self, mock_get: Mock, transient_status: int):
+        """A transient status error for a candidate does not prevent fallback to an older dump."""
+        mock_html = """
+        <a href="20230501/">20230501/</a>
+        <a href="20230601/">20230601/</a>
+        """
+        completed_dump = {
+            "jobs": {
+                "articlesmultistreamdump": {
+                    "status": "done",
+                    "files": {"enwiki-20230501-pages-articles-multistream1.xml.bz2": {}},
+                }
+            }
+        }
+
+        def mock_get_side_effect(url: str, **kwargs) -> Mock:  # noqa: ARG001
+            response = Mock(status_code=200)
+            if url == "https://dumps.wikimedia.org/enwiki":
+                response.content = mock_html.encode("utf-8")
+            elif url == "https://dumps.wikimedia.org/enwiki/20230601/dumpstatus.json":
+                response.status_code = transient_status
+                response.raise_for_status.side_effect = requests.HTTPError(
+                    f"{transient_status} Server Error", response=response
+                )
+            elif url == "https://dumps.wikimedia.org/enwiki/20230501/dumpstatus.json":
+                response.content = json.dumps(completed_dump).encode("utf-8")
+            else:
+                error_msg = f"Unexpected URL: {url}"
+                raise ValueError(error_msg)
+            return response
+
+        mock_get.side_effect = mock_get_side_effect
+
+        urls = WikipediaUrlGenerator(language="en").generate_urls()
+
+        assert urls == [
+            "https://dumps.wikimedia.org/enwiki/20230501/"
+            "enwiki-20230501-pages-articles-multistream1.xml.bz2"
+        ]
+
     @patch("requests.get")
     def test_get_latest_dump_date_http_error(self, mock_get: Mock):
         """HTTP errors from the dump index are surfaced instead of parsed as HTML."""