Building a basic multiprocess worker in python

I’ve recently been working quite a bit on a worker application in python. This worker picks up tasks over a gRPC service executes those tasks, and publishes the results via gRPC. Because of python’s Global Interpreter Lock (GIL), threads don’t help for CPU intensive workloads. Instead, I’m using multiprocessing to run additional processes per worker and get better CPU utilization at the cost of using a ton of memory. When a python process opens a child process it can also share state with the child process. The most interesting of these are synchronization and communication primitives like Event, Queue, Pipe, and Lock. At a high level the worker’s workflow is:

Theoretically, N is limited by how much memory and cores you have. In practice though, I’ve found I had diminishing results with N over 16. More than that, and I found my processes were wasting a lot of memory, and not getting much additional throughput as multiprocessing is using sockets under the hood to pass messages between processes, and there is serialization and locking overhead with those constructs.

Fetching work

While my worker fetches tasks via a gRPC API, the source of work can be a variety of things, HTTP requests, kafka message, MQ message etc. We’ll call this a TaskClient. Our worker uses the client to fetch one or more tasks from its underlying data stream. For this prototype, we’ll start with something very simple:

Show Plain Text
  1. class StubClient:
  2.     def __init__(self, count: int = 100):
  3.         self._remaining = count
  4.         self._updates: dict[str, str] = {}
  5.  
  6.     def fetch_task(self):
  7.         # This could be an RPC call that returns a more complex
  8.         # object. I've used a simple dictionary as a prototype
  9.         if self._remaining <= 0:
  10.             return None
  11.         self._remaining -= 1
  12.         return {
  13.             "task_id": uuid.uuid4().hex,
  14.             "name": "echo",
  15.             "payload": f"task # {self._remaining + 1}"
  16.         }
  17.  
  18.     def update_task(self, task_id: str, status: str):
  19.         # This could be an RPC call that returns a result object
  20.         self._updates[task_id] = status

This is just a stub implementation, but this could easily have a gRPC, HTTP, or SQL backend. Next we have the bulk of our logic the WorkerManager. This class has a few responsibilities:

1. Using the client to fetch tasks and publish results.
2. Manage the pool of worker processes.
3. Handle shutdown.

Our WorkerManager needs a constructor and a way to start its processes up:

Show Plain Text
  1. class WorkerManager:
  2.     def __init__(self, client: StubClient, process_count: int = NUM_PROCESS, queue_size: int = 10) -> None:
  3.         self._client = client
  4.         self._process_count = process_count
  5.  
  6.         # One queue to send work to the processes, and another to receive results.
  7.         self._process_work: multiprocessing.Queue = multiprocessing.Queue(maxsize=queue_size)
  8.         self._process_results: multiprocessing.Queue = multiprocessing.Queue(maxsize=queue_size)
  9.  
  10.         # We could use multiprocessing.ProcessPool, but I learned more doing it the hard way.
  11.         self._processes: list[multiprocessing.Process] = []
  12.  
  13.         # We want to publish results on a thread, and keep our main thread for fetching tasks.
  14.         self._result_thread: threading.Thread | None = None
  15.  
  16.         # We'll use this to shut down our result_thread later.
  17.         self._shutdown = threading.Event()
  18.  
  19.     def start_processes(self) -> None:
  20.         # Check existing workers for aliveness, and spawn more minions as required.
  21.         processes = [p for p in self._processes if p.is_alive()]
  22.         if len(processes) == self._process_count:
  23.             return
  24.  
  25.         while len(processes) < self._process_count:
  26.             # I'll cover run_task later
  27.             p = multiprocessing.Process(target=run_task, args=(self._process_work, self._process_results))
  28.             processes.append(p)
  29.             p.start()
  30.  
  31.         self._processes = processes

I’m using a bunch multiprocessing objects here. The Queue instances let us communicate with our worker processes. These queues should always be bounded. Eventually appending to them will block which provides some natural backpressure and helps keep memory usage bounded. Once work has been fetched, it needs to be appended to then process_work queue, which each of our worker processes consumes from. We add the following to WorkerManager

Show Plain Text
  1. def _fetch_tasks(self) -> None:
  2.         task = self._client.fetch_task()
  3.         if not task:
  4.             # Backoff and wait for task to arrive
  5.             time.sleep(1)
  6.             return
  7.  
  8.         # Could fail with queue.Full
  9.         self._process_work.put(task)

We also need to start our result thread, have a run loop, and shutdown for our worker:

Show Plain Text
  1. def _start_send_result_thread(self) -> None:
  2.         # This function will run in a thread. While threads in python
  3.         # aren't great when you want concurrent CPU work, they are ok on IO
  4.         # tasks like network.
  5.         def send_result() -> None:
  6.             while True:
  7.                 try:
  8.                     # Read results from worker processes.
  9.                     result = self._process_results.get(timeout=1.0)
  10.                     self._client.update_task(result["task_id"], result["status"])
  11.                 except queue.Empty:
  12.                     # We have drained the queue and we've been shutdown.
  13.                     if self._shutdown.is_set():
  14.                         return
  15.                     continue
  16.  
  17.         self._result_thread = threading.Thread(
  18.             name="send-result", target=send_result, daemon=True
  19.         )
  20.         self._result_thread.start()
  21.  
  22.     def run(self):
  23.         self.start_processes()
  24.         self._start_send_result_thread()
  25.  
  26.         # When the process gets SIGINT/SIGTERM try to shutdown cleanly.
  27.         atexit.register(lambda: self.shutdown())
  28.  
  29.         while True:
  30.             self._fetch_tasks()
  31.             # Simulate latency with sleep
  32.             time.sleep(0.05)
  33.  
  34.             # To handle if a processes dies, while not possible in this prototype,
  35.             # it could happen as workers could oom, or crash.
  36.             self.start_processes()
  37.  
  38.     def shutdown(self) -> None:
  39.         # Trigger shutdown to synchronize result thread
  40.         self._shutdown.set()
  41.  
  42.         # Terminate and wait for all the child processes to die.
  43.         for p in self._processes:
  44.             p.terminate()
  45.             p.join()
  46.  
  47.         # Wait for the result thread to shutdown.
  48.         if self._result_thread:
  49.             self._result_thread.join(timeout=5)

One thing I like about this approach is that we have separated our IO work (fetching tasks and sending results) from our CPU/IO workload (executing tasks). With fetching and result sending in separate threads, we do our best to execute keep the CPU busy. If fetch/update operations are expensive we have a lot of flexibility on batching as well. From my experience if you have long running tasks, fetching them individually is totally fine. In higher throughput scenarios, or where IO is expensive batching can work out better as you can fetch blocks of work, and publish batches of results.

Executing work

The next stage of our worker is executing tasks. Up until now, our run_task function has been a TODO. What we need
is a worker loop that can consume from our process_work queue, and publish results to process_results. Ideally, we look up tasks by name from a registry, and execute them. Task functions contain arbitrary application logic that our worker is executing as background ‘tasks’ or ‘jobs’. My simple implementation, without the task lookup or error handling looks like:

Show Plain Text
  1. def run_task(
  2.     process_work: multiprocessing.Queue,
  3.     process_results: multiprocessing.Queue,
  4. ) -> None:
  5.     # This function needs to be importable from the module scope
  6.     # because of how multiprocessing works.
  7.     while True:
  8.         # Read from the queue in the manager process.
  9.         try:
  10.             task = process_work.get_nowait()
  11.         except queue.Empty:
  12.             continue
  13.  
  14.         # TODO look up tasks from a map of registered tasks
  15.         # and call the user defined function. Catch errors
  16.         # and provide retry behavior or mark tasks failed.
  17.         do_work(task)
  18.  
  19.         result = {
  20.             "task_id": task["task_id"],
  21.             "status": "complete"
  22.         }
  23.         # Publish results back to the main process
  24.         process_results.put(result)
  25.  
  26.  
  27. # Our 'task' to be executed. Ideally we use the task name
  28. # to locate the appropriate task and parse parameters.
  29. def do_work(task):
  30.     print(f"Ran task {task}")

While, I’ve only implemented the most basic solution. In a more robust implementation, I would handle more task outcomes. Ideally tasks complete without errors and are complete. If the task function does raise an error, it is a failure. Within failure, there are failures that we can retry again later and see if we get a successful result. We can detect those scenarios from the task definition and return a retry status.

You can run the worker demo with:

Show Plain Text
  1. def main():
  2.     client = StubClient()
  3.     manager = WorkerManager(client)
  4.     manager.run()
  5.  
  6.  
  7. if __name__ == "__main__":
  8.     main()

After implementing a task registry system, I would also recommend adding task timeouts using with SIGALRM. A system like this also benefits from metrics and logging/tracing so that you can understand how your workloads are operating.

Comments

There are no comments, be the first!

Have your say: