-
Notifications
You must be signed in to change notification settings - Fork 25
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add helpers for launching MPI jobs on various platforms #37
base: main
Are you sure you want to change the base?
Conversation
How should we handle the Also, in order for the MPI tests to run, the CI will need to set the |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for working on this. A few comments below.
@@ -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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
.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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It might be useful to allow arguments here. Not every user might know about functools.partial
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Would it make sense to create some rudimentary auto-detection facility (which can be overridden by a documented env var)?
|
||
# {{{ MPI executors | ||
|
||
class MPIExecutorParams: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Dataclass maybe?
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_test
decorator (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.