From ff968dc2179a2060b08c40745df133e63d3ea0ea Mon Sep 17 00:00:00 2001 From: Evan Morris Date: Fri, 4 Sep 2026 17:50:37 -0400 Subject: [PATCH] prevent ack timeout from dropping callback ids --- shepherd_utils/config.py | 4 + .../aragorn/test_aragorn_lookup_submit.py | 130 ++++++++++++++++++ tests/unit/test_bte_lookup_branches.py | 102 ++++++++++++++ workers/aragorn_lookup/worker.py | 45 ++++-- workers/bte_lookup/worker.py | 46 +++++-- 5 files changed, 311 insertions(+), 16 deletions(-) create mode 100644 tests/unit/aragorn/test_aragorn_lookup_submit.py diff --git a/shepherd_utils/config.py b/shepherd_utils/config.py index d3526a7..29d0905 100644 --- a/shepherd_utils/config.py +++ b/shepherd_utils/config.py @@ -81,6 +81,10 @@ class Settings(BaseSettings): # byte count; 0 (or unparseable) disables the limit. callback_max_request_size: str = "0" kg_retrieval_url: str = "http://host.docker.internal:8080/asyncquery" + # How long a lookup worker waits for the retrieval service to ACK an + # /asyncquery submission. This bounds only the acknowledgement, the + # results arrive later via callback. + kg_retrieval_submit_timeout: float = 100.0 sync_kg_retrieval_url: str = "http://host.docker.internal:8080/query" kg_rehydrate_url: str = "http://host.docker.internal:8080/rehydrate" default_data_tier: int = 0 diff --git a/tests/unit/aragorn/test_aragorn_lookup_submit.py b/tests/unit/aragorn/test_aragorn_lookup_submit.py new file mode 100644 index 0000000..914db19 --- /dev/null +++ b/tests/unit/aragorn/test_aragorn_lookup_submit.py @@ -0,0 +1,130 @@ +"""Tests for how ``workers.aragorn_lookup.worker`` handles the /asyncquery +submission ACK. + +A creative query fans out to ~20 submissions at once and the retrieval +service has been seen taking 20s+ to ACK the burst. A submission whose ACK +times out has still been received -- the service runs it and POSTs the +callback later -- so its callback id must survive, or the callback handler +rejects the results with a 500 and the lookup waits out its deadline for +nothing. +""" + +import json +import logging + +import httpx +import pytest + +from workers.aragorn_lookup import worker as al +from workers.aragorn_lookup.worker import ( + AsyncResponse, + aragorn_lookup, + run_async_lookup, +) + +logger = logging.getLogger(__name__) + + +def _make_task(): + return [ + "test", + { + "query_id": "qid", + "response_id": "rid", + "workflow": json.dumps([{"id": "aragorn.lookup"}]), + "log_level": "20", + "otel": json.dumps({}), + }, + ] + + +INFERRED_MSG = { + "message": { + "query_graph": { + "nodes": { + "a": {"ids": ["X:1"], "categories": ["biolink:Drug"]}, + "b": {"categories": ["biolink:Disease"]}, + }, + "edges": { + "e0": { + "subject": "a", + "object": "b", + "knowledge_type": "inferred", + "predicates": ["biolink:treats"], + } + }, + } + }, + "parameters": {"timeout": 5}, +} + + +@pytest.mark.asyncio +async def test_run_async_lookup_read_timeout_is_in_flight(redis_mock, mocker): + mocker.patch.object(al, "add_callback_id", new_callable=mocker.AsyncMock) + + client = mocker.Mock() + client.timeout.read = 100.0 + client.post = mocker.AsyncMock(side_effect=httpx.ReadTimeout("")) + + out = await run_async_lookup(client, {"message": {}}, "qid", logger) + assert out.success is False + assert out.in_flight is True + assert "ReadTimeout" in out.error + assert "100.0s" in out.error + + +@pytest.mark.asyncio +async def test_run_async_lookup_connect_error_is_not_in_flight(redis_mock, mocker): + mocker.patch.object(al, "add_callback_id", new_callable=mocker.AsyncMock) + + client = mocker.Mock() + client.post = mocker.AsyncMock(side_effect=httpx.ConnectError("boom")) + + out = await run_async_lookup(client, {"message": {}}, "qid", logger) + assert out.success is False + assert out.in_flight is False + assert "ConnectError: boom" == out.error + + +@pytest.mark.asyncio +async def test_aragorn_lookup_keeps_in_flight_callback_ids(redis_mock, mocker): + """Only submissions that definitely failed lose their callback id.""" + mocker.patch.object( + al, "get_message", new_callable=mocker.AsyncMock, return_value=INFERRED_MSG + ) + mocker.patch.object( + al, + "expand_aragorn_query", + return_value=[ + {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"}, + {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"}, + {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"}, + ], + ) + mocker.patch.object( + al, + "run_async_lookup", + new_callable=mocker.AsyncMock, + side_effect=[ + AsyncResponse(status_code=200, success=True, callback_id="ok-cb"), + AsyncResponse( + status_code=500, + success=False, + callback_id="slow-ack-cb", + error="ReadTimeout", + in_flight=True, + ), + AsyncResponse( + status_code=500, success=False, callback_id="failed-cb", error="x" + ), + ], + ) + mock_remove = mocker.patch.object( + al, "remove_callback_id", new_callable=mocker.AsyncMock + ) + mocker.patch.object( + al, "get_running_callbacks", new_callable=mocker.AsyncMock, return_value=[] + ) + await aragorn_lookup(_make_task(), logger) + mock_remove.assert_awaited_once_with("failed-cb", logger) diff --git a/tests/unit/test_bte_lookup_branches.py b/tests/unit/test_bte_lookup_branches.py index 5233eca..3987e5d 100644 --- a/tests/unit/test_bte_lookup_branches.py +++ b/tests/unit/test_bte_lookup_branches.py @@ -90,6 +90,39 @@ async def test_run_async_lookup_returns_500_when_post_raises(redis_mock, mocker) assert out.success is False assert out.status_code == 500 assert "boom" in out.error + assert out.in_flight is False + + +@pytest.mark.asyncio +async def test_run_async_lookup_read_timeout_is_in_flight(redis_mock, mocker): + """A read timeout means the submission was sent but not ACKed: the service + may still deliver the callback, so the response is flagged in_flight.""" + mocker.patch.object(btel, "add_callback_id", new_callable=mocker.AsyncMock) + + client = mocker.Mock() + client.timeout.read = 100.0 + client.post = mocker.AsyncMock(side_effect=httpx.ReadTimeout("")) + + out = await run_async_lookup(client, {"message": {}}, "qid", logger) + assert out.success is False + assert out.in_flight is True + # httpx timeouts stringify to nothing; the error should still say what happened. + assert "ReadTimeout" in out.error + assert "100.0s" in out.error + + +@pytest.mark.asyncio +async def test_run_async_lookup_connect_timeout_is_not_in_flight(redis_mock, mocker): + """A connect timeout never reached the service, so nothing is coming back.""" + mocker.patch.object(btel, "add_callback_id", new_callable=mocker.AsyncMock) + + client = mocker.Mock() + client.post = mocker.AsyncMock(side_effect=httpx.ConnectTimeout("")) + + out = await run_async_lookup(client, {"message": {}}, "qid", logger) + assert out.success is False + assert out.in_flight is False + assert "ConnectTimeout" in out.error @pytest.mark.asyncio @@ -220,6 +253,75 @@ async def test_bte_lookup_inferred_removes_failed_callback_ids(redis_mock, mocke mock_remove.assert_awaited_once_with("failed-cb", logger) +@pytest.mark.asyncio +async def test_bte_lookup_inferred_keeps_in_flight_callback_ids(redis_mock, mocker): + """A submission that timed out waiting for the ACK keeps its callback id: + the retrieval service may still POST the results, and removing the id + would make the callback handler reject them with a 500.""" + inferred_msg = { + "message": { + "query_graph": { + "nodes": { + "a": {"ids": ["X:1"], "categories": ["biolink:Drug"]}, + "b": {"categories": ["biolink:Disease"]}, + }, + "edges": { + "e0": { + "subject": "a", + "object": "b", + "knowledge_type": "inferred", + "predicates": ["biolink:treats"], + } + }, + } + }, + "parameters": {"timeout": 5}, + } + mocker.patch.object( + btel, + "get_message", + new_callable=mocker.AsyncMock, + return_value=inferred_msg, + ) + mocker.patch.object( + btel, + "expand_bte_query", + return_value=[ + {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"}, + {"message": {"query_graph": {}}, "parameters": {}, "submitter": "t"}, + ], + ) + mocker.patch.object( + btel, + "run_async_lookup", + new_callable=mocker.AsyncMock, + side_effect=[ + AsyncResponse( + status_code=500, + success=False, + callback_id="slow-ack-cb", + error="ReadTimeout", + in_flight=True, + ), + AsyncResponse( + status_code=500, success=False, callback_id="failed-cb", error="x" + ), + ], + ) + mock_remove = mocker.patch.object( + btel, "remove_callback_id", new_callable=mocker.AsyncMock + ) + mocker.patch.object( + btel, + "get_running_callbacks", + new_callable=mocker.AsyncMock, + return_value=[], + ) + await bte_lookup(_make_task(), logger) + # Only the genuinely failed submission is dropped. + mock_remove.assert_awaited_once_with("failed-cb", logger) + + @pytest.mark.asyncio async def test_bte_lookup_inferred_logs_exception_responses(redis_mock, mocker): """An exception in ``asyncio.gather`` (return_exceptions=True) is logged diff --git a/workers/aragorn_lookup/worker.py b/workers/aragorn_lookup/worker.py index ec19a06..f942f7c 100644 --- a/workers/aragorn_lookup/worker.py +++ b/workers/aragorn_lookup/worker.py @@ -84,6 +84,11 @@ class AsyncResponse: success: bool callback_id: str error: Optional[str] = None + # True when the submission was sent but we never heard the ACK. The + # retrieval service may well be running the lookup and will POST the + # callback later, so the callback id must be kept or the results will be + # rejected on arrival. False means the request never reached the service. + in_flight: bool = False async def run_async_lookup( @@ -119,13 +124,25 @@ async def run_async_lookup( success=response.status_code == 200, callback_id=callback_id, ) + except httpx.ReadTimeout as e: + # The request went out and the service just didn't ACK in time. + # It has the query and the callback URL, so treat it as still + # running rather than failed. + span.record_exception(e) + return AsyncResponse( + status_code=500, + success=False, + callback_id=callback_id, + error=f"{type(e).__name__}: no ACK within {client.timeout.read}s", + in_flight=True, + ) except Exception as e: span.record_exception(e) return AsyncResponse( status_code=500, success=False, callback_id=callback_id, - error=str(e), + error=f"{type(e).__name__}: {e}", ) @@ -168,7 +185,9 @@ async def aragorn_lookup(task, logger: logging.Logger): ) with tracer.start_as_current_span("aragorn.lookup") as span: span.set_attribute("callback.id", callback_id) - async with httpx.AsyncClient(timeout=100) as client: + async with httpx.AsyncClient( + timeout=settings.kg_retrieval_submit_timeout + ) as client: await client.post( settings.kg_retrieval_url, json=message, @@ -180,7 +199,9 @@ async def aragorn_lookup(task, logger: logging.Logger): requests = [] # send all messages to lookup service - async with httpx.AsyncClient(timeout=20) as client: + async with httpx.AsyncClient( + timeout=settings.kg_retrieval_submit_timeout + ) as client: for expanded_message in expanded_messages: requests.append( run_async_lookup(client, expanded_message, query_id, logger) @@ -196,12 +217,20 @@ async def aragorn_lookup(task, logger: logging.Logger): f"Failed to do lookup and unable to remove callback id: {response}" ) elif isinstance(response, AsyncResponse): - if not response.success: - logger.error( - f"[{response.callback_id}] Failed to do lookup, " - f"removing callback id: {response.error}" + if response.success: + continue + if response.in_flight: + logger.warning( + f"[{response.callback_id}] Lookup submitted but not " + f"acknowledged, waiting for its callback anyway: " + f"{response.error}" ) - await remove_callback_id(response.callback_id, logger) + continue + logger.error( + f"[{response.callback_id}] Failed to do lookup, " + f"removing callback id: {response.error}" + ) + await remove_callback_id(response.callback_id, logger) else: logger.error( f"Failed to do lookup and unable to remove callback id: {response}" diff --git a/workers/bte_lookup/worker.py b/workers/bte_lookup/worker.py index 5333bd6..707da19 100644 --- a/workers/bte_lookup/worker.py +++ b/workers/bte_lookup/worker.py @@ -84,6 +84,11 @@ class AsyncResponse: success: bool callback_id: str error: Optional[str] = None + # True when the submission was sent but we never heard the ACK. The + # retrieval service may well be running the lookup and will POST the + # callback later, so the callback id must be kept or the results will be + # rejected on arrival. False means the request never reached the service. + in_flight: bool = False async def run_async_lookup( @@ -119,12 +124,25 @@ async def run_async_lookup( success=response.status_code == 200, callback_id=callback_id, ) + except httpx.ReadTimeout as e: + # The request went out and the service just didn't ACK in time. + # It has the query and the callback URL, so treat it as still + # running rather than failed. + span.record_exception(e) + return AsyncResponse( + status_code=500, + success=False, + callback_id=callback_id, + error=f"{type(e).__name__}: no ACK within {client.timeout.read}s", + in_flight=True, + ) except Exception as e: + span.record_exception(e) return AsyncResponse( status_code=500, success=False, callback_id=callback_id, - error=str(e), + error=f"{type(e).__name__}: {e}", ) @@ -161,7 +179,9 @@ async def bte_lookup(task, logger: logging.Logger): ) with tracer.start_as_current_span("bte.lookup") as span: span.set_attribute("callback.id", callback_id) - async with httpx.AsyncClient(timeout=100) as client: + async with httpx.AsyncClient( + timeout=settings.kg_retrieval_submit_timeout + ) as client: await client.post( settings.kg_retrieval_url, json=message, @@ -173,7 +193,9 @@ async def bte_lookup(task, logger: logging.Logger): logger.info(f"Expanded to {len(expanded_messages)} messages") requests = [] # send all messages to retriever - async with httpx.AsyncClient(timeout=20) as client: + async with httpx.AsyncClient( + timeout=settings.kg_retrieval_submit_timeout + ) as client: for expanded_message in expanded_messages: requests.append( run_async_lookup(client, expanded_message, query_id, logger) @@ -189,12 +211,20 @@ async def bte_lookup(task, logger: logging.Logger): f"Failed to do lookup and unable to remove callback id: {response}" ) elif isinstance(response, AsyncResponse): - if not response.success: - logger.error( - f"[{response.callback_id}] Failed to do lookup, " - f"removing callback id: {response.error}" + if response.success: + continue + if response.in_flight: + logger.warning( + f"[{response.callback_id}] Lookup submitted but not " + f"acknowledged, waiting for its callback anyway: " + f"{response.error}" ) - await remove_callback_id(response.callback_id, logger) + continue + logger.error( + f"[{response.callback_id}] Failed to do lookup, " + f"removing callback id: {response.error}" + ) + await remove_callback_id(response.callback_id, logger) else: logger.error( f"Failed to do lookup and unable to remove callback id: {response}"