diff --git a/config/config.default.yaml b/config/config.default.yaml index ca850515..313fa406 100644 --- a/config/config.default.yaml +++ b/config/config.default.yaml @@ -44,6 +44,7 @@ job: ttl_max: 604800 # Time in seconds after which job state is cleared, regardless of last touch log: log_to_mongo: true # Persist logs in MongoDB. + mongo_level: INFO # Minimum level of logs persisted to MongoDB. mongo_ttl: 604800 # Time in seconds for a log to persist. redis: host: localhost @@ -65,6 +66,7 @@ mongo: shutdown_timeout: 3 # Time in seconds to wait for serialize task to finish. compression_level: 10 # zstd level used to compress stored job documents. Higher means greater compression, at greater CPU cost. max_stored_doc_bytes: 134217728 # Compressed job-docs at or above this size are not persisted. + stored_success_proportion: 0.1 # Fraction of succeeded response bodies to persist. Failures are always kept. tier0: backend: gandalf backend_infores: infores:dogpark-tier0 diff --git a/pyproject.toml b/pyproject.toml index c2fa5931..66dfa182 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,6 +139,7 @@ launch = 'retriever' JOB__LOOKUP__TIER1_TIMEOUT=-1 \ JOB__LOOKUP__TIER2_TIMEOUT=-1 \ JOB__METAKG__TIMEOUT=-1 \ + MONGO__STORED_SUCCESS_PROPORTION=1.0 \ WORKERS=1 \ retriever """ diff --git a/src/retriever/config/general.py b/src/retriever/config/general.py index bbd93f19..76750f77 100644 --- a/src/retriever/config/general.py +++ b/src/retriever/config/general.py @@ -105,6 +105,14 @@ class MongoSettings(BaseModel): description="Compressed job-docs at or above this size are not persisted." ), ] = 128 * 1024**2 + stored_success_proportion: Annotated[ + float, + Field( + ge=0, + le=1, + description="Fraction of succeeded response bodies to persist. Failures are always kept.", + ), + ] = 0.1 class TelemetrySettings(BaseModel): @@ -231,6 +239,9 @@ class LogSettings(BaseModel): """Settings for log handling.""" log_to_mongo: Annotated[bool, Field(description="Persist logs in MongoDB.")] = True + mongo_level: Annotated[ + LogLevel, Field(description="Minimum level of logs persisted to MongoDB.") + ] = "INFO" mongo_ttl: Annotated[ int, Field(description="Time in seconds for a log to persist.") ] = 604_800 diff --git a/src/retriever/query.py b/src/retriever/query.py index 35018c87..755ecc80 100644 --- a/src/retriever/query.py +++ b/src/retriever/query.py @@ -260,14 +260,14 @@ async def in_progress_payload( ) -> tuple[HTTPStatus, dict[str, Any]]: """Return an AsyncQueryStatusResponse for in-progress or missing job. - "Running" if the row exists or there's log evidence, "Error" if the - fetch raised, otherwise "Not Found". + "Running" if the row exists, "Error" if the fetch raised, otherwise "Not + Found". Status polling returns empty log list for performance. Logs can be obtained + when query completes, or looking at /logs. """ - logs = await job_logs(job_id) - if exists or len(logs) > 0: + if exists: return HTTPStatus.OK, { "status": "Running", - "logs": logs, + "logs": [], "description": "Job is running.", } if error is not None: @@ -286,7 +286,11 @@ async def in_progress_payload( async def terminal_status_payload( - job_id: str, request: Request, status_doc: JobStatus + job_id: str, + request: Request, + status_doc: JobStatus, + *, + include_logs: bool = False, ) -> dict[str, Any]: """Build an AsyncQueryStatusResponse for a terminal job. @@ -307,7 +311,7 @@ async def terminal_status_payload( return { "status": to_async_lifecycle(job_status), "description": description, - "logs": await job_logs(job_id), + "logs": await job_logs(job_id) if include_logs else [], "response_url": f"{request.base_url}response/{job_id}", } @@ -362,7 +366,9 @@ async def get_job_response( # If job is abandoned, respond with the query if status_doc.get("abandoned"): - payload = await terminal_status_payload(job_id, request, status_doc) + payload = await terminal_status_payload( + job_id, request, status_doc, include_logs=True + ) if job is not None and job.get("doc") is not None: query = unpack_doc(job) payload = { @@ -372,7 +378,9 @@ async def get_job_response( return HTTPStatus.OK, payload if job is None or job.get("doc") is None: - return HTTPStatus.OK, await terminal_status_payload(job_id, request, status_doc) + return HTTPStatus.OK, await terminal_status_payload( + job_id, request, status_doc, include_logs=True + ) response = unpack_doc(job) return HTTPStatus.OK, { diff --git a/src/retriever/utils/logs.py b/src/retriever/utils/logs.py index 4ffeec8b..6b427654 100644 --- a/src/retriever/utils/logs.py +++ b/src/retriever/utils/logs.py @@ -229,7 +229,7 @@ def mongo_sink(message: loguru.Message) -> None: serialize=True, enqueue=True, filter=lambda record: not record["extra"].get("no_mongo_log", False), - level=CONFIG.log_level, + level=CONFIG.log.mongo_level, ) diff --git a/src/retriever/utils/mongo.py b/src/retriever/utils/mongo.py index 1df0e780..0070d760 100644 --- a/src/retriever/utils/mongo.py +++ b/src/retriever/utils/mongo.py @@ -1623,7 +1623,7 @@ async def get_logs( """Return a generator of a filtered set of logs.""" query = _build_log_query(start, end, level) if job_id is not None: - query["extra.job_id"] = {"$regex": job_id} + query["extra.job_id"] = job_id async for document in self._yield_logs(query): yield document @@ -1747,6 +1747,16 @@ async def wrapup(self) -> None: self.client.close() +def _store_job_success_response(job_id: str) -> bool: + """Whether to persist this succeeded response body. + + Deterministic and uniform over `job_id` (a random uuid4 hex), so the two + writes an async job makes agree and never churn a stored blob. + `stored_success_proportion` of 0.0 keeps nothing, 1.0 keeps everything. + """ + return int(job_id[:8], 16) / 0x1_0000_0000 < CONFIG.mongo.stored_success_proportion + + class MongoOutage(Exception): """Raised by `MongoQueue.put` when MongoClient is down; callers attach a warning.""" @@ -1811,6 +1821,13 @@ async def job_state(self, batch: list[QueryState | ResponseState]) -> None: f"Job {state['job_id']} {doc_type} is {size} bytes (>= {CONFIG.mongo.max_stored_doc_bytes} cap); delivering without storing it.", no_mongo_log=True, ) + elif ( + doc_type == "response" + and cast("ResponseState", state)["status"] in TERMINAL_SUCCESS + and not _store_job_success_response(state["job_id"]) + ): + # Don't store (but keep the query) + store_doc = False elif size > GRIDFS_INLINE_LIMIT: doc_ref = await self.client.offload_doc_blob(state["job_id"], blob) # pyright: ignore[reportArgumentType] size>0 implies blob is not None ops.append( diff --git a/tests/test_query.py b/tests/test_query.py index ac861d1d..6e638477 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -5,6 +5,7 @@ relevant query function. """ +from http import HTTPStatus from types import SimpleNamespace from unittest.mock import AsyncMock, Mock, patch @@ -187,3 +188,97 @@ async def test_async_lookup_reapplies_data_tier_tag(): call.args and call.args[0].get("data_tier") == 2 for call in set_tags.call_args_list ) + + +def _status_request() -> SimpleNamespace: + """A stand-in Request exposing only the `base_url` these payloads read.""" + return SimpleNamespace(base_url="http://test/") + + +def _mongo_stub( + *, + status_doc: dict[str, object] | None = None, + job_doc: dict[str, object] | None = None, +) -> Mock: + """A MongoClient() whose status/doc reads return the given fixtures.""" + mongo = Mock() + mongo.get_job_status = AsyncMock(return_value=status_doc) + mongo.get_job_doc = AsyncMock(return_value=job_doc) + return mongo + + +@pytest.mark.asyncio +async def test_asyncquery_status_terminal_returns_empty_logs_without_scan(): + """`/asyncquery_status` on a terminal job returns [] logs and never scans log_dump.""" + mongo = _mongo_stub(status_doc={"status": "Success", "description": None}) + with ( + patch.object(query_module, "MongoClient", return_value=mongo), + patch.object( + query_module, "job_logs", AsyncMock(return_value=["scanned"]) + ) as jl, + ): + code, payload = await query_module.get_job_status("abc123", _status_request()) + + assert code == HTTPStatus.OK + assert payload["logs"] == [] + jl.assert_not_called() + + +@pytest.mark.asyncio +async def test_status_in_progress_returns_empty_logs_without_scan(): + """A Running job reports Running with [] logs and no log_dump scan.""" + mongo = _mongo_stub(status_doc={"status": "Running", "description": None}) + with ( + patch.object(query_module, "MongoClient", return_value=mongo), + patch.object( + query_module, "job_logs", AsyncMock(return_value=["scanned"]) + ) as jl, + ): + code, payload = await query_module.get_job_status("abc123", _status_request()) + + assert code == HTTPStatus.OK + assert payload["status"] == "Running" + assert payload["logs"] == [] + jl.assert_not_called() + + +@pytest.mark.asyncio +async def test_response_completed_stored_uses_blob_logs_without_scan(): + """`/response` for a stored completed job serves the blob's logs, no scan.""" + mongo = _mongo_stub( + status_doc={"status": "Success", "description": None}, + job_doc={"job_id": "abc123", "doc": b"blob"}, + ) + blob = {"message": {}, "logs": ["from-blob"]} + with ( + patch.object(query_module, "MongoClient", return_value=mongo), + patch.object(query_module, "unpack_doc", return_value=blob), + patch.object( + query_module, "job_logs", AsyncMock(return_value=["scanned"]) + ) as jl, + ): + code, payload = await query_module.get_job_response("abc123", _status_request()) + + assert code == HTTPStatus.OK + assert payload["logs"] == ["from-blob"] + jl.assert_not_called() + + +@pytest.mark.asyncio +async def test_response_completed_not_stored_falls_back_to_scan(): + """`/response` for a completed-but-unstored job falls back to a log_dump scan.""" + mongo = _mongo_stub( + status_doc={"status": "Success", "description": None}, + job_doc=None, # sampled out: no stored body + ) + with ( + patch.object(query_module, "MongoClient", return_value=mongo), + patch.object( + query_module, "job_logs", AsyncMock(return_value=["scanned"]) + ) as jl, + ): + code, payload = await query_module.get_job_response("abc123", _status_request()) + + assert code == HTTPStatus.OK + assert payload["logs"] == ["scanned"] + jl.assert_awaited_once() diff --git a/tests/test_utils_mongo_gridfs_live.py b/tests/test_utils_mongo_gridfs_live.py index cd89d5f1..af22be64 100644 --- a/tests/test_utils_mongo_gridfs_live.py +++ b/tests/test_utils_mongo_gridfs_live.py @@ -39,6 +39,7 @@ def _response_state(job_id: str, payload: bytes) -> ResponseState: aux_graphs=0, results=1, status="Success", + event_time=datetime.now().astimezone(), ) @@ -47,6 +48,16 @@ async def _write(state: ResponseState) -> None: await MongoQueue().job_state([state]) +@pytest.fixture(autouse=True) +def _keep_all_successes(monkeypatch: pytest.MonkeyPatch) -> None: + """These tests exercise storage mechanics, not sampling. + + Force every succeeded body to be stored so the store/GridFS assertions are + deterministic; the sampling tests below re-monkeypatch the proportion down. + """ + monkeypatch.setattr(CONFIG.mongo, "stored_success_proportion", 1.0) + + @pytest.mark.asyncio async def test_large_blob_round_trips_via_gridfs(test_mongo: MongoClient) -> None: # noqa: F811 """A blob over the inline limit is spilled to GridFS and hydrated back intact.""" @@ -242,6 +253,7 @@ async def test_dehydrated_response_stores_state_but_no_body( aux_graphs=0, results=5, status="Success", + event_time=datetime.now().astimezone(), dehydrated=True, ) ] @@ -286,6 +298,7 @@ async def test_dehydrated_flag_recorded_at_enqueue(test_mongo: MongoClient) -> N status="Running", worker_pid=1234, worker_started_at=datetime.now().astimezone(), + event_time=datetime.now().astimezone(), dehydrated=True, ) ] @@ -300,3 +313,93 @@ async def test_dehydrated_flag_recorded_at_enqueue(test_mongo: MongoClient) -> N raw = await docs.find_one({"job_id": job_id}) assert raw is not None assert raw.get("doc") == b"compressed-query-bytes" + + +@pytest.mark.asyncio +async def test_sampled_out_success_not_stored( + test_mongo: MongoClient, # noqa: F811 + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A succeeded response sampled out keeps its status but stores no body.""" + monkeypatch.setattr(CONFIG.mongo, "stored_success_proportion", 0.0) + + job_id = uuid.uuid4().hex + await _write(_response_state(job_id, b"body-that-should-not-persist")) + + status, docs = test_mongo.get_job_collection() + status_doc = await status.find_one({"job_id": job_id}) + assert status_doc is not None + assert status_doc["status"] == "Success" # state is still recorded + + raw = await docs.find_one({"job_id": job_id}) + assert raw is not None + assert raw.get("doc") is None + assert raw.get("doc_ref") is None + assert ( + await test_mongo._doc_blob_files().count_documents({"metadata.job_id": job_id}) + == 0 + ) + + +@pytest.mark.asyncio +async def test_sampled_out_success_clears_initial_query_blob( + test_mongo: MongoClient, # noqa: F811 + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A sampled-out success clears the initial query blob, so it's never served back.""" + monkeypatch.setattr(CONFIG.mongo, "stored_success_proportion", 0.0) + + job_id = uuid.uuid4().hex + # The initial query write (unconditional for sync + async) stores the query. + await MongoQueue().job_state( + [ + QueryState( + job_id=job_id, + query=b"compressed-query-bytes", + job_timeout=30.0, + submitter="tester", + data_tier=0, + is_async=True, + qnodes=2, + qedges=1, + qpaths=0, + status="Running", + worker_pid=1234, + worker_started_at=datetime.now().astimezone(), + event_time=datetime.now().astimezone(), + ) + ] + ) + _, docs = test_mongo.get_job_collection() + assert (await docs.find_one({"job_id": job_id}) or {}).get("doc") is not None + + # The sampled-out response must drop that prior query blob. + await _write(_response_state(job_id, b"response-body-dropped")) + + raw = await docs.find_one({"job_id": job_id}) + assert raw is not None + assert raw.get("doc") is None + assert raw.get("doc_ref") is None + + job = await test_mongo.get_job_doc(job_id) + assert job is not None + assert job.get("doc") is None + + +@pytest.mark.asyncio +async def test_failure_always_stored_despite_zero_proportion( + test_mongo: MongoClient, # noqa: F811 + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Failures are always persisted; success-only sampling never drops them.""" + monkeypatch.setattr(CONFIG.mongo, "stored_success_proportion", 0.0) + + job_id = uuid.uuid4().hex + state = _response_state(job_id, b"failure-body") + state["status"] = "Failed" + await _write(state) + + _, docs = test_mongo.get_job_collection() + raw = await docs.find_one({"job_id": job_id}) + assert raw is not None + assert raw.get("doc") == b"failure-body"