Add guards against loading pkl files (#9048)

* Add guards against loading pkl files

* remove extra comment

* add tests and modify old tests

* add import and fix tests

* Update test to allow loading of pickled models with a safety flag

* Change from dangersouly_allow_pickle to allow_pickle, remove env var, and suggest saving with module.save(x.json)

* fix extra whitespace

* fix test warning

Author: isaacbmiller

docs/docs/tutorials/games/index.ipynb

@@ -746,14 +746,21 @@
     "If you want to load and use the agent program, you can do that as follows."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "> **⚠️ Security Warning:** Loading `.pkl` files can execute arbitrary code and may be dangerous. Only save and load pickle files from trusted sources in secure environments. Consider using JSON format when possible for safer serialization."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
     "loaded = Agent()\n",
-    "loaded.load('finetuned_4o_mini_001.pkl')"
+    "loaded.load('finetuned_4o_mini_001.pkl', allow_pickle=True)"
    ]
   }
  ],

docs/docs/tutorials/saving/index.md

@@ -38,6 +38,9 @@ compiled_dspy_program.save("./dspy_program/program.json", save_program=False)
 
 To save the state of your program to a pickle file:
 
+!!! danger "Security Warning: Pickle Files Can Execute Arbitrary Code"
+    Loading `.pkl` files can execute arbitrary code and may be dangerous. Only load pickle files from trusted sources in secure environments. **Prefer using `.json` files whenever possible**. If you must use pickle files, ensure you trust the source and use the `allow_pickle=True` parameter when loading.
+
 ```python
 compiled_dspy_program.save("./dspy_program/program.pkl", save_program=False)
 ```
@@ -57,9 +60,12 @@ assert str(compiled_dspy_program.signature) == str(loaded_dspy_program.signature
 
 Or load the state from a pickle file:
 
+!!! danger "Security Warning"
+    Remember to use `allow_pickle=True` when loading pickle files, and only load from trusted sources.
+
 ```python
 loaded_dspy_program = dspy.ChainOfThought("question -> answer") # Recreate the same program.
-loaded_dspy_program.load("./dspy_program/program.pkl")
+loaded_dspy_program.load("./dspy_program/program.pkl", allow_pickle=True)
 
 assert len(compiled_dspy_program.demos) == len(loaded_dspy_program.demos)
 for original_demo, loaded_demo in zip(compiled_dspy_program.demos, loaded_dspy_program.demos):
@@ -70,6 +76,9 @@ assert str(compiled_dspy_program.signature) == str(loaded_dspy_program.signature
 
 ## Whole Program Saving
 
+!!! warning "Security Notice: Whole Program Saving Uses Pickle"
+    Whole program saving uses `cloudpickle` for serialization, which has the same security risks as pickle files. Only load programs from trusted sources in secure environments.
+
 Starting from `dspy>=2.6.0`, DSPy supports saving the whole program, including the architecture and the state. This feature
 is powered by `cloudpickle`, which is a library for serializing and deserializing Python objects.
 

dspy/primitives/base_module.py

@@ -202,7 +202,8 @@ def save(self, path, save_program=False, modules_to_serialize=None):
             if not path.exists():
                 # Create the directory (and any parent directories)
                 path.mkdir(parents=True)
-
+            logger.warning("Loading untrusted .pkl files can run arbitrary code, which may be dangerous. To avoid "
+                          'this, prefer saving using json format using module.save("module.json").')
             try:
                 modules_to_serialize = modules_to_serialize or []
                 for module in modules_to_serialize:
@@ -233,26 +234,34 @@ def save(self, path, save_program=False, modules_to_serialize=None):
                     "with `.pkl`, or saving the whole program by setting `save_program=True`."
                 )
         elif path.suffix == ".pkl":
+            logger.warning("Loading untrusted .pkl files can run arbitrary code, which may be dangerous. To avoid "
+                          'this, prefer saving using json format using module.save("module.json").')
             state = self.dump_state(json_mode=False)
             state["metadata"] = metadata
             with open(path, "wb") as f:
                 cloudpickle.dump(state, f)
         else:
             raise ValueError(f"`path` must end with `.json` or `.pkl` when `save_program=False`, but received: {path}")
 
-    def load(self, path):
+    def load(self, path, allow_pickle=False):
         """Load the saved module. You may also want to check out dspy.load, if you want to
         load an entire program, not just the state for an existing program.
 
         Args:
             path (str): Path to the saved state file, which should be a .json or a .pkl file
+            allow_pickle (bool): If True, allow loading .pkl files, which can run arbitrary code.
+                This is dangerous and should only be used if you are sure about the source of the file and in a trusted environment.
         """
         path = Path(path)
 
         if path.suffix == ".json":
             with open(path, "rb") as f:
                 state = orjson.loads(f.read())
         elif path.suffix == ".pkl":
+            if not allow_pickle:
+                raise ValueError("Loading .pkl files can run arbitrary code, which may be dangerous. Prefer "
+                                 "saving with .json files if possible. Set `allow_pickle=True` "
+                                 "if you are sure about the source of the file and in a trusted environment.")
             with open(path, "rb") as f:
                 state = cloudpickle.load(f)
         else:

dspy/utils/saving.py

@@ -24,17 +24,21 @@ def get_dependency_versions():
     }
 
 
-def load(path: str) -> "Module":
+def load(path: str, allow_pickle: bool = False) -> "Module":
     """Load saved DSPy model.
 
     This method is used to load a saved DSPy model with `save_program=True`, i.e., the model is saved with cloudpickle.
 
     Args:
         path (str): Path to the saved model.
+        allow_pickle (bool): Whether to allow loading the model with pickle. This is dangerous and should only be used if you are sure you trust the source of the model.
 
     Returns:
         The loaded model, a `dspy.Module` instance.
     """
+    if not allow_pickle:
+        raise ValueError("Loading with pickle is not allowed. Please set `allow_pickle=True` if you are sure you trust the source of the model.")
+
     path = Path(path)
     if not path.exists():
         raise FileNotFoundError(f"The path '{path}' does not exist.")

tests/predict/test_predict.py

@@ -253,7 +253,7 @@ def test_lm_field_after_dump_and_load_state(tmp_path, filename):
     assert file_path.exists()
 
     loaded_predict = dspy.Predict("q->a")
-    loaded_predict.load(file_path)
+    loaded_predict.load(file_path, allow_pickle=True)
 
     assert original_predict.dump_state() == loaded_predict.dump_state()
 

tests/primitives/test_base_module.py

@@ -123,7 +123,7 @@ def dummy_metric(example, pred, trace=None):
     compiled_cot.save(save_path)
 
     new_cot = dspy.ChainOfThought(MySignature)
-    new_cot.load(save_path)
+    new_cot.load(save_path, allow_pickle=True)
 
     assert str(new_cot.predict.signature) == str(compiled_cot.predict.signature)
     assert new_cot.predict.demos == compiled_cot.predict.demos
@@ -162,7 +162,7 @@ def forward(self, q):
 
         # Test the loading fails without using `modules_to_serialize`
         with pytest.raises(ModuleNotFoundError):
-            dspy.load(tmp_path)
+            dspy.load(tmp_path, allow_pickle=True)
 
         sys.path.insert(0, str(tmp_path))
         import custom_module
@@ -179,7 +179,7 @@ def forward(self, q):
         sys.path.remove(str(tmp_path))
         del custom_module
 
-        loaded_module = dspy.load(tmp_path)
+        loaded_module = dspy.load(tmp_path, allow_pickle=True)
         assert loaded_module.cot.predict.signature == cot.cot.predict.signature
 
     finally:
@@ -223,12 +223,16 @@ def emit(self, record):
         # Mock version during load
         with patch("dspy.primitives.base_module.get_dependency_versions", return_value=load_versions):
             loaded_predict = dspy.Predict("question->answer")
-            loaded_predict.load(save_path)
+            loaded_predict.load(save_path, allow_pickle=True)
 
-        # Assert warnings were logged, and one warning for each mismatched dependency.
-        assert len(handler.messages) == 3
+        # Assert warnings were logged: 1 for pickle loading + 3 for version mismatches
+        assert len(handler.messages) == 4
 
-        for msg in handler.messages:
+        # First message is about pickle loading
+        assert ".pkl" in handler.messages[0]
+
+        # Rest are version mismatch warnings
+        for msg in handler.messages[1:]:
             assert "There is a mismatch of" in msg
 
         # Verify the model still loads correctly despite version mismatches

tests/utils/test_saving.py

@@ -14,7 +14,7 @@ def test_save_predict(tmp_path):
     assert (tmp_path / "metadata.json").exists()
     assert (tmp_path / "program.pkl").exists()
 
-    loaded_predict = dspy.load(tmp_path)
+    loaded_predict = dspy.load(tmp_path, allow_pickle=True)
     assert isinstance(loaded_predict, dspy.Predict)
 
     assert predict.signature == loaded_predict.signature
@@ -29,7 +29,7 @@ def __init__(self):
     model = CustomModel()
     model.save(tmp_path, save_program=True)
 
-    loaded_model = dspy.load(tmp_path)
+    loaded_model = dspy.load(tmp_path, allow_pickle=True)
     assert isinstance(loaded_model, CustomModel)
 
     assert len(model.predictors()) == len(loaded_model.predictors())
@@ -51,7 +51,7 @@ class MySignature(dspy.Signature):
     predict.signature = predict.signature.with_instructions("You are a helpful assistant.")
     predict.save(tmp_path, save_program=True)
 
-    loaded_predict = dspy.load(tmp_path)
+    loaded_predict = dspy.load(tmp_path, allow_pickle=True)
     assert isinstance(loaded_predict, dspy.Predict)
 
     assert predict.signature == loaded_predict.signature
@@ -77,7 +77,7 @@ def dummy_metric(example, pred, trace=None):
     compiled_predict = optimizer.compile(predict, trainset=trainset)
     compiled_predict.save(tmp_path, save_program=True)
 
-    loaded_predict = dspy.load(tmp_path)
+    loaded_predict = dspy.load(tmp_path, allow_pickle=True)
     assert compiled_predict.demos == loaded_predict.demos
     assert compiled_predict.signature == loaded_predict.signature
 
@@ -115,7 +115,7 @@ def emit(self, record):
 
         # Mock version during load
         with patch("dspy.utils.saving.get_dependency_versions", return_value=load_versions):
-            loaded_predict = dspy.load(tmp_path)
+            loaded_predict = dspy.load(tmp_path, allow_pickle=True)
 
         # Assert warnings were logged, and one warning for each mismatched dependency.
         assert len(handler.messages) == 3
@@ -131,3 +131,45 @@ def emit(self, record):
         # Clean up: restore original level and remove handler
         logger.setLevel(original_level)
         logger.removeHandler(handler)
+
+
+def test_pickle_loading_requires_explicit_permission(tmp_path):
+    """Test that loading pickle files requires explicit permission."""
+    predict = dspy.Predict("question->answer")
+    predict.save(tmp_path, save_program=True)
+
+    # Should fail without dangerously_allow_pickle
+    with pytest.raises(ValueError, match="Loading with pickle is not allowed"):
+        dspy.load(tmp_path)
+
+    # Should succeed with dangerously_allow_pickle
+    loaded_predict = dspy.load(tmp_path, allow_pickle=True)
+    assert isinstance(loaded_predict, dspy.Predict)
+
+
+def test_pkl_file_loading_requires_explicit_permission(tmp_path):
+    """Test that loading .pkl files requires explicit permission."""
+    predict = dspy.Predict("question->answer")
+    pkl_path = tmp_path / "model.pkl"
+    predict.save(pkl_path)
+
+    # Should fail without allow_pickle
+    new_predict = dspy.Predict("question->answer")
+    with pytest.raises(ValueError, match="Loading .pkl files can run arbitrary code"):
+        new_predict.load(pkl_path)
+
+    # Should succeed with allow_pickle
+    new_predict.load(pkl_path, allow_pickle=True)
+    assert new_predict.dump_state() == predict.dump_state()
+
+
+def test_json_file_loading_works_without_permission(tmp_path):
+    """Test that loading .json files works without explicit permission."""
+    predict = dspy.Predict("question->answer")
+    json_path = tmp_path / "model.json"
+    predict.save(json_path)
+
+    # Should succeed without allow_pickle
+    new_predict = dspy.Predict("question->answer")
+    new_predict.load(json_path)
+    assert new_predict.dump_state() == predict.dump_state()