Using PyO3 to get around slow multiprocessing "wait" calls

[matroscoe]

I am wondering if there is a way to use PyO3 and rust to get around very slow multiprocessing wait calls. The processes that are doing the wait are doing it on a Python Barrier object. I do NOT need safe memory access or synchronization for stored data I do need the spawned processes to stay temporally in sync with each other. I was hoping I could leverage PyO3 to perform some crazy basic IPC (time sync signal only) between the spawned processes. Does this approach make any sense it is based mostly on rust being generally faster than Python so looking to see if I can leverage it for this rater odd need for process synchronization.

[adamreichold]

Are you sure the performance issue is how Barrier is implemented and not due to the communication pattern itself? From a first guess, I would not expect manually re-implementing the necessary parts of Barrier's IPC manually using Rust to improve the speed of the operation itself. (It is most likely already partially implemented in C and you would probably end up issuing similar system calls.)

If you say that the processes do not exchange data but only need to synchronize, that sounds like it would be particularly likely that the overhead of synchronization is high no matter which mechanism is applied. Could you provide a small reproducer so that the design problem could be discussed here? From what you have written so far, it does not sound likely that replacing the IPC mechanism with something written in Rust would make a difference.

[matroscoe]

I can it is attached here as a runnable python script (well txt because GH won't let me post .py) and will post it below for ease of discussion. I have also attached a performance log captured with py-side that can be opened in the amazing speedscope tool.
base-example (profile as .svg)
base-example.txt

In this example case, the performance is NOT as bad as when I run in the actual software so I need to track down why. However, I still see upwards of 40% of the execution time is spent on wait when looking at Process 51862 in left heavy mode. (11.28 s for run, ~8.1s for all the various waits). The amount of wait on 51863 is MUCH higher and closer to what I see in my code. Any insights you have on what I could be doing wrong or if I could leverage rust to work around the sync slowness would be appreciated!

import multiprocessing
from threading import Thread
from multiprocessing import Process, Barrier, Value
from loguru import logger
from time import sleep
import psutil


class SyncControl(Thread):

    def __init__(self, sync_var: Barrier):
        super().__init__()
        self.__sync_var: Barrier = sync_var
        self.__continue_running = Value('b', False)
        self.__timeout_count: int = 0

    def stop(self):
        logger.warning("Stopping SyncControl")
        self.__continue_running.value = False

    def run(self) -> None:
        logger.info("Starting SyncControl")
        self.__continue_running.value = True

        while self.__continue_running.value:

            try:
                self.__sync_var.wait(timeout=1)

            except:
                logger.error("BrokenBarrierError")
                self.__timeout_count += 1

                if self.__timeout_count <= 3:
                    continue
                else:
                    self.stop()

        logger.warning("SyncControl has finished")


class LongRunningProcess(Process):

    def __init__(self, sync_var: Barrier, cpu_core: int):
        super().__init__()
        self.__sync_var: Barrier = sync_var
        self.__cpu_core: int = cpu_core
        self.__continue_running = Value('b', False)
        self.__timeout_count: int = 0

    def start(self):
        super().start()
        process = psutil.Process()
        logger.warning(f"Core {self.__cpu_core} - Current CPU affinity: {process.cpu_affinity()}")
        process.cpu_affinity([self.__cpu_core])
        logger.warning(f"Core {self.__cpu_core} - Updated CPU affinity: {process.cpu_affinity()}")

    def stop(self):
        logger.warning(f"Core {self.__cpu_core} - Stopping LongRunningProcess")
        self.__continue_running.value = False

    def run(self) -> None:
        logger.info(f"Core {self.__cpu_core} -  LongRunningProcess")
        self.__continue_running.value = True

        while self.__continue_running.value:

            try:
                self.__sync_var.wait(timeout=1)

                self.__perform_logic()

            except:
                logger.error(f"Core {self.__cpu_core} - BrokenBarrierError")
                self.__timeout_count += 1

                if self.__timeout_count <= 3:
                    continue
                else:
                    self.stop()

        logger.warning(f"Core {self.__cpu_core} - LongRunningProcess has finished")

    def __perform_logic(self):
        # debugging check to make sure it is doing stuff:
        # logger.debug(f"Core {self.__cpu_core} - I am doing work")

        # The sleep time here is representative of the time the work the LongRunningProcess takes (~50μs)
        sleep(0.00005)


if __name__ == "__main__":

    sync_var: Barrier = Barrier(3)

    sync_control: SyncControl = SyncControl(sync_var=sync_var)

    os_cores: list[int] = [3, 4]
    long_running_processes: list[LongRunningProcess] = []

    # Create LongRunningProcesses:
    for core in os_cores:
        long_running_processes.append(LongRunningProcess(sync_var=sync_var, cpu_core=core))

    # Start LongRunningProcesses:
    for temp in long_running_processes:
        temp.start()

    # Start the StepControl:
    sync_control.start()

    # Let it do some "work":
    sleep(30)

    # Stop the LongRunningProcesses:
    for temp in long_running_processes:
        temp.stop()
        temp.join()

    # Stop the StepControl:
    sync_control.stop()
    sync_control.join()

    logger.info("Example is done!")

[adamreichold]

From my point of view, your problem is to have a group of processes (maybe one per CPU core) which need to rendezvous every 50µs. The profile you captured does not seem to include the C code executed by the multiprocessing implementation but I suspect that the time spent in those wait calls is mostly spent blocking in the system calls instead of executing Python code.

For example, when I run the above code, I see only about 15% CPU utilization from those processes which means 85% of the time is spent off-CPU waiting to become scheduled again. Put differently, with your current communication pattern, you could waste another 85% of a CPU executing Python byte code before that would become your bottleneck. (This is likely too simplistic as wasting time before actually blocking on the system calls affect all processes, not just the ones faster than the slowest one which determines the actual rendezvous time.)

I therefore think your problem is not Python being slow and just re-implementing the synchronization mechanism in Rust would not help. You could go port the whole program to Rust, use multi-threading and exercise the resulting fine-control over the synchronization mechanism. But I suspect that even this would not give you orders of magnitude improvements (which is what porting CPU-bound Python code to Rust often yields).

I think the best thing to improve your program would be to do more work before having to synchronize or at least further partition the workers, i.e. not have every process synchronize with every other. If and how that can be done is most likely specific to your problem and is a common issue when writing parallel simulations. Often it is preferable to do more work (redundantly or speculatively) if it means avoiding communication for a bit longer. Sometimes, it is also important to make the workload as balanced as possible as the slowest process will determine a rendezvous like the one you described. These issues are also mostly independent of the implementation language.

(A specific technology that is often used in these settings and therefore tuned to this is MPI. You could try that via the mpi4py package and an MPI implementation tuned for your system, probably OpenMPI if running on your personal computer. But as written above, I would not expect orders of magnitude improvements from swapping out multiprocessing.Barrier for mpi4py.MPI.Comm.barrier.)

[matroscoe]

Thanks, I was thinking that could be the case and I increased the amount of "work" (sleep) that was occurring, and the amount of time spent waiting decreased dramatically. In the end, the system won't be running at 50us update rates it is just what my current testing allowed me to see. Your very well-written and detailed answers have helped out a ton!

Not that this is the right forum but I have to ask anyway is there a general idea of a maximum update rate I could expect to get Python to realistically be able to synchronize with (1ms? 5ms?) if the CPU usage is higher? At that point, if I start to hit a performance bottleneck and cannot solve it by adding more CPUs would that be the point to consider a re-implementation of some of the more "expensive" code in a more performant language such as rust?

[davidhewitt]

Not that this is the right forum but I have to ask anyway is there a general idea of a maximum update rate I could expect to get Python to realistically be able to synchronize with (1ms? 5ms?) if the CPU usage is higher?

I don't think there is a rule of thumb that would just give you this answer. Benchmark it at different update rates and see what works best!