pyiron / pympipool

up-scale python functions for high performance computing
https://pympipool.readthedocs.io
BSD 3-Clause "New" or "Revised" License
11 stars 3 forks source link

Add option to specify resources during submit() call rather than on the executor level #293

Closed jan-janssen closed 1 month ago

jan-janssen commented 1 month ago

This pull request introduces a new class of interfaces PyFluxStepExecutor, PyMPIStepExecutor and PySlurmStepExecutor. Each of these allow to submit a function by specifying the resource requirements on the function rather than on the executor.

Example:

import numpy as np
from pympipool.mpi import PyMPIStepExecutor

def calc(i, j, k):
    from mpi4py import MPI
    size = MPI.COMM_WORLD.Get_size()
    rank = MPI.COMM_WORLD.Get_rank()
    return np.array([i, j, k]), size, rank

with PyMPIStepExecutor(max_cores=2) as p:
    fs = p.submit(calc, 2, j=4, k=3, resource_dict={"cores": 2})
    print(fs.result())

Output:

>>> [(array([2, 4, 3]), 2, 0), (array([2, 4, 3]), 2, 1)]
mgt16-LANL commented 1 month ago

I think there's a good chance this would work. When I have a chance I want to download and try. I think the thing that really stands out to me from my initial pass was the ability to do multi-core MPI "jobs" locally. That's SUPER useful for local dev.