Harden FRAMES corpus fetching against Wikimedia rate limits
This commit is contained in:
Yiorgis Gozadinos
parent
152d57d6e2
commit
c00ccb7322
2 changed files with 158 additions and 11 deletions
|
|
@ -10,6 +10,7 @@ import ast
|
|||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
|
|
@ -27,6 +28,9 @@ from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
USER_AGENT = "haiku.rag-evaluations (https://github.com/ggozad/haiku.rag)"
|
||||
FETCH_ATTEMPTS = 3
|
||||
THROTTLE_SECONDS = 1.0
|
||||
RATE_LIMIT_BACKOFF_SECONDS = 60.0
|
||||
|
||||
|
||||
def load_frames_test() -> Dataset:
|
||||
|
|
@ -153,6 +157,13 @@ def _fetch_article_page(
|
|||
return response.text, "html", revid, title
|
||||
|
||||
|
||||
def _backoff_seconds(error: Exception, attempt: int) -> float:
|
||||
if isinstance(error, httpx.HTTPStatusError) and error.response.status_code == 429:
|
||||
retry_after = error.response.headers.get("retry-after")
|
||||
return float(retry_after) if retry_after else RATE_LIMIT_BACKOFF_SECONDS
|
||||
return 5.0 * attempt
|
||||
|
||||
|
||||
def fetch_article(
|
||||
uri: str, cache_dir: Path, client: httpx.Client | None
|
||||
) -> dict[str, Any] | None:
|
||||
|
|
@ -170,17 +181,24 @@ def fetch_article(
|
|||
return row
|
||||
|
||||
assert client is not None
|
||||
try:
|
||||
title = unquote(urlsplit(uri).path[len("/wiki/") :])
|
||||
if title.startswith("Category:"):
|
||||
content, format, revid = _fetch_category_page(
|
||||
urlsplit(uri).netloc, title, client
|
||||
)
|
||||
else:
|
||||
content, format, revid, title = _fetch_article_page(uri, client)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to fetch {uri}: {e}")
|
||||
return None
|
||||
title = unquote(urlsplit(uri).path[len("/wiki/") :])
|
||||
# Wikimedia throttles sustained bot traffic; pace uncached fetches.
|
||||
time.sleep(THROTTLE_SECONDS)
|
||||
for attempt in range(1, FETCH_ATTEMPTS + 1):
|
||||
try:
|
||||
if title.startswith("Category:"):
|
||||
content, format, revid = _fetch_category_page(
|
||||
urlsplit(uri).netloc, title, client
|
||||
)
|
||||
else:
|
||||
content, format, revid, title = _fetch_article_page(uri, client)
|
||||
break
|
||||
except Exception as e:
|
||||
if attempt == FETCH_ATTEMPTS:
|
||||
logger.warning(f"Failed to fetch {uri}: {e}")
|
||||
return None
|
||||
logger.info(f"Retrying {uri} after error: {e}")
|
||||
time.sleep(_backoff_seconds(e, attempt))
|
||||
|
||||
row: dict[str, Any] = {
|
||||
"uri": uri,
|
||||
|
|
@ -228,6 +246,11 @@ def load_frames_corpus() -> list[dict[str, Any]]:
|
|||
if index % 100 == 0:
|
||||
logger.info(f"Fetched {index}/{len(uris)} articles")
|
||||
logger.info(f"Fetched {len(rows)}/{len(uris)} articles")
|
||||
if len(rows) < len(uris):
|
||||
raise RuntimeError(
|
||||
f"Fetched only {len(rows)}/{len(uris)} FRAMES articles; "
|
||||
"refusing to build a partial corpus. Re-run to resume from cache."
|
||||
)
|
||||
_cached_corpus = rows
|
||||
return _cached_corpus
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from evaluations.datasets.frames import (
|
||||
FETCH_ATTEMPTS,
|
||||
build_frames_case,
|
||||
fetch_article,
|
||||
map_frames_document,
|
||||
|
|
@ -564,3 +567,124 @@ class TestFrames:
|
|||
content = Path(row["path"]).read_text()
|
||||
assert "1908 Summer Olympics" in content
|
||||
assert "2012 Summer Olympics" in content
|
||||
|
||||
def test_fetch_article_retries_transient_failures(
|
||||
self, tmp_path: Path, monkeypatch
|
||||
) -> None:
|
||||
sleeps: list[float] = []
|
||||
monkeypatch.setattr(
|
||||
"evaluations.datasets.frames.time.sleep", lambda s: sleeps.append(s)
|
||||
)
|
||||
|
||||
class FlakyResponse:
|
||||
text = "<html><body>ok</body></html>"
|
||||
headers = {"etag": 'W/"42/uuid"'}
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
class FlakyClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def get(self, url: str, params: dict | None = None) -> FlakyResponse:
|
||||
self.calls += 1
|
||||
if self.calls < 3:
|
||||
raise OSError("connection reset")
|
||||
return FlakyResponse()
|
||||
|
||||
client = FlakyClient()
|
||||
row = fetch_article(
|
||||
"https://en.wikipedia.org/wiki/Capybara",
|
||||
tmp_path,
|
||||
client=client, # ty: ignore[invalid-argument-type]
|
||||
)
|
||||
assert row is not None
|
||||
assert row["revid"] == "42"
|
||||
assert client.calls == 3
|
||||
# One throttle sleep before fetching plus one backoff per failure.
|
||||
assert len(sleeps) == 3
|
||||
|
||||
def test_fetch_article_honors_retry_after_on_rate_limit(
|
||||
self, tmp_path: Path, monkeypatch
|
||||
) -> None:
|
||||
import httpx
|
||||
|
||||
sleeps: list[float] = []
|
||||
monkeypatch.setattr(
|
||||
"evaluations.datasets.frames.time.sleep", lambda s: sleeps.append(s)
|
||||
)
|
||||
request = httpx.Request("GET", "https://en.wikipedia.org/x")
|
||||
|
||||
class OkResponse:
|
||||
text = "<html><body>ok</body></html>"
|
||||
headers = {"etag": 'W/"42/uuid"'}
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
class RateLimitedClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def get(self, url: str, params: dict | None = None) -> OkResponse:
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
raise httpx.HTTPStatusError(
|
||||
"429 too many requests",
|
||||
request=request,
|
||||
response=httpx.Response(
|
||||
429, headers={"retry-after": "13"}, request=request
|
||||
),
|
||||
)
|
||||
return OkResponse()
|
||||
|
||||
row = fetch_article(
|
||||
"https://en.wikipedia.org/wiki/Capybara",
|
||||
tmp_path,
|
||||
client=RateLimitedClient(), # ty: ignore[invalid-argument-type]
|
||||
)
|
||||
assert row is not None
|
||||
assert 13.0 in sleeps
|
||||
|
||||
def test_fetch_article_gives_up_after_max_attempts(
|
||||
self, tmp_path: Path, monkeypatch
|
||||
) -> None:
|
||||
monkeypatch.setattr("evaluations.datasets.frames.time.sleep", lambda s: None)
|
||||
|
||||
class DeadClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def get(self, url: str, params: dict | None = None):
|
||||
self.calls += 1
|
||||
raise OSError("connection reset")
|
||||
|
||||
client = DeadClient()
|
||||
row = fetch_article(
|
||||
"https://en.wikipedia.org/wiki/Capybara",
|
||||
tmp_path,
|
||||
client=client, # ty: ignore[invalid-argument-type]
|
||||
)
|
||||
assert row is None
|
||||
assert client.calls == FETCH_ATTEMPTS
|
||||
|
||||
def test_load_corpus_raises_on_partial_fetch(self, monkeypatch) -> None:
|
||||
import evaluations.datasets.frames as frames
|
||||
|
||||
monkeypatch.setattr(frames, "_cached_corpus", None)
|
||||
monkeypatch.setattr(
|
||||
frames,
|
||||
"load_frames_test",
|
||||
lambda: [
|
||||
{
|
||||
"wiki_links": "['https://en.wikipedia.org/wiki/A', "
|
||||
"'https://en.wikipedia.org/wiki/B']"
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
frames, "fetch_article", lambda uri, cache_dir, client: None
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="0/2"):
|
||||
frames.load_frames_corpus()
|
||||
|
|
|
|||
Loading…
Reference in a new issue