Fix NaN/Infinity serialization in save_to_json_file (#4144)

ax/api/client.py

@@ -1057,7 +1057,7 @@ def save_to_json_file(self, filepath: str = "ax_client_snapshot.json") -> None:
         to a .json file by the given path.
         """
         with open(filepath, "w+") as file:
-            file.write(json.dumps(self._to_json_snapshot()))
+            file.write(json.dumps(self._to_json_snapshot(), allow_nan=False))
             logger.debug(
                 f"Saved JSON-serialized state of optimization to `{filepath}`."
             )

ax/api/tests/test_client.py

@@ -5,8 +5,12 @@
 
 # pyre-strict
 
+import json
+import math
 import random
+import tempfile
 from collections.abc import Mapping
+from pathlib import Path
 from typing import Any
 from unittest import mock
 
@@ -1556,6 +1560,41 @@ def test_json_storage(self) -> None:
             str(client._generation_strategy), str(other_client._generation_strategy)
         )
 
+    def test_json_file_storage_serializes_non_finite_floats(self) -> None:
+        client = Client()
+        client.configure_experiment(
+            parameters=[
+                RangeParameterConfig(name="x1", parameter_type="float", bounds=(-1, 1))
+            ],
+            name="foo",
+        )
+        client._experiment._properties["map_key_infos"] = [
+            {"key": "step", "default_value": np.float64("nan")}
+        ]
+        client._experiment._properties["inf_property"] = float("inf")
+        client._experiment._properties["negative_inf_property"] = float("-inf")
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            filepath = Path(tmpdir) / "ax_client_snapshot.json"
+            client.save_to_json_file(filepath=str(filepath))
+            serialized = filepath.read_text()
+            json.loads(
+                serialized,
+                parse_constant=lambda constant: self.fail(
+                    f"Invalid JSON constant found: {constant}"
+                ),
+            )
+            self.assertNotIn("NaN", serialized)
+            self.assertNotIn("Infinity", serialized)
+            self.assertNotIn("-Infinity", serialized)
+
+            other_client = Client.load_from_json_file(filepath=str(filepath))
+
+        properties = other_client._experiment._properties
+        self.assertTrue(math.isnan(properties["map_key_infos"][0]["default_value"]))
+        self.assertEqual(properties["inf_property"], float("inf"))
+        self.assertEqual(properties["negative_inf_property"], float("-inf"))
+
     def test_sql_storage(self) -> None:
         init_test_engine_and_session_factory(force_init=True)
         client = Client(storage_config=StorageConfig())

ax/storage/json_store/decoder.py

@@ -190,16 +190,22 @@ def object_from_json(
                 result[key] = _object_from_json(v)
             return result
 
+        if "value" in object_json:
+            object_json = {
+                **object_json,
+                "value": _object_from_json(object_json["value"]),
+            }
+
         _type = object_json.pop("__type")
 
-        if _type == "datetime":
+        if _type == "float":
+            return float(object_json["value"])
+        elif _type == "datetime":
             return datetime.datetime.strptime(
                 object_json["value"], "%Y-%m-%d %H:%M:%S.%f"
             )
         elif _type == "OrderedDict":
-            return OrderedDict(
-                [(k, _object_from_json(v)) for k, v in object_json["value"]]
-            )
+            return OrderedDict(object_json["value"])
         elif _type == "DataFrame":
             # Need dtype=False, otherwise infers arm_names like "4_1"
             # should be int 41

ax/storage/json_store/decoders.py

@@ -304,6 +304,7 @@ def class_from_json(json: dict[str, Any]) -> type[Any]:
 
 def tensor_from_json(json: dict[str, Any]) -> torch.Tensor:
     try:
+        value = json["value"]
         device = (
             assert_is_instance(
                 torch_type_from_str(
@@ -315,7 +316,7 @@ def tensor_from_json(json: dict[str, Any]) -> torch.Tensor:
             else torch.device("cpu")
         )
         return torch.tensor(
-            json["value"],
+            value,
             dtype=assert_is_instance(
                 torch_type_from_str(
                     identifier=json["dtype"]["value"], type_name="dtype"

ax/storage/json_store/encoder.py

@@ -9,6 +9,7 @@
 import dataclasses
 import datetime
 import enum
+import math
 from collections import OrderedDict
 from collections.abc import Callable
 from functools import partial
@@ -75,6 +76,8 @@ def object_to_json(
     if _type in encoder_registry:
         obj_dict = encoder_registry[_type](obj)
         return {k: _object_to_json(v) for k, v in obj_dict.items()}
+    if isinstance(obj, (float, np.floating)) and not np.isfinite(obj):
+        return {"__type": "float", "value": str(float(obj))}
     # Python built-in types + `typing` module types
     if _type in (str, int, float, bool, type(None)):
         return obj
@@ -115,11 +118,23 @@ def object_to_json(
     elif issubclass(_type, enum.Enum):
         return {"__type": _type.__name__, "name": obj.name}
     elif _type is np.ndarray or issubclass(_type, np.ndarray):
-        return {"__type": _type.__name__, "value": obj.tolist()}
+        value = obj.tolist()
+        if np.issubdtype(obj.dtype, np.floating) and not np.isfinite(obj).all():
+            value = _object_to_json(value)
+        return {"__type": _type.__name__, "value": value}
     elif _type is set:
-        return {"__type": _type.__name__, "value": list(obj)}
+        raw_value = list(obj)
+        value = [numpy_type_to_python_type(x) for x in raw_value]
+        if any(
+            isinstance(x, (float, np.floating)) and not math.isfinite(x) for x in obj
+        ) or any(type(x) not in (str, int, float, bool, type(None)) for x in value):
+            value = _object_to_json(raw_value)
+        return {"__type": _type.__name__, "value": value}
     elif _type is torch.Tensor:
-        return tensor_to_dict(obj=obj)
+        obj_dict = tensor_to_dict(obj=obj)
+        if obj.is_floating_point() and not torch.isfinite(obj).all():
+            obj_dict["value"] = _object_to_json(obj_dict["value"])
+        return obj_dict
     elif _type.__module__ == "torch":
         # Torch does not support saving to string, so save to buffer first
         return {"__type": f"torch_{_type.__name__}", "value": torch_type_to_str(obj)}

ax/storage/json_store/tests/test_json_store.py

@@ -8,6 +8,7 @@
 
 import dataclasses
 import json
+import math
 import os
 import tempfile
 from collections import OrderedDict
@@ -673,6 +674,112 @@ def __post_init__(self, doesnt_serialize: None) -> None:
         self.assertEqual(recovered.not_a_field, 1)
         self.assertEqual(obj, recovered)
 
+    def test_EncodeDecode_non_finite_floats(self) -> None:
+        obj = {
+            "python_nan": float("nan"),
+            "numpy_nan": np.float64("nan"),
+            "python_inf": float("inf"),
+            "python_negative_inf": float("-inf"),
+            "nested": [np.float32("nan"), {"inf": np.float64("inf")}],
+            "ndarray": np.array([[1.0, np.nan], [np.inf, -np.inf]]),
+            "tensor": torch.tensor(
+                [[1.0, float("nan")], [float("inf"), -float("inf")]]
+            ),
+            "set": {float("inf"), float("-inf")},
+        }
+
+        obj_json = object_to_json(
+            obj,
+            encoder_registry=CORE_ENCODER_REGISTRY,
+            class_encoder_registry=CORE_CLASS_ENCODER_REGISTRY,
+        )
+        serialized = json.dumps(obj_json, allow_nan=False)
+        json.loads(
+            serialized,
+            parse_constant=lambda constant: self.fail(
+                f"Invalid JSON constant found: {constant}"
+            ),
+        )
+        self.assertNotIn("NaN", serialized)
+        self.assertNotIn("Infinity", serialized)
+        self.assertNotIn("-Infinity", serialized)
+
+        recovered = object_from_json(
+            obj_json,
+            decoder_registry=CORE_DECODER_REGISTRY,
+            class_decoder_registry=CORE_CLASS_DECODER_REGISTRY,
+        )
+        self.assertTrue(math.isnan(recovered["python_nan"]))
+        self.assertTrue(math.isnan(recovered["numpy_nan"]))
+        self.assertEqual(recovered["python_inf"], float("inf"))
+        self.assertEqual(recovered["python_negative_inf"], float("-inf"))
+        self.assertTrue(math.isnan(recovered["nested"][0]))
+        self.assertEqual(recovered["nested"][1]["inf"], float("inf"))
+        self.assertTrue(np.isnan(recovered["ndarray"][0, 1]))
+        self.assertEqual(recovered["ndarray"][1, 0], float("inf"))
+        self.assertEqual(recovered["ndarray"][1, 1], float("-inf"))
+        self.assertTrue(torch.isnan(recovered["tensor"][0, 1]))
+        self.assertTrue(torch.isinf(recovered["tensor"][1, 0]))
+        self.assertTrue(torch.isinf(recovered["tensor"][1, 1]))
+        self.assertLess(recovered["tensor"][1, 1].item(), 0)
+        self.assertEqual(recovered["set"], {float("inf"), float("-inf")})
+
+    def test_EncodeDecode_zero_dimensional_non_finite_arrays_and_tensors(self) -> None:
+        obj = {
+            "ndarray": np.array(np.nan),
+            "tensor": torch.tensor(float("nan")),
+        }
+
+        obj_json = object_to_json(
+            obj,
+            encoder_registry=CORE_ENCODER_REGISTRY,
+            class_encoder_registry=CORE_CLASS_ENCODER_REGISTRY,
+        )
+        serialized = json.dumps(obj_json, allow_nan=False)
+        json.loads(
+            serialized,
+            parse_constant=lambda constant: self.fail(
+                f"Invalid JSON constant found: {constant}"
+            ),
+        )
+        self.assertNotIn("NaN", serialized)
+
+        recovered = object_from_json(
+            obj_json,
+            decoder_registry=CORE_DECODER_REGISTRY,
+            class_decoder_registry=CORE_CLASS_DECODER_REGISTRY,
+        )
+        self.assertTrue(np.isnan(recovered["ndarray"]))
+        self.assertTrue(torch.isnan(recovered["tensor"]))
+
+    def test_Encode_finite_arrays_and_tensors_unchanged(self) -> None:
+        ndarray = np.array([[1.0, 2.0], [3.0, 4.0]])
+        tensor = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64, device=torch.device("cpu")
+        )
+
+        self.assertEqual(
+            object_to_json(
+                ndarray,
+                encoder_registry=CORE_ENCODER_REGISTRY,
+                class_encoder_registry=CORE_CLASS_ENCODER_REGISTRY,
+            ),
+            {"__type": "ndarray", "value": [[1.0, 2.0], [3.0, 4.0]]},
+        )
+        self.assertEqual(
+            object_to_json(
+                tensor,
+                encoder_registry=CORE_ENCODER_REGISTRY,
+                class_encoder_registry=CORE_CLASS_ENCODER_REGISTRY,
+            ),
+            {
+                "__type": "Tensor",
+                "value": [[1.0, 2.0], [3.0, 4.0]],
+                "dtype": {"__type": "torch_dtype", "value": "torch.float64"},
+                "device": {"__type": "torch_device", "value": "cpu"},
+            },
+        )
+
     def test_EncodeDecodeTorchTensor(self) -> None:
         x = torch.tensor(
             [[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64, device=torch.device("cpu")