fix: avoid 500 errors when job parameters cannot be stored

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

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."""

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."

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,

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")})

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]:

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)

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(