feat/support parent mutliprocessing fan out (#36)

* support parent mutliprocessing fan out

* Remove mention of num_processes from pinecone connector

* Create new implementation of BatchPipelineStep

* fix method signatures on fsspec connectors

* Remove --num-processes from pinecone dest ingest test

* fix tidy

* fix tidy

* fix import error

Author: rbiseck3

CHANGELOG.md

@@ -1,4 +1,8 @@
-## 0.0.7-dev0
+## 0.0.7-dev1
+
+### Enhancements
+
+* **support sharing parent multiprocessing for uploaders** If an uploader needs to fan out it's process using multiprocessing, support that using the parent pipeline approach rather than handling it explicitly by the connector logic.  
 
 ## 0.0.6
 

test_e2e/dest/pinecone.sh

@@ -9,7 +9,6 @@ OUTPUT_FOLDER_NAME=s3-pinecone-dest
 OUTPUT_DIR=$SCRIPT_DIR/structured-output/$OUTPUT_FOLDER_NAME
 WORK_DIR=$SCRIPT_DIR/workdir/$OUTPUT_FOLDER_NAME
 max_processes=${MAX_PROCESSES:=$(python3 -c "import os; print(os.cpu_count())")}
-writer_processes=$(((max_processes - 1) > 1 ? (max_processes - 1) : 2))
 
 if [ -z "$PINECONE_API_KEY" ]; then
   echo "Skipping Pinecone ingest test because PINECONE_API_KEY env var is not set."
@@ -101,8 +100,7 @@ PYTHONPATH=. ./unstructured_ingest/main.py \
   pinecone \
   --api-key "$PINECONE_API_KEY" \
   --index-name "$PINECONE_INDEX" \
-  --batch-size 80 \
-  --num-processes "$writer_processes"
+  --batch-size 80
 
 # It can take some time for the index to catch up with the content that was written, this check between 10s sleeps
 # to give it that time process the writes. Will timeout after checking for a minute.

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.0.7-dev0"  # pragma: no cover
+__version__ = "0.0.7-dev1"  # pragma: no cover

unstructured_ingest/v2/interfaces/uploader.py

@@ -1,4 +1,4 @@
-from abc import ABC, abstractmethod
+from abc import ABC
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, TypeVar
@@ -31,9 +31,14 @@ class Uploader(BaseProcess, BaseConnector, ABC):
     def is_async(self) -> bool:
         return False
 
-    @abstractmethod
-    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
-        pass
+    def is_batch(self) -> bool:
+        return False
+
+    def run_batch(self, contents: list[UploadContent], **kwargs: Any) -> None:
+        raise NotImplementedError()
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        raise NotImplementedError()
 
     async def run_async(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         return self.run(contents=[UploadContent(path=path, file_data=file_data)], **kwargs)

unstructured_ingest/v2/pipeline/interfaces.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 import multiprocessing as mp
-from abc import ABC
+from abc import ABC, abstractmethod
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from functools import wraps
@@ -12,7 +12,7 @@
 from tqdm import tqdm
 from tqdm.asyncio import tqdm as tqdm_asyncio
 
-from unstructured_ingest.v2.interfaces import BaseProcess, ProcessorConfig
+from unstructured_ingest.v2.interfaces import BaseProcess, ProcessorConfig, Uploader
 from unstructured_ingest.v2.logger import logger, make_default_logger
 
 BaseProcessT = TypeVar("BaseProcessT", bound=BaseProcess)
@@ -167,3 +167,27 @@ async def run_async(self, _fn: Optional[Callable] = None, **kwargs: Any) -> Opti
     @property
     def cache_dir(self) -> Path:
         return Path(self.context.work_dir) / self.identifier
+
+
+@dataclass
+class BatchPipelineStep(PipelineStep, ABC):
+    process: Uploader
+
+    @timed
+    def __call__(self, iterable: Optional[iterable_input] = None) -> Any:
+        if self.context.mp_supported and self.process.is_batch():
+            return self.run_batch(contents=iterable)
+        super().__call__(iterable=iterable)
+
+    @abstractmethod
+    def _run_batch(self, contents: iterable_input, **kwargs) -> Any:
+        pass
+
+    def run_batch(self, contents: iterable_input, **kwargs) -> Any:
+        try:
+            return self._run_batch(contents=contents, **kwargs)
+        except Exception as e:
+            self.context.status[self.identifier] = {"step_error": str(e)}
+            if self.context.raise_on_error:
+                raise e
+            return None

unstructured_ingest/v2/pipeline/pipeline.py

@@ -4,7 +4,7 @@
 from time import time
 from typing import Any, Optional, Union
 
-from unstructured_ingest.v2.interfaces import ProcessorConfig
+from unstructured_ingest.v2.interfaces import ProcessorConfig, Uploader
 from unstructured_ingest.v2.logger import logger, make_default_logger
 from unstructured_ingest.v2.pipeline.steps.chunk import Chunker, ChunkStep
 from unstructured_ingest.v2.pipeline.steps.download import DownloaderT, DownloadStep
@@ -14,7 +14,7 @@
 from unstructured_ingest.v2.pipeline.steps.partition import Partitioner, PartitionStep
 from unstructured_ingest.v2.pipeline.steps.stage import UploadStager, UploadStageStep
 from unstructured_ingest.v2.pipeline.steps.uncompress import Uncompressor, UncompressStep
-from unstructured_ingest.v2.pipeline.steps.upload import Uploader, UploadStep
+from unstructured_ingest.v2.pipeline.steps.upload import UploadStep
 from unstructured_ingest.v2.processes.chunker import ChunkerConfig
 from unstructured_ingest.v2.processes.connector_registry import (
     ConnectionConfig,

unstructured_ingest/v2/pipeline/steps/upload.py

@@ -4,9 +4,9 @@
 from typing import Callable, Optional, TypedDict
 
 from unstructured_ingest.v2.interfaces import FileData
-from unstructured_ingest.v2.interfaces.uploader import UploadContent, Uploader
+from unstructured_ingest.v2.interfaces.uploader import UploadContent
 from unstructured_ingest.v2.logger import logger
-from unstructured_ingest.v2.pipeline.interfaces import PipelineStep, iterable_input, timed
+from unstructured_ingest.v2.pipeline.interfaces import BatchPipelineStep
 
 STEP_ID = "upload"
 
@@ -17,8 +17,7 @@ class UploadStepContent(TypedDict):
 
 
 @dataclass
-class UploadStep(PipelineStep):
-    process: Uploader
+class UploadStep(BatchPipelineStep):
     identifier: str = STEP_ID
 
     def __str__(self):
@@ -34,25 +33,12 @@ def __post_init__(self):
             f"connection configs: {connection_config}"
         )
 
-    def process_whole(self, iterable: iterable_input):
-        self.run(contents=iterable)
-
-    @timed
-    def __call__(self, iterable: iterable_input):
-        logger.info(
-            f"Calling {self.__class__.__name__} " f"with {len(iterable)} docs",  # type: ignore
-        )
-        if self.process.is_async():
-            self.process_async(iterable=iterable)
-        else:
-            self.process_whole(iterable=iterable)
-
-    def _run(self, fn: Callable, contents: list[UploadStepContent]):
+    def _run_batch(self, contents: list[UploadStepContent]) -> None:
         upload_contents = [
             UploadContent(path=Path(c["path"]), file_data=FileData.from_file(c["file_data_path"]))
             for c in contents
         ]
-        fn(contents=upload_contents)
+        self.process.run_batch(contents=upload_contents)
 
     async def _run_async(self, path: str, file_data_path: str, fn: Optional[Callable] = None):
         fn = fn or self.process.run_async

unstructured_ingest/v2/processes/connectors/astradb.py

@@ -14,7 +14,6 @@
     AccessConfig,
     ConnectionConfig,
     FileData,
-    UploadContent,
     Uploader,
     UploaderConfig,
     UploadStager,
@@ -139,13 +138,9 @@ def get_collection(self) -> "AstraDBCollection":
         )
         return astra_db_collection
 
-    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
-        elements_dict = []
-        for content in contents:
-            with open(content.path) as elements_file:
-                elements = json.load(elements_file)
-                elements_dict.extend(elements)
-
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        with path.open("r") as file:
+            elements_dict = json.load(file)
         logger.info(
             f"writing {len(elements_dict)} objects to destination "
             f"collection {self.upload_config.collection_name}"

unstructured_ingest/v2/processes/connectors/azure_cognitive_search.py

@@ -12,7 +12,7 @@
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
-    UploadContent,
+    FileData,
     Uploader,
     UploaderConfig,
     UploadStager,
@@ -192,14 +192,9 @@ def precheck(self) -> None:
     def write_dict_wrapper(self, elements_dict):
         return self.write_dict(elements_dict=elements_dict)
 
-    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
-
-        elements_dict = []
-        for content in contents:
-            with open(content.path) as elements_file:
-                elements = json.load(elements_file)
-                elements_dict.extend(elements)
-
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        with path.open("r") as file:
+            elements_dict = json.load(file)
         logger.info(
             f"writing document batches to destination"
             f" endpoint at {str(self.connection_config.endpoint)}"

unstructured_ingest/v2/processes/connectors/chroma.py

@@ -15,7 +15,6 @@
     AccessConfig,
     ConnectionConfig,
     FileData,
-    UploadContent,
     Uploader,
     UploaderConfig,
     UploadStager,
@@ -186,13 +185,9 @@ def prepare_chroma_list(chunk: tuple[dict[str, Any]]) -> dict[str, list[Any]]:
         )
         return chroma_dict
 
-    def run(self, contents: list[UploadContent], **kwargs: Any) -> None:
-
-        elements_dict = []
-        for content in contents:
-            with open(content.path) as elements_file:
-                elements = json.load(elements_file)
-                elements_dict.extend(elements)
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        with path.open("r") as file:
+            elements_dict = json.load(file)
 
         logger.info(
             f"writing {len(elements_dict)} objects to destination "