Commit 99105f1c authored by Pierre Augier's avatar Pierre Augier
Browse files

Start to work on MPI

parent 9b4238c9eec6
......@@ -81,6 +81,7 @@ from .util import (
from .annotation import make_signatures_from_typehinted_func
from .compat import open
from . import mpi
modules = {}
......@@ -249,10 +250,10 @@ class CachedJIT:
mod = _get_module_cachedjit(index_frame)
mod.cachedjit_functions[func_name] = self
module_name = mod.module_name
print(f"Make new function to replace {func_name} ({module_name})")
path_pythran = path_cachedjit / module_name.replace(".", os.path.sep)
path_pythran.mkdir(parents=True, exist_ok=True)
if mpi.rank == 0:
path_pythran.mkdir(parents=True, exist_ok=True)
path_pythran = (path_pythran / func_name).with_suffix(".py")
path_pythran_header = path_pythran.with_suffix(".pythran")
......@@ -267,7 +268,7 @@ class CachedJIT:
src = None
if has_to_write:
if has_to_write and mpi.rank == 0:
import_lines = [
line.split("# pythran ")[1]
for line in mod.get_source().split("\n")
......@@ -298,6 +299,8 @@ class CachedJIT:
with open(path_pythran) as file:
src = file.read()
# FIXME MPI all processes have to get these 2 variables
# hash from src (to produce the extension name)
hex_src = make_hex(src)
......
"""Minimalist MPI module
========================
"""
try:
from fluiddyn.util import mpi as _mpi
except ImportError:
try:
from mpi4py import MPI
except ImportError:
rank = 0
else:
comm = MPI.COMM_WORLD
nb_proc = comm.size
rank = comm.Get_rank()
else:
rank = _mpi.rank
nb_proc = _mpi.nb_proc
if nb_proc > 1:
comm = _mpi.comm
def barrier():
if nb_proc > 1:
comm.barrier()
\ No newline at end of file
......@@ -32,6 +32,7 @@ import sysconfig
import hashlib
from .compat import open, implementation
from . import mpi
ext_suffix = sysconfig.get_config_var("EXT_SUFFIX") or ".so"
......@@ -56,6 +57,8 @@ class SchedulerPopen:
deltat = 0.2
def __init__(self, parallel=True):
if mpi.rank > 0:
return
self.processes = []
if parallel:
self.limit_nb_processes = max(1, multiprocessing.cpu_count() // 2)
......@@ -64,6 +67,8 @@ class SchedulerPopen:
def launch_popen(self, words_command, cwd=None, parallel=True):
"""Launch a program (blocking if too many processes launched)"""
if mpi.rank > 0:
return
if parallel:
limit = self.limit_nb_processes
......@@ -86,11 +91,14 @@ class SchedulerPopen:
def wait_for_all_extensions(self):
"""Wait until all compilation processes are done"""
while self.processes:
time.sleep(self.deltat)
self.processes = [
process for process in self.processes if process.poll() is None
]
if mpi.rank == 0:
while self.processes:
time.sleep(self.deltat)
self.processes = [
process for process in self.processes if process.poll() is None
]
mpi.barrier()
scheduler = SchedulerPopen()
......@@ -103,18 +111,28 @@ def wait_for_all_extensions():
def name_ext_from_path_pythran(path_pythran):
"""Return an extension name given the path of a Pythran file"""
if path_pythran.exists():
with open(path_pythran) as file:
src = file.read()
if mpi.rank == 0:
if path_pythran.exists():
with open(path_pythran) as file:
src = file.read()
else:
src = ""
name = path_pythran.stem + "_" + make_hex(src) + ext_suffix_short
else:
src = ""
name = None
if mpi.nb_proc > 1:
name = mpi.comm.bcast(name, root=0)
return path_pythran.stem + "_" + make_hex(src) + ext_suffix_short
return name
def compile_pythran_files(
paths: Iterable[Path], str_pythran_flags: str, parallel=True
):
if mpi.rank > 0:
return
pythran_flags = str_pythran_flags.strip().split()
......@@ -135,6 +153,8 @@ def compile_pythran_file(
xsimd=True,
openmp=False,
):
if mpi.rank > 0:
return
if not isinstance(path, Path):
path = Path(path)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment