60 lines
1.6 KiB
Python
60 lines
1.6 KiB
Python
'''
|
|
Integration test: spawning tractor actors from an MPI process.
|
|
|
|
When a parent is launched via ``mpirun``, Open MPI sets
|
|
``OMPI_*`` env vars that bind ``MPI_Init`` to the ``orted``
|
|
daemon. Tractor children inherit those env vars, so if
|
|
``inherit_parent_main=True`` (the default)
|
|
the child re-executes ``__main__``, re-imports ``mpi4py``, and
|
|
``MPI_Init_thread`` fails because the child was never spawned by
|
|
``orted``::
|
|
|
|
getting local rank failed
|
|
--> Returned value No permission (-17) instead of ORTE_SUCCESS
|
|
|
|
Passing ``inherit_parent_main=False`` and placing RPC functions in a
|
|
separate importable module (``_child``) avoids the re-import
|
|
entirely.
|
|
|
|
Usage::
|
|
|
|
mpirun --allow-run-as-root -np 1 python -m \
|
|
examples.integration.mpi4py.inherit_parent_main
|
|
|
|
'''
|
|
|
|
from mpi4py import MPI
|
|
|
|
import os
|
|
import trio
|
|
import tractor
|
|
|
|
from ._child import child_fn
|
|
|
|
|
|
async def main() -> None:
|
|
'''
|
|
Spawn an MPI-safe child without replaying the parent main.
|
|
|
|
'''
|
|
rank: int = MPI.COMM_WORLD.Get_rank()
|
|
pid: int = os.getpid()
|
|
print(f'[parent] rank={rank} pid={pid}', flush=True)
|
|
|
|
an: tractor.ActorNursery
|
|
async with tractor.open_nursery(start_method='trio') as an:
|
|
portal: tractor.Portal = await an.start_actor(
|
|
'mpi-child',
|
|
enable_modules=[child_fn.__module__],
|
|
# Without this the child replays __main__, which
|
|
# re-imports mpi4py and crashes on MPI_Init.
|
|
inherit_parent_main=False,
|
|
)
|
|
result: str = await portal.run(child_fn)
|
|
print(f'[parent] got: {result}', flush=True)
|
|
await portal.cancel_actor()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
trio.run(main)
|