From 62bef41a9577685b346b4fd56838163446b0266e Mon Sep 17 00:00:00 2001 From: SiddharthSingh Date: Wed, 8 Jul 2026 15:56:15 +0530 Subject: [PATCH 1/2] Fix NaN serialization in save_to_json_file (#4144) --- ax/api/client.py | 2 +- ax/api/tests/test_client.py | 39 ++++++++++++++ ax/storage/json_store/decoder.py | 8 +-- ax/storage/json_store/decoders.py | 7 ++- ax/storage/json_store/encoder.py | 8 +-- .../json_store/tests/test_json_store.py | 51 +++++++++++++++++++ 6 files changed, 107 insertions(+), 8 deletions(-) diff --git a/ax/api/client.py b/ax/api/client.py index 0de96370b3d..1b7ec8efb3c 100644 --- a/ax/api/client.py +++ b/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}`." ) diff --git a/ax/api/tests/test_client.py b/ax/api/tests/test_client.py index 2162c16d221..c67a59e4b03 100644 --- a/ax/api/tests/test_client.py +++ b/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()) diff --git a/ax/storage/json_store/decoder.py b/ax/storage/json_store/decoder.py index 89ed3879a38..3764ba25ae6 100644 --- a/ax/storage/json_store/decoder.py +++ b/ax/storage/json_store/decoder.py @@ -192,7 +192,9 @@ def object_from_json( _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" ) @@ -206,7 +208,7 @@ def object_from_json( # pyrefly: ignore [no-matching-overload] return pd.read_json(StringIO(object_json["value"]), dtype=False) elif _type == "ndarray": - return np.array(object_json["value"]) + return np.array(_object_from_json(object_json["value"])) elif _type == "Tensor": return tensor_from_json(json=object_json) elif _type.startswith("torch"): @@ -219,7 +221,7 @@ def object_from_json( list_surrogate_json=object_json, **vars(registry_kwargs) ) elif _type == "set": - return set(object_json["value"]) + return set(_object_from_json(object_json["value"])) # Used for decoding classes (not objects). elif _type in class_decoder_registry: return class_decoder_registry[_type](object_json) diff --git a/ax/storage/json_store/decoders.py b/ax/storage/json_store/decoders.py index 1d7fd35a181..d30c20522fd 100644 --- a/ax/storage/json_store/decoders.py +++ b/ax/storage/json_store/decoders.py @@ -304,6 +304,11 @@ def class_from_json(json: dict[str, Any]) -> type[Any]: def tensor_from_json(json: dict[str, Any]) -> torch.Tensor: try: + value = json["value"] + if isinstance(value, list): + from ax.storage.json_store.decoder import object_from_json + + value = object_from_json(object_json=value) device = ( assert_is_instance( torch_type_from_str( @@ -315,7 +320,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" diff --git a/ax/storage/json_store/encoder.py b/ax/storage/json_store/encoder.py index 18a0130178c..c38a3012917 100644 --- a/ax/storage/json_store/encoder.py +++ b/ax/storage/json_store/encoder.py @@ -75,6 +75,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 +117,11 @@ 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()} + return {"__type": _type.__name__, "value": _object_to_json(obj.tolist())} elif _type is set: - return {"__type": _type.__name__, "value": list(obj)} + return {"__type": _type.__name__, "value": _object_to_json(list(obj))} elif _type is torch.Tensor: - return tensor_to_dict(obj=obj) + return {k: _object_to_json(v) for k, v in tensor_to_dict(obj=obj).items()} 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)} diff --git a/ax/storage/json_store/tests/test_json_store.py b/ax/storage/json_store/tests/test_json_store.py index b2a04dc4226..e0d6b6fb9ae 100644 --- a/ax/storage/json_store/tests/test_json_store.py +++ b/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,56 @@ 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_EncodeDecodeTorchTensor(self) -> None: x = torch.tensor( [[1.0, 2.0], [3.0, 4.0]], dtype=torch.float64, device=torch.device("cpu") From 67948da8da8dd7a8d0e7c9867ad2382085404241 Mon Sep 17 00:00:00 2001 From: SiddharthSingh Date: Wed, 15 Jul 2026 23:12:25 +0530 Subject: [PATCH 2/2] Address review: fix circular import, avoid unnecessary per-element conversion --- ax/storage/json_store/decoder.py | 14 +++-- ax/storage/json_store/decoders.py | 4 -- ax/storage/json_store/encoder.py | 19 ++++++- .../json_store/tests/test_json_store.py | 56 +++++++++++++++++++ 4 files changed, 81 insertions(+), 12 deletions(-) diff --git a/ax/storage/json_store/decoder.py b/ax/storage/json_store/decoder.py index 3764ba25ae6..7e96085b11e 100644 --- a/ax/storage/json_store/decoder.py +++ b/ax/storage/json_store/decoder.py @@ -190,6 +190,12 @@ 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 == "float": @@ -199,16 +205,14 @@ def object_from_json( 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 # pyrefly: ignore [no-matching-overload] return pd.read_json(StringIO(object_json["value"]), dtype=False) elif _type == "ndarray": - return np.array(_object_from_json(object_json["value"])) + return np.array(object_json["value"]) elif _type == "Tensor": return tensor_from_json(json=object_json) elif _type.startswith("torch"): @@ -221,7 +225,7 @@ def object_from_json( list_surrogate_json=object_json, **vars(registry_kwargs) ) elif _type == "set": - return set(_object_from_json(object_json["value"])) + return set(object_json["value"]) # Used for decoding classes (not objects). elif _type in class_decoder_registry: return class_decoder_registry[_type](object_json) diff --git a/ax/storage/json_store/decoders.py b/ax/storage/json_store/decoders.py index d30c20522fd..0d8ce50368c 100644 --- a/ax/storage/json_store/decoders.py +++ b/ax/storage/json_store/decoders.py @@ -305,10 +305,6 @@ def class_from_json(json: dict[str, Any]) -> type[Any]: def tensor_from_json(json: dict[str, Any]) -> torch.Tensor: try: value = json["value"] - if isinstance(value, list): - from ax.storage.json_store.decoder import object_from_json - - value = object_from_json(object_json=value) device = ( assert_is_instance( torch_type_from_str( diff --git a/ax/storage/json_store/encoder.py b/ax/storage/json_store/encoder.py index c38a3012917..b1908c243b4 100644 --- a/ax/storage/json_store/encoder.py +++ b/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 @@ -117,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": _object_to_json(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": _object_to_json(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 {k: _object_to_json(v) for k, v in tensor_to_dict(obj=obj).items()} + 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)} diff --git a/ax/storage/json_store/tests/test_json_store.py b/ax/storage/json_store/tests/test_json_store.py index e0d6b6fb9ae..9d568558136 100644 --- a/ax/storage/json_store/tests/test_json_store.py +++ b/ax/storage/json_store/tests/test_json_store.py @@ -724,6 +724,62 @@ def test_EncodeDecode_non_finite_floats(self) -> None: 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")