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 json
|
||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
|
import time
|
||||||
from collections.abc import Mapping
|
from collections.abc import Mapping
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
@ -27,6 +28,9 @@ from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
USER_AGENT = "haiku.rag-evaluations (https://github.com/ggozad/haiku.rag)"
|
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:
|
def load_frames_test() -> Dataset:
|
||||||
|
|
@ -153,6 +157,13 @@ def _fetch_article_page(
|
||||||
return response.text, "html", revid, title
|
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(
|
def fetch_article(
|
||||||
uri: str, cache_dir: Path, client: httpx.Client | None
|
uri: str, cache_dir: Path, client: httpx.Client | None
|
||||||
) -> dict[str, Any] | None:
|
) -> dict[str, Any] | None:
|
||||||
|
|
@ -170,17 +181,24 @@ def fetch_article(
|
||||||
return row
|
return row
|
||||||
|
|
||||||
assert client is not None
|
assert client is not None
|
||||||
try:
|
title = unquote(urlsplit(uri).path[len("/wiki/") :])
|
||||||
title = unquote(urlsplit(uri).path[len("/wiki/") :])
|
# Wikimedia throttles sustained bot traffic; pace uncached fetches.
|
||||||
if title.startswith("Category:"):
|
time.sleep(THROTTLE_SECONDS)
|
||||||
content, format, revid = _fetch_category_page(
|
for attempt in range(1, FETCH_ATTEMPTS + 1):
|
||||||
urlsplit(uri).netloc, title, client
|
try:
|
||||||
)
|
if title.startswith("Category:"):
|
||||||
else:
|
content, format, revid = _fetch_category_page(
|
||||||
content, format, revid, title = _fetch_article_page(uri, client)
|
urlsplit(uri).netloc, title, client
|
||||||
except Exception as e:
|
)
|
||||||
logger.warning(f"Failed to fetch {uri}: {e}")
|
else:
|
||||||
return None
|
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] = {
|
row: dict[str, Any] = {
|
||||||
"uri": uri,
|
"uri": uri,
|
||||||
|
|
@ -228,6 +246,11 @@ def load_frames_corpus() -> list[dict[str, Any]]:
|
||||||
if index % 100 == 0:
|
if index % 100 == 0:
|
||||||
logger.info(f"Fetched {index}/{len(uris)} articles")
|
logger.info(f"Fetched {index}/{len(uris)} articles")
|
||||||
logger.info(f"Fetched {len(rows)}/{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
|
_cached_corpus = rows
|
||||||
return _cached_corpus
|
return _cached_corpus
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,9 @@
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
from evaluations.datasets.frames import (
|
from evaluations.datasets.frames import (
|
||||||
|
FETCH_ATTEMPTS,
|
||||||
build_frames_case,
|
build_frames_case,
|
||||||
fetch_article,
|
fetch_article,
|
||||||
map_frames_document,
|
map_frames_document,
|
||||||
|
|
@ -564,3 +567,124 @@ class TestFrames:
|
||||||
content = Path(row["path"]).read_text()
|
content = Path(row["path"]).read_text()
|
||||||
assert "1908 Summer Olympics" in content
|
assert "1908 Summer Olympics" in content
|
||||||
assert "2012 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