diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index c3f090d94f8..ea867c8a90a 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -49,6 +49,9 @@ async def _provider_job_cursor( ) -> str | None: if after is None or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return after + decoded: Final = is_base64_encoded_unified_file_id(after) + if not decoded: + return after await validate_managed_id_requirement( resource_id=after, resource_kind="fine-tuning job", @@ -56,8 +59,7 @@ async def _provider_job_cursor( managed_files_obj=managed_files_obj, require_owned_resource=True, ) - decoded: Final = is_base64_encoded_unified_file_id(after) - return get_batch_id_from_unified_batch_id(decoded) if decoded else after + return get_batch_id_from_unified_batch_id(decoded) async def _owned_jobs(response: object, user_api_key_dict: UserAPIKeyAuth, managed_files_obj: object) -> object: @@ -118,19 +120,28 @@ async def _collect_owned_job_pages( user_api_key_dict: UserAPIKeyAuth, managed_files_obj: object, ) -> object: + loop: Final = asyncio.get_running_loop() + deadline: Final = loop.time() + _JOB_LIST_SCAN_TIMEOUT + + async def owned_page(cursor: str | None) -> _OwnedJobPage: + return await _owned_job_page(await next_page(cursor), user_api_key_dict, managed_files_obj) + state: tuple[str | None, _OwnedJobPage] = ( # rebind-ok: advance through provider pages without recursion None, await _owned_job_page(response, user_api_key_dict, managed_files_obj), ) for _ in range(_JOB_LIST_MAX_PAGES - 1): - if state[1].next_cursor is None or state[1].next_cursor == state[0]: + if state[1].next_cursor is None or state[1].next_cursor == state[0] or loop.time() >= deadline: + break + try: + state = ( + state[1].next_cursor, + await asyncio.wait_for(owned_page(state[1].next_cursor), timeout=deadline - loop.time()), + ) + except asyncio.TimeoutError: break - state = ( - state[1].next_cursor, - await _owned_job_page(await next_page(state[1].next_cursor), user_api_key_dict, managed_files_obj), - ) if state[1].next_cursor is not None and state[1].next_cursor != state[0]: - raise HTTPException(status_code=503, detail="Fine-tuning job listing exceeded its page budget") + return {**state[1].response, "has_more": True, "last_id": state[1].next_cursor} return {**state[1].response, "has_more": False} if state[1].next_cursor is not None else state[1].response @@ -142,13 +153,7 @@ async def _owned_job_pages( ) -> object: if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return response - try: - return await asyncio.wait_for( - _collect_owned_job_pages(response, next_page, user_api_key_dict, managed_files_obj), - timeout=_JOB_LIST_SCAN_TIMEOUT, - ) - except asyncio.TimeoutError as e: - raise HTTPException(status_code=504, detail="Fine-tuning job listing timed out") from e + return await _collect_owned_job_pages(response, next_page, user_api_key_dict, managed_files_obj) def set_fine_tuning_config(config): diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index c8be3ea7c6e..5876f3ec32e 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -189,21 +189,33 @@ async def _collect_managed_file_list_pages( user_api_key_dict: UserAPIKeyAuth, before: bool, ) -> VectorStoreFileListResponse: + loop: Final = asyncio.get_running_loop() + deadline: Final = loop.time() + _FILE_LIST_SCAN_TIMEOUT + + async def owned_page(cursor: str) -> _ManagedFilePage: + return await _managed_file_list_page(await next_page(cursor), managed_files_obj, user_api_key_dict, before) + state: tuple[str | None, _ManagedFilePage] = ( # rebind-ok: advance a bounded page scan None, await _managed_file_list_page(response, managed_files_obj, user_api_key_dict, before), ) for _ in range(_FILE_LIST_MAX_PAGES - 1): - if state[1].next_cursor is None or state[1].next_cursor == state[0]: + if state[1].next_cursor is None or state[1].next_cursor == state[0] or loop.time() >= deadline: + break + try: + state = ( + state[1].next_cursor, + await asyncio.wait_for(owned_page(state[1].next_cursor), timeout=deadline - loop.time()), + ) + except asyncio.TimeoutError: break - state = ( - state[1].next_cursor, - await _managed_file_list_page( - await next_page(state[1].next_cursor), managed_files_obj, user_api_key_dict, before - ), - ) if state[1].next_cursor is not None and state[1].next_cursor != state[0]: - raise HTTPException(status_code=503, detail="Vector store file listing exceeded its page budget") + return { + **state[1].response, + "has_more": True, + "first_id": state[1].next_cursor if before else None, + "last_id": state[1].next_cursor if not before else None, + } return {**state[1].response, "has_more": False} if state[1].next_cursor is not None else state[1].response @@ -216,13 +228,7 @@ async def _with_managed_file_list_pages( ) -> VectorStoreFileListResponse: if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return await _with_managed_file_list_ids(response, managed_files_obj, user_api_key_dict) - try: - return await asyncio.wait_for( - _collect_managed_file_list_pages(response, next_page, managed_files_obj, user_api_key_dict, before), - timeout=_FILE_LIST_SCAN_TIMEOUT, - ) - except asyncio.TimeoutError as e: - raise HTTPException(status_code=504, detail="Vector store file listing timed out") from e + return await _collect_managed_file_list_pages(response, next_page, managed_files_obj, user_api_key_dict, before) async def _update_request_data_with_managed_file_id( diff --git a/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py index 320fb617a22..831da6f4aa7 100644 --- a/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py +++ b/tests/unit/proxy/fine_tuning_endpoints/test_endpoints.py @@ -10,6 +10,7 @@ Each test asserts BOTH that the request is rejected AND that every downstream pr seam stayed untouched, so a guard that raises after the provider call would still fail. """ +import asyncio import base64 from contextlib import ExitStack from collections.abc import Mapping, Sequence @@ -341,14 +342,14 @@ PROVIDER_JOBS_PAGE: Final = { } -async def _list(custom_llm_provider: str | None): +async def _list(custom_llm_provider: str | None, *, after: str | None = None, limit: int | None = None): return await endpoints.list_fine_tuning_jobs( request=FakeRequest(), fastapi_response=Response(), custom_llm_provider=custom_llm_provider, target_model_names=None, - after=None, - limit=None, + after=after, + limit=limit, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), ) @@ -433,10 +434,8 @@ async def test_admin_job_cursor_keeps_provider_value(cursor: str | None) -> None @pytest.mark.asyncio -async def test_job_cursor_requires_a_managed_reference() -> None: - with pytest.raises(HTTPException) as exc: - await endpoints._provider_job_cursor(RAW_JOB_ID, UserAPIKeyAuth(api_key="test"), None) - assert exc.value.status_code == 403 +async def test_job_list_cursor_accepts_a_provider_page_boundary() -> None: + assert await endpoints._provider_job_cursor(RAW_JOB_ID, UserAPIKeyAuth(api_key="test"), None) == RAW_JOB_ID @pytest.mark.asyncio @@ -460,12 +459,71 @@ async def test_job_listing_continues_until_an_owned_page(seams) -> None: @pytest.mark.asyncio -async def test_job_listing_stops_at_its_page_budget() -> None: +async def test_job_listing_returns_a_resumable_page_at_its_budget() -> None: async def next_page(cursor: str | None) -> object: assert cursor is not None return {"data": [{"id": cursor + "-next"}], "has_more": True} fetch: Final = AsyncMock(side_effect=next_page) + assert await endpoints._owned_job_pages( + {"data": [{"id": "job-first"}], "has_more": True}, + fetch, + UserAPIKeyAuth(api_key="test"), + _OwnedJobIds({}), + ) == {"data": [], "has_more": True, "last_id": "job-first" + "-next" * (endpoints._JOB_LIST_MAX_PAGES - 1)} + assert fetch.await_count == endpoints._JOB_LIST_MAX_PAGES - 1 + + +@pytest.mark.asyncio +async def test_job_listing_returns_a_resumable_page_at_its_deadline(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(endpoints, "_JOB_LIST_SCAN_TIMEOUT", 0) + + async def next_page(cursor: str | None) -> object: + pytest.fail("An expired scan must not request another page") + + assert await endpoints._owned_job_pages( + {"data": [{"id": "job-first"}], "has_more": True}, next_page, UserAPIKeyAuth(api_key="test"), _OwnedJobIds({}) + ) == {"data": [], "has_more": True, "last_id": "job-first"} + + +@pytest.mark.asyncio +@respx.mock +async def test_job_listing_resumes_after_an_empty_scan(seams) -> None: + job: Final = _job().model_dump() + route: Final = respx.get(f"{FINE_TUNING_API_BASE}/fine_tuning/jobs").mock( + side_effect=[ + *( + httpx.Response(200, json={"data": [{**job, "id": f"ftjob-page-{index}"}], "has_more": True}) + for index in range(endpoints._JOB_LIST_MAX_PAGES) + ), + httpx.Response(200, json={"data": [job], "has_more": False}), + ] + ) + seams.logging.get_proxy_hook.return_value = _OwnedJobIds({RAW_JOB_ID: _unified_job_id()}) + provider: Final = {"custom_llm_provider": "openai", "api_key": "sk-provider", "api_base": FINE_TUNING_API_BASE} + with patch.object(endpoints, "fine_tuning_config", [provider]): + page: Final = await _list(custom_llm_provider="openai", limit=1) + assert page == {"data": [], "has_more": True, "last_id": f"ftjob-page-{endpoints._JOB_LIST_MAX_PAGES - 1}"} + following: Final = await _list(custom_llm_provider="openai", after=page["last_id"], limit=1) + assert following == {"data": [{**job, "id": _unified_job_id()}], "has_more": False} + assert route.calls[-1].request.url.params["after"] == page["last_id"] + assert route.calls[-1].request.url.params["limit"] == "1" + + +@pytest.mark.asyncio +async def test_job_scan_timeout_keeps_the_last_completed_page() -> None: + fetch: Final = AsyncMock(side_effect=asyncio.TimeoutError) + assert await endpoints._owned_job_pages( + {"data": [{"id": "job-first"}], "has_more": True}, + fetch, + UserAPIKeyAuth(api_key="test"), + _OwnedJobIds({}), + ) == {"data": [], "has_more": True, "last_id": "job-first"} + + +@pytest.mark.asyncio +async def test_job_scan_preserves_an_upstream_error() -> None: + fetch: Final = AsyncMock(side_effect=HTTPException(status_code=502, detail="upstream unavailable")) with pytest.raises(HTTPException) as exc: await endpoints._owned_job_pages( {"data": [{"id": "job-first"}], "has_more": True}, @@ -473,21 +531,4 @@ async def test_job_listing_stops_at_its_page_budget() -> None: UserAPIKeyAuth(api_key="test"), _OwnedJobIds({}), ) - assert exc.value.status_code == 503 - assert exc.value.detail == "Fine-tuning job listing exceeded its page budget" - assert fetch.await_count == endpoints._JOB_LIST_MAX_PAGES - 1 - - -@pytest.mark.asyncio -async def test_job_listing_reports_an_expired_scan_deadline(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(endpoints, "_JOB_LIST_SCAN_TIMEOUT", 0) - - async def next_page(cursor: str | None) -> object: - pytest.fail("An expired scan must not request another page") - - with pytest.raises(HTTPException) as exc: - await endpoints._owned_job_pages( - {"data": [], "has_more": False}, next_page, UserAPIKeyAuth(api_key="test"), _OwnedJobIds({}) - ) - assert exc.value.status_code == 504 - assert exc.value.detail == "Fine-tuning job listing timed out" + assert exc.value.status_code == 502 diff --git a/tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py b/tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py index e782c0873f3..abef6967741 100644 --- a/tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py +++ b/tests/unit/proxy/vector_store_files_endpoints/test_endpoints.py @@ -9,6 +9,7 @@ A raw or forged managed-looking file id has no ownership row, so without the gua is attached to a vector store or read back under shared provider credentials. """ +import asyncio import base64 from collections.abc import Mapping, Sequence from copy import deepcopy @@ -316,28 +317,75 @@ async def test_admin_file_listing_preserves_a_provider_page() -> None: @pytest.mark.asyncio -async def test_file_listing_reports_its_page_budget() -> None: +@pytest.mark.parametrize("before", [False, True]) +async def test_file_listing_resumes_after_its_page_budget(before: bool) -> None: async def next_page(cursor: str) -> VectorStoreFileListResponse: + if cursor.count("-next") == endpoints._FILE_LIST_MAX_PAGES - 1: + return {"data": [_vector_store_file_row(RAW_FILE_ID)], "has_more": False} return {"data": [_vector_store_file_row(cursor + "-next")], "has_more": True} fetch: Final = AsyncMock(side_effect=next_page) resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={})) + page: Final = await endpoints._with_managed_file_list_pages( + {"data": [_vector_store_file_row("file-first")], "has_more": True}, fetch, resolver, CALLER, before + ) + cursor: Final = "file-first" + "-next" * (endpoints._FILE_LIST_MAX_PAGES - 1) + assert page == { + "data": [], + "has_more": True, + "first_id": cursor if before else None, + "last_id": None if before else cursor, + } + assert fetch.await_count == endpoints._FILE_LIST_MAX_PAGES - 1 + query: Final = _with_provider_file_id_cursors({"before" if before else "after": cursor}) + assert query == {"before" if before else "after": cursor} + following: Final = await fetch(query["before" if before else "after"]) + owned_resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={RAW_FILE_ID: _unified_file_id()})) + assert await endpoints._with_managed_file_list_pages(following, fetch, owned_resolver, CALLER, before) == { + "data": [_vector_store_file_row(_unified_file_id())], + "has_more": False, + "first_id": _unified_file_id(), + "last_id": _unified_file_id(), + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("before", [False, True]) +async def test_file_listing_returns_a_resumable_page_at_its_deadline( + monkeypatch: pytest.MonkeyPatch, before: bool +) -> None: + monkeypatch.setattr(endpoints, "_FILE_LIST_SCAN_TIMEOUT", 0) + fetch: Final = AsyncMock() + resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={})) + assert await endpoints._with_managed_file_list_pages( + {"data": [_vector_store_file_row("file-first")], "has_more": True}, fetch, resolver, CALLER, before + ) == { + "data": [], + "has_more": True, + "first_id": "file-first" if before else None, + "last_id": None if before else "file-first", + } + fetch.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_file_scan_timeout_keeps_the_last_completed_page() -> None: + resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={})) + fetch: Final = AsyncMock(side_effect=asyncio.TimeoutError) + assert await endpoints._with_managed_file_list_pages( + {"data": [_vector_store_file_row("file-first")], "has_more": True}, fetch, resolver, CALLER + ) == {"data": [], "has_more": True, "first_id": None, "last_id": "file-first"} + + +@pytest.mark.asyncio +async def test_file_scan_preserves_an_upstream_error() -> None: + resolver: Final = ManagedFileIdResolverStub(AsyncMock(return_value={})) + fetch: Final = AsyncMock(side_effect=HTTPException(status_code=502, detail="upstream unavailable")) with pytest.raises(HTTPException) as exc: await endpoints._with_managed_file_list_pages( {"data": [_vector_store_file_row("file-first")], "has_more": True}, fetch, resolver, CALLER ) - assert exc.value.status_code == 503 - assert fetch.await_count == endpoints._FILE_LIST_MAX_PAGES - 1 - - -@pytest.mark.asyncio -async def test_file_listing_reports_an_expired_deadline(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(endpoints, "_FILE_LIST_SCAN_TIMEOUT", 0) - fetch: Final = AsyncMock() - with pytest.raises(HTTPException) as exc: - await endpoints._with_managed_file_list_pages({"data": [], "has_more": False}, fetch, None, CALLER) - assert exc.value.status_code == 504 - fetch.assert_not_awaited() + assert exc.value.status_code == 502 @pytest.mark.asyncio