diff --git a/miles/utils/workers/rpc/common/serialization.py b/miles/utils/workers/rpc/common/serialization.py index 63aef8ba391..8df09b68749 100644 --- a/miles/utils/workers/rpc/common/serialization.py +++ b/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: diff --git a/tests/fast/utils/workers/e2e/e2e_worker.py b/tests/fast/utils/workers/e2e/e2e_worker.py index 328f223cb4a..f0e510383c0 100644 --- a/tests/fast/utils/workers/e2e/e2e_worker.py +++ b/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 diff --git a/tests/fast/utils/workers/e2e/test_complex_types.py b/tests/fast/utils/workers/e2e/test_complex_types.py index 5ecf8a04ec7..847efdb1a75 100644 --- a/tests/fast/utils/workers/e2e/test_complex_types.py +++ b/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.""" diff --git a/tests/fast/utils/workers/rpc/common/test_serialization.py b/tests/fast/utils/workers/rpc/common/test_serialization.py index 1b28f6c0588..0fe3d53996e 100644 --- a/tests/fast/utils/workers/rpc/common/test_serialization.py +++ b/tests/fast/utils/workers/rpc/common/test_serialization.py @@ -2,6 +2,7 @@ import datetime import enum import json +import math import uuid from decimal import Decimal from pathlib import PurePosixPath @@ -43,6 +44,11 @@ class Point: y: int +@dataclasses.dataclass +class Reading: + value: float + + class Blob(StrictBaseModel): data: bytes @@ -52,6 +58,9 @@ class Options(TypedDict): label: str +NON_FINITE_FLOAT_TAG = "__miles_non_finite_float__" + + _ROUNDTRIP_CASES = [ ("int", int, 42), ("negative_int", int, -7), @@ -103,6 +112,24 @@ class Options(TypedDict): ] +_NON_FINITE_CASES = [ + ("float_nan", float, math.nan), + ("float_inf", float, math.inf), + ("float_negative_inf", float, -math.inf), + ("optional_float_nan", float | None, math.nan), + ("optional_float_inf", float | None, math.inf), + ("any_nan", Any, math.nan), + ("any_negative_inf", Any, -math.inf), + ("bare_dict_nan", dict, {"a": math.nan, "b": 1.0}), + ("dict_str_float_inf", dict[str, float], {"a": math.inf}), + ("list_of_float_nan", list[float], [math.nan, 1.0, -math.inf]), + ("tuple_of_float_nan_and_inf", tuple[float, float], (math.nan, math.inf)), + ("model_field_nan", Inner, Inner(name="x", score=math.nan)), + ("model_field_negative_inf", Inner, Inner(name="x", score=-math.inf)), + ("dataclass_field_inf", Reading, Reading(value=math.inf)), + ("deeply_nested_nan", dict[str, list[dict[str, float]]], {"a": [{"b": math.nan}]}), +] + _BYTES_CASES = [ ("non_utf8_bytes", bytes, b"\x00\x80\xff"), ("empty_bytes", bytes, b""), @@ -121,6 +148,14 @@ def _through_the_wire(payload: Any) -> Any: return json.loads(json.dumps(payload)) +def _through_the_strict_wire(payload: Any) -> Any: + return json.loads(json.dumps(payload, allow_nan=False)) + + +def _equal_including_non_finite_floats(actual: Any, expected: Any) -> bool: + return repr(actual) == repr(expected) + + @pytest.mark.parametrize("case", _ROUNDTRIP_CASES, ids=[case[0] for case in _ROUNDTRIP_CASES]) class TestWireRoundtrip: def test_result_survives_the_wire_as_the_declared_type(self, case): @@ -233,6 +268,161 @@ def test_encoded_payloads_are_plain_json_types(self): assert json.dumps(serializer.encode_result(value)) +@pytest.mark.parametrize("case", _NON_FINITE_CASES, ids=[case[0] for case in _NON_FINITE_CASES]) +class TestNonFiniteFloatRoundtrip: + def test_non_finite_result_survives_the_wire(self, case): + """Every non-finite float result comes back from json as the same non-finite float.""" + _, annotation, value = case + serializer = _serializer(annotation) + decoded = serializer.decode_result(_through_the_strict_wire(serializer.encode_result(value))) + assert _equal_including_non_finite_floats(decoded, value) + + def test_non_finite_argument_survives_the_wire(self, case): + """Every non-finite float argument reaches the worker as the same non-finite float.""" + _, annotation, value = case + serializer = _serializer(annotation) + decoded = serializer.decode_query(_through_the_strict_wire(serializer.encode_query({"payload": value}))) + assert _equal_including_non_finite_floats(decoded["payload"], value) + + def test_encoded_non_finite_result_is_strict_json(self, case): + """Encoding a non-finite float result yields a payload json accepts without allow_nan.""" + _, annotation, value = case + serializer = _serializer(annotation) + assert json.dumps(serializer.encode_result(value), allow_nan=False) + + def test_encoded_non_finite_argument_is_strict_json(self, case): + """Encoding a non-finite float argument yields a payload json accepts without allow_nan.""" + _, annotation, value = case + serializer = _serializer(annotation) + assert json.dumps(serializer.encode_query({"payload": value}), allow_nan=False) + + +class TestNonFiniteFloatValues: + def test_nan_result_decodes_to_nan(self): + """A NaN result decodes back to a real NaN float rather than None or a string.""" + serializer = _serializer(float) + decoded = serializer.decode_result(_through_the_strict_wire(serializer.encode_result(float("nan")))) + assert isinstance(decoded, float) + assert math.isnan(decoded) + + def test_infinity_result_decodes_to_positive_infinity(self): + """An infinity result decodes back to positive infinity rather than None or a string.""" + serializer = _serializer(float) + decoded = serializer.decode_result(_through_the_strict_wire(serializer.encode_result(float("inf")))) + assert decoded == math.inf + + def test_negative_infinity_result_decodes_to_negative_infinity(self): + """A negative infinity result decodes back to negative infinity rather than None or a string.""" + serializer = _serializer(float) + decoded = serializer.decode_result(_through_the_strict_wire(serializer.encode_result(float("-inf")))) + assert decoded == -math.inf + + def test_nan_result_under_any_annotation_decodes_to_nan(self): + """A NaN result declared as Any keeps its float identity instead of becoming a marker dict.""" + serializer = _serializer(Any) + decoded = serializer.decode_result(_through_the_strict_wire(serializer.encode_result(math.nan))) + assert isinstance(decoded, float) + assert math.isnan(decoded) + + def test_nan_result_inside_a_model_field_decodes_to_nan(self): + """A NaN nested in a model field decodes back to NaN on the revived model.""" + serializer = _serializer(Inner) + decoded = serializer.decode_result( + _through_the_strict_wire(serializer.encode_result(Inner(name="x", score=math.nan))) + ) + assert isinstance(decoded, Inner) + assert math.isnan(decoded.score) + + def test_infinity_result_inside_a_dataclass_field_decodes_to_infinity(self): + """An infinity nested in a dataclass field decodes back to infinity on the revived dataclass.""" + serializer = _serializer(Reading) + decoded = serializer.decode_result(_through_the_strict_wire(serializer.encode_result(Reading(value=math.inf)))) + assert isinstance(decoded, Reading) + assert decoded.value == math.inf + + def test_nan_argument_decodes_to_nan(self): + """A NaN argument reaches the worker as a real NaN float.""" + serializer = _serializer(float) + decoded = serializer.decode_query(_through_the_strict_wire(serializer.encode_query({"payload": math.nan}))) + assert math.isnan(decoded["payload"]) + + def test_infinity_argument_decodes_to_infinity(self): + """An infinity argument reaches the worker as positive infinity.""" + serializer = _serializer(float) + decoded = serializer.decode_query(_through_the_strict_wire(serializer.encode_query({"payload": math.inf}))) + assert decoded["payload"] == math.inf + + def test_nan_argument_inside_a_dict_decodes_to_nan(self): + """A NaN nested inside a dict argument reaches the worker as a real NaN float.""" + serializer = _serializer(dict[str, float]) + decoded = serializer.decode_query( + _through_the_strict_wire(serializer.encode_query({"payload": {"a": math.nan}})) + ) + assert math.isnan(decoded["payload"]["a"]) + + def test_nan_argument_inside_a_model_field_decodes_to_nan(self): + """A NaN nested in a model argument field reaches the worker as a real NaN float.""" + serializer = _serializer(Inner) + decoded = serializer.decode_query( + _through_the_strict_wire(serializer.encode_query({"payload": Inner(name="x", score=math.nan)})) + ) + assert math.isnan(decoded["payload"].score) + + +@pytest.mark.parametrize("value", [0.0, -0.0, 1e308, -1e308, 1.5]) +class TestFiniteFloatsAreUntouched: + def test_finite_float_result_is_encoded_as_a_plain_float(self, value): + """A finite float result encodes to the same plain float without any marker.""" + serializer = _serializer(float) + encoded = serializer.encode_result(value) + assert isinstance(encoded, float) + assert NON_FINITE_FLOAT_TAG not in json.dumps(encoded) + + def test_finite_float_result_survives_the_wire_unchanged(self, value): + """A finite float result comes back with the exact same value and sign.""" + serializer = _serializer(float) + decoded = serializer.decode_result(_through_the_strict_wire(serializer.encode_result(value))) + assert _equal_including_non_finite_floats(decoded, value) + + def test_finite_float_argument_is_encoded_without_any_marker(self, value): + """A finite float argument encodes without any marker key on the wire.""" + serializer = _serializer(float) + assert NON_FINITE_FLOAT_TAG not in json.dumps(serializer.encode_query({"payload": value})) + + +class TestMarkerLookalikes: + @pytest.mark.parametrize("value", ["nan", "inf", "-inf"]) + def test_string_that_looks_like_a_token_stays_a_string(self, value): + """A string spelled like a non-finite token stays a string instead of being revived as a float.""" + serializer = _serializer(str) + decoded = serializer.decode_result(_through_the_strict_wire(serializer.encode_result(value))) + assert decoded == value + + def test_token_string_inside_a_dict_stays_a_string(self): + """A token-shaped string nested in a dict stays a string rather than becoming a float.""" + serializer = _serializer(dict[str, str]) + decoded = serializer.decode_result(_through_the_strict_wire(serializer.encode_result({"a": "nan"}))) + assert decoded == {"a": "nan"} + + def test_marker_shaped_dict_with_extra_keys_stays_a_dict(self): + """A dict carrying the marker key alongside other keys decodes as a plain dict.""" + serializer = _serializer(dict) + payload = {NON_FINITE_FLOAT_TAG: "nan", "extra": 1} + assert serializer.decode_result(payload) == payload + + def test_marker_shaped_dict_with_unknown_token_stays_a_dict(self): + """A dict carrying the marker key with an unknown token decodes as a plain dict.""" + serializer = _serializer(dict) + payload = {NON_FINITE_FLOAT_TAG: "not-a-token"} + assert serializer.decode_result(payload) == payload + + def test_marker_shaped_dict_with_extra_keys_stays_a_dict_as_an_argument(self): + """A marker-shaped argument dict with extra keys reaches the worker as a plain dict.""" + serializer = _serializer(dict) + payload = {NON_FINITE_FLOAT_TAG: "nan", "extra": 1} + assert serializer.decode_query({"payload": payload})["payload"] == payload + + @pytest.mark.parametrize("case", _BYTES_CASES, ids=[case[0] for case in _BYTES_CASES]) class TestBytesOnTheWire: def test_bytes_result_survives_the_wire(self, case): @@ -280,6 +470,36 @@ def test_non_utf8_bytes_nested_in_a_model_are_refused_loudly(self): class TestRejectedPayloads: + def test_result_containing_the_reserved_marker_key_is_rejected(self): + """A result dict already using the reserved marker key fails loudly instead of being ambiguous.""" + serializer = _serializer(dict) + with pytest.raises(ValueError, match=NON_FINITE_FLOAT_TAG): + serializer.encode_result({NON_FINITE_FLOAT_TAG: "nan"}) + + def test_result_using_the_reserved_marker_key_alongside_others_is_rejected(self): + """The reserved key is refused even when it shares a dict with unrelated keys.""" + serializer = _serializer(dict) + with pytest.raises(ValueError, match=NON_FINITE_FLOAT_TAG): + serializer.encode_result({NON_FINITE_FLOAT_TAG: "nan", "extra": 1}) + + def test_result_using_the_reserved_marker_key_with_an_unknown_token_is_rejected(self): + """The reserved key is refused whatever value it carries, not only recognized tokens.""" + serializer = _serializer(dict) + with pytest.raises(ValueError, match=NON_FINITE_FLOAT_TAG): + serializer.encode_result({NON_FINITE_FLOAT_TAG: "not-a-token"}) + + def test_nested_result_containing_the_reserved_marker_key_is_rejected(self): + """A result carrying the reserved marker key deep inside is rejected too.""" + serializer = _serializer(dict) + with pytest.raises(ValueError): + serializer.encode_result({"outer": [{NON_FINITE_FLOAT_TAG: "inf"}]}) + + def test_argument_containing_the_reserved_marker_key_is_rejected(self): + """An argument dict already using the reserved marker key fails loudly instead of being ambiguous.""" + serializer = _serializer(dict) + with pytest.raises(ValueError): + serializer.encode_query({"payload": {NON_FINITE_FLOAT_TAG: "nan"}}) + def test_unserializable_result_is_rejected(self): """A result that is not json encodable fails rather than being silently coerced.""" serializer = _serializer(int)