shmem4py: High-Performance One-Sided Communication for Python Applications

Python is one of the most widely used programming languages in computing today. It owes its popularity to being easy to learn and use, as well as the vast ecosystem of available libraries. In particular, Python has become the de facto language in data science and machine learning. Packages such as NumPy [22], SciPy [51], Pandas [34], Matplotlib [26], and Project Jupyter [30] have been significant contributors to such success. While Python itself is not particularly high-performant due to its interpreted nature, Python applications that utilize libraries written in lower-level languages such as C, C++, and Fortran are able to support excellent performance in a range of scenarios. More recently, just-in-time (JIT) compilation frameworks have made it even more straightforward to achieve performance without leaving the comfort of Python syntax. Python JIT frameworks include Numba [33], JAX [5], Theano [48], and CuPy [39].

Being a language conceived in the 1990s, initial design choices in the default and most popular implementation, CPython, hindered its capability to adapt to the multicore era. CPython's global interpreter lock (GIL) prevented multithreaded applications from taking full advantage of CPUs with multiple computing cores [4]. Due to this limitation, Python has always favored a message-passing and process-based approach to parallel computing.

Since 2008, the Python standard library has provided the multiprocessing package [36] to assist developers with process-based parallelism within a single compute node. More recently, in 2011, the Python standard library added concurrent.futures [41] to make task-based parallelism even easier. Both of these solutions, however, are based on standard operating system facilities for process management and interprocess communication, and therefore are limited to execution within a single compute node.

For distributed memory parallelism across compute node boundaries, Python supports a range of models with various levels of abstraction. Towards the bottom of the stack, the standard networking interfaces available through the socket module from the Python standard library is a feasible, although very cumbersome, approach. The Python wrapper to ZeroMQ [23] is close to socket programming in spirit but with a much richer feature set and a simpler interface. In the scientific community, MPI is the de facto standard for message passing on distributed memory architectures. The most popular Python MPI wrapper is mpi4py [12], which allows for both low-level communication of array-like data and higher-level message passing involving arbitrary Python objects. In the higher level end, a much larger number of solutions is available. Those typically follow the Single Program Multiple Data (SPMD) approach with interprocess communication mostly hidden from the user. Among these packages, we mention Python wrappers for distributed math libraries like Global Arrays [32] and PETSc [13], distributed-memory backends for NumPy [3, 11, 25, 37, 38], and task-based frameworks like Dask [9, 43], Ray [35], and the mpi4py.futures package [44].

In this paper, we describe shmem4py, a framework for using the OpenSHMEM API [6] within Python applications. OpenSHMEM resembles MPI in some ways but specializes in one-sided communication and focuses on features that have native support in high-performance networks. Like mpi4py, shmem4py is a lower-level abstraction. Users are required to invoke communication and synchronization calls explicitly. shmem4py provides the backbone communication capabilities and targets the development of higher-level parallel applications and frameworks.

shmem4py intends to fulfill at least three goals:

1. exposing the high-performance communication capability of OpenSHMEM libraries to Python applications;
   
2. enabling OpenSHMEM users to do rapid prototyping in Python, before switching to lower-level languages (usually C); and
   
3. providing a foundation for building a Python interface to NVSHMEM [24, 40] and related GPU-centric communication libraries.
   

Goal 1 is straightforward; if OpenSHMEM provides better performance than MPI for a particular communication pattern, then it is useful to make it accessible in Python. The need for goal 2 is unclear, but we have found that designing parallel algorithms with Python and mpi4py is much more productive than implementing them from scratch in C/C++ or Fortran. Goal 3 was one of the inspirations for starting the shmem4py effort. NVSHMEM (and its counterparts on other GPU platforms) is arguably the best abstraction for GPU-to-GPU communication. Due to the strong interest in Python and multi-GPU computing for machine learning applications, having direct access to this communication API would allow for fast prototyping of new communication algorithms with direct control over the low-level protocol. We hope for shmem4py to expose NVSHMEM to Python applications in the future.

The latest OpenSHMEM specification defines library functions, constants, variables, and language bindings for the C and C++ programming languages. The first design decision in implementing shmem4py involved choosing, from the many options available, the approach to access OpenSHMEM functionalities in Python. To achieve this, shmem4py uses C Foreign Function Interface for Python (CFFI) [42]. We have considered it an appropriate tool to access a C API with the peculiarities of OpenSHMEM. By providing CFFI with a C header file listing all the types, constants and functions from the OpenSHMEM C API, a Python extension module can be automatically generated. This low-level extension module contains large sets of functions named after the OpenSHMEM operation and the data types they operate on. Higher-level pure Python code then builds a name string out of operations and data types and queries the low-level CFFI extension module. An alternative implementation of shmem4py could have used Cython. However, Cython was designed as a Python language extension and not an automatic wrapper generator. Due to the peculiarities of the OpenSHMEM API, a Cython-based implementation of shmem4py would require large portions of code devoted to function call dispatching based on data types. Nonetheless, those issues could be easily addressed by using code generation with a template engine similar to that used to implement NumPy internals. The choice of CFFI over Cython raises some concerns regarding performance. The runtime data type to function name translation occurs in pure Python code and therefore introduces some minor latency that contributes to overall Python function call overhead. Had we decided to implement shmem4py in Cython, this translation would occur in compiled C code and be significantly faster. Nevertheless, we would expect such overhead to be noticeable only in latency-sensitive applications involving the communication of very small messages. The benchmarks we present in this paper include such scenarios.

In general lines, shmem4py is implemented with a similar philosophy as mpi4py. Rather than hiding low-level APIs, shmem4py exposes as many features as possible to allow for maximum flexibility and provide value to the experienced developer. Rather than building new ad-hoc interfaces, shmem4py tries to stay close to the original OpenSHMEM specification. Nonetheless, the shmem4py interface cannot possibly be a one-to-one C-to-Python translation. The reason for this is twofold: the C and Python programming languages have fundamental differences, and Python users would expect a set of features that adjust to the many language idioms and the existing software ecosystem. OpenSHMEM C API functions have to be invoked with buffer pointers and sizes, while data types are encoded in each function name. On the other hand, shmem4py is NumPy-centric. All shmem4py functions accept NumPy array objects as arguments. Memory addresses, buffer sizes and data types can be determined by introspection of array objects to then dispatch the appropriate low-level OpenSHMEM C routine. This approach allows for a convenient type-generic interface: the Python shmem4py functions operating on NumPy arrays runtime-translate data types to the corresponding low-level OpenSHMEM function names. OpenSHMEM symmetric memory allocations are exposed to Python through NumPy array objects. NumPy memory management is automatic and relies on the Python garbage collector (GC), which executes in a non-deterministic way by design. Oppositely, OpenSHMEM memory management is explicit and collective: allowing the Python GC to perform symmetric memory deallocations could lead to a deadlock state. Therefore, shmem4py requires users to handle symmetric memory deallocations explicitly. The strong reliance on NumPy does not hinder interoperability: NumPy can easily wrap arbitrary memory buffers within its ndarray object type. In fact, such a feature is the one that made possible the current implementation of shmem4py.

As an example of how NumPy is integrated into shmem4py, we sketch a simplified implementation of the shmem.zeros() function, which mimics the interface of the numpy.zeros() function. shmem.zeros() wraps the OpenSHMEM shmem_calloc routine and instantiates a NumPy array object wrapping the symmetric memory buffer.

The shmem4py package supports all major implementations of the OpenSHMEM specification:

• Cray OpenSHMEMX
  
• Open MPI OpenSHMEM (OSHMEM)
  
• Open Source Software Solutions (OSSS) OpenSHMEM
  
• OpenSHMEM Implementation on MPI v2 (OSHMPI)
  
• Sandia OpenSHMEM (SOS)
  

The following C program uses the OpenSHMEM 1.5 API and showcases the use of the SHMEM_FCOLLECT collective.

The corresponding Python program using shmem4py follows. Note the automatic, implicit handling of initialization and finalization of the OpenSHMEM runtime, as well as the NumPy-based exposure of symmetric memory buffers.

We evaluate shmem4py in two dimensions. First, we compare the overhead of Python versus lower-level languages. We achieve this by comparing shmem4py versus existing implementations of the same communication pattern in C on a single node. For some of the benchmarks, we add an implementation in Python with critical kernels compiled with Numba to further highlight the performance difference between interpreted and compiled languages when it comes to arithmetic. In the second dimension, we compare shmem4py to mpi4py within a single node and on a supercomputer. Here, we expect similar conclusions to a comparison of OpenSHMEM to MPI in other languages, although it is useful to verify such a hypothesis.

First, we compare the performance of four OpenSHMEM backends using a PingPong test implemented in Python. This benchmark emphasizes exclusively on communication performance and allows us to limit further analysis to the most performant OpenSHMEM backend. After that, we repeat the PingPong test, this time contrasting the performance between C vs. Python and MPI vs. OpenSHMEM. The goal of this study is to determine the overhead of wrapping C library calls in Python for both MPI and OpenSHMEM, as well as to measure the achievable bandwidth using both frameworks. Subsequently, we benchmark three applications from the Parallel Research Kernels suite [15, 49, 50]. These benchmarks include more complicated communication patterns as well as local computation. We use them to characterize performance in more realistic application scenarios.

6.1 PingPong - OpenSHMEM backends

In this experiment, we compare shmem4py using four different OpenSHMEM backends. We report the highest bandwidth observed during 100 repetitions of the benchmark. For convenience, we run the benchmarks within Podman containers. For validation purposes, we also compare to one backend running natively on the host operating system.

The results of the experiment are shown in Fig. 1. Sandia OpenSHMEM (SOS) underperforms for large messages. Open MPI SHMEM (OSHMEM), on the other hand, performs worse than other backends for messages smaller than ≈ 2 MiB. OSHMPI, running both natively and inside a Podman container, performs almost identically, and its achieved performance is among the highest regardless of the message size. Similar results can be reproduced using C implementations, i.e., the discrepancies point to differences between OpenSHMEM backends and not shmem4py. Based on these results, we use OSHMPI for the remaining single-node experiments.

6.2 PingPong - C vs. Python and MPI vs. OpenSHMEM

In this experiment, we compare the PingPong maximum achievable bandwidth repeating most of the setup of Section 6.1. This time, we fix the OpenSHMEM backend to OSHMPI. We test the following variants: C+MPI, C+OpenSHMEM, Python+mpi4py, and Python+shmem4py. Again, we report the fastest of 100 iterations.

The results are shown in Fig. 2. As in earlier work analyzing mpi4py performance, there is no significant difference between C and Python implementations for large message sizes [14]. OpenSHMEM behaves similarly. Interestingly, we notice mpi4py outperforms C+MPI for some message sizes. Our analysis points to cache effects caused by different memory allocation strategies (NumPy's use of madvise vs. C's use of plain malloc). For small messages, we see increased latency when using communication libraries through Python wrappers. The ratio is about the same for both mpi4py and shmem4py.

Overall, we note a rather large performance advantage of OpenSHMEM over MPI for large messages. For this intra-node communication case, such an outcome is expected. OpenSHMEM naturally involves a single buffer copy from the source memory space to the target shared memory segment. On the other hand, the MPI backend in use (MPICH) implements the send-receive pair with a double copy using an intermediate shared memory buffer.² Note that an implementation based on MPI-3 Remote Memory Access (RMA) features is expected to achieve similar performance as OpenSHMEM.

It is also worth noting that the OpenSHMEM implementation gets very close to the practically achievable memory bandwidth. The dashed line in Fig. 2 shows the fastest aggregate memcpy bandwidth measured over 100 repetitions using two processes and large memory buffer sizes of 500 MB.

6.3 Parallel Research Kernels

In this section, we present the results obtained running three applications from the Parallel Research Kernels suite, as briefly described in Section 5.1. In an inner loop, we execute each kernel for 4 iterations and measure the average performance of the last 3 iterations. In an outer loop, each kernel runs 5 times, and we report the fastest of the inner averages.

As shown in Fig. 3, 4, and 5, the trends are similar across all three kernels. Generally, kernels written in C perform the best. Pure Python implementations may be up to four orders of magnitude slower than C. This huge performance gap is typical of pure Python codes performing arithmetic operations in inner loops. The introduction of Numba to JIT compile the triple nested loop containing the update to the stencil values increases the performance dramatically. This can be seen in Fig. 3 – for the rightmost point, the gap between Python (388 MFLOPS/s) and C (128 GFLOPS/s) is huge, however using Numba immediately increases Python's performance to 79 GFLOPS/s. Similar observations can be made when analyzing the results of the Synch_p2p benchmark (Fig. 4).

Overall, we do not notice large differences between the performance achieved using MPI vs. OpenSHMEM. MPI is slightly faster for kernels implemented in C, partly due to PRK tests being a good match for MPI operations. Moreover, the OpenSHMEM implementations in C are based on the MPI ones and thus may not take advantage of the different programming model OpenSHMEM favors. When using Python, either with or without Numba, the results are mixed. Synch_p2p is a perfect use case for MPI send-receive, and thus it is not surprising to see MPI winning there. With the Stencil benchmark, there is a small advantage for OpenSHMEM, which is consistent with the expectation that it has a lower overhead relative to MPI. With the Transpose benchmark, the behavior is nearly identical, which is expected since both implementations use the all-to-all collective operation.

The results shown for the Transpose benchmark in Fig. 5 are the least remarkable and should be considered as a validation case. Python and C implementations perform very similarly, and Numba did contribute to increasing the performance (results not included). This shows that the Transpose operation, which is the local computation, is implemented efficiently in NumPy. As NumPy's transpose operation is already implemented in C, there is no room for improvement from Numba's JIT compiler. MPI and OpenSHMEM also performed similarly. Both implementations use an all-to-all collective operation, and in the particular case of OSHMPI, the underlying MPI's all-to-all operation is ultimately used.

In this section, we present multi-node experiments performed on a Cray XC40 supercomputer using the same three Parallel Research Kernels as in Section 6. We use the OpenSHMEM backend recommended on this system, that is, the vendor-provided Cray OpenSHMEMX. We do not compare programming languages as the performance implications were already covered in our single-node experiments. Instead, we focus the analysis on comparing mpi4py and shmem4py.

All experiments are run using 32 processes per node, that is, fully subscribing all the available cores. In order to account for job placement on the supercomputer, all benchmarks are repeated in 10 SLURM jobs. Within each job, both mpi4py and shmem4py variants of the code are executed. PRK kernels are then executed 11 times, and the average time over the last 10 iterations is considered. We then report the fastest average obtained in any job for a given number of processes. The Transpose benchmark result is missing the 32,768 process data point, as it requires a matrix with a number of rows and columns divisible by the number of processes.

The results are presented in Fig. 6. We observe that the results vary between benchmarks. For Transpose, the OpenSHMEM implementation typically outperforms the one based on MPI. For 8,192 processes, we note the largest difference with an aggregate bandwidth of 619 GB/s (shmem4py) compared to 314 GB/s (mpi4py). This suggests that the all-to-all collective operation in Cray OpenSHMEMX outperforms the one in Cray MPICH. On the other hand, in the Synch_p2p benchmark, MPI has a large advantage over OpenSHMEM (7.06 GFLOPS/s vs. 1.64 GFLOPS/s for 32,768 processes). Arguably, MPI is not only faster, but it also provides a better communication model for the task. For the Stencil benchmark, the difference was the smallest and did not exceed 11.8%.

In this work, we introduced shmem4py - a Python wrapper for the OpenSHMEM standard specification. We briefly described the landscape of distributed memory parallelism in Python and shmem4py’s niche next to mpi4py. We then discussed the design choices and shmem4py’s implementation details. We also showed an example contrasting the use of OpenSHMEM features from both C and Python. In the second part of the paper, we analyzed the performance of shmem4py on both a workstation and a supercomputer. There, we compared the performance of different OpenSHMEM backends. Additionally, we discussed the overhead of using Python wrappers such as shmem4py and mpi4py vs. accessing the frameworks directly from C. A number of benchmarks from the PRK suite provided a baseline to compare the performance of shmem4py and mpi4py. However, the benchmarks used do not take advantage of the unique properties of OpenSHMEM, nor do they compare OpenSHMEM to MPI RMA (one-sided), because these issues are well understood outside of the Python language context. Our experiments primarily serve to validate the design, to ensure that the overheads from the Python wrapper are negligible, and that common algorithm patterns are straightforward to implement.

Overall, we obtained mixed results. We were able to identify scenarios in which OpenSHMEM was up to 2.0 × faster than MPI. In other cases, MPI was able to outperform OpenSHMEM by a factor of 4.3 ×. Overall, the performance of OpenSHMEM and MPI and their respective Python wrappers is on par, and the differences are mostly unremarkable. Nonetheless, we consider the outcome of these experiments as positive and ultimately interesting for the practitioner community at large. As both frameworks perform similarly, the choice of a programming model can be made purely on the basis of what suits the application better.