ThreadPoolExecutor with stateful workers

Question

I'm working with a Backend class which spawns a subprocess to perform the CPU-bound work. I have no control over that class and basically the only way of interaction is to create an instance backend = Backend() and submit work via backend.run(data) (this in turn submits the work to the subprocess and blocks until completion). Because these computations take quite some time, I'd like to perform them in parallel. Since the Backend class already spawns its own subprocess to perform the actual work, this appears to be an IO-bound situation.

So I thought about using multiple threads, each of which uses its own Backend instance. I could create these threads manually and connect them via queues. The following is an example implementation with some Backend mock class:

import os
import pty
from queue import Queue
from subprocess import PIPE, Popen
from threading import Thread


class Backend:
    def __init__(self):
        f, g = pty.openpty()
        self.process = Popen(
            ['bash'],  # example program
            text=True, bufsize=1, stdin=PIPE, stdout=g)
        self.write = self.process.stdin.write
        self.read = os.fdopen(f).readline

    def __enter__(self):
        self.write('sleep 2\n')  # startup work
        return self

    def __exit__(self, *exc):
        self.process.stdin.close()
        self.process.kill()

    def run(self, x):
        self.write(f'sleep {x} && echo "ok"\n')  # perform work
        return self.read().strip()


class Worker(Thread):
    def __init__(self, inq, outq, **kwargs):
        super().__init__(**kwargs)
        self.inq = inq
        self.outq = outq

    def run(self):
        with Backend() as backend:
            while True:
                data = self.inq.get()
                result = backend.run(data)
                self.outq.put((data, result))


task_queue = Queue()
result_queue = Queue()

n_workers = 3
threads = [Worker(task_queue, result_queue, daemon=True) for _ in range(n_workers)]
for thread in threads:
    thread.start()

data = [2]*7
for x in data:
    task_queue.put(x)

for _ in data:
    print(f'Result ready: {result_queue.get()}')

Since the Backend needs to perform some work at startup, I don't want to create a new instance for each task. Hence each Worker creates one Backend instance for its whole life cycle. It's also important that each of the workers has its own backend, so they won't interfere with each other.

Now here's the question: Can I also use concurrent.futures.ThreadPoolExecutor to accomplish this? It looks like the Executor.map method would be the right candidate, but I can't figure out how to ensure that each worker receives its own instance of Backend (which needs to be persistent between tasks).

Answer 1

The state of worker threads can be saved in the global namespace, e.g. as a dict. Then threading.current_thread can be used to save/load the state for each of the workers. contextlib.ExitStack can be used to handle Backend appropriately as a context manager.

from concurrent.futures import ThreadPoolExecutor
from contextlib import ExitStack
import os
import pty
from subprocess import PIPE, Popen
import threading


class Backend:
    ...


backends = {}
exit_stack = ExitStack()


def init_backend():
    backends[threading.current_thread()] = exit_stack.enter_context(Backend())


def compute(data):
    return data, backends[threading.current_thread()].run(data)


with exit_stack:
    with ThreadPoolExecutor(max_workers=3, initializer=init_backend) as executor:
        for result in executor.map(compute, [2]*7):
            print(f'Result ready: {result}')

Answer 2

Building on the existing answer from a_guest, you can avoid the usage of globals by leveraging Thread-Local Data with threading.local.

Starting with an example class of:

class Backend:
    def __enter__(self):
        logging.info("Initializing")
        return self

    def __exit__(self, *args):
        logging.info("Exiting")

    def work(self, val):
        logging.info(f"Running {val}")
        return val

We can create a thread local wrapper for it like so:

class BackendContainer:
    def __init__(self):
        self._stack = ExitStack()
        self._locals = threading.local()

    def __enter__(self):
        return self

    def initialize(self):
        self._locals.backend = self._stack.enter_context(Backend())

    def __exit__(self, *args):
        self._stack.close()

    def work(self, val):
        return self._locals.backend.work(val)

Usage:

logging.basicConfig(level=logging.INFO, format="%(threadName)s: %(message)s")

with BackendContainer() as b, ThreadPoolExecutor(3, initializer=b.initialize) as ex:
    for val in ex.map(b.work, range(8)):
        pass

Output:

ThreadPoolExecutor-0_0: Initializing
ThreadPoolExecutor-0_0: Running 0
ThreadPoolExecutor-0_1: Initializing
ThreadPoolExecutor-0_0: Running 1
ThreadPoolExecutor-0_1: Running 2
ThreadPoolExecutor-0_2: Initializing
ThreadPoolExecutor-0_0: Running 3
ThreadPoolExecutor-0_1: Running 4
ThreadPoolExecutor-0_2: Running 5
ThreadPoolExecutor-0_0: Running 6
ThreadPoolExecutor-0_1: Running 7
MainThread: Exiting
MainThread: Exiting
MainThread: Exiting

You could also choose to avoid the initializer method by making the access to backend lazily created:

class BackendContainer:
    def __init__(self):
        self._stack = ExitStack()
        self._locals = threading.local()

    def __enter__(self):
        return self

    @property
    def backend(self):
        if not hasattr(self._locals, "backend"):
            self._locals.backend = self._stack.enter_context(Backend())
        return self._locals.backend

    def __exit__(self, *args):
        self._stack.close()

    def work(self, val):
        return self.backend.work(val)


logging.basicConfig(level=logging.INFO, format="%(threadName)s: %(message)s")

with BackendContainer() as b, ThreadPoolExecutor(3) as ex:
    for val in ex.map(b.work, range(8)):
        pass