Support non-finite floats across the wire

miles/utils/workers/rpc/common/serialization.py

@@ -1,12 +1,18 @@
 from __future__ import annotations
 
 import dataclasses
+import math
 from typing import Any
 
 import typing_extensions
 from pydantic import BaseModel, ConfigDict, TypeAdapter, create_model
 
-_WIRE_CONFIG = ConfigDict(ser_json_bytes="base64", val_json_bytes="base64")
+NON_FINITE_FLOAT_TAG = "__miles_non_finite_float__"
+
+_WIRE_CONFIG = ConfigDict(ser_json_inf_nan="constants", ser_json_bytes="base64", val_json_bytes="base64")
+
+_TOKEN_BY_FLOAT = {math.inf: "inf", -math.inf: "-inf"}
+_FLOAT_BY_TOKEN = {"nan": math.nan, "inf": math.inf, "-inf": -math.inf}
 
 
 @dataclasses.dataclass(frozen=True)
@@ -23,16 +29,41 @@ def create(cls, *, query_model_name: str, query_fields: dict[str, Any], result_a
         return cls(query_model=query_model, result_adapter=TypeAdapter(result_annotation, config=result_config))
 
     def encode_query(self, kwargs: dict[str, Any]) -> dict[str, Any]:
-        return self.query_model(**kwargs).model_dump(mode="json")
+        return _NonFiniteFloatCodec.encode(self.query_model(**kwargs).model_dump(mode="json"))
 
     def decode_query(self, query: dict[str, Any]) -> dict[str, Any]:
-        return dict(self.query_model(**query))
+        return dict(self.query_model(**_NonFiniteFloatCodec.decode(query)))
 
     def encode_result(self, result: Any) -> Any:
-        return self.result_adapter.dump_python(result, mode="json")
+        return _NonFiniteFloatCodec.encode(self.result_adapter.dump_python(result, mode="json"))
 
     def decode_result(self, payload: Any) -> Any:
-        return self.result_adapter.validate_python(payload)
+        return self.result_adapter.validate_python(_NonFiniteFloatCodec.decode(payload))
+
+
+class _NonFiniteFloatCodec:
+    @classmethod
+    def encode(cls, value: Any) -> Any:
+        if isinstance(value, float) and not math.isfinite(value):
+            return {NON_FINITE_FLOAT_TAG: _TOKEN_BY_FLOAT.get(value, "nan")}
+        if isinstance(value, dict):
+            if NON_FINITE_FLOAT_TAG in value:
+                raise ValueError(f"rpc payloads must not contain the reserved key {NON_FINITE_FLOAT_TAG!r}")
+            return {key: cls.encode(item) for key, item in value.items()}
+        if isinstance(value, list):
+            return [cls.encode(item) for item in value]
+        return value
+
+    @classmethod
+    def decode(cls, value: Any) -> Any:
+        if isinstance(value, dict):
+            token = value.get(NON_FINITE_FLOAT_TAG)
+            if len(value) == 1 and token in _FLOAT_BY_TOKEN:
+                return _FLOAT_BY_TOKEN[token]
+            return {key: cls.decode(item) for key, item in value.items()}
+        if isinstance(value, list):
+            return [cls.decode(item) for item in value]
+        return value
 
 
 def _carries_own_config(annotation: Any) -> bool:

tests/fast/utils/workers/e2e/e2e_worker.py

@@ -40,6 +40,11 @@ class Point:
     y: int
 
 
+class Metric(StrictBaseModel):
+    name: str
+    value: float
+
+
 class Event(StrictBaseModel):
     tag: str
     phase: str
@@ -151,6 +156,27 @@ async def demo_model_list(self, items: list[Item]) -> list[Item]:
     async def demo_bytes(self, blob: bytes) -> bytes:
         return blob
 
+    async def demo_nan_result(self) -> float:
+        return float("nan")
+
+    async def demo_float(self, value: float) -> float:
+        return value
+
+    async def demo_optional_float(self, value: float | None) -> float | None:
+        return value
+
+    async def demo_float_metrics(self) -> dict:
+        return {"loss": float("nan"), "grad_norm": float("inf"), "lr": -float("inf"), "step": 3.0}
+
+    async def demo_float_list(self, values: list[float]) -> list[float]:
+        return values
+
+    async def demo_metric_model(self, metric: Metric) -> Metric:
+        return metric
+
+    async def report_float_repr(self, value: float) -> str:
+        return repr(value)
+
     async def demo_bytes_list(self, blobs: list[bytes]) -> list[bytes]:
         return blobs
 

tests/fast/utils/workers/e2e/test_complex_types.py

@@ -1,11 +1,12 @@
 import datetime
+import math
 import uuid
 from decimal import Decimal
 
 import pytest
 from pydantic import ValidationError
 
-from tests.fast.utils.workers.e2e.e2e_worker import Colour, Item, Nested, Point
+from tests.fast.utils.workers.e2e.e2e_worker import Colour, Item, Metric, Nested, Point
 
 
 class TestComplexArgumentsAndResults:
@@ -51,6 +52,15 @@ async def test_aware_datetime_keeps_its_offset(self, handle):
         assert isinstance(result, datetime.datetime)
         assert result == value
 
+    async def test_non_utc_datetime_keeps_its_offset(self, handle):
+        """A datetime with a non-zero utc offset arrives with that same offset, not normalised to utc."""
+        tzinfo = datetime.timezone(datetime.timedelta(hours=5, minutes=30))
+        value = datetime.datetime(2026, 7, 27, 12, 30, tzinfo=tzinfo)
+        result = await handle.demo_datetime(when=value)
+        assert result == value
+        assert result.utcoffset() == datetime.timedelta(hours=5, minutes=30)
+        assert (result.hour, result.minute) == (12, 30)
+
     async def test_naive_datetime_stays_naive(self, handle):
         """A naive datetime does not acquire a timezone on the way through."""
         value = datetime.datetime(2026, 7, 27, 12, 30)
@@ -134,6 +144,49 @@ async def test_none_result_roundtrips(self, handle):
         assert await handle.demo_none_result() is None
 
 
+class TestNonFiniteFloats:
+    async def test_nan_result_arrives_as_nan(self, handle):
+        """A NaN result reaches the caller as NaN instead of being silently nulled."""
+        assert math.isnan(await handle.demo_nan_result())
+
+    @pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf, 0.0, 1.5])
+    async def test_float_argument_and_result_roundtrip(self, handle, value: float):
+        """Every float, finite or not, survives both directions unchanged."""
+        result = await handle.demo_float(value=value)
+        assert repr(result) == repr(value)
+
+    @pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf, None])
+    async def test_optional_float_keeps_nan_distinct_from_none(self, handle, value: float | None):
+        """A NaN under an optional annotation is not confused with None."""
+        result = await handle.demo_optional_float(value=value)
+        assert repr(result) == repr(value)
+
+    async def test_non_finite_floats_survive_a_loose_dict_result(self, handle):
+        """Non-finite floats inside an untyped dict result are not collapsed to None."""
+        metrics = await handle.demo_float_metrics()
+        assert math.isnan(metrics["loss"])
+        assert metrics["grad_norm"] == math.inf
+        assert metrics["lr"] == -math.inf
+        assert metrics["step"] == 3.0
+
+    async def test_non_finite_floats_survive_a_list_argument(self, handle):
+        """A list mixing finite and non-finite floats roundtrips element by element."""
+        values = [math.nan, math.inf, -math.inf, 0.0, -0.0, 1e308]
+        result = await handle.demo_float_list(values=values)
+        assert [repr(item) for item in result] == [repr(item) for item in values]
+
+    async def test_non_finite_float_survives_a_model_field(self, handle):
+        """A NaN inside a model field roundtrips as NaN."""
+        result = await handle.demo_metric_model(metric=Metric(name="loss", value=math.nan))
+        assert result.name == "loss"
+        assert math.isnan(result.value)
+
+    async def test_worker_receives_the_non_finite_argument(self, handle):
+        """The worker is handed a real NaN, not None and not the string 'nan'."""
+        assert await handle.report_float_repr(value=math.nan) == "nan"
+        assert await handle.report_float_repr(value=-math.inf) == "-inf"
+
+
 class TestWorkerSideTypes:
     async def test_worker_receives_model_instances(self, handle):
         """The worker is handed revived models, not the raw json dicts."""