Add helpers for launching MPI jobs on various platforms #37
Conversation
|
How should we handle the Also, in order for the MPI tests to run, the CI will need to set the |
| @@ -29,3 +31,184 @@ def run_with_mpi_ranks(py_script, ranks, callable_, args=(), kwargs=None): | |||
| check_call(["mpirun", "-np", str(ranks), | |||
| sys.executable, py_script, "--mpi-relaunch", callable_and_args], | |||
| env=newenv) | |||
|
|
|||
|
|
|||
| # {{{ MPI executors | |||
There was a problem hiding this comment.
Could you link this into the documentation? (There might not be an MPI section yet, please create one. :)
| from functools import partial | ||
|
|
||
|
|
||
| def get_test_mpi_executor(): |
There was a problem hiding this comment.
Do these tests run as part of CI? If not I'd like them to.
|
|
||
|
|
||
| class MPIExecutor(metaclass=abc.ABCMeta): | ||
| """Base class for a general MPI launcher.""" |
There was a problem hiding this comment.
Add
.. automethod:: xxxy
to the the class docstring for each method.
| """Calls *func* with MPI. Note: *func* must be picklable.""" | ||
| import pickle | ||
| calling_code = ('import sys; import pickle; pickle.loads(bytes.fromhex("' | ||
| + pickle.dumps(func).hex() + '"))()') |
There was a problem hiding this comment.
.hex() is not an efficient encoding. Base64 is better.
| if exit_code != 0: | ||
| raise MPIExecError(exit_code) | ||
|
|
||
| def call(self, func, exec_params=None): |
There was a problem hiding this comment.
It might be useful to allow arguments here. Not every user might know about functools.partial.
There was a problem hiding this comment.
Is there a convention for passing in the argument list? As a dict? Or something else?
| raise MPIExecError(exit_code) | ||
|
|
||
| def call(self, func, exec_params=None): | ||
| """Calls *func* with MPI. Note: *func* must be picklable.""" |
There was a problem hiding this comment.
Document exec_params. (Here and elsewhere, by reference "See :meth:\...`` for exec_params" OK.)
| "slurm": SlurmMPIExecutor, | ||
| "lclsf": LCLSFMPIExecutor | ||
| } | ||
| return type_name_map[executor_type_name]() |
There was a problem hiding this comment.
Would it make sense to create some rudimentary auto-detection facility (which can be overridden by a documented env var)?
|
|
||
| # {{{ MPI executors | ||
|
|
||
| class MPIExecutorParams: |
Sure. You can hack it in yourself in |
For the CI? No big deal IMO. Just install it. But pytools itself should not declare a hard dependency on mpi4py. Also, maybe this stuff wants to be a separate package. What do you think? |
I made an attempt: link |
This has been generalized a little more from what you saw last week, mainly to make it easier to support functionality that exists on some platforms but not on others (like GPUs). I think it's now at a point where we can start using it and add on to it as needed.
The meshmode usage is here. I ended up deciding not to do the
@mpi_testdecorator (though I did get it working). I think it would just cause confusion, since there are things you can do in a normal test that you can't do in an MPI-launched test (pytest.skip()for example).I tested this on my Mac, quartz, and lassen; it behaves as expected on all three.