Repository: explosion/spaCy
Pull request: Avoid pickling `Doc` inputs passed to `Language.pipe()`
## Description
<!--- Use this section to describe your changes. If your changes required
testing, include information about the testing environment and the tests you
ran. If your test fixes a bug reported in an issue, don't forget to include the
issue number. If your PR is still a work in progress, that's totally fine – just
include a note to let us know. -->
A user [reported](https://github.com/explosion/spaCy/issues/10839) significantly slower performance when calling `Language.pipe()` with `as_tuples=True` (as compared to `as_tuples=False`). Closer investigation revealed the following:
- When `as_tuples` is set to `True`, the inputs passed as `texts` are preemptively converted to `Doc` objects (to store their respective context data in the `_context` attribute). Then, `pipe()` is called again with the augmented inputs.
- When multiprocessing is enabled (`n_process > 1`), `Doc` objects passed to `pipe()` are pickled as-is and sent to child processes along with their shared `Vocab` and `Vectors` instances.
- This incurs a significant serialization/deserialization overhead, which is further exacerbated by small batch sizes.

This PR address the issue through the following changes:
- Serialize `Doc` objects in the input to byte arrays using `Doc.to_bytes()` when multiprocessing is enabled. The overhead of pickling a byte array is negligible.
- In child processes, deserialize incoming byte arrays to their correspond `Doc` representations before executing the pipeline.
- Ensure that the `_context` attribute is serialized when calling `Doc.to_bytes/to_dict()` in order to prevent the loss of context information.

### Relevant issues
- #10839

### Types of change
<!-- What type of change does your PR cover? Is it a bug fix, an enhancement
or new feature, or a change to the documentation? -->
Performance enhancement

## Checklist
<!--- Before you submit the PR, go over this checklist and make sure you can
tick off all the boxes. [] -> [x] -->
- [x] I confirm that I have the right to submit this contribution under the project's MIT license.
- [x] I ran the tests, and all new and existing tests passed.
- [x] My changes don't require a change to the documentation, or if they do, I've added all required information.


Problem statement:
Pipe() process time is very slow with as_tuples attribute
When I set `as_tuples=True` attribute to the `.pipe` method, the process time is very very slow.

Results without tuples :
```
0:00:02.997644 Starting pipe
text processed
[...]
text processed
Pipe finished in 0:00:13.833113
```

Results with tuples :
```
0:00:02.879235 Starting pipe
text processed
[...]
text processed
Pipe finished in 0:02:19.793344
```

## How to reproduce the behaviour
Here a code sample :
```
import datetime

import spacy


TIME = datetime.datetime.now()


def time_profiler():
    global TIME
    diff = datetime.datetime.now() - TIME
    TIME = datetime.datetime.now()
    return diff


if __name__ == "__main__":
    nlp = spacy.load("fr_core_news_lg")
    # text_to_processed = nlp.pipe(
    #     [f"<p>Testing without tuple</p>" * 100] * 50,
    #     n_process=2, batch_size=1
    # )
    text_to_processed = nlp.pipe(
        [(f"<p>Testing tuple</p>" * 100, -1)] * 50,
        n_process=2, batch_size=1, as_tuples=True
    )
    print(time_profiler(), "Starting pipe")
    # for text_processed in text_to_processed:
    #     print("text processed")
    for text_processed, _id in text_to_processed:
        print("text processed", _id)
    print("Pipe finished in", time_profiler())
```

## Your Environment
- **spaCy version:** 3.3.0
- **Platform:** macOS-12.4-arm64-arm-64bit
- **Python version:** 3.9.11
- **Pipelines:** fr_core_news_lg (3.3.0)

Hints from issue discussion:
Thanks for the report and the code sample! I'm able to reproduce the slow down on my end (on Linux) as well. Looks like the code spends an inordinate amount of time serializing and deserializing data during IPC (inter-process communication).

Commit message: Remove superfluous comma

Patch to review:
diff --git a/spacy/errors.py b/spacy/errors.py
index c82ffe8824b..8df5edee809 100644
--- a/spacy/errors.py
+++ b/spacy/errors.py
@@ -919,6 +919,7 @@ class Errors(metaclass=ErrorsWithCodes):
     E1035 = ("Token index {i} out of bounds ({length})")
     E1036 = ("Cannot index into NoneNode")
     E1037 = ("Invalid attribute value '{attr}'.")
+    E1038 = ("Expected a string, 'Doc', or 'bytes' as input, but got: {type}.")
 
 
 # Deprecated model shortcuts, only used in errors and warnings
diff --git a/spacy/language.py b/spacy/language.py
index faca1f25883..97bb8692e91 100644
--- a/spacy/language.py
+++ b/spacy/language.py
@@ -1090,16 +1090,21 @@ def make_doc(self, text: str) -> Doc:
             )
         return self.tokenizer(text)
 
-    def _ensure_doc(self, doc_like: Union[str, Doc]) -> Doc:
-        """Create a Doc if need be, or raise an error if the input is not a Doc or a string."""
+    def _ensure_doc(self, doc_like: Union[str, Doc, bytes]) -> Doc:
+        """Create a Doc if need be, or raise an error if the input is not
+        a Doc, string, or a byte array (generated by Doc.to_bytes())."""
         if isinstance(doc_like, Doc):
             return doc_like
         if isinstance(doc_like, str):
             return self.make_doc(doc_like)
-        raise ValueError(Errors.E866.format(type=type(doc_like)))
+        if isinstance(doc_like, bytes):
+            return Doc(self.vocab).from_bytes(doc_like)
+        raise ValueError(Errors.E1038.format(type=type(doc_like)))
 
-    def _ensure_doc_with_context(self, doc_like: Union[str, Doc], context: Any) -> Doc:
-        """Create a Doc if need be and add as_tuples context, or raise an error if the input is not a Doc or a string."""
+    def _ensure_doc_with_context(
+        self, doc_like: Union[str, Doc, bytes], context: _AnyContext
+    ) -> Doc:
+        """Call _ensure_doc to generate a Doc and set its context object."""
         doc = self._ensure_doc(doc_like)
         doc._context = context
         return doc
@@ -1519,7 +1524,6 @@ def pipe(  # noqa: F811
 
         DOCS: https://spacy.io/api/language#pipe
         """
-        # Handle texts with context as tuples
         if as_tuples:
             texts = cast(Iterable[Tuple[Union[str, Doc], _AnyContext]], texts)
             docs_with_contexts = (
@@ -1597,8 +1601,21 @@ def _multiprocessing_pipe(
         n_process: int,
         batch_size: int,
     ) -> Iterator[Doc]:
+        def prepare_doc(
+            texts: Iterable[Union[str, Doc]]
+        ) -> Iterable[Tuple[Union[str, bytes], _AnyContext]]:
+            # Serialize Doc inputs to bytes to avoid incurring pickling
+            # overhead when they are passed to child processes. Also yield
+            # any context objects they might have separately (as they are not serialized).
+            for doc_like in texts:
+                if isinstance(doc_like, Doc):
+                    yield (doc_like.to_bytes(), cast(_AnyContext, doc_like._context))
+                else:
+                    yield (doc_like, cast(_AnyContext, None))
+
+        serialized_texts_with_ctx = prepare_doc(texts)  # type: ignore
         # raw_texts is used later to stop iteration.
-        texts, raw_texts = itertools.tee(texts)
+        texts, raw_texts = itertools.tee(serialized_texts_with_ctx)  # type: ignore
         # for sending texts to worker
         texts_q: List[mp.Queue] = [mp.Queue() for _ in range(n_process)]
         # for receiving byte-encoded docs from worker
@@ -1618,7 +1635,13 @@ def _multiprocessing_pipe(
         procs = [
             mp.Process(
                 target=_apply_pipes,
-                args=(self._ensure_doc, pipes, rch, sch, Underscore.get_state()),
+                args=(
+                    self._ensure_doc_with_context,
+                    pipes,
+                    rch,
+                    sch,
+                    Underscore.get_state(),
+                ),
             )
             for rch, sch in zip(texts_q, bytedocs_send_ch)
         ]
@@ -1631,15 +1654,15 @@ def _multiprocessing_pipe(
             recv.recv() for recv in cycle(bytedocs_recv_ch)
         )
         try:
-            for i, (_, (byte_doc, byte_context, byte_error)) in enumerate(
+            for i, (_, (byte_doc, context, error)) in enumerate(
                 zip(raw_texts, byte_tuples), 1
             ):
                 if byte_doc is not None:
                     doc = Doc(self.vocab).from_bytes(byte_doc)
-                    doc._context = byte_context
+                    doc._context = context
                     yield doc
-                elif byte_error is not None:
-                    error = srsly.msgpack_loads(byte_error)
+                elif error is not None:
+                    error = srsly.msgpack_loads(error)
                     self.default_error_handler(
                         None, None, None, ValueError(Errors.E871.format(error=error))
                     )
@@ -2163,7 +2186,7 @@ def _copy_examples(examples: Iterable[Example]) -> List[Example]:
 
 
 def _apply_pipes(
-    ensure_doc: Callable[[Union[str, Doc]], Doc],
+    ensure_doc: Callable[[Union[str, Doc, bytes], _AnyContext], Doc],
     pipes: Iterable[Callable[..., Iterator[Doc]]],
     receiver,
     sender,
@@ -2184,17 +2207,19 @@ def _apply_pipes(
     Underscore.load_state(underscore_state)
     while True:
         try:
-            texts = receiver.get()
-            docs = (ensure_doc(text) for text in texts)
+            texts_with_ctx = receiver.get()
+            docs = (
+                ensure_doc(doc_like, context) for doc_like, context in texts_with_ctx
+            )
             for pipe in pipes:
                 docs = pipe(docs)  # type: ignore[arg-type, assignment]
             # Connection does not accept unpickable objects, so send list.
             byte_docs = [(doc.to_bytes(), doc._context, None) for doc in docs]
-            padding = [(None, None, None)] * (len(texts) - len(byte_docs))
+            padding = [(None, None, None)] * (len(texts_with_ctx) - len(byte_docs))
             sender.send(byte_docs + padding)  # type: ignore[operator]
         except Exception:
             error_msg = [(None, None, srsly.msgpack_dumps(traceback.format_exc()))]
-            padding = [(None, None, None)] * (len(texts) - 1)
+            padding = [(None, None, None)] * (len(texts_with_ctx) - 1)
             sender.send(error_msg + padding)
 
 
diff --git a/spacy/tests/doc/test_pickle_doc.py b/spacy/tests/doc/test_pickle_doc.py
index 738a751a043..28cb66714d6 100644
--- a/spacy/tests/doc/test_pickle_doc.py
+++ b/spacy/tests/doc/test_pickle_doc.py
@@ -5,11 +5,9 @@
 def test_pickle_single_doc():
     nlp = Language()
     doc = nlp("pickle roundtrip")
-    doc._context = 3
     data = pickle.dumps(doc, 1)
     doc2 = pickle.loads(data)
     assert doc2.text == "pickle roundtrip"
-    assert doc2._context == 3
 
 
 def test_list_of_docs_pickles_efficiently():
diff --git a/spacy/tokens/doc.pyx b/spacy/tokens/doc.pyx
index d25247b13bb..4d44fd3870d 100644
--- a/spacy/tokens/doc.pyx
+++ b/spacy/tokens/doc.pyx
@@ -1737,18 +1737,17 @@ cdef int [:,:] _get_lca_matrix(Doc doc, int start, int end):
 def pickle_doc(doc):
     bytes_data = doc.to_bytes(exclude=["vocab", "user_data", "user_hooks"])
     hooks_and_data = (doc.user_data, doc.user_hooks, doc.user_span_hooks,
-                      doc.user_token_hooks, doc._context)
+                      doc.user_token_hooks)
     return (unpickle_doc, (doc.vocab, srsly.pickle_dumps(hooks_and_data), bytes_data))
 
 
 def unpickle_doc(vocab, hooks_and_data, bytes_data):
-    user_data, doc_hooks, span_hooks, token_hooks, _context = srsly.pickle_loads(hooks_and_data)
+    user_data, doc_hooks, span_hooks, token_hooks = srsly.pickle_loads(hooks_and_data)
 
     doc = Doc(vocab, user_data=user_data).from_bytes(bytes_data, exclude=["user_data"])
     doc.user_hooks.update(doc_hooks)
     doc.user_span_hooks.update(span_hooks)
     doc.user_token_hooks.update(token_hooks)
-    doc._context = _context
     return doc
 
 


Code review comments:
diff_hunk: @@ -1597,8 +1601,21 @@ def _multiprocessing_pipe(
         n_process: int,
         batch_size: int,
     ) -> Iterator[Doc]:
+        def prepare_doc(
path: spacy/language.py
text: @user1:
I'm still not happy with my own naming suggestion. `prepare_input`? `split_context`? 

(But I care a lot less when the method is here.)

@author:
I'll go for `prepare_input`. The comment inside the function should describe the intent more concretely, in any event.