diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index afafbe5f9c..eb107db1c4 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -895,6 +895,9 @@ jobs: env: EXTRA: cpu HF_TOKEN: ${{ secrets.HF_TOKEN }} + NMP_FILES_HF_RETRY_ATTEMPTS: "7" + NMP_FILES_HF_RETRY_INITIAL_DELAY_SECONDS: "1" + NMP_FILES_HF_RETRY_MAX_DELAY_SECONDS: "30" PYTEST_WORKERS: "4" _TYPER_FORCE_DISABLE_TERMINAL: "1" - name: Upload test artifacts @@ -1059,6 +1062,9 @@ jobs: E2E_SERVICES_LOG_DIR: ${{ runner.temp }}/e2e-services-logs NGC_API_KEY: ${{ secrets.NGC_REGISTRY_READ_TOKEN }} HF_TOKEN: ${{ secrets.HF_TOKEN }} + NMP_FILES_HF_RETRY_ATTEMPTS: "7" + NMP_FILES_HF_RETRY_INITIAL_DELAY_SECONDS: "1" + NMP_FILES_HF_RETRY_MAX_DELAY_SECONDS: "30" - name: Dump server logs if: always() run: | diff --git a/docs/set-up/config-reference.mdx b/docs/set-up/config-reference.mdx index 434afdbf4f..39a2b7510b 100644 --- a/docs/set-up/config-reference.mdx +++ b/docs/set-up/config-reference.mdx @@ -188,6 +188,12 @@ files: file_lock_ttl_seconds: 300 # Maximum concurrent downloads during cache warming | default: 3 cache_warming_max_concurrent: 3 + # Maximum Hugging Face request attempts for transient failures. | default: 4 + hf_retry_attempts: 4 + # Initial Hugging Face retry delay in seconds before exponential backoff. | default: 0.5 + hf_retry_initial_delay_seconds: 0.5 + # Maximum Hugging Face retry delay in seconds. | default: 5.0 + hf_retry_max_delay_seconds: 5.0 ``` ### `inference_gateway` diff --git a/services/core/files/src/nmp/core/files/app/backends/huggingface.py b/services/core/files/src/nmp/core/files/app/backends/huggingface.py index 2bf7356a5e..eb316fc78f 100644 --- a/services/core/files/src/nmp/core/files/app/backends/huggingface.py +++ b/services/core/files/src/nmp/core/files/app/backends/huggingface.py @@ -6,13 +6,15 @@ from __future__ import annotations import logging +from collections.abc import AsyncIterator, Callable from dataclasses import dataclass, field -from typing import ( - AsyncIterator, -) +from datetime import datetime, timezone +from email.utils import parsedate_to_datetime +from typing import TypeVar import aiohttp -from anyio import to_thread +import httpx +from anyio import sleep, to_thread from huggingface_hub import HfApi, get_hf_file_metadata, hf_hub_url from huggingface_hub.utils import ( EntryNotFoundError, @@ -29,6 +31,7 @@ from nmp.core.files.app.external_hosts import validate_external_host from nmp.core.files.app.http_session import get_http_session from nmp.core.files.app.streaming import download_url_streaming +from nmp.core.files.config import FilesConfig, files_config from nmp.core.files.exceptions import ( NotFoundError, StorageAccessError, @@ -39,6 +42,8 @@ logger = logging.getLogger(__name__) +_T = TypeVar("_T") + class HuggingfaceBackendError(StorageBackendError): """Raised when there's issues talking to Huggingface.""" @@ -101,6 +106,224 @@ def raise_for_hf_status( raise HuggingfaceBackendError(f"HTTP {status_code}{context}") +def _map_hf_http_error(exc: HfHubHTTPError) -> Exception: + """Map HuggingFace Hub HTTP exceptions into storage backend exceptions.""" + if exc.response is not None: + try: + raise_for_hf_status( + exc.response.status_code, + dict(exc.response.headers), + str(exc.response.url), + ) + except ( + HuggingfaceAccessError, + HuggingfaceConfigError, + HuggingfaceUnavailableError, + HuggingfaceBackendError, + ) as mapped: + return mapped + return HuggingfaceBackendError(f"HuggingFace API error: {exc}") + + +def _ratelimit_reset_seconds(headers: dict[str, str]) -> float | None: + """Parse HuggingFace's RateLimit t= field.""" + raw_value = headers.get("RateLimit") or headers.get("ratelimit") + if not raw_value: + return None + + for part in raw_value.split(";"): + key, separator, value = part.strip().partition("=") + if key != "t" or not separator: + continue + try: + return max(0.0, float(value.split(",", 1)[0].strip().strip('"'))) + except ValueError: + continue + return None + + +def _retry_after_seconds(headers: dict[str, str] | None, status_code: int | None = None) -> float | None: + """Parse Retry-After, or HF's 429 RateLimit reset, into seconds.""" + if not headers: + return None + + raw_value = headers.get("Retry-After") or headers.get("retry-after") + if raw_value: + try: + return max(0.0, float(raw_value)) + except ValueError: + pass + + try: + retry_at = parsedate_to_datetime(raw_value) + except (TypeError, ValueError): + pass + else: + if retry_at.tzinfo is None: + retry_at = retry_at.replace(tzinfo=timezone.utc) + return max(0.0, (retry_at - datetime.now(timezone.utc)).total_seconds()) + + if status_code == 429: + return _ratelimit_reset_seconds(headers) + return None + + +async def _sleep_before_retry( + *, + operation: str, + attempt: int, + retry_config: FilesConfig, + headers: dict[str, str] | None, + status_code: int | None = None, + error: Exception, +) -> None: + """Sleep according to Retry-After or exponential backoff before retrying.""" + retry_after = _retry_after_seconds(headers, status_code=status_code) + delay = ( + min(retry_after, retry_config.hf_retry_max_delay_seconds) + if retry_after is not None + else min( + retry_config.hf_retry_initial_delay_seconds * (2 ** (attempt - 1)), + retry_config.hf_retry_max_delay_seconds, + ) + ) + logger.warning( + "Transient HuggingFace error during %s; retrying attempt %s/%s after %.2fs: %s", + operation, + attempt + 1, + retry_config.hf_retry_attempts, + delay, + error, + ) + await sleep(delay) + + +async def _run_hf_request(operation: str, request: Callable[[], _T]) -> _T: + """Run a blocking HuggingFace Hub request with transient retry handling.""" + retry_config = files_config() + for attempt in range(1, retry_config.hf_retry_attempts + 1): + try: + return await to_thread.run_sync(request) + except EntryNotFoundError: + raise + except HfHubHTTPError as exc: + mapped = _map_hf_http_error(exc) + if isinstance(mapped, HuggingfaceUnavailableError) and attempt < retry_config.hf_retry_attempts: + headers = dict(exc.response.headers) if exc.response is not None else None + status_code = exc.response.status_code if exc.response is not None else None + await _sleep_before_retry( + operation=operation, + attempt=attempt, + retry_config=retry_config, + headers=headers, + status_code=status_code, + error=mapped, + ) + continue + raise mapped from exc + except httpx.RequestError as exc: + mapped = HuggingfaceUnavailableError(f"Network error during {operation}: {exc}") + if attempt < retry_config.hf_retry_attempts: + await _sleep_before_retry( + operation=operation, + attempt=attempt, + retry_config=retry_config, + headers=None, + status_code=None, + error=mapped, + ) + continue + raise mapped from exc + + raise AssertionError("unreachable in _run_hf_request") + + +def _map_hf_download_error( + *, + path: str, + download_url: str, + exc: aiohttp.ClientError, +) -> tuple[Exception, dict[str, str] | None]: + """Map aiohttp download errors and preserve response headers for retry delays.""" + if isinstance(exc, aiohttp.ClientResponseError): + response_headers = dict(exc.headers) if exc.headers else None + try: + raise_for_hf_status(exc.status, response_headers, download_url) + except ( + HuggingfaceAccessError, + HuggingfaceConfigError, + HuggingfaceUnavailableError, + HuggingfaceBackendError, + ) as mapped: + return mapped, response_headers + return HuggingfaceBackendError(f"HTTP error downloading file {path}: {exc.status}"), response_headers + + return HuggingfaceUnavailableError(f"Network error downloading file {path}: {exc}"), None + + +async def _map_hf_download_stream_errors( + stream: AsyncIterator[bytes], + *, + path: str, + download_url: str, +) -> AsyncIterator[bytes]: + """Translate stream errors after retrying is no longer safe.""" + try: + async for chunk in stream: + yield chunk + except aiohttp.ClientError as exc: + mapped, _ = _map_hf_download_error(path=path, download_url=download_url, exc=exc) + raise mapped from exc + + +async def _stream_hf_download_with_retries( + *, + path: str, + download_url: str, + session: aiohttp.ClientSession, + headers: dict[str, str] | None, + byte_range: ByteRange | None, + chunk_size: int, +) -> AsyncIterator[bytes]: + """Stream a download, retrying only errors before the first chunk is emitted.""" + retry_config = files_config() + for attempt in range(1, retry_config.hf_retry_attempts + 1): + stream = download_url_streaming( + url=download_url, + session=session, + headers=dict(headers) if headers else None, + byte_range=byte_range, + chunk_size=chunk_size, + ) + + try: + first_chunk = await anext(stream) + except StopAsyncIteration: + return + except aiohttp.ClientError as exc: + mapped, response_headers = _map_hf_download_error(path=path, download_url=download_url, exc=exc) + if isinstance(mapped, HuggingfaceUnavailableError) and attempt < retry_config.hf_retry_attempts: + status_code = exc.status if isinstance(exc, aiohttp.ClientResponseError) else None + await _sleep_before_retry( + operation="download file", + attempt=attempt, + retry_config=retry_config, + headers=response_headers, + status_code=status_code, + error=mapped, + ) + continue + raise mapped from exc + + # Once bytes have reached the caller, retrying would replay the file prefix. + yield first_chunk + async for chunk in _map_hf_download_stream_errors(stream, path=path, download_url=download_url): + yield chunk + return + + raise AssertionError("unreachable in _stream_hf_download_with_retries") + + @dataclass class HuggingfaceStorageImpl(StorageImpl): config: HuggingfaceStorageConfig @@ -126,28 +349,20 @@ async def resolve_config(self) -> HuggingfaceStorageConfig: Raises: HuggingfaceConfigError: If the repository or revision is not found. """ - try: - info = await to_thread.run_sync( - lambda: self._api.repo_info( - repo_id=self.config.repo_id, - repo_type=self.config.repo_type, - revision=self.config.revision, - ) - ) - return self.config.model_copy( - update={ - "original_revision": self.config.revision, - "revision": info.sha, - } - ) - except HfHubHTTPError as exc: - if exc.response is not None: - raise_for_hf_status( - exc.response.status_code, - dict(exc.response.headers), - str(exc.response.url), - ) - raise HuggingfaceBackendError(f"HuggingFace API error: {exc}") from exc + info = await _run_hf_request( + "resolve repository revision", + lambda: self._api.repo_info( + repo_id=self.config.repo_id, + repo_type=self.config.repo_type, + revision=self.config.revision, + ), + ) + return self.config.model_copy( + update={ + "original_revision": self.config.revision, + "revision": info.sha, + } + ) def _get_download_url(self, filepath: str) -> str: """Generate a download URL for a file in the Huggingface repo.""" @@ -162,14 +377,18 @@ def _get_download_url(self, filepath: str) -> str: async def _get_hf_file_metadata(self, filepath: str): """Get file metadata from Huggingface for a specific file.""" url = self._get_download_url(filepath) - return await to_thread.run_sync(lambda: get_hf_file_metadata(url=url, token=self.secrets.get("token"))) + return await _run_hf_request( + "get file metadata", + lambda: get_hf_file_metadata(url=url, token=self.secrets.get("token")), + ) async def list_files(self, path: str | None = None) -> list[FileInfo]: """List files in the Huggingface repository.""" try: # list_repo_tree returns RepoFile and RepoFolder objects # We filter for files only (items with size attribute) - items = await to_thread.run_sync( + items = await _run_hf_request( + "list repository tree", lambda: list( self._api.list_repo_tree( repo_id=self.config.repo_id, @@ -178,7 +397,7 @@ async def list_files(self, path: str | None = None) -> list[FileInfo]: path_in_repo=path, recursive=True, ) - ) + ), ) except EntryNotFoundError: # list_repo_tree expects a directory path. If path points to a file @@ -192,14 +411,6 @@ async def list_files(self, path: str | None = None) -> list[FileInfo]: # Neither a directory nor a file - return empty list return [] return [] - except HfHubHTTPError as exc: - if exc.response is not None: - raise_for_hf_status( - exc.response.status_code, - dict(exc.response.headers), - str(exc.response.url), - ) - raise HuggingfaceBackendError(f"HuggingFace API error: {exc}") from exc file_infos = [] for item in items: @@ -225,17 +436,12 @@ async def get_file(self, path: str) -> FileInfo: url = self._get_download_url(path) try: - metadata = await to_thread.run_sync(lambda: get_hf_file_metadata(url=url, token=self.secrets.get("token"))) + metadata = await _run_hf_request( + "get file metadata", + lambda: get_hf_file_metadata(url=url, token=self.secrets.get("token")), + ) except EntryNotFoundError as exc: raise NotFoundError(f"File '{path}' not found in {self.config.repo_id}@{self.config.revision}") from exc - except HfHubHTTPError as exc: - if exc.response is not None: - raise_for_hf_status( - exc.response.status_code, - dict(exc.response.headers), - str(exc.response.url), - ) - raise HuggingfaceBackendError(f"HuggingFace API error: {exc}") from exc return FileInfo(path=path, size=metadata.size) @@ -255,20 +461,15 @@ async def download(self, path: str, byte_range: ByteRange | None) -> AsyncIterat async def _download() -> AsyncIterator[bytes]: session = get_http_session() - try: - async for chunk in download_url_streaming( - url=download_url, - session=session, - headers=headers if headers else None, - byte_range=byte_range, - chunk_size=self.config.read_chunk_size, - ): - yield chunk - except aiohttp.ClientResponseError as exc: - response_headers = dict(exc.headers) if exc.headers else None - raise_for_hf_status(exc.status, response_headers, download_url) - except aiohttp.ClientError as exc: - raise HuggingfaceBackendError(f"Network error downloading file {path}") from exc + async for chunk in _stream_hf_download_with_retries( + path=path, + download_url=download_url, + session=session, + headers=headers if headers else None, + byte_range=byte_range, + chunk_size=self.config.read_chunk_size, + ): + yield chunk return _download() @@ -285,12 +486,13 @@ async def validate_storage(self): """ validate_external_host(self.config.endpoint) try: - repo_info = await to_thread.run_sync( + repo_info = await _run_hf_request( + "validate repository", lambda: self._api.repo_info( repo_id=self.config.repo_id, repo_type=self.config.repo_type, revision=self.config.revision, - ) + ), ) # Verify we can actually download files by checking a file's metadata. @@ -299,14 +501,13 @@ async def validate_storage(self): sibling = repo_info.siblings[0] await self._get_hf_file_metadata(sibling.rfilename) - except HfHubHTTPError as exc: - if exc.response is not None: - raise_for_hf_status( - exc.response.status_code, - dict(exc.response.headers), - str(exc.response.url), - ) - raise HuggingfaceBackendError(f"HuggingFace API error: {exc}") from exc + except ( + HuggingfaceAccessError, + HuggingfaceConfigError, + HuggingfaceUnavailableError, + HuggingfaceBackendError, + ): + raise except Exception as exc: raise HuggingfaceBackendError( f"Failed to access Huggingface repository {self.config.repo_id}@{self.config.revision}" diff --git a/services/core/files/src/nmp/core/files/config.py b/services/core/files/src/nmp/core/files/config.py index de8a549a82..0dc1596a0e 100644 --- a/services/core/files/src/nmp/core/files/config.py +++ b/services/core/files/src/nmp/core/files/config.py @@ -61,6 +61,24 @@ def validate_allowed_external_hosts(self) -> "FilesConfig": description="Maximum concurrent downloads during cache warming", ) + hf_retry_attempts: int = Field( + default=4, + ge=1, + description="Maximum Hugging Face request attempts for transient failures.", + ) + + hf_retry_initial_delay_seconds: float = Field( + default=0.5, + ge=0.0, + description="Initial Hugging Face retry delay in seconds before exponential backoff.", + ) + + hf_retry_max_delay_seconds: float = Field( + default=5.0, + ge=0.0, + description="Maximum Hugging Face retry delay in seconds.", + ) + # TODO(v2): CONFIG @cache diff --git a/services/core/files/tests/test_files_config.py b/services/core/files/tests/test_files_config.py index 569355511c..d2b3f54e1e 100644 --- a/services/core/files/tests/test_files_config.py +++ b/services/core/files/tests/test_files_config.py @@ -67,3 +67,55 @@ def test_allowed_external_hosts_missing_netloc_raises(self) -> None: """Test that URL without scheme/netloc raises ValueError from model validator.""" with pytest.raises(ValueError, match="must be a valid URL"): FilesConfig(allowed_external_hosts="https://valid.com,not-a-url") + + +class TestFilesConfigHuggingFaceRetries: + """Tests for Hugging Face retry configuration.""" + + def test_hf_retry_defaults(self) -> None: + """Test Hugging Face retry default values.""" + config = FilesConfig() + assert config.hf_retry_attempts == 4 + assert config.hf_retry_initial_delay_seconds == 0.5 + assert config.hf_retry_max_delay_seconds == 5.0 + + def test_hf_retry_from_env_vars(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test NMP_FILES_HF_RETRY_* env vars set Hugging Face retry config.""" + monkeypatch.setenv("NMP_FILES_HF_RETRY_ATTEMPTS", "7") + monkeypatch.setenv("NMP_FILES_HF_RETRY_INITIAL_DELAY_SECONDS", "1") + monkeypatch.setenv("NMP_FILES_HF_RETRY_MAX_DELAY_SECONDS", "30") + + config = FilesConfig() + + assert config.hf_retry_attempts == 7 + assert config.hf_retry_initial_delay_seconds == 1.0 + assert config.hf_retry_max_delay_seconds == 30.0 + + def test_hf_retry_from_yaml(self) -> None: + """Test files.hf_retry_* values load from YAML settings.""" + settings = { + "files": { + "hf_retry_attempts": 6, + "hf_retry_initial_delay_seconds": 2.5, + "hf_retry_max_delay_seconds": 20, + }, + } + + config = Configuration.global_settings_to_service_config(settings, FilesConfig) + + assert config.hf_retry_attempts == 6 + assert config.hf_retry_initial_delay_seconds == 2.5 + assert config.hf_retry_max_delay_seconds == 20.0 + + def test_hf_retry_invalid_values_raise(self) -> None: + """Test invalid Hugging Face retry values fail validation.""" + from pydantic import ValidationError + + with pytest.raises(ValidationError): + FilesConfig(hf_retry_attempts=0) + + with pytest.raises(ValidationError): + FilesConfig(hf_retry_initial_delay_seconds=-1) + + with pytest.raises(ValidationError): + FilesConfig(hf_retry_max_delay_seconds=-1) diff --git a/services/core/files/tests/test_huggingface_backend.py b/services/core/files/tests/test_huggingface_backend.py index 629e1135fb..aae83198f0 100644 --- a/services/core/files/tests/test_huggingface_backend.py +++ b/services/core/files/tests/test_huggingface_backend.py @@ -3,12 +3,13 @@ """Tests for Huggingface storage backend.""" -from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, Mock, call, patch import aiohttp import httpx import pytest from nmp.common.api.common import SecretRef +from nmp.common.config import Configuration from nmp.core.files.app.backends.base import ByteRange from nmp.core.files.app.backends.factory import storage_impl_factory from nmp.core.files.app.backends.huggingface import ( @@ -18,8 +19,15 @@ HuggingfaceStorageConfig, HuggingfaceStorageImpl, HuggingfaceUnavailableError, + _retry_after_seconds, raise_for_hf_status, ) +from nmp.core.files.config import FilesConfig, files_config + + +def _clear_files_config_cache() -> None: + Configuration.clear_cache() + files_config.cache_clear() def _hf_http_error_without_response(message: str): @@ -32,6 +40,17 @@ def _hf_http_error_without_response(message: str): return error +def _hf_http_error(status_code: int, message: str = "HuggingFace error", headers: dict[str, str] | None = None): + """Create a Hugging Face HTTP error with a mock response.""" + from huggingface_hub.utils import HfHubHTTPError + + response = Mock() + response.status_code = status_code + response.headers = headers or {} + response.url = "https://huggingface.co/test-org/test-repo" + return HfHubHTTPError(message, response=response) + + @pytest.fixture def mock_httpx_response(): """Create a mock httpx.Response for Huggingface exceptions.""" @@ -151,6 +170,26 @@ async def test_list_files_no_filter(hf_config, mock_hf_api, hf_secrets_empty): ) +async def test_list_files_network_error_retries_then_succeeds(hf_config, mock_hf_api, hf_secrets_empty): + """Transient HuggingFace transport errors are retried.""" + mock_file = Mock() + mock_file.path = "file1.txt" + mock_file.size = 100 + mock_hf_api.list_repo_tree.side_effect = [ + httpx.ConnectError("Connection reset by peer"), + [mock_file], + ] + + impl = HuggingfaceStorageImpl(hf_config, hf_secrets_empty) + + with patch("nmp.core.files.app.backends.huggingface.sleep", new_callable=AsyncMock) as mock_sleep: + files = await impl.list_files() + + assert files[0].path == "file1.txt" + assert mock_hf_api.list_repo_tree.call_count == 2 + mock_sleep.assert_awaited_once_with(0.5) + + async def test_list_files_with_path_filter(hf_config, mock_hf_api, hf_secrets_empty): """Test listing files with path filter.""" mock_file1 = Mock() @@ -424,16 +463,84 @@ async def mock_chunks(): raise aiohttp.ClientError("Connection failed") yield # pragma: no cover - mock_stream.return_value = mock_chunks() + mock_stream.side_effect = lambda *args, **kwargs: mock_chunks() impl = HuggingfaceStorageImpl(hf_config, hf_secrets_empty) download_iter = await impl.download("test.txt", None) - with pytest.raises(HuggingfaceBackendError) as exc_info: + with ( + patch("nmp.core.files.app.backends.huggingface.sleep", new_callable=AsyncMock) as mock_sleep, + pytest.raises(HuggingfaceUnavailableError) as exc_info, + ): async for _ in download_iter: pass assert "Network error" in str(exc_info.value) + retry_attempts = FilesConfig().hf_retry_attempts + assert mock_stream.call_count == retry_attempts + assert mock_sleep.await_count == retry_attempts - 1 + + +async def test_download_retries_reuse_http_session(hf_config, hf_secrets_empty): + """Retry attempts reuse the backend's pooled HTTP session.""" + with ( + patch("nmp.core.files.app.backends.huggingface.hf_hub_url") as mock_url, + patch("nmp.core.files.app.backends.huggingface.get_http_session") as mock_get_session, + patch("nmp.core.files.app.backends.huggingface.download_url_streaming") as mock_stream, + ): + mock_url.return_value = "https://huggingface.co/test-org/test-repo/resolve/main/test.txt" + session = Mock() + mock_get_session.return_value = session + + async def failing_chunks(): + raise aiohttp.ClientError("Connection failed") + yield # pragma: no cover + + mock_stream.side_effect = lambda *args, **kwargs: failing_chunks() + + impl = HuggingfaceStorageImpl(hf_config, hf_secrets_empty) + download_iter = await impl.download("test.txt", None) + + with ( + patch("nmp.core.files.app.backends.huggingface.sleep", new_callable=AsyncMock), + pytest.raises(HuggingfaceUnavailableError), + ): + async for _ in download_iter: + pass + + assert mock_get_session.call_count == 1 + assert {call.kwargs["session"] for call in mock_stream.call_args_list} == {session} + + +async def test_download_does_not_retry_after_yielding_chunk(hf_config, hf_secrets_empty): + """A streaming failure after yielding bytes is propagated without replaying content.""" + with ( + patch("nmp.core.files.app.backends.huggingface.hf_hub_url") as mock_url, + patch("nmp.core.files.app.backends.huggingface.download_url_streaming") as mock_stream, + ): + mock_url.return_value = "https://huggingface.co/test-org/test-repo/resolve/main/test.txt" + + async def chunks_then_error(): + yield b"chunk1" + raise aiohttp.ClientError("Connection dropped") + + mock_stream.side_effect = lambda *args, **kwargs: chunks_then_error() + + impl = HuggingfaceStorageImpl(hf_config, hf_secrets_empty) + download_iter = await impl.download("test.txt", None) + + chunks = [] + with ( + patch("nmp.core.files.app.backends.huggingface.sleep", new_callable=AsyncMock) as mock_sleep, + pytest.raises(HuggingfaceUnavailableError) as exc_info, + ): + async for chunk in download_iter: + chunks.append(chunk) + + assert "Network error" in str(exc_info.value) + assert chunks == [b"chunk1"] + mock_stream.assert_called_once() + mock_sleep.assert_not_awaited() async def test_validate_storage_success(hf_config, mock_hf_api, hf_secrets_empty): @@ -585,25 +692,94 @@ async def test_get_file_gated_repo_error(hf_config, mock_httpx_response, hf_secr async def test_get_file_rate_limit_error(hf_config, hf_secrets_empty): """Test get_file when rate limited by HuggingFace.""" - from huggingface_hub.utils import HfHubHTTPError - - # Create a mock response with 429 status code - mock_response = Mock() - mock_response.status_code = 429 - mock_response.headers = {} - mock_response.url = "https://huggingface.co/test-org/test-repo/test.txt" - - mock_error = HfHubHTTPError("Rate limited", response=mock_response) + mock_error = _hf_http_error(429, "Rate limited", headers={"Retry-After": "0"}) with patch("nmp.core.files.app.backends.huggingface.get_hf_file_metadata") as mock_metadata: mock_metadata.side_effect = mock_error impl = HuggingfaceStorageImpl(hf_config, hf_secrets_empty) - with pytest.raises(HuggingfaceUnavailableError) as exc_info: + with ( + patch("nmp.core.files.app.backends.huggingface.sleep", new_callable=AsyncMock) as mock_sleep, + pytest.raises(HuggingfaceUnavailableError) as exc_info, + ): await impl.get_file("test.txt") assert "Rate limited" in str(exc_info.value) + retry_attempts = FilesConfig().hf_retry_attempts + assert mock_metadata.call_count == retry_attempts + assert mock_sleep.await_count == retry_attempts - 1 + + +async def test_get_file_rate_limit_uses_retry_env_override(hf_config, hf_secrets_empty, monkeypatch): + """HF retry attempts and exponential delays can be increased by environment.""" + monkeypatch.setenv("NMP_FILES_HF_RETRY_ATTEMPTS", "3") + monkeypatch.setenv("NMP_FILES_HF_RETRY_INITIAL_DELAY_SECONDS", "2") + monkeypatch.setenv("NMP_FILES_HF_RETRY_MAX_DELAY_SECONDS", "3") + _clear_files_config_cache() + + try: + mock_error = _hf_http_error(429, "Rate limited") + + with patch("nmp.core.files.app.backends.huggingface.get_hf_file_metadata") as mock_metadata: + mock_metadata.side_effect = mock_error + + impl = HuggingfaceStorageImpl(hf_config, hf_secrets_empty) + + with ( + patch("nmp.core.files.app.backends.huggingface.sleep", new_callable=AsyncMock) as mock_sleep, + pytest.raises(HuggingfaceUnavailableError), + ): + await impl.get_file("test.txt") + + assert mock_metadata.call_count == 3 + mock_sleep.assert_has_awaits([call(2.0), call(3.0)]) + finally: + _clear_files_config_cache() + + +def test_retry_after_seconds_uses_hf_ratelimit_header_for_429(): + """HF's RateLimit header exposes t= for 429s.""" + assert _retry_after_seconds({"RateLimit": '"resolvers";r=0;t=7'}, status_code=429) == 7.0 + + +def test_retry_after_seconds_ignores_hf_ratelimit_header_for_non_429(): + """RateLimit headers can appear on non-429 responses, so don't sleep on t=.""" + assert _retry_after_seconds({"RateLimit": '"resolvers";r=0;t=7'}, status_code=503) is None + + +def test_retry_after_seconds_prefers_retry_after_header(): + assert ( + _retry_after_seconds( + {"Retry-After": "2", "RateLimit": '"resolvers";r=0;t=7'}, + status_code=429, + ) + == 2.0 + ) + + +async def test_get_file_rate_limit_uses_hf_ratelimit_delay(hf_config, hf_secrets_empty, monkeypatch): + """429 retries use HF's RateLimit reset when Retry-After is absent.""" + monkeypatch.setenv("NMP_FILES_HF_RETRY_ATTEMPTS", "2") + monkeypatch.setenv("NMP_FILES_HF_RETRY_MAX_DELAY_SECONDS", "10") + _clear_files_config_cache() + + try: + mock_error = _hf_http_error(429, "Rate limited", headers={"RateLimit": '"resolvers";r=0;t=7'}) + + with patch("nmp.core.files.app.backends.huggingface.get_hf_file_metadata") as mock_metadata: + mock_metadata.side_effect = [mock_error, Mock(size=1234)] + + impl = HuggingfaceStorageImpl(hf_config, hf_secrets_empty) + + with patch("nmp.core.files.app.backends.huggingface.sleep", new_callable=AsyncMock) as mock_sleep: + file_info = await impl.get_file("test.txt") + + assert file_info.path == "test.txt" + assert file_info.size == 1234 + mock_sleep.assert_awaited_once_with(7.0) + finally: + _clear_files_config_cache() async def test_validate_storage_gated_repo_error(hf_config, mock_hf_api, mock_httpx_response, hf_secrets_empty): @@ -750,6 +926,29 @@ async def test_resolve_config_tag_to_sha(mock_hf_api, hf_secrets_empty): assert resolved_config.revision == "def789abc123456" +async def test_resolve_config_retries_rate_limit_then_succeeds(mock_hf_api, hf_secrets_empty): + """Transient HuggingFace rate limits are retried during revision resolution.""" + mock_repo_info = Mock() + mock_repo_info.sha = "abc123def456789" + mock_hf_api.repo_info.side_effect = [ + _hf_http_error(429, "Rate limited", headers={"Retry-After": "0"}), + mock_repo_info, + ] + + impl = HuggingfaceStorageImpl( + HuggingfaceStorageConfig(repo_id="test-org/test-repo", repo_type="model", revision="main"), + hf_secrets_empty, + ) + + with patch("nmp.core.files.app.backends.huggingface.sleep", new_callable=AsyncMock) as mock_sleep: + resolved_config = await impl.resolve_config() + + assert resolved_config.original_revision == "main" + assert resolved_config.revision == "abc123def456789" + assert mock_hf_api.repo_info.call_count == 2 + mock_sleep.assert_awaited_once_with(0.0) + + async def test_resolve_config_repo_not_found(hf_config, mock_hf_api, mock_httpx_response, hf_secrets_empty): """Test resolve_config raises HuggingfaceConfigError when repo not found.""" from huggingface_hub.utils import HfHubHTTPError diff --git a/services/core/inference-gateway/tests/unit/test_openai_router.py b/services/core/inference-gateway/tests/unit/test_openai_router.py index 24746f02db..1d7bd5749e 100644 --- a/services/core/inference-gateway/tests/unit/test_openai_router.py +++ b/services/core/inference-gateway/tests/unit/test_openai_router.py @@ -1123,9 +1123,18 @@ def test_virtual_model_proxy_mock_provider_keeps_qualified_body_model(app: FastA MOCK_RESPONSE_MAP_HEADER, MOCK_SERVED_MODELS_HEADER, ) - from nmp.core.inference_gateway.config import config as igw_config - mocker.patch.object(igw_config, "mock_provider_prefix", "igw-mock-") + # Patch the function that is actually called at dispatch time. Patching the + # module-level ``config`` instance is unreliable: if the module was first + # imported while a ``create_test_client`` context was active, ``config`` + # captures that transient override instance; once the context exits and + # ``Configuration.clear_overrides()`` runs, ``_get_mock_provider_prefix`` + # resolves through ``Configuration.get_service_config`` to the @cache-d + # baseline instance (a *different* object), so the attribute patch is lost. + mocker.patch( + "nmp.core.inference_gateway.api.mock_provider.utils._get_mock_provider_prefix", + return_value="igw-mock-", + ) workspace = "vm-mock-ws" served_name = "echo-model" @@ -1175,9 +1184,11 @@ def test_virtual_model_proxy_streaming_mock_provider_runs_response_middleware(ap MOCK_RESPONSE_MAP_HEADER, MOCK_SERVED_MODELS_HEADER, ) - from nmp.core.inference_gateway.config import config as igw_config - mocker.patch.object(igw_config, "mock_provider_prefix", "igw-mock-") + mocker.patch( + "nmp.core.inference_gateway.api.mock_provider.utils._get_mock_provider_prefix", + return_value="igw-mock-", + ) workspace = "vm-mock-ws" served_name = "stream-model"