diff --git a/.github/workflows/notebook-tests.yml b/.github/workflows/notebook-tests.yml index 8ef3c22..53e1dbb 100644 --- a/.github/workflows/notebook-tests.yml +++ b/.github/workflows/notebook-tests.yml @@ -15,7 +15,7 @@ permissions: jobs: run-openmpi: - runs-on: ubuntu-latest + runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v3 @@ -26,7 +26,7 @@ jobs: - name: Install dependencies run: | sudo apt-get update - sudo apt-get install -y openmpi-bin + sudo apt-get install -y openmpi-bin libopenmpi-dev python -m pip install --upgrade pip pip install . pip install nbformat mpi4py @@ -37,7 +37,7 @@ jobs: run-local: - runs-on: ubuntu-latest + runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v3 @@ -56,7 +56,7 @@ jobs: run-dask-cluster: - runs-on: ubuntu-latest + runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v3 diff --git a/examples/osu_bw_cupy.py b/examples/osu_bw_cupy.py new file mode 100644 index 0000000..f0bdb46 --- /dev/null +++ b/examples/osu_bw_cupy.py @@ -0,0 +1,89 @@ +# Adapted from https://github.com/mpi4py/mpi4py/blob/master/demo/osu_bw.py +# and https://mpi4py.readthedocs.io/en/stable/tutorial.html#cuda-aware-mpi-python-gpu-arrays +# http://mvapich.cse.ohio-state.edu/benchmarks/ + +from mpi4py import MPI +import cupy as cp +import os + + +def osu_bw( + BENCHMARH = "MPI G2G Bandwidth Test", + skip = 10, + loop = 100, + window_size = 64, + skip_large = 2, + loop_large = 20, + window_size_large = 64, + large_message_size = 8192, + MAX_MSG_SIZE = 1<<22, + ): + + comm = MPI.COMM_WORLD + myid = comm.Get_rank() + numprocs = comm.Get_size() + + cp.cuda.Device(myid).use() + + if numprocs != 2: + if myid == 0: + errmsg = "This test requires exactly two processes" + else: + errmsg = None + raise SystemExit(errmsg) + + s_buf = cp.arange(MAX_MSG_SIZE, dtype='i') + r_buf = cp.empty_like(s_buf) + cp.cuda.get_current_stream().synchronize() + + if myid == 0: + print ('# %s' % (BENCHMARH,)) + if myid == 0: + print ('# %-8s%20s' % ("Size [B]", "Bandwidth [MB/s]")) + + message_sizes = [2**i for i in range(30)] + for size in message_sizes: + if size > MAX_MSG_SIZE: + break + if size > large_message_size: + skip = skip_large + loop = loop_large + window_size = window_size_large + + iterations = list(range(loop+skip)) + window_sizes = list(range(window_size)) + requests = [MPI.REQUEST_NULL] * window_size + # + comm.Barrier() + if myid == 0: + s_msg = [s_buf, size, MPI.BYTE] + r_msg = [r_buf, 4, MPI.BYTE] + for i in iterations: + if i == skip: + t_start = MPI.Wtime() + for j in window_sizes: + requests[j] = comm.Isend(s_msg, 1, 100) + MPI.Request.Waitall(requests) + comm.Recv(r_msg, 1, 101) + + t_end = MPI.Wtime() + elif myid == 1: + s_msg = [s_buf, 4, MPI.BYTE] + r_msg = [r_buf, size, MPI.BYTE] + for i in iterations: + for j in window_sizes: + requests[j] = comm.Irecv(r_msg, 0, 100) + MPI.Request.Waitall(requests) + comm.Send(s_msg, 0, 101) + # + if myid == 0: + MB = size / 1e6 * loop * window_size + s = t_end - t_start + print ('%-10d%20.2f' % (size, MB/s)) + + + cp.allclose(s_buf, r_buf) + + +if __name__ == '__main__': + osu_bw() diff --git a/examples/osu_bw_mpi4py-cupy.ipynb b/examples/osu_bw_mpi4py-cupy.ipynb new file mode 100644 index 0000000..22474ba --- /dev/null +++ b/examples/osu_bw_mpi4py-cupy.ipynb @@ -0,0 +1,135 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "86b84834-b048-4832-bd49-f49c29d4fdfb", + "metadata": { + "tags": [] + }, + "source": [ + "## OSU G2G Bandwidth Benchmark with MPI4Py\n", + "In this example we use [IPCMagic](https://github.com/eth-cscs/ipcluster_magic/tree/master) to run a test from the [OSU Bandwidth benchmark](http://mvapich.cse.ohio-state.edu/benchmarks/) with MPI4Py from a Jupyter notebook.\n", + "Using [this example](https://mpi4py.readthedocs.io/en/stable/tutorial.html#cuda-aware-mpi-python-gpu-arrays), we adapted the [osu_bw.py](https://github.com/mpi4py/mpi4py/blob/d0228f0397403ff73d8f41d90d97b411efda6128/demo/osu_bw.py) script from the MPI4Py repository so it uses arrays allocated on the GPU." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b55135ce-b328-4d73-b498-d8369ead8380", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import ipcmagic" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "64b6492d-1dbd-441d-95fa-7eb1717e0bdb", + "metadata": {}, + "outputs": [], + "source": [ + "# os.environ['MPICH_RDMA_ENABLED_CUDA'] = '1' # Enable direct communication between GPUs\n", + "os.environ['MPICH_GPU_SUPPORT_ENABLED'] = '1'" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "79928bb7-8091-458f-b438-b66e1af23a8b", + "metadata": {}, + "outputs": [], + "source": [ + "%ipcluster --version" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a7d237d6-f385-49c6-9bad-323ba3c2d7a6", + "metadata": {}, + "outputs": [], + "source": [ + "%ipcluster start -n 2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cbaea7a2-4171-438d-8ad9-3a00411fd007", + "metadata": {}, + "outputs": [], + "source": [ + "# Disable IPyParallel's progress bar\n", + "%pxconfig --progress-after -1" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12b975be-ceed-412d-9117-b3d8b7904549", + "metadata": {}, + "outputs": [], + "source": [ + "%%px\n", + "import socket\n", + "\n", + "socket.gethostname()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d81d6c89-d339-45ac-aff7-88a11558ce15", + "metadata": {}, + "outputs": [], + "source": [ + "%%px\n", + "from osu_bw_cupy import osu_bw" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8ebf8f80-37b1-4d74-a911-33c1a32a461d", + "metadata": {}, + "outputs": [], + "source": [ + "%%px\n", + "osu_bw()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1cb656f2-2367-405a-9173-d79fe53281a0", + "metadata": {}, + "outputs": [], + "source": [ + "%ipcluster stop" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "test env 2", + "language": "python", + "name": "test-env-2" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}