Repository: mlflow/mlflow

Issue:
[BUG] Langchain: `TypeError: cannot pickle`
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Willingness to contribute

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

### MLflow version

2.7.1


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: databricks cluster
- **Python version**: 3.10.12
- **yarn version, if running the dev UI**: NA


### Describe the problem

I've created a langchain agent as follow:

```python
from langchain.llms import OpenAI
from langchain.agents import create_spark_dataframe_agent

df = spark.read.csv("/databricks-datasets/COVID/coronavirusdataset/Region.csv", header=True, inferSchema=True)
agent = create_spark_dataframe_agent(llm=OpenAI(temperature=0), df=df, verbose=True)
agent.run("How many rows are there?") # this works fine

import mlflow

with mlflow.start_run():
  logged_model = mlflow.langchain.log_model(agent, "langchain_model") # this fails with PicklingError: Can't pickle <built-in function input>: it's not the same object as builtins.input
````

The exception comes from:
```
python3.10/site-packages/mlflow/langchain/__init__.py:573, in _save_model(model, path, loader_fn, persist_dir)
    571     tools_data_path = os.path.join(path, _TOOLS_DATA_FILE_NAME)
    572     with open(tools_data_path, "wb") as f:
--> 573         cloudpickle.dump(model.tools, f)
    574     model_data_kwargs[_TOOLS_DATA_KEY] = _TOOLS_DATA_FILE_NAME
    575 else:
```

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
System information: Linux #49~20.04.1-Ubuntu SMP Wed Jul 12 12:44:56 UTC 2023
Python version: 3.10.12
MLflow version: 2.7.1
MLflow module location: /local_disk0/.ephemeral_nfs/envs/pythonEnv-cbf388f4-5594-493c-bba2-9053f5df6f7d/lib/python3.10/site-packages/mlflow/__init__.py
Tracking URI: databricks
Registry URI: databricks
Databricks runtime version: 14.0.x-cpu-ml-scala2.12
MLflow environment variables: 
  MLFLOW_CONDA_HOME: /databricks/conda
  MLFLOW_GATEWAY_URI: databricks
  MLFLOW_PYTHON_EXECUTABLE: /databricks/spark/scripts/mlflow_python.sh
  MLFLOW_TRACKING_URI: databricks
MLflow dependencies: 
  Flask: 2.2.5
  Jinja2: 3.1.2
  aiohttp: 3.8.5
  alembic: 1.12.0
  azure-storage-file-datalake: 12.12.0
  boto3: 1.24.28
  click: 8.0.4
  cloudpickle: 2.2.1
  databricks-cli: 0.17.7
  docker: 6.1.3
  entrypoints: 0.4
  fastapi: 0.98.0
  gitpython: 3.1.27
  google-cloud-storage: 2.10.0
  gunicorn: 20.1.0
  importlib-metadata: 4.11.3
  markdown: 3.4.1
  matplotlib: 3.7.0
  numpy: 1.23.5
  packaging: 22.0
  pandas: 1.5.3
  protobuf: 4.24.0
  psutil: 5.9.0
  pyarrow: 8.0.0
  pydantic: 1.10.6
  pytz: 2022.7
  pyyaml: 6.0
  querystring-parser: 1.2.4
  requests: 2.28.1
  scikit-learn: 1.1.1
  scipy: 1.10.0
  sqlalchemy: 1.4.39
  sqlparse: 0.4.2
  uvicorn: 0.23.2
  virtualenv: 20.16.7
  watchfiles: 0.20.0
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
import os
os.environ["OPENAI_API_KEY"] = "..."

from langchain.llms import OpenAI
from langchain.agents import create_spark_dataframe_agent

df = spark.read.csv("/databricks-datasets/COVID/coronavirusdataset/Region.csv", header=True, inferSchema=True)
agent = create_spark_dataframe_agent(llm=OpenAI(temperature=0), df=df, verbose=True)
agent.run("How many rows are there?") # this works fine

import mlflow

with mlflow.start_run():
  logged_model = mlflow.langchain.log_model(agent, "langchain_model") # this fails with PicklingError: Can't pickle <built-in function input>: it's not the same object as builtins.input

```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
---------------------------------------------------------------------------
PicklingError                             Traceback (most recent call last)
File <command-4172266444168703>, line 4
      1 import mlflow
      3 with mlflow.start_run():
----> 4   logged_model = mlflow.langchain.log_model(agent, "langchain_model")

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-cbf388f4-5594-493c-bba2-9053f5df6f7d/lib/python3.10/site-packages/mlflow/langchain/__init__.py:535, in log_model(lc_model, artifact_path, conda_env, code_paths, registered_model_name, signature, input_example, await_registration_for, pip_requirements, extra_pip_requirements, metadata, loader_fn, persist_dir)
    532     output_schema = Schema(output_columns)
    533     signature = ModelSignature(input_schema, output_schema)
--> 535 return Model.log(
    536     artifact_path=artifact_path,
    537     flavor=mlflow.langchain,
    538     registered_model_name=registered_model_name,
    539     lc_model=lc_model,
    540     conda_env=conda_env,
    541     code_paths=code_paths,
    542     signature=signature,
    543     input_example=input_example,
    544     await_registration_for=await_registration_for,
    545     pip_requirements=pip_requirements,
    546     extra_pip_requirements=extra_pip_requirements,
    547     metadata=metadata,
    548     loader_fn=loader_fn,
    549     persist_dir=persist_dir,
    550 )

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-cbf388f4-5594-493c-bba2-9053f5df6f7d/lib/python3.10/site-packages/mlflow/models/model.py:579, in Model.log(cls, artifact_path, flavor, registered_model_name, await_registration_for, metadata, **kwargs)
    573 if (
    574     (tracking_uri == "databricks" or get_uri_scheme(tracking_uri) == "databricks")
    575     and kwargs.get("signature") is None
    576     and kwargs.get("input_example") is None
    577 ):
    578     _logger.warning(_LOG_MODEL_MISSING_SIGNATURE_WARNING)
--> 579 flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
    580 mlflow.tracking.fluent.log_artifacts(local_path, mlflow_model.artifact_path)
    581 try:

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-cbf388f4-5594-493c-bba2-9053f5df6f7d/lib/python3.10/site-packages/mlflow/langchain/__init__.py:285, in save_model(lc_model, path, conda_env, code_paths, mlflow_model, signature, input_example, pip_requirements, extra_pip_requirements, metadata, loader_fn, persist_dir)
    282 if metadata is not None:
    283     mlflow_model.metadata = metadata
--> 285 model_data_kwargs = _save_model(lc_model, path, loader_fn, persist_dir)
    287 pyfunc.add_to_model(
    288     mlflow_model,
    289     loader_module="mlflow.langchain",
   (...)
    293     **model_data_kwargs,
    294 )
    295 flavor_conf = {
    296     _MODEL_TYPE_KEY: lc_model.__class__.__name__,
    297     **model_data_kwargs,
    298 }

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-cbf388f4-5594-493c-bba2-9053f5df6f7d/lib/python3.10/site-packages/mlflow/langchain/__init__.py:573, in _save_model(model, path, loader_fn, persist_dir)
    571     tools_data_path = os.path.join(path, _TOOLS_DATA_FILE_NAME)
    572     with open(tools_data_path, "wb") as f:
--> 573         cloudpickle.dump(model.tools, f)
    574     model_data_kwargs[_TOOLS_DATA_KEY] = _TOOLS_DATA_FILE_NAME
    575 else:

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-cbf388f4-5594-493c-bba2-9053f5df6f7d/lib/python3.10/site-packages/cloudpickle/cloudpickle_fast.py:57, in dump(obj, file, protocol, buffer_callback)
     45 def dump(obj, file, protocol=None, buffer_callback=None):
     46     """Serialize obj as bytes streamed into file
     47 
     48     protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to
   (...)
     53     compatibility with older versions of Python.
     54     """
     55     CloudPickler(
     56         file, protocol=protocol, buffer_callback=buffer_callback
---> 57     ).dump(obj)

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-cbf388f4-5594-493c-bba2-9053f5df6f7d/lib/python3.10/site-packages/cloudpickle/cloudpickle_fast.py:632, in CloudPickler.dump(self, obj)
    630 def dump(self, obj):
    631     try:
--> 632         return Pickler.dump(self, obj)
    633     except RuntimeError as e:
    634         if "recursion" in e.args[0]:

PicklingError: Can't pickle <built-in function input>: it's not the same object as builtins.input
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/gateway`: AI Gateway service, Gateway client APIs, third-party Gateway integrations
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [X] `integrations/databricks`: Databricks integrations


Discussion:
Seems like there's some sort of global ref object in the agent's implementation in LangChain that is causing a pickle issue. @liangz1 was this agent ever tested or supported?
Is this conditional logic target for this AgentExecutor the right path for this agent type? https://github.com/mlflow/mlflow/blob/55180340a6a3c460b5e3084a9eb2f2fa834c94f8/mlflow/langchain/__init__.py#L576-L580 
@sunishsheth2009 has more context for agents. Could you share insights?
<!-- assign-maintainer -->
@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.

Resolving patch:
diff --git a/mlflow/langchain/__init__.py b/mlflow/langchain/__init__.py
--- a/mlflow/langchain/__init__.py
+++ b/mlflow/langchain/__init__.py
@@ -571,8 +571,14 @@ def _save_model(model, path, loader_fn, persist_dir):
 
         if model.tools:
             tools_data_path = os.path.join(path, _TOOLS_DATA_FILE_NAME)
-            with open(tools_data_path, "wb") as f:
-                cloudpickle.dump(model.tools, f)
+            try:
+                with open(tools_data_path, "wb") as f:
+                    cloudpickle.dump(model.tools, f)
+            except Exception as e:
+                raise mlflow.MlflowException(
+                    "Error when attempting to pickle the AgentExecutor tools. "
+                    "This model likely does not support serialization."
+                ) from e
             model_data_kwargs[_TOOLS_DATA_KEY] = _TOOLS_DATA_FILE_NAME
         else:
             raise mlflow.MlflowException.invalid_parameter_value(


Test patch:
diff --git a/tests/langchain/test_langchain_model_export.py b/tests/langchain/test_langchain_model_export.py
--- a/tests/langchain/test_langchain_model_export.py
+++ b/tests/langchain/test_langchain_model_export.py
@@ -12,6 +12,7 @@
 import pytest
 import transformers
 from langchain import SQLDatabase
+from langchain.agents import AgentType, initialize_agent
 from langchain.chains import (
     APIChain,
     ConversationChain,
@@ -32,6 +33,7 @@
 from langchain.prompts import PromptTemplate
 from langchain.requests import TextRequestsWrapper
 from langchain.text_splitter import CharacterTextSplitter
+from langchain.tools import Tool
 from langchain.vectorstores import FAISS
 from langchain_experimental.sql import SQLDatabaseChain
 from packaging import version
@@ -736,3 +738,29 @@ def test_unsupported_class():
     ):
         with mlflow.start_run():
             mlflow.langchain.log_model(llm, "fake_llm")
+
+
+def test_agent_with_unpicklable_tools(tmp_path):
+    tmp_file = tmp_path / "temp_file.txt"
+    with open(tmp_file, mode="w") as temp_file:
+        # files that aren't opened for reading cannot be pickled
+        tools = [
+            Tool.from_function(
+                func=lambda: temp_file,
+                name="Write 0",
+                description="If you need to write 0 to a file",
+            )
+        ]
+        agent = initialize_agent(
+            llm=OpenAI(temperature=0), tools=tools, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION
+        )
+
+        with pytest.raises(
+            MlflowException,
+            match=(
+                "Error when attempting to pickle the AgentExecutor tools. "
+                "This model likely does not support serialization."
+            ),
+        ):
+            with mlflow.start_run():
+                mlflow.langchain.log_model(agent, "unpicklable_tools")