ProcessPoolExecutor in apply endpoint causes `tesseract run` to hang indefinitely

[nmheim]

Description

See title.

Steps to reproduce

When building a tesseract with the following setup

# tesseract_config.yaml
name: "reproducer"
version: "0.0.1"

# tesseract_api.py

from concurrent.futures import ProcessPoolExecutor
from pydantic import BaseModel


class InputSchema(BaseModel):
    pass


class OutputSchema(BaseModel):
    pass


def preprocess_fn(data_id: int):
    return data_id


def apply(inputs: InputSchema):

    data_ids = list(range(10))

    pool = ProcessPoolExecutor()
    futures = []

    for idx in data_ids:
        x = pool.submit(preprocess_fn, idx)
        futures.append(x)
        print(idx, "submitted")

    for f in futures:
        res = f.result()
        print(res, "done")

    return OutputSchema()

Running via the command below never finishes

tesseract run reproducer apply '{"inputs":{}}' --output-path outputs

Logs

OS

Mac

Tesseract version

1.0.0

[nmheim]

I just tried adding a threading.Lock in the LogPipe.__exit__ method to make sure only one thread is closing the same file descriptor... something along the lines of

class LogPipe(threading.Thread):
    daemon = True

    def __init__(self, *sinks: Callable) -> None:
        """Initialize the LogPipe with the given logging level."""
        super().__init__()
        self._sinks = sinks
        self._fd_read, self._fd_write = os.pipe()
        # note that we need closefd=False for some reason
        self._pipe_reader = os.fdopen(self._fd_read, closefd=False)
        self._captured_lines = []

        self._lock = threading.Lock()
        self._closed = False

    def __exit__(self, *args: Any) -> None:
        """Close the pipe and join the thread."""
        with self._lock:
            if not self._closed:
                os.close(self._fd_write)
                self.join(timeout=1)
                os.close(self._fd_read)
                self._pipe_reader.close()
                self._closed = True

running this via

tesseract-runtime apply '{"inputs":{}}'

which interestingly resulted in

PicklingError: Can't pickle <function preprocess_fn at 0x104b85ea0>: import of module
'tesseract_api' failed

Full stack trace

_RemoteTraceback:
"""
Traceback (most recent call last):
  File
"/Users/niklas/.local/share/uv/python/cpython-3.10.16-macos-aarch64-none/lib/python3.10/multiproc
essing/queues.py", line 244, in _feed
    obj = _ForkingPickler.dumps(obj)
  File
"/Users/niklas/.local/share/uv/python/cpython-3.10.16-macos-aarch64-none/lib/python3.10/multiproc
essing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
_pickle.PicklingError: Can't pickle <function preprocess_fn at 0x104b85ea0>: import of module
'tesseract_api' failed
"""

The above exception was the direct cause of the following exception:

╭────────────────────────────── Traceback (most recent call last) ──────────────────────────────╮
│ /Users/niklas/repos/pasteur/tesseract-core/tesseract_core/runtime/cli.py:447 in command_func  │
│                                                                                               │
│   444 │   │                                                                                   │
│   445 │   │   def command_func(payload: str):                                                 │
│   446 │   │   │   parsed_payload = _parse_payload(payload)                                    │
│ ❱ 447 │   │   │   return _callback_wrapper(payload=parsed_payload)                            │
│   448 │   else:                                                                               │
│   449 │   │                                                                                   │
│   450 │   │   def command_func():                                                             │
│                                                                                               │
│ ╭───────────── locals ─────────────╮                                                          │
│ │ parsed_payload = {'inputs': {}}  │                                                          │
│ │        payload = '{"inputs":{}}' │                                                          │
│ ╰──────────────────────────────────╯                                                          │
│                                                                                               │
│ /Users/niklas/repos/pasteur/tesseract-core/tesseract_core/runtime/cli.py:414 in               │
│ _callback_wrapper                                                                             │
│                                                                                               │
│   411 │   │   │   Path(output_path).mkdir(parents=True, exist_ok=True)                        │
│   412 │   │                                                                                   │
│   413 │   │   with start_run(base_dir=output_path):                                           │
│ ❱ 414 │   │   │   result = user_function(**user_function_args)                                │
│   415 │   │                                                                                   │
│   416 │   │   result = output_to_bytes(result, output_format, output_path)                    │
│   417                                                                                         │
│                                                                                               │
│ ╭──────────────────────────────────── locals ────────────────────────────────────╮            │
│ │             config = RuntimeConfig(                                            │            │
│ │                      │   api_path=PosixPath('tesseract_api.py'),               │            │
│ │                      │   name='Tesseract',                                     │            │
│ │                      │   description='',                                       │            │
│ │                      │   version='unknown',                                    │            │
│ │                      │   debug=False,                                          │            │
│ │                      │   input_path='.',                                       │            │
│ │                      │   output_path='.',                                      │            │
│ │                      │   output_format='json',                                 │            │
│ │                      │   output_file='',                                       │            │
│ │                      │   mlflow_tracking_uri=''                                │            │
│ │                      )                                                         │            │
│ │         input_path = '.'                                                       │            │
│ │             kwargs = {'payload': {'inputs': {}}}                               │            │
│ │         out_stream = <_io.TextIOWrapper name=3 mode='w' encoding='UTF-8'>      │            │
│ │        out_stream_ = <_io.TextIOWrapper name=3 mode='w' encoding='UTF-8'>      │            │
│ │        output_file = ''                                                        │            │
│ │      output_format = 'json'                                                    │            │
│ │        output_path = '.'                                                       │            │
│ │            payload = {'inputs': {}}                                            │            │
│ │ user_function_args = {'payload': ApplyInputSchema(inputs=Apply_InputSchema())} │            │
│ ╰────────────────────────────────────────────────────────────────────────────────╯            │
│                                                                                               │
│ /Users/niklas/repos/pasteur/tesseract-core/tesseract_core/runtime/core.py:155 in apply        │
│                                                                                               │
│   152 │   @assemble_docstring(api_module.apply)                                               │
│   153 │   def apply(payload: ApplyInputSchema) -> ApplyOutputSchema:                          │
│   154 │   │   """Apply the Tesseract to the input data."""                                    │
│ ❱ 155 │   │   out = api_module.apply(payload.inputs)                                          │
│   156 │   │   if isinstance(out, api_module.OutputSchema):                                    │
│   157 │   │   │   out = out.model_dump()                                                      │
│   158 │   │   return ApplyOutputSchema.model_validate(out)                                    │
│                                                                                               │
│ ╭───────────────────────────────────────── locals ──────────────────────────────────────────╮ │
│ │ api_module = <module 'tesseract_api' from                                                 │ │
│ │              '/Users/niklas/repos/pasteur/tesseract-core/examples/reproducer/tesseract_a… │ │
│ │    payload = ApplyInputSchema(inputs=Apply_InputSchema())                                 │ │
│ ╰───────────────────────────────────────────────────────────────────────────────────────────╯ │
│                                                                                               │
│ /Users/niklas/repos/pasteur/tesseract-core/examples/reproducer/tesseract_api.py:36 in apply   │
│                                                                                               │
│   33 │   │   print(idx, "submitted")                                                          │
│   34 │                                                                                        │
│   35 │   for f in futures:                                                                    │
│ ❱ 36 │   │   res = f.result()                                                                 │
│   37 │   │   print(res, "done")                                                               │
│   38 │                                                                                        │
│   39 │   return OutputSchema()                                                                │
│                                                                                               │
│ ╭───────────────────────────────────── locals ──────────────────────────────────────╮         │
│ │ data_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]                                         │         │
│ │        f = <Future at 0x1082a4c70 state=finished raised PicklingError>            │         │
│ │  futures = [                                                                      │         │
│ │            │   <Future at 0x1082a4c70 state=finished raised PicklingError>,       │         │
│ │            │   <Future at 0x1082a5960 state=finished raised PicklingError>,       │         │
│ │            │   <Future at 0x1082a5ab0 state=finished raised PicklingError>,       │         │
│ │            │   <Future at 0x1082a5ea0 state=finished raised PicklingError>,       │         │
│ │            │   <Future at 0x1082a63b0 state=finished raised PicklingError>,       │         │
│ │            │   <Future at 0x1082a5ff0 state=finished raised PicklingError>,       │         │
│ │            │   <Future at 0x1082a6710 state=finished raised PicklingError>,       │         │
│ │            │   <Future at 0x1082a6800 state=finished raised PicklingError>,       │         │
│ │            │   <Future at 0x1082a6aa0 state=finished raised PicklingError>,       │         │
│ │            │   <Future at 0x1082a6d40 state=finished raised PicklingError>        │         │
│ │            ]                                                                      │         │
│ │      idx = 9                                                                      │         │
│ │   inputs = Apply_InputSchema()                                                    │         │
│ │     pool = <concurrent.futures.process.ProcessPoolExecutor object at 0x108266500> │         │
│ │        x = <Future at 0x1082a6d40 state=finished raised PicklingError>            │         │
│ ╰───────────────────────────────────────────────────────────────────────────────────╯         │
│                                                                                               │
│ /Users/niklas/.local/share/uv/python/cpython-3.10.16-macos-aarch64-none/lib/python3.10/concur │
│ rent/futures/_base.py:451 in result                                                           │
│                                                                                               │
│   448 │   │   │   │   if self._state in [CANCELLED, CANCELLED_AND_NOTIFIED]:                  │
│   449 │   │   │   │   │   raise CancelledError()                                              │
│   450 │   │   │   │   elif self._state == FINISHED:                                           │
│ ❱ 451 │   │   │   │   │   return self.__get_result()                                          │
│   452 │   │   │   │                                                                           │
│   453 │   │   │   │   self._condition.wait(timeout)                                           │
│   454                                                                                         │
│                                                                                               │
│ ╭──── locals ────╮                                                                            │
│ │    self = None │                                                                            │
│ │ timeout = None │                                                                            │
│ ╰────────────────╯                                                                            │
│                                                                                               │
│ /Users/niklas/.local/share/uv/python/cpython-3.10.16-macos-aarch64-none/lib/python3.10/concur │
│ rent/futures/_base.py:403 in __get_result                                                     │
│                                                                                               │
│   400 │   def __get_result(self):                                                             │
│   401 │   │   if self._exception:                                                             │
│   402 │   │   │   try:                                                                        │
│ ❱ 403 │   │   │   │   raise self._exception                                                   │
│   404 │   │   │   finally:                                                                    │
│   405 │   │   │   │   # Break a reference cycle with the exception in self._exception         │
│   406 │   │   │   │   self = None                                                             │
│                                                                                               │
│ ╭── locals ───╮                                                                               │
│ │ self = None │                                                                               │
│ ╰─────────────╯                                                                               │
│                                                                                               │
│ /Users/niklas/.local/share/uv/python/cpython-3.10.16-macos-aarch64-none/lib/python3.10/multip │
│ rocessing/queues.py:244 in _feed                                                              │
│                                                                                               │
│   241 │   │   │   │   │   │   │   return                                                      │
│   242 │   │   │   │   │   │                                                                   │
│   243 │   │   │   │   │   │   # serialize the data before acquiring the lock                  │
│ ❱ 244 │   │   │   │   │   │   obj = _ForkingPickler.dumps(obj)                                │
│   245 │   │   │   │   │   │   if wacquire is None:                                            │
│   246 │   │   │   │   │   │   │   send_bytes(obj)                                             │
│   247 │   │   │   │   │   │   else:                                                           │
│                                                                                               │
│ ╭───────────────────────────────────────── locals ──────────────────────────────────────────╮ │
│ │     bpopleft = <built-in method popleft of collections.deque object at 0x108253340>       │ │
│ │       buffer = deque()                                                                    │ │
│ │ ignore_epipe = True                                                                       │ │
│ │     nacquire = <built-in method acquire of _thread.lock object at 0x105917300>            │ │
│ │     notempty = <Condition(<unlocked _thread.lock object at 0x105917300>, 1)>              │ │
│ │     nrelease = <built-in method release of _thread.lock object at 0x105917300>            │ │
│ │        nwait = <bound method Condition.wait of <Condition(<unlocked _thread.lock object   │ │
│ │                at 0x105917300>, 1)>>                                                      │ │
│ │          obj = <concurrent.futures.process._CallItem object at 0x108266c80>               │ │
│ │      onerror = <bound method _SafeQueue._on_queue_feeder_error of                         │ │
│ │                <concurrent.futures.process._SafeQueue object at 0x108267790>>             │ │
│ │    queue_sem = <BoundedSemaphore(value=unknown, maxvalue=12)>                             │ │
│ │ reader_close = <bound method _ConnectionBase.close of                                     │ │
│ │                <multiprocessing.connection.Connection object at 0x1082677f0>>             │ │
│ │   send_bytes = <bound method _ConnectionBase.send_bytes of                                │ │
│ │                <multiprocessing.connection.Connection object at 0x108267850>>             │ │
│ │     sentinel = <object object at 0x10745d3e0>                                             │ │
│ │     wacquire = <built-in method acquire of _multiprocessing.SemLock object at             │ │
│ │                0x1050276b0>                                                               │ │
│ │     wrelease = <built-in method release of _multiprocessing.SemLock object at             │ │
│ │                0x1050276b0>                                                               │ │
│ │    writelock = <Lock(owner=unknown)>                                                      │ │
│ │ writer_close = <bound method _ConnectionBase.close of                                     │ │
│ │                <multiprocessing.connection.Connection object at 0x108267850>>             │ │
│ ╰───────────────────────────────────────────────────────────────────────────────────────────╯ │
│                                                                                               │
│ /Users/niklas/.local/share/uv/python/cpython-3.10.16-macos-aarch64-none/lib/python3.10/multip │
│ rocessing/reduction.py:51 in dumps                                                            │
│                                                                                               │
│    48 │   @classmethod                                                                        │
│    49 │   def dumps(cls, obj, protocol=None):                                                 │
│    50 │   │   buf = io.BytesIO()                                                              │
│ ❱  51 │   │   cls(buf, protocol).dump(obj)                                                    │
│    52 │   │   return buf.getbuffer()                                                          │
│    53 │                                                                                       │
│    54 │   loads = pickle.loads                                                                │
│                                                                                               │
│ ╭──────────────────────────────── locals ─────────────────────────────────╮                   │
│ │      buf = <_io.BytesIO object at 0x1082a2f20>                          │                   │
│ │      obj = <concurrent.futures.process._CallItem object at 0x1082a5720> │                   │
│ │ protocol = None                                                         │                   │
│ ╰─────────────────────────────────────────────────────────────────────────╯                   │
╰───────────────────────────────────────────────────────────────────────────────────────────────╯
PicklingError: Can't pickle <function preprocess_fn at 0x104b85ea0>: import of module
'tesseract_api' failed

Replacing pool.submit(preprocess_fn, idx) with e.g. pool.submit(np.identity, idx) makes both tesseract-runtime and tesseract run succeed. so it seems like pickling our apply endpoint is also sometimes a problem.

[dionhaefner]

Thanks for the report + repro @nmheim, I'll do some hacking on this.

[dionhaefner]

I'm not sure whether we can support the reproducer as written. It seems like our internal machinery modifies tesseract_api.py to the point where function objects aren't identical pre- and post-pickling.

In the meantime, this workaround works for me:

# tesseract_config.yaml
name: "reproducer"
version: "0.0.1"

build_config:
  package_data:
    - ["funceroni.py", "funceroni.py"]

# funceroni.py
def preprocess_fn(data_id: int):
    return data_id

# tesseract_api.py

from concurrent.futures import ProcessPoolExecutor
from pydantic import BaseModel

import os
import sys
sys.path.append(os.path.dirname(__file__))

from funceroni import preprocess_fn

class InputSchema(BaseModel):
    pass


class OutputSchema(BaseModel):
    pass


def apply(inputs: InputSchema):
    pool = ProcessPoolExecutor()

    data_ids = list(range(10))

    futures = []
    with pool as executor:
        for idx in data_ids:
            x = executor.submit(preprocess_fn, idx)
            futures.append(x)
            print(idx, "submitted")

        for f in futures:
            res = f.result()
            print(res, "done")

    return OutputSchema()

Would you be able to verify that @nmheim ?

[nmheim]

Cool, thanks a lot for looking into this! So tesseract-runtime apply does not throw the pickling error anymore, but it still hangs for me:

$ tesseract-runtime apply '{"inputs":{}}'
0 submitted
1 submitted
2 submitted
3 submitted
4 submitted
5 submitted
6 submitted
7 submitted
8 submitted
9 submitted
0 done
1 done
2 done
3 done
4 done
5 done
6 done
7 done
8 done
9 done
^C^C^C^C^C^C^C^C^C^C^C^C^C^C^C^C^C^C^C^C^C^C^C

With tesseract run reproduce apply it seems to work though

$ tesseract run reproducer apply '{"inputs":{}}'
0 submitted
1 submitted
2 submitted
3 submitted
4 submitted
5 submitted
6 submitted
7 submitted
8 submitted
9 submitted
0 done
1 done
2 done
3 done
4 done
5 done
6 done
7 done
8 done
9 done
{}%  

[dionhaefner]

Thanks for testing. Is that with main or #392?

[dionhaefner]

Either way we should probably add a sharp edge to the docs about this.

[nmheim]

Thanks for testing. Is that with main or #392?

Tested on main. On #392 its working fine 👍