apache / apache/hamilton

`MultiProcessingExecutor` pickle'ing error when using `extract_fields`

Open
#1,285 2 comments 0 reactions 0 assignees View on GitHub
triage
Dominant language
Jupyter Notebook
Stars
2.6k
Forks
213
PR merge metrics
No merged PRs in 30d

Description

# Current behavior
`MultiProcessingExecutor` pickle'ing error when using `extract_fields`.

## Stack Traces
```Traceback (most recent call last):
File "c:\codebase\api\venv\Lib\site-packages\hamilton\execution\executors.py", line 204, in get_state
self.future.result()
File "C:\Program Files\Python312\Lib\concurrent\futures\_base.py", line 449, in result
return self.__get_result()
^^^^^^^^^^^^^^^^^^^
File "C:\Program Files\Python312\Lib\concurrent\futures\_base.py", line 401, in __get_result
raise self._exception
File "C:\Program Files\Python312\Lib\multiprocessing\queues.py", line 244, in _feed
obj = _ForkingPickler.dumps(obj)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Program Files\Python312\Lib\multiprocessing\reduction.py", line 51, in dumps
cls(buf, protocol).dump(obj)
AttributeError: Can't pickle local object 'extract_fields.transform_node..dict_generator'
```
## Screenshots

![Image](https://github.com/user-attachments/assets/db122492-4710-4e65-a0a5-d3a635145b65)

## Steps to replicate behavior
```python
import pandas as pd
from typing import Dict, List, Tuple

from hamilton.htypes import Parallelizable, Collect
from hamilton.function_modifiers import extract_fields

def _load_full_dataset() -> pd.DataFrame:
url = "https://raw.githubusercontent.com/mwaskom/seaborn-data/master/penguins.csv"
return pd.read_csv(url)

def kind(kinds: List[str]) -> Parallelizable[str]:
for kind in kinds:
yield kind

def raw_kind(kind: str) -> pd.DataFrame:
return df[df.species == kind]

@extract_fields(dict(raw_df=pd.DataFrame, agg_df=pd.DataFrame))
def all_kinds(raw_kind: Collect[pd.DataFrame]) -> Dict[str, pd.DataFrame]:
raw = pd.concat(raw_kind)
return {
"raw_df": raw,
"agg_df": raw.groupby("species")[
["bill_length_mm", "bill_depth_mm", "flipper_length_mm", "body_mass_g"]
].mean().reset_index(),
}

def all_data(raw_df: pd.DataFrame, agg_df: pd.DataFrame) -> Tuple[pd.DataFrame]:
return raw_df, agg_df

if __name__ == "__main__":
from hamilton import driver, settings
from hamilton.execution import executors

import __main__

def debug_parallellism(mode: str) -> Tuple:
if mode == "local":
executor = executors.SynchronousLocalTaskExecutor()
elif mode=="multithreading":
executor = executors.MultiThreadingExecutor(max_tasks=4)
elif mode=="multiprocessing":
executor = executors.MultiProcessingExecutor(max_tasks=6)

config = {settings.ENABLE_POWER_USER_MODE: True, "kinds": species}
dr = (
driver.Builder()
.with_modules(__main__)
.with_config(config)
.enable_dynamic_execution(allow_experimental_mode=True)
.with_remote_executor(executor)
.build()
)

dag = dr.visualize_execution(["all_data"], f"parallel-penguins-{mode}-dag.png", bypass_validation=True)
results = dr.execute(["all_data"])
return dag, results


df = _load_full_dataset()
species = df.species.unique().tolist()

modes = ["local", "multithreading", "multiprocessing"]
for mode in modes:
try:
dag, results = debug_parallellism(mode)
except Exception as e:
print(f"Error in {mode=}: {e}")
```

## Library & System Information
Microsoft Windows [Version 10.0.22621.4890]
Python 3.12
`sf-hamilton==1.87.0`

# Expected behavior
I was hoping this would work since the `SynchronousLocalTaskExecutor` works as expected.
I tried `MultiThreadingExecutor` for my use case but it hangs longer than the sync executor which is why I am trying the `MultiProcessingExecutor`

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.