diff --git a/openevolve/utils/trace_export_utils.py b/openevolve/utils/trace_export_utils.py index 967f82d71..a24d4c683 100644 --- a/openevolve/utils/trace_export_utils.py +++ b/openevolve/utils/trace_export_utils.py @@ -80,7 +80,7 @@ def export_traces_json( logger.info(f"Exported {len(traces)} traces to {output_path}") -# Names the metadata attributes that were JSON-encoded, so a list or dict can be told apart +# Names the attributes that were JSON-encoded, so a list or dict can be told apart # from a string that merely looks like JSON _HDF5_JSON_KEYS_ATTR = "__json_encoded_keys__" @@ -129,6 +129,7 @@ def export_traces_hdf5( for i, trace in enumerate(traces): trace_dict = trace.to_dict() if hasattr(trace, "to_dict") else trace trace_group = traces_group.create_group(f"trace_{i:06d}") + json_encoded_keys = [] for key, value in trace_dict.items(): if value is None: @@ -137,6 +138,7 @@ def export_traces_hdf5( if isinstance(value, dict): # Store dictionaries as JSON strings in attributes trace_group.attrs[key] = json.dumps(value) + json_encoded_keys.append(key) elif isinstance(value, list): # Convert lists to numpy arrays if possible try: @@ -145,6 +147,7 @@ def export_traces_hdf5( except (ValueError, TypeError): # Fall back to JSON for complex lists trace_group.attrs[key] = json.dumps(value) + json_encoded_keys.append(key) elif isinstance(value, str): # Store strings as attributes trace_group.attrs[key] = value @@ -154,6 +157,9 @@ def export_traces_hdf5( else: # Store other types as JSON trace_group.attrs[key] = json.dumps(value) + json_encoded_keys.append(key) + + trace_group.attrs[_HDF5_JSON_KEYS_ATTR] = json.dumps(json_encoded_keys) logger.info(f"Exported {len(traces)} traces to {output_path}") @@ -284,12 +290,21 @@ def load_traces_hdf5(input_path: Union[str, Path]) -> tuple[List[Dict[str, Any]] for trace_name in sorted(traces_group.keys()): trace_group = traces_group[trace_name] trace_dict = {} + json_encoded_keys = trace_group.attrs.get(_HDF5_JSON_KEYS_ATTR) + if json_encoded_keys is not None: + json_encoded_keys = set(json.loads(json_encoded_keys)) # Load attributes for key in trace_group.attrs: + if key == _HDF5_JSON_KEYS_ATTR: + continue value = trace_group.attrs[key] - # Try to parse JSON strings - if isinstance(value, str) and (value.startswith("{") or value.startswith("[")): + if json_encoded_keys is not None: + trace_dict[key] = json.loads(value) if key in json_encoded_keys else value + # Preserve the heuristic for files written without an encoding key list. + elif isinstance(value, str) and ( + value.startswith("{") or value.startswith("[") + ): try: trace_dict[key] = json.loads(value) except json.JSONDecodeError: diff --git a/tests/test_trace_export_hdf5.py b/tests/test_trace_export_hdf5.py index 93c55f83f..0fda06dd3 100644 --- a/tests/test_trace_export_hdf5.py +++ b/tests/test_trace_export_hdf5.py @@ -4,6 +4,7 @@ import unittest from pathlib import Path +from openevolve.evolution_trace import EvolutionTrace from openevolve.utils.trace_export_utils import export_traces, load_traces try: @@ -65,5 +66,49 @@ def test_files_without_a_key_list_load_as_before(self): self.assertEqual(loaded["broken"], "{not json") +@unittest.skipIf(h5py is None, "h5py is not installed") +class TestHdf5TraceRoundTrip(unittest.TestCase): + def test_json_looking_code_and_response_stay_strings(self): + trace = EvolutionTrace( + iteration=1, + timestamp=123.0, + parent_id="parent", + child_id="child", + parent_metrics={"score": 1.0}, + child_metrics={"score": 2.0}, + parent_code="[1, 2]", + child_code='{"answer": 42}', + llm_response='{"code": "result = 42"}', + child_changes_description='["improved score"]', + artifacts={"nested": [{"value": 1}]}, + ) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "trace.h5" + export_traces([trace], path, format="hdf5") + traces, _ = load_traces(path) + self.assertEqual(traces, [trace.to_dict()]) + + def test_numeric_and_json_encoded_lists_round_trip(self): + trace = {"scores": [1.0, 2.0], "labels": ["alpha"], "nested": [{"a": [1]}]} + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "trace.h5" + export_traces([trace], path, format="hdf5") + traces, _ = load_traces(path) + self.assertEqual(traces, [trace]) + + def test_trace_without_a_key_list_loads_as_before(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "legacy.h5" + with h5py.File(path, "w") as f: + attrs = f.create_group("traces/trace_000000").attrs + attrs["metrics"] = '{"score": 1}' + attrs["labels"] = '["alpha"]' + attrs["broken"] = "{not json" + traces, _ = load_traces(path) + self.assertEqual( + traces, [{"metrics": {"score": 1}, "labels": ["alpha"], "broken": "{not json"}] + ) + + if __name__ == "__main__": unittest.main()