dosubot[bot]
commented
2 weeks ago Hey there, @melya! I'm here to help you troubleshoot the issue you're experiencing with the ingestion pipeline. I can assist with bug fixes, answer your questions, and guide you on becoming a contributor while you wait for a human maintainer. Let's dive into resolving this issue!
To fix the AttributeError: Can't pickle local object 'split_by_sentence_tokenizer.<locals>.split' when running the ingestion pipeline with num_workers > 1 using the SentenceWindowNodeParser, you need to ensure that the function or object being passed to the multiprocessing pool is pickleable. This error typically occurs because local functions or lambdas cannot be pickled.
Here's how you can refactor the code to resolve this issue:
- Move the local function
splitoutside of thesplit_by_sentence_tokenizerfunction. - Ensure that any objects or functions passed to the
ProcessPoolExecutorare defined at the top level of the module.
Here is an example of how you can refactor the code:
import os
import asyncio
from concurrent.futures import ProcessPoolExecutor
from functools import partial
from typing import List, Sequence, Any, Optional
from llama_index.core import SimpleDirectoryReader, IngestionPipeline, Document, BaseNode
from llama_index.core.node_parser import SentenceWindowNodeParser
from llama_index.core.ingestion.cache import IngestionCache
from llama_index.core.ingestion.transform import TransformComponent
# Define the split function at the top level
def split(text):
# Your split logic here
pass
def split_by_sentence_tokenizer(text):
# Use the top-level split function
return split(text)
# Ensure arun_transformations_wrapper is also at the top level
def arun_transformations_wrapper(
nodes: List[BaseNode],
transformations: Sequence[TransformComponent],
in_place: bool = True,
cache: Optional[IngestionCache] = None,
cache_collection: Optional[str] = None,
**kwargs: Any,
) -> List[BaseNode]:
"""Wrapper for async run_transformation. To be used in loop.run_in_executor
within a ProcessPoolExecutor.
"""
loop = asyncio.new_event_loop()
nodes = loop.run_until_complete(
arun_transformations(
nodes=nodes,
transformations=transformations,
in_place=in_place,
cache=cache,
cache_collection=cache_collection,
**kwargs,
)
)
loop.close()
return nodes
# In the run method, ensure the ProcessPoolExecutor uses the top-level function
def run(
self,
show_progress: bool = False,
documents: Optional[List[Document]] = None,
nodes: Optional[List[BaseNode]] = None,
cache_collection: Optional[str] = None,
in_place: bool = True,
store_doc_text: bool = True,
num_workers: Optional[int] = None,
**kwargs: Any,
) -> Sequence[BaseNode]:
input_nodes = self._prepare_inputs(documents, nodes)
if num_workers and num_workers > 1:
if num_workers > os.cpu_count():
warnings.warn(
"Specified num_workers exceed number of CPUs in the system. "
"Setting `num_workers` down to the maximum CPU count."
)
loop = asyncio.get_event_loop()
with ProcessPoolExecutor(max_workers=num_workers) as p:
node_batches = self._node_batcher(
num_batches=num_workers, nodes=nodes_to_run
)
tasks = [
loop.run_in_executor(
p,
partial(
arun_transformations_wrapper,
transformations=self.transformations,
in_place=in_place,
cache=self.cache if not self.disable_cache else None,
cache_collection=cache_collection,
),
batch,
)
for batch in node_batches
]
result: List[List[BaseNode]] = await asyncio.gather(*tasks)
nodes = reduce(lambda x, y: x + y, result, [])
else:
nodes = await arun_transformations(
nodes_to_run,
self.transformations,
show_progress=show_progress,
cache=self.cache if not self.disable_cache else None,
cache_collection=cache_collection,
in_place=in_place,
**kwargs,
)
if self.vector_store is not None:
await self.vector_store.async_add(
[n for n in nodes if n.embedding is not None]
)
return nodes
# Example usage
num_cores = os.cpu_count()
num_workers = min(4, num_cores)
director_reader = SimpleDirectoryReader(
input_dir="./data"
)
documents = director_reader.load_data()
pipeline = IngestionPipeline(
transformations=[
SentenceWindowNodeParser.from_defaults(
window_size=6,
window_metadata_key="window",
original_text_metadata_key="original_text",
),
],
)
pipeline.run(documents=documents, show_progress=True, num_workers=num_workers)By moving the split function and arun_transformations_wrapper to the top level, they become pickleable, and the ProcessPoolExecutor can serialize them correctly [1][2][3][4].
To continue talking to Dosu, mention @dosu.
melya
houtaroy
Bug Description
When attempting to run the ingestion pipeline with more than one worker (
num_workers > 1), the process fails when using theSentenceWindowNodeParser.Version
0.10.47
Steps to Reproduce
Here is code which reproduces the bug
Relevant Logs/Tracbacks