Parallel execution: Can't pickle

[dmpetrov]

Description

The code below fails in parallel mode (works with no parallelizm).

We had this issue several times 🙁 - iterative/dvcx#1620

import os

from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage, ChatCompletionResponse

from datachain.lib.dc import DataChain, Column

MistralResponse = pydantic_to_feature(ChatCompletionResponse)

PROMPT = "Was this bot dialog successful? Describe the 'result' as 'Yes' or 'No' in a short JSON"

model = "mistral-large-latest"
api_key = os.environ["MISTRAL_API_KEY"]

chain = (
    DataChain.from_storage("gs://datachain-demo/chatbot-KiT/")
    .limit(5)
    .settings(cache=True, parallel=4)
    .map(
        mistral_response=lambda file: MistralResponse(
            **MistralClient(api_key=api_key).chat(
                    model=model,
                    response_format={"type": "json_object"},
                    messages=[
                        ChatMessage(role="user", content=f"{PROMPT}: {file.get_value()}")
                    ],
                ).dict()
        ),
        output=MistralResponse
    )
    .save()
)

print(chain.select("mistral_response").collect())

Error:

...
  File "/usr/local/Cellar/python@3.11/3.11.0/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
    ^^^^^^^^^^^^
  File "/Users/dmitry/venv/venv-py311/lib/python3.11/site-packages/dill/_dill.py", line 1869, in save_type
    StockPickler.save_global(pickler, obj, name=obj_name)
  File "/usr/local/Cellar/python@3.11/3.11.0/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 1071, in save_global
    raise PicklingError(
_pickle.PicklingError: Can't pickle <class 'datachain.lib.feature_utils.FunctionCall'>: it's not found as datachain.lib.feature_utils.FunctionCall

Version Info

0.1.dev3+gb377a5b
Python 3.9.16

3.11 also fails. `main` branch - the same.

[dmpetrov]

CC @amritghimire

[dmpetrov]

note, this code uses dynamic schema

[amritghimire]

At this point with all this issue, we should probably plan for moving away from pydantic issue if possible in my opinion. I am not sure how to identify the dynamic object during parsing stage if it is feature class or not

[amritghimire]

cc. @dtulga Were you able to figure out any alternatives to make these pickle friendly?

[dmpetrov]

If that helps - all dynamic Feature classes are registered in the cache in feature_utils.py:

feature_cache: dict[type[BaseModel], type[Feature]] = {}

[dmpetrov]

@dtulga Were you able to figure out any alternatives to make these pickle friendly?

That would be great if we can find some solution.

Also, I was thinking... we can try re-creating all dynamic Feature classes in each of the thread/process. I'll try this on the weekends.

[dmpetrov]

we can try re-creating all dynamic Feature classes in each of the thread/process. I'll try this on the weekends.

Well... re-creating the objects should not be a problem. However, it's not clear how to delete the objects before deserialization 😅 Cleaning up the cache is not enough - the objects can be created in user code.

Is there a way to specify for pickle/dill not to serialize particular type of objects (subclasses of Feature)?

[amritghimire]

Is there a way to specify for pickle/dill not to serialize particular type of objects (subclasses of Feature)?

I dont think it is possible implicitly but we can determine it and make sure it is not pickled when forming the udf_info that is passed as pickled.

[dmpetrov]

Closed by #45