Embracing Dynamism: Control State Serialization for High-Performance Python

Python's ease of use has driven its adoption in data science, machine learning, and increasingly, high-performance computing (HPC), but its performance lags due to its dynamic nature. While many efforts accelerate Python by restricting its features, this work leverages Python's dynamism to accelerate it. We introduce control state serialization to Python through Sauerkraut, a library that captures the complete execution state (call stack, instruction pointers, operand stacks, local variables, global context) of running functions, complementing existing data serialization. Sauerkraut enables snapshots of function execution to be serialized, transferred, and resumed later or elsewhere. Sauerkraut is compatible with off-the-shelf Python installations. We demonstrate its utility by building Kombucha, a partial MPI implementation that provides general-purpose load balancing by migrating virtualized MPI ranks between processes using Sauerkraut to transfer control state. Evaluations show Sauerkraut adds minimal overhead over standard serialization, and Kombucha achieves significant speedups (up to 2.21x for a CPU Particle-in-Cell code and 1.43x for a GPU Jacobi3D code with synthetic imbalance) with minimal application code changes. Control state serialization opens new avenues for performance optimization in Python, including load balancing, checkpoint/restart, and replay debugging, enhancing Python's suitability for the HPC community.