diff --git a/diracx-core/src/diracx/core/exceptions.py b/diracx-core/src/diracx/core/exceptions.py index 4ac333787..082a2b0a6 100644 --- a/diracx-core/src/diracx/core/exceptions.py +++ b/diracx-core/src/diracx/core/exceptions.py @@ -3,6 +3,7 @@ __all__ = [ "AuthorizationError", "DiracError", + "DocumentUpsertError", "IAMClientError", "IAMServerError", "InvalidCredentialsError", @@ -56,6 +57,10 @@ class InvalidQueryError(DiracError): """It was not possible to build a valid database query from the given input.""" +class DocumentUpsertError(DiracError): + """The backend rejected a document upsert, e.g. because it cannot be indexed.""" + + class TokenNotFoundError(DiracError): def __init__(self, jti: str, detail: str = ""): self.jti: str = jti diff --git a/diracx-core/src/diracx/core/models/job.py b/diracx-core/src/diracx/core/models/job.py index d577e6c07..79d4c4935 100644 --- a/diracx-core/src/diracx/core/models/job.py +++ b/diracx-core/src/diracx/core/models/job.py @@ -5,10 +5,11 @@ from __future__ import annotations +import math from enum import StrEnum -from typing import Literal +from typing import Any, Literal -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, field_validator, model_validator from .types import UTCDatetime @@ -21,12 +22,16 @@ class InsertedJob(BaseModel): class HeartbeatData(BaseModel, extra="forbid"): - load_average: float | None = Field(None, alias="LoadAverage") - memory_used: float | None = Field(None, alias="MemoryUsed") - vsize: float | None = Field(None, alias="Vsize") - available_disk_space: float | None = Field(None, alias="AvailableDiskSpace") - cpu_consumed: float | None = Field(None, alias="CPUConsumed") - wall_clock_time: float | None = Field(None, alias="WallClockTime") + load_average: float | None = Field(None, alias="LoadAverage", allow_inf_nan=False) + memory_used: float | None = Field(None, alias="MemoryUsed", allow_inf_nan=False) + vsize: float | None = Field(None, alias="Vsize", allow_inf_nan=False) + available_disk_space: float | None = Field( + None, alias="AvailableDiskSpace", allow_inf_nan=False + ) + cpu_consumed: float | None = Field(None, alias="CPUConsumed", allow_inf_nan=False) + wall_clock_time: float | None = Field( + None, alias="WallClockTime", allow_inf_nan=False + ) standard_output: str | None = Field(None, alias="StandardOutput") @@ -36,6 +41,18 @@ class JobCommand(BaseModel): arguments: str | None = None +def _ensure_finite_numbers(value: Any, path: str) -> None: + """Raise ValueError if a (possibly nested) value contains NaN or infinity.""" + if isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"{path}: non-finite numbers are not supported") + if isinstance(value, dict): + for key, item in value.items(): + _ensure_finite_numbers(item, f"{path}.{key}") + elif isinstance(value, (list, tuple)): + for i, item in enumerate(value): + _ensure_finite_numbers(item, f"{path}[{i}]") + + class JobParameters(BaseModel, populate_by_name=True, extra="allow"): """Some of the most important parameters that can be set for a job.""" @@ -73,6 +90,17 @@ def convert_cpu_fields_to_int(cls, v): return int(v) return v + @model_validator(mode="after") + def validate_extra_fields_are_json_safe(self): + """Reject extra field values which cannot be represented in strict JSON. + + Python's JSON parser accepts NaN and (-)Infinity so such values survive + request parsing, but OpenSearch rejects documents containing them. + """ + for name, value in (self.model_extra or {}).items(): + _ensure_finite_numbers(value, name) + return self + class JobAttributes(BaseModel, populate_by_name=True, extra="forbid"): """All the attributes that can be set for a job.""" diff --git a/diracx-core/src/diracx/core/models/pilot.py b/diracx-core/src/diracx/core/models/pilot.py index a325ca044..58fee08c6 100644 --- a/diracx-core/src/diracx/core/models/pilot.py +++ b/diracx-core/src/diracx/core/models/pilot.py @@ -67,7 +67,10 @@ class PilotMetadata(BaseModel, populate_by_name=True, extra="forbid"): None, alias="Status", description="Current pilot status." ) benchmark: float | None = Field( - None, alias="BenchMark", description="Pilot benchmark value." + None, + alias="BenchMark", + allow_inf_nan=False, + description="Pilot benchmark value.", ) destination_site: str | None = Field( None, alias="DestinationSite", max_length=128, description="Destination site." diff --git a/diracx-db/src/diracx/db/os/utils.py b/diracx-db/src/diracx/db/os/utils.py index be31d224d..f7db62b12 100644 --- a/diracx-db/src/diracx/db/os/utils.py +++ b/diracx-db/src/diracx/db/os/utils.py @@ -10,8 +10,9 @@ from typing import Any, Self from opensearchpy import AsyncOpenSearch +from opensearchpy.exceptions import RequestError -from diracx.core.exceptions import InvalidQueryError +from diracx.core.exceptions import DocumentUpsertError, InvalidQueryError from diracx.core.extensions import DiracEntryPoint, select_from_extension from diracx.core.settings import FactorySettings from diracx.db.exceptions import DBUnavailableError @@ -183,12 +184,24 @@ async def create_index_template(self) -> None: async def upsert(self, vo: str, doc_id: int, document: Any) -> None: index_name = self.index_name(vo, doc_id) - response = await self.client.update( - index=index_name, - id=doc_id, - body={"doc": document, "doc_as_upsert": True}, - params=dict(retry_on_conflict=10), - ) + try: + response = await self.client.update( + index=index_name, + id=doc_id, + body={"doc": document, "doc_as_upsert": True}, + params=dict(retry_on_conflict=10), + ) + except RequestError as e: + logger.error( + "Failed to upsert document %s in index %s: %s (document: %r)", + doc_id, + index_name, + e.info, + document, + ) + raise DocumentUpsertError( + f"Failed to upsert document {doc_id} in {self.__class__.__name__}: {e.error}" + ) from e logger.debug( "Upserted document %s in index %s with response: %s", doc_id, diff --git a/diracx-db/tests/opensearch/test_upsert.py b/diracx-db/tests/opensearch/test_upsert.py new file mode 100644 index 000000000..a81566135 --- /dev/null +++ b/diracx-db/tests/opensearch/test_upsert.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import pytest + +from diracx.core.exceptions import DocumentUpsertError +from diracx.testing.osdb import DummyOSDB + + +async def test_upsert_valid_document(dummy_opensearch_db: DummyOSDB): + """Sanity check that a well-formed document can be upserted.""" + await dummy_opensearch_db.upsert("dummyvo", 1, {"IntField": 1234}) + await dummy_opensearch_db.client.indices.refresh( + index=f"{dummy_opensearch_db.index_prefix}*" + ) + results = await dummy_opensearch_db.search( + None, [{"parameter": "IntField", "operator": "eq", "value": "1234"}], [] + ) + assert len(results) == 1 + + +async def test_upsert_unparsable_document_raises(dummy_opensearch_db: DummyOSDB): + """NaN survives Python JSON serialization but OpenSearch rejects it. + + This must surface as a DocumentUpsertError rather than an unhandled + RequestError, and the offending document must be logged. + """ + with pytest.raises(DocumentUpsertError, match="Failed to upsert document"): + await dummy_opensearch_db.upsert("dummyvo", 2, {"IntField": float("nan")}) diff --git a/diracx-logic/src/diracx/logic/jobs/status.py b/diracx-logic/src/diracx/logic/jobs/status.py index 38d0339d3..f6c77f347 100644 --- a/diracx-logic/src/diracx/logic/jobs/status.py +++ b/diracx-logic/src/diracx/logic/jobs/status.py @@ -19,6 +19,7 @@ ) from diracx.core.config import Config +from diracx.core.exceptions import DiracError from diracx.core.models import ( HeartbeatData, JobAttributes, @@ -577,41 +578,54 @@ async def add_heartbeat( if result["Status"] in [JobStatus.MATCHED, JobStatus.STALLED] } - async with TaskGroup() as tg: - if status_changes: - tg.create_task( - set_job_statuses( - status_changes=status_changes, - config=config, - job_db=job_db, - job_logging_db=job_logging_db, - task_queue_db=task_queue_db, - job_parameters_db=job_parameters_db, + try: + async with TaskGroup() as tg: + if status_changes: + tg.create_task( + set_job_statuses( + status_changes=status_changes, + config=config, + job_db=job_db, + job_logging_db=job_logging_db, + task_queue_db=task_queue_db, + job_parameters_db=job_parameters_db, + ) ) - ) - - if other_ids := set(data) - set(status_changes): - # If there are no status changes, we still need to update the heartbeat time - heartbeat_updates = { - job_id: {"HeartBeatTime": utcnow()} for job_id in other_ids - } - tg.create_task(job_db.set_job_attributes(heartbeat_updates)) - os_data_by_job_id: defaultdict[int, dict[str, Any]] = defaultdict(dict) - for job_id, job_data in data.items(): - sql_data = {} - for key, value in job_data.model_dump( - by_alias=True, exclude_defaults=True - ).items(): - if key in job_db.heartbeat_fields: - sql_data[key] = value - else: - os_data_by_job_id[job_id][key] = value - - if sql_data: - tg.create_task(job_db.add_heartbeat_data(job_id, sql_data)) - - await _insert_parameters(os_data_by_job_id, job_parameters_db, job_db) + if other_ids := set(data) - set(status_changes): + # If there are no status changes, we still need to update the heartbeat time + heartbeat_updates = { + job_id: {"HeartBeatTime": utcnow()} for job_id in other_ids + } + tg.create_task(job_db.set_job_attributes(heartbeat_updates)) + + os_data_by_job_id: defaultdict[int, dict[str, Any]] = defaultdict(dict) + for job_id, job_data in data.items(): + sql_data = {} + for key, value in job_data.model_dump( + by_alias=True, exclude_defaults=True + ).items(): + if key in job_db.heartbeat_fields: + sql_data[key] = value + else: + os_data_by_job_id[job_id][key] = value + + if sql_data: + tg.create_task(job_db.add_heartbeat_data(job_id, sql_data)) + + await _insert_parameters(os_data_by_job_id, job_parameters_db, job_db) + except* DiracError as eg: + # Re-raise a DiracError directly rather than the surrounding + # ExceptionGroup so callers can catch it by exception type + raise _first_leaf(eg) from eg + + +def _first_leaf(eg: BaseExceptionGroup) -> BaseException: + """Return the first non-group exception contained in an exception group.""" + exc: BaseException = eg + while isinstance(exc, BaseExceptionGroup): + exc = exc.exceptions[0] + return exc async def _insert_parameters( @@ -641,11 +655,16 @@ async def _insert_parameters( job_id_to_vo = {int(x["JobID"]): str(x["VO"]) for x in job_vos} # Upsert the parameters into the JobParametersDB # TODO: can we do a bulk upsert instead - async with TaskGroup() as tg: - for job_id, job_params in updates.items(): - tg.create_task( - job_parameters_db.upsert(job_id_to_vo[job_id], job_id, job_params) - ) + try: + async with TaskGroup() as tg: + for job_id, job_params in updates.items(): + tg.create_task( + job_parameters_db.upsert(job_id_to_vo[job_id], job_id, job_params) + ) + except* DiracError as eg: + # Re-raise a DiracError directly rather than the surrounding + # ExceptionGroup so callers can catch it by exception type + raise _first_leaf(eg) from eg async def get_job_commands(job_ids: Iterable[int], job_db: JobDB) -> list[JobCommand]: diff --git a/diracx-logic/tests/jobs/test_status.py b/diracx-logic/tests/jobs/test_status.py index f5c482df7..e4dfbb5d7 100644 --- a/diracx-logic/tests/jobs/test_status.py +++ b/diracx-logic/tests/jobs/test_status.py @@ -5,6 +5,7 @@ import pytest +from diracx.core.exceptions import DocumentUpsertError from diracx.core.models import JobMetaData from diracx.db.os.job_parameters import JobParametersDB as RealJobParametersDB from diracx.db.sql.job.db import JobDB @@ -158,3 +159,27 @@ async def test_patch_metadata_updates_attributes_and_parameters( assert prow["does_not_exist"] == "unknown" assert "UserPriority" not in prow assert "HeartBeatTime" not in prow + + +@pytest.mark.asyncio +async def test_upsert_failure_propagates_as_bare_dirac_error( + job_db: JobDB, + job_parameters_db: _MockJobParametersDB, + valid_job_id: int, + monkeypatch: pytest.MonkeyPatch, +): + """A DiracError raised while upserting job parameters must propagate as is. + + The TaskGroup wraps failures in an ExceptionGroup, which callers cannot + catch by exception type; the logic layer must collapse it. + """ + + async def failing_upsert(vo, doc_id, document): + raise DocumentUpsertError("failed to parse field [doc]") + + monkeypatch.setattr(job_parameters_db, "upsert", failing_upsert) + + updates = {valid_job_id: JobMetaData.model_validate({"SomeParameter": "value"})} + with pytest.raises(DocumentUpsertError): + async with job_db: + await set_job_parameters_or_attributes(updates, job_db, job_parameters_db) diff --git a/diracx-routers/src/diracx/routers/factory.py b/diracx-routers/src/diracx/routers/factory.py index 7ca979ad2..bf4ab3430 100644 --- a/diracx-routers/src/diracx/routers/factory.py +++ b/diracx-routers/src/diracx/routers/factory.py @@ -6,6 +6,7 @@ import inspect import logging +import math from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable, Sequence from functools import partial from http import HTTPStatus @@ -16,7 +17,7 @@ from cachetools import TTLCache from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request from fastapi.dependencies.models import Dependant -from fastapi.exception_handlers import request_validation_exception_handler +from fastapi.encoders import jsonable_encoder from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, Response @@ -443,6 +444,17 @@ def route_unavailable_error_hander(request: Request, exc: DBUnavailableError): ) +def _replace_non_finite(obj): + """Replace NaN and infinity with their repr as they cannot be serialized to JSON.""" + if isinstance(obj, float) and not math.isfinite(obj): + return repr(obj) + if isinstance(obj, dict): + return {k: _replace_non_finite(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_replace_non_finite(v) for v in obj] + return obj + + async def validation_error_handler(request: Request, exc: RequestValidationError): logger_422.warning( "Got validation error: %s in %s %s with body %r", @@ -460,7 +472,13 @@ async def validation_error_handler(request: Request, exc: RequestValidationError # } # }, ) - return await request_validation_exception_handler(request, exc) + # The rejected input is echoed in the error detail and may contain values + # which cannot be represented in strict JSON, such as NaN + detail = _replace_non_finite(jsonable_encoder(exc.errors())) + return JSONResponse( + status_code=HTTPStatus.UNPROCESSABLE_ENTITY, + content={"detail": detail}, + ) def find_dependents( diff --git a/diracx-routers/tests/jobs/test_heartbeat_commands.py b/diracx-routers/tests/jobs/test_heartbeat_commands.py index 0804e15c2..01aedf843 100644 --- a/diracx-routers/tests/jobs/test_heartbeat_commands.py +++ b/diracx-routers/tests/jobs/test_heartbeat_commands.py @@ -24,6 +24,23 @@ ) +def test_heartbeat_rejects_non_finite_values( + normal_user_client: TestClient, valid_job_id: int +): + """Non-finite floats survive Python JSON parsing but cannot be stored. + + They used to be forwarded to the database backends, which reject them, + resulting in an internal server error. + """ + # Send the raw body as httpx itself refuses to serialize NaN + r = normal_user_client.patch( + "/api/jobs/heartbeat", + content=f'{{"{valid_job_id}": {{"Vsize": NaN}}}}', + headers={"Content-Type": "application/json"}, + ) + assert r.status_code == 422, r.text + + def test_heartbeat(frozen_time, normal_user_client: TestClient, valid_job_id: int): search_body = { "search": [{"parameter": "JobID", "operator": "eq", "value": valid_job_id}] diff --git a/diracx-routers/tests/jobs/test_status.py b/diracx-routers/tests/jobs/test_status.py index d84a8868a..863c6ad93 100644 --- a/diracx-routers/tests/jobs/test_status.py +++ b/diracx-routers/tests/jobs/test_status.py @@ -899,6 +899,28 @@ def test_patch_metadata(normal_user_client: TestClient, valid_job_id: int): assert r.json()[0]["UserPriority"] == 2 +@pytest.mark.parametrize( + "raw_value", + ["NaN", "Infinity", "-Infinity", '{"nested": [NaN]}'], +) +def test_patch_metadata_rejects_non_finite_values( + normal_user_client: TestClient, valid_job_id: int, raw_value: str +): + """Non-finite floats survive Python JSON parsing but cannot be stored. + + They used to be forwarded to OpenSearch, which rejects them, resulting + in an internal server error. + """ + # Send the raw body as httpx itself refuses to serialize NaN + r = normal_user_client.patch( + "/api/jobs/metadata", + content=f'{{"{valid_job_id}": {{"SomeParameter": {raw_value}}}}}', + headers={"Content-Type": "application/json"}, + ) + assert r.status_code == 422, r.text + assert "non-finite" in r.text + + def test_diracx_476(normal_user_client: TestClient, valid_job_id: int): """Test fix for https://github.com/DIRACGrid/diracx/issues/476.""" inner_payload = {"Status": JobStatus.FAILED.value, "MinorStatus": "Payload failed"} diff --git a/diracx-routers/tests/pilots/test_management.py b/diracx-routers/tests/pilots/test_management.py index 0b4e447a3..44b61b201 100644 --- a/diracx-routers/tests/pilots/test_management.py +++ b/diracx-routers/tests/pilots/test_management.py @@ -95,6 +95,26 @@ async def test_update_pilot_metadata_applies_partial_fields(normal_test_client): assert by_stamp["stamp_m2"]["BenchMark"] == 0.0 # untouched +async def test_update_pilot_metadata_rejects_non_finite_benchmark(normal_test_client): + """Non-finite floats survive Python JSON parsing but cannot be stored. + + They used to be forwarded to the database, which rejects them, resulting + in an internal server error. + """ + r = normal_test_client.post( + "/api/pilots/", json={"pilot_stamp": "stamp_nan", "vo": MAIN_VO} + ) + assert r.status_code == 201 + + # Send the raw body as httpx itself refuses to serialize NaN + r = normal_test_client.patch( + "/api/pilots/metadata", + content='{"stamp_nan": {"BenchMark": NaN}}', + headers={"Content-Type": "application/json"}, + ) + assert r.status_code == 422, r.text + + async def test_update_pilot_metadata_unknown_stamp_returns_404(normal_test_client): r = normal_test_client.patch( "/api/pilots/metadata", diff --git a/docs/dev/reference/coding-conventions.md b/docs/dev/reference/coding-conventions.md index 7def05a2f..d15087ad7 100644 --- a/docs/dev/reference/coding-conventions.md +++ b/docs/dev/reference/coding-conventions.md @@ -83,6 +83,38 @@ delay = datetime.datetime.now() + datetime.timedelta(hours=1)