Repository: stanfordnlp/dspy
Pull request: Fix save predict model with LM
This PR enables the serialization of dspy.LM to fix the issue of model saving.
Resolve #7852

Problem statement:
[Bug] save() is broken when predict has lm instance
### What happened?

save() is broken when predict has lm assigned. See code below for reproducing.

To fix it, we can serialize the model name, then at loading time, reconstruct the lm instance by calling `dspy.LM(saved_model_name)`

### Steps to reproduce

```
>>> import dspy
>>> predict = dspy.Predict("q->a")
>>> predict.lm = dspy.LM("openai/gpt-4o-mini")
>>> predict.save("test_lm_saving.json", save_program=False)
```

Raises error:

```
>>> predict.save("test_lm_saving.json", save_program=False)
Traceback (most recent call last):
  File "/Users/chen.qian/Documents/mlflow_team/dspy/dspy/primitives/module.py", line 218, in save
    f.write(ujson.dumps(state, indent=2))
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: <dspy.clients.lm.LM object at 0x10563c2c0> is not JSON serializable

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/chen.qian/Documents/mlflow_team/dspy/dspy/primitives/module.py", line 220, in save
    raise RuntimeError(
RuntimeError: Failed to save state to test_lm_saving.json with error: <dspy.clients.lm.LM object at 0x10563c2c0> is not JSON serializable. Your DSPy program may contain non json-serializable objects, please consider saving the state in .pkl by using `path` ending with `.pkl`, or saving the whole program by setting `save_program=True`.
```

### DSPy version

2.6.6

Hints from issue discussion:

Commit message: fix test_lm_after_dump_and_load_state

Signed-off-by: TomuHirata <tomu.hirata@gmail.com>

Patch to review:
diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py
index 6f5468b84e..617f532230 100644
--- a/dspy/clients/lm.py
+++ b/dspy/clients/lm.py
@@ -246,6 +246,10 @@ def copy(self, **kwargs):
                 new_instance.kwargs[key] = value
 
         return new_instance
+    
+    def dump_state(self):
+        # TODO: callbacks cannot be saved. We should consider how to save callbacks.
+        return { k: v for k, v in self.__dict__.items() if k not in ["provider", "callbacks", "kwargs"] }
 
 
 def request_cache(maxsize: Optional[int] = None):
diff --git a/dspy/predict/predict.py b/dspy/predict/predict.py
index 03bc44d869..1c51a2ab71 100644
--- a/dspy/predict/predict.py
+++ b/dspy/predict/predict.py
@@ -7,6 +7,7 @@
 from dspy.primitives.program import Module
 from dspy.signatures.signature import ensure_signature
 from dspy.utils.callback import with_callbacks
+from dspy.clients.lm import LM
 
 
 class Predict(Module, Parameter):
@@ -24,7 +25,7 @@ def reset(self):
         self.demos = []
 
     def dump_state(self):
-        state_keys = ["lm", "traces", "train"]
+        state_keys = ["traces", "train"]
         state = {k: getattr(self, k) for k in state_keys}
 
         state["demos"] = []
@@ -38,6 +39,7 @@ def dump_state(self):
             state["demos"].append(demo)
 
         state["signature"] = self.signature.dump_state()
+        state["lm"] = self.lm.dump_state() if self.lm else None
         return state
 
     def load_state(self, state):
@@ -49,13 +51,14 @@ def load_state(self, state):
         Returns:
             self: Returns self to allow method chaining
         """
-        excluded_keys = ["signature", "extended_signature"]
+        excluded_keys = ["signature", "extended_signature", "lm"]
         for name, value in state.items():
             # `excluded_keys` are fields that go through special handling.
             if name not in excluded_keys:
                 setattr(self, name, value)
                     
         self.signature = self.signature.load_state(state["signature"])
+        self.lm = LM(**state["lm"]) if state["lm"] else None
 
         if "extended_signature" in state: # legacy, up to and including 2.5, for CoT.
             raise NotImplementedError("Loading extended_signature is no longer supported in DSPy 2.6+")
diff --git a/tests/predict/test_predict.py b/tests/predict/test_predict.py
index 0ae2c000d3..c68785232c 100644
--- a/tests/predict/test_predict.py
+++ b/tests/predict/test_predict.py
@@ -34,11 +34,18 @@ def test_reset_method():
 
 def test_lm_after_dump_and_load_state():
     predict_instance = Predict("input -> output")
-    predict_instance.lm = "lm_state"
+    lm = dspy.LM(
+        model="openai/gpt-4o-mini", 
+        model_type="chat",
+        temperature=1,
+        max_tokens=100,
+        num_retries=10,
+    )
+    predict_instance.lm = lm
     dumped_state = predict_instance.dump_state()
     new_instance = Predict("input -> output")
     new_instance.load_state(dumped_state)
-    assert new_instance.lm == "lm_state"
+    assert new_instance.lm.dump_state() == lm.dump_state()
 
 
 def test_call_method():
@@ -211,6 +218,28 @@ class CustomSignature2(dspy.Signature):
     new_instance.load(file_path)
     assert new_instance.signature.dump_state() == original_instance.signature.dump_state()
 
+@pytest.mark.parametrize("filename", ["model.json", "model.pkl"])
+def test_lm_field_after_dump_and_load_state(tmp_path, filename):
+    file_path = tmp_path / filename
+    lm = dspy.LM(
+        model="openai/gpt-4o-mini", 
+        model_type="chat",
+        temperature=1,
+        max_tokens=100,
+        num_retries=10,
+    )
+    original_predict = dspy.Predict("q->a")
+    original_predict.lm = lm
+
+    original_predict.save(file_path)
+
+    assert file_path.exists()
+
+    loaded_predict = dspy.Predict("q->a")
+    loaded_predict.load(file_path)
+
+    assert original_predict.dump_state() == loaded_predict.dump_state()
+
 
 def test_forward_method():
     program = Predict("question -> answer")


Code review comments:
diff_hunk: @@ -246,6 +246,10 @@ def copy(self, **kwargs):
                 new_instance.kwargs[key] = value
 
         return new_instance
+    
+    def dump_state(self):
+        # TODO: callbacks cannot be saved. We should consider how to save callbacks.
+        return { k: v for k, v in self.__dict__.items() if k not in ["provider", "callbacks", "kwargs"] }
path: dspy/clients/lm.py
text: @user1:
let's explicitly list out the attributes to save for clarity.

diff_hunk: @@ -246,6 +246,10 @@ def copy(self, **kwargs):
                 new_instance.kwargs[key] = value
 
         return new_instance
+    
+    def dump_state(self):
+        # TODO: callbacks cannot be saved. We should consider how to save callbacks.
path: dspy/clients/lm.py
text: @user1:
good catch! actually callbacks doesn't belong to the state, it's part of runtime config, we don't need to save it with this `dump_state`

diff_hunk: @@ -34,11 +34,18 @@ def test_reset_method():
 
 def test_lm_after_dump_and_load_state():
     predict_instance = Predict("input -> output")
-    predict_instance.lm = "lm_state"
+    lm = dspy.LM(
+        model="openai/gpt-4o-mini", 
+        model_type="chat",
+        temperature=1,
+        max_tokens=100,
+        num_retries=10,
+    )
+    predict_instance.lm = lm
     dumped_state = predict_instance.dump_state()
     new_instance = Predict("input -> output")
     new_instance.load_state(dumped_state)
-    assert new_instance.lm == "lm_state"
+    assert new_instance.lm.dump_state() == lm.dump_state()
path: tests/predict/test_predict.py
text: @user1:
minor comment - here we are assuming `dump_state()` is working correctly, let's add a few `assert` check to verify the dumped state contains the desired fields with the correct value.

@author:
Good call. Added the assertion here and also added a unit test for the LM side.