From 058ced37ffd2cf8e6f478379c6fec71896a14bbb Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Thu, 4 Jun 2026 09:47:23 +0100 Subject: [PATCH 1/3] fix(retrievers): fix 429 silently truncating embedding downloads The embedding-parts download loop in download_mtrag_embeddings caught urllib.error.HTTPError to detect end-of-parts (404 = no more files). Since HTTPError covers all HTTP errors, a transient 429 or 5xx was treated as end-of-data, silently producing a truncated embedding set with no error raised. Fix: only break on 404; all other HTTP errors now propagate. Also adds _urlretrieve_with_retry to retry transient errors (429, 5xx) with exponential back-off before propagating. Closes #1198 Assisted-by: Claude Code Signed-off-by: Nigel Jones --- mellea/formatters/granite/retrievers/util.py | 36 ++++++++++++++++---- 1 file changed, 30 insertions(+), 6 deletions(-) diff --git a/mellea/formatters/granite/retrievers/util.py b/mellea/formatters/granite/retrievers/util.py index e5552fba7..e73c730f9 100644 --- a/mellea/formatters/granite/retrievers/util.py +++ b/mellea/formatters/granite/retrievers/util.py @@ -5,6 +5,7 @@ # Standard import os import pathlib +import time import urllib.error import urllib.request import zipfile @@ -19,6 +20,21 @@ 'Please install them with: pip install "mellea[granite_retriever]"' ) from e +_RETRYABLE_HTTP_CODES = (429, 500, 502, 503, 504) + + +def _urlretrieve_with_retry(url: str, target: str, max_attempts: int = 3) -> None: + """Download *url* to *target*, retrying on transient HTTP errors (429, 5xx).""" + for attempt in range(1, max_attempts + 1): + try: + urllib.request.urlretrieve(url, target) + return + except urllib.error.HTTPError as exc: + if exc.code in _RETRYABLE_HTTP_CODES and attempt < max_attempts: + time.sleep(2**attempt) + continue + raise + def download_mtrag_corpus(target_dir: str, corpus_name: str) -> pathlib.Path: """Download a corpus file from the MTRAG benchmark if the file hasn't already present. @@ -33,6 +49,8 @@ def download_mtrag_corpus(target_dir: str, corpus_name: str) -> pathlib.Path: Raises: ValueError: If `corpus_name` is not one of the supported corpus names. + urllib.error.HTTPError: On transient HTTP errors (429, 5xx) after all + retry attempts are exhausted. """ corpus_names = ("cloud", "clapnq", "fiqa", "govt") if corpus_name not in corpus_names: @@ -44,7 +62,7 @@ def download_mtrag_corpus(target_dir: str, corpus_name: str) -> pathlib.Path: f"https://github.com/IBM/mt-rag-benchmark/raw/refs/heads/main/" f"corpora/{corpus_name}.jsonl.zip" ) - urllib.request.urlretrieve(source_url, target_file) + _urlretrieve_with_retry(source_url, str(target_file)) return target_file @@ -99,7 +117,9 @@ def read_mtrag_corpus(corpus_file: str | pathlib.Path) -> pa.Table: return t -def download_mtrag_embeddings(embedding_name: str, corpus_name: str, target_dir: str): +def download_mtrag_embeddings( + embedding_name: str, corpus_name: str, target_dir: str +) -> None: """Download precomputed embeddings for a corpus in the MTRAG benchmark. Args: @@ -114,6 +134,8 @@ def download_mtrag_embeddings(embedding_name: str, corpus_name: str, target_dir: ValueError: If `corpus_name` is not one of the supported corpus names, or if no precomputed embeddings are found for the given corpus and embedding model combination. + urllib.error.HTTPError: On transient HTTP errors (429, 5xx) after all + retry attempts are exhausted. """ corpus_names = ("cloud", "clapnq", "fiqa", "govt") if corpus_name not in corpus_names: @@ -134,11 +156,13 @@ def download_mtrag_embeddings(embedding_name: str, corpus_name: str, target_dir: ) target_file = target_root / parquet_file_name try: - urllib.request.urlretrieve(source_url, target_file) + _urlretrieve_with_retry(source_url, str(target_file)) part_num += 1 - except urllib.error.HTTPError: - # Found all the parts; flow through - break + except urllib.error.HTTPError as exc: + if exc.code == 404: + # Found all the parts; flow through + break + raise # 429/5xx propagate rather than silently truncating if part_num == 1: raise ValueError( From 5c025868eea3151c9844c4bda6e360e30c975c9f Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Thu, 4 Jun 2026 11:39:50 +0100 Subject: [PATCH 2/3] test(formatters): add unit tests for _urlretrieve_with_retry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Covers success, per-retryable-code retry, non-retryable 404 fast-fail, exhausted retries, and max_attempts=1 — all mocked so no network I/O. Assisted-by: Claude Code Signed-off-by: Nigel Jones --- .../granite/test_retrievers_util.py | 104 ++++++++++++++++++ 1 file changed, 104 insertions(+) create mode 100644 test/formatters/granite/test_retrievers_util.py diff --git a/test/formatters/granite/test_retrievers_util.py b/test/formatters/granite/test_retrievers_util.py new file mode 100644 index 000000000..72a44f2ec --- /dev/null +++ b/test/formatters/granite/test_retrievers_util.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for mellea.formatters.granite.retrievers.util._urlretrieve_with_retry.""" + +import http.client +import urllib.error +from unittest.mock import call, patch + +import pytest + +pytest.importorskip( + "pyarrow", reason="pyarrow not installed — install mellea[granite_retriever]" +) + +from mellea.formatters.granite.retrievers.util import ( + _urlretrieve_with_retry, # type: ignore[reportAttributeAccessIssue] +) + + +def _http_error(code: int) -> urllib.error.HTTPError: + return urllib.error.HTTPError( + url="http://example.com", + code=code, + msg=str(code), + hdrs=http.client.HTTPMessage(), + fp=None, + ) + + +def test_success_first_attempt(): + with ( + patch("urllib.request.urlretrieve") as mock_retrieve, + patch("time.sleep") as mock_sleep, + ): + _urlretrieve_with_retry("http://example.com/f", "/tmp/f") + + mock_retrieve.assert_called_once_with("http://example.com/f", "/tmp/f") + mock_sleep.assert_not_called() + + +def test_retries_on_429_succeeds_second_attempt(): + with ( + patch("urllib.request.urlretrieve") as mock_retrieve, + patch("time.sleep") as mock_sleep, + ): + mock_retrieve.side_effect = [_http_error(429), None] + _urlretrieve_with_retry("http://example.com/f", "/tmp/f") + + assert mock_retrieve.call_count == 2 + mock_sleep.assert_called_once_with(2) # 2**1 + + +@pytest.mark.parametrize("code", [500, 502, 503, 504]) +def test_retries_on_5xx(code: int): + with ( + patch("urllib.request.urlretrieve") as mock_retrieve, + patch("time.sleep") as mock_sleep, + ): + mock_retrieve.side_effect = [_http_error(code), None] + _urlretrieve_with_retry("http://example.com/f", "/tmp/f") + + assert mock_retrieve.call_count == 2 + mock_sleep.assert_called_once_with(2) + + +def test_non_retryable_raises_immediately(): + with ( + patch("urllib.request.urlretrieve") as mock_retrieve, + patch("time.sleep") as mock_sleep, + ): + mock_retrieve.side_effect = _http_error(404) + with pytest.raises(urllib.error.HTTPError) as exc_info: + _urlretrieve_with_retry("http://example.com/f", "/tmp/f") + + assert exc_info.value.code == 404 + mock_retrieve.assert_called_once() + mock_sleep.assert_not_called() + + +def test_raises_after_all_retries_exhausted(): + with ( + patch("urllib.request.urlretrieve") as mock_retrieve, + patch("time.sleep") as mock_sleep, + ): + mock_retrieve.side_effect = _http_error(429) + with pytest.raises(urllib.error.HTTPError) as exc_info: + _urlretrieve_with_retry("http://example.com/f", "/tmp/f") + + assert exc_info.value.code == 429 + assert mock_retrieve.call_count == 3 # default max_attempts=3 + assert mock_sleep.call_args_list == [call(2), call(4)] # 2**1, 2**2 + + +def test_max_attempts_one_no_retry(): + with ( + patch("urllib.request.urlretrieve") as mock_retrieve, + patch("time.sleep") as mock_sleep, + ): + mock_retrieve.side_effect = _http_error(429) + with pytest.raises(urllib.error.HTTPError): + _urlretrieve_with_retry("http://example.com/f", "/tmp/f", max_attempts=1) + + mock_retrieve.assert_called_once() + mock_sleep.assert_not_called() From 22cc6f4ec4a514a89447414d9c5ad844e8a68749 Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Thu, 4 Jun 2026 11:53:58 +0100 Subject: [PATCH 3/3] fix(retrievers): propagate HTTP errors instead of retrying in library code Replace _urlretrieve_with_retry (which embedded application-level retry/sleep logic) with direct urllib.request.urlretrieve calls. Error handling is the caller's responsibility per project conventions. The only behaviour change: download_mtrag_embeddings now breaks only on 404 (end-of-parts signal) and propagates all other HTTP errors (429, 5xx) rather than silently treating them as end-of-data and returning a truncated result. Assisted-by: Claude Code Signed-off-by: Nigel Jones --- mellea/formatters/granite/retrievers/util.py | 27 +--- .../granite/test_retrievers_util.py | 121 +++++++++--------- 2 files changed, 63 insertions(+), 85 deletions(-) diff --git a/mellea/formatters/granite/retrievers/util.py b/mellea/formatters/granite/retrievers/util.py index e73c730f9..383ace898 100644 --- a/mellea/formatters/granite/retrievers/util.py +++ b/mellea/formatters/granite/retrievers/util.py @@ -5,7 +5,6 @@ # Standard import os import pathlib -import time import urllib.error import urllib.request import zipfile @@ -20,21 +19,6 @@ 'Please install them with: pip install "mellea[granite_retriever]"' ) from e -_RETRYABLE_HTTP_CODES = (429, 500, 502, 503, 504) - - -def _urlretrieve_with_retry(url: str, target: str, max_attempts: int = 3) -> None: - """Download *url* to *target*, retrying on transient HTTP errors (429, 5xx).""" - for attempt in range(1, max_attempts + 1): - try: - urllib.request.urlretrieve(url, target) - return - except urllib.error.HTTPError as exc: - if exc.code in _RETRYABLE_HTTP_CODES and attempt < max_attempts: - time.sleep(2**attempt) - continue - raise - def download_mtrag_corpus(target_dir: str, corpus_name: str) -> pathlib.Path: """Download a corpus file from the MTRAG benchmark if the file hasn't already present. @@ -49,8 +33,7 @@ def download_mtrag_corpus(target_dir: str, corpus_name: str) -> pathlib.Path: Raises: ValueError: If `corpus_name` is not one of the supported corpus names. - urllib.error.HTTPError: On transient HTTP errors (429, 5xx) after all - retry attempts are exhausted. + urllib.error.HTTPError: On any HTTP error downloading the corpus file. """ corpus_names = ("cloud", "clapnq", "fiqa", "govt") if corpus_name not in corpus_names: @@ -62,7 +45,7 @@ def download_mtrag_corpus(target_dir: str, corpus_name: str) -> pathlib.Path: f"https://github.com/IBM/mt-rag-benchmark/raw/refs/heads/main/" f"corpora/{corpus_name}.jsonl.zip" ) - _urlretrieve_with_retry(source_url, str(target_file)) + urllib.request.urlretrieve(source_url, str(target_file)) return target_file @@ -134,8 +117,8 @@ def download_mtrag_embeddings( ValueError: If `corpus_name` is not one of the supported corpus names, or if no precomputed embeddings are found for the given corpus and embedding model combination. - urllib.error.HTTPError: On transient HTTP errors (429, 5xx) after all - retry attempts are exhausted. + urllib.error.HTTPError: On any HTTP error other than 404 (which signals + end-of-parts). """ corpus_names = ("cloud", "clapnq", "fiqa", "govt") if corpus_name not in corpus_names: @@ -156,7 +139,7 @@ def download_mtrag_embeddings( ) target_file = target_root / parquet_file_name try: - _urlretrieve_with_retry(source_url, str(target_file)) + urllib.request.urlretrieve(source_url, str(target_file)) part_num += 1 except urllib.error.HTTPError as exc: if exc.code == 404: diff --git a/test/formatters/granite/test_retrievers_util.py b/test/formatters/granite/test_retrievers_util.py index 72a44f2ec..3633dde1e 100644 --- a/test/formatters/granite/test_retrievers_util.py +++ b/test/formatters/granite/test_retrievers_util.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for mellea.formatters.granite.retrievers.util._urlretrieve_with_retry.""" +"""Unit tests for mellea.formatters.granite.retrievers.util download helpers.""" import http.client import urllib.error -from unittest.mock import call, patch +from unittest.mock import patch import pytest @@ -12,8 +12,9 @@ "pyarrow", reason="pyarrow not installed — install mellea[granite_retriever]" ) -from mellea.formatters.granite.retrievers.util import ( - _urlretrieve_with_retry, # type: ignore[reportAttributeAccessIssue] +from mellea.formatters.granite.retrievers.util import ( # type: ignore[reportAttributeAccessIssue] + download_mtrag_corpus, + download_mtrag_embeddings, ) @@ -27,78 +28,72 @@ def _http_error(code: int) -> urllib.error.HTTPError: ) -def test_success_first_attempt(): - with ( - patch("urllib.request.urlretrieve") as mock_retrieve, - patch("time.sleep") as mock_sleep, - ): - _urlretrieve_with_retry("http://example.com/f", "/tmp/f") +# --------------------------------------------------------------------------- +# download_mtrag_corpus +# --------------------------------------------------------------------------- - mock_retrieve.assert_called_once_with("http://example.com/f", "/tmp/f") - mock_sleep.assert_not_called() +def test_corpus_raises_on_invalid_name(tmp_path): + with pytest.raises(ValueError, match="cloud"): + download_mtrag_corpus(str(tmp_path), "invalid_corpus") -def test_retries_on_429_succeeds_second_attempt(): - with ( - patch("urllib.request.urlretrieve") as mock_retrieve, - patch("time.sleep") as mock_sleep, - ): - mock_retrieve.side_effect = [_http_error(429), None] - _urlretrieve_with_retry("http://example.com/f", "/tmp/f") - assert mock_retrieve.call_count == 2 - mock_sleep.assert_called_once_with(2) # 2**1 +def test_corpus_skips_download_if_file_exists(tmp_path): + target = tmp_path / "cloud.jsonl.zip" + target.write_bytes(b"dummy") + with patch("urllib.request.urlretrieve") as mock_retrieve: + download_mtrag_corpus(str(tmp_path), "cloud") + mock_retrieve.assert_not_called() -@pytest.mark.parametrize("code", [500, 502, 503, 504]) -def test_retries_on_5xx(code: int): - with ( - patch("urllib.request.urlretrieve") as mock_retrieve, - patch("time.sleep") as mock_sleep, - ): - mock_retrieve.side_effect = [_http_error(code), None] - _urlretrieve_with_retry("http://example.com/f", "/tmp/f") - - assert mock_retrieve.call_count == 2 - mock_sleep.assert_called_once_with(2) +def test_corpus_downloads_when_file_missing(tmp_path): + with patch("urllib.request.urlretrieve") as mock_retrieve: + download_mtrag_corpus(str(tmp_path), "fiqa") + mock_retrieve.assert_called_once() + assert "fiqa" in mock_retrieve.call_args[0][0] -def test_non_retryable_raises_immediately(): - with ( - patch("urllib.request.urlretrieve") as mock_retrieve, - patch("time.sleep") as mock_sleep, - ): - mock_retrieve.side_effect = _http_error(404) +def test_corpus_propagates_http_error(tmp_path): + with patch("urllib.request.urlretrieve", side_effect=_http_error(429)): with pytest.raises(urllib.error.HTTPError) as exc_info: - _urlretrieve_with_retry("http://example.com/f", "/tmp/f") + download_mtrag_corpus(str(tmp_path), "cloud") + assert exc_info.value.code == 429 - assert exc_info.value.code == 404 - mock_retrieve.assert_called_once() - mock_sleep.assert_not_called() +# --------------------------------------------------------------------------- +# download_mtrag_embeddings — 404 breaks cleanly, other errors propagate +# --------------------------------------------------------------------------- -def test_raises_after_all_retries_exhausted(): - with ( - patch("urllib.request.urlretrieve") as mock_retrieve, - patch("time.sleep") as mock_sleep, - ): - mock_retrieve.side_effect = _http_error(429) - with pytest.raises(urllib.error.HTTPError) as exc_info: - _urlretrieve_with_retry("http://example.com/f", "/tmp/f") - assert exc_info.value.code == 429 - assert mock_retrieve.call_count == 3 # default max_attempts=3 - assert mock_sleep.call_args_list == [call(2), call(4)] # 2**1, 2**2 +def test_embeddings_404_on_first_part_raises_value_error(tmp_path): + """A 404 on part_001 means no embeddings exist — ValueError, not HTTPError.""" + with patch("urllib.request.urlretrieve", side_effect=_http_error(404)): + with pytest.raises(ValueError, match="No precomputed embeddings"): + download_mtrag_embeddings("model", "cloud", str(tmp_path)) -def test_max_attempts_one_no_retry(): - with ( - patch("urllib.request.urlretrieve") as mock_retrieve, - patch("time.sleep") as mock_sleep, - ): - mock_retrieve.side_effect = _http_error(429) - with pytest.raises(urllib.error.HTTPError): - _urlretrieve_with_retry("http://example.com/f", "/tmp/f", max_attempts=1) +def test_embeddings_stops_after_404(tmp_path): + """Downloads part_001, then 404 on part_002 → returns cleanly with one part.""" + call_count = 0 - mock_retrieve.assert_called_once() - mock_sleep.assert_not_called() + def side_effect(_url, dest): + nonlocal call_count + call_count += 1 + if call_count == 1: + open(dest, "wb").close() # create the file + return + raise _http_error(404) + + with patch("urllib.request.urlretrieve", side_effect=side_effect): + download_mtrag_embeddings("model", "cloud", str(tmp_path)) + + assert call_count == 2 + + +@pytest.mark.parametrize("code", [429, 500, 502, 503, 504]) +def test_embeddings_propagates_non_404_errors(tmp_path, code): + """429/5xx must propagate — never silently treated as end-of-parts.""" + with patch("urllib.request.urlretrieve", side_effect=_http_error(code)): + with pytest.raises(urllib.error.HTTPError) as exc_info: + download_mtrag_embeddings("model", "cloud", str(tmp_path)) + assert exc_info.value.code == code