diff --git a/dspy/primitives/base_module.py b/dspy/primitives/base_module.py index a583368467..f11ccaa801 100644 --- a/dspy/primitives/base_module.py +++ b/dspy/primitives/base_module.py @@ -230,8 +230,9 @@ def save(self, path, save_program=False, modules_to_serialize=None): state = self.dump_state() state["metadata"] = metadata try: + serialized_state = orjson.dumps(state, option=orjson.OPT_INDENT_2 | orjson.OPT_APPEND_NEWLINE) with open(path, "wb") as f: - f.write(orjson.dumps(state, option=orjson.OPT_INDENT_2 | orjson.OPT_APPEND_NEWLINE)) + f.write(serialized_state) except Exception as e: raise RuntimeError( f"Failed to save state to {path} with error: {e}. Your DSPy program may contain non " diff --git a/tests/primitives/test_module.py b/tests/primitives/test_module.py index 0545a50184..2ba4e62735 100644 --- a/tests/primitives/test_module.py +++ b/tests/primitives/test_module.py @@ -26,6 +26,28 @@ def test_module_initialization(): assert module._compiled is False, "Module _compiled attribute should be False upon initialization" +@pytest.mark.parametrize("existing", [False, True]) +def test_json_save_serialization_failure_preserves_checkpoint(tmp_path, existing): + program = dspy.Predict("question -> answer") + path = tmp_path / "checkpoint.json" + previous = None + if existing: + program.save(path) + previous = path.read_bytes() + + program.traces.append({"unserializable": object()}) + with pytest.raises(RuntimeError, match="Failed to save state"): + program.save(path) + + if existing: + assert path.read_bytes() == previous + restored = dspy.Predict("question -> answer") + restored.load(path) + assert restored.traces == [] + else: + assert not path.exists() + + def test_named_predictors(): module = HopModule() named_preds = module.named_predictors()