MPIPool() object can't be passed through multiple functions

Question

I am not very familiar with MPI or schwimmbad, but I'm trying to edit this code that I've inherited and I can't seem to get it to work. The original code uses mpi4py and schwimmbad to pass an MPIPool () instance to a function (from another package, so this part can't be changed). It can also be run with or without MPI installed (i.e. it just runs in series if MPI is not there). For example:

if __name__=='__main__':
    try:
        import mpi4py
        from mpi4py import MPI
        from schwimmbad import MPIPool

        mpi4py.rc.threads = False
        mpi4py.rc.recv_mprobe = False

        comm = MPI.COMM_WORLD
        size = comm.Get_size()

        withmpi = comm.Get_size() > 1
    except ImportError as e:
        print('Failed to start MPI; are mpi4py and schwimmbad installed? Proceeding without MPI.')
        print(f'Message: {e}')
        withmpi = False
    
    # .... other code here .... #
    
    if withmpi:
        with MPIPool() as pool:
            if not pool.is_master():
                pool.wait()
                sys.exit(0)
            do_important_function(other_args, pool)

This code works fine with or without MPI. However, now I want to add a run() function so that I can call this code from another script. For example:

run(other_args, pool):
    # ... other code here ... #
    do_important_function(other_args, pool)

if __name__=='__main__':
    # try/except code here
    if withmpi:
        with MPIPool() as pool:
            if not pool.is_master():
                pool.wait()
                sys.exit(0)
        run(other_args, pool)

This is breaking with a pretty unclear error message (YOUR APPLICATION TERMINATED WITH THE EXIT STRING: Killed (signal 9)). Can anyone tell if there's something wrong with the way I've set this up, or is there possibly a different error that's causing things to go wrong?

I've tried putting the if not pool.is_master() lines in the other function, tried moving the with MPIPool() as pool portion, etc., but no configuration I've tried so far seems to work. The schwimmbad site has one single example of how to use MPIPool() and it's not the scenario I'm trying to work with, so I'm not sure what to do from here. I don't think I technically have to use MPI or schwimmbad at all as long as I pass a multiprocessing pool to the right function, so if you can think of a better way to do this without them feel free to let me know.

Thanks in advance for any pointers!