-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathmpiimport.py
66 lines (51 loc) · 2 KB
/
mpiimport.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import os
import imp
import mpiimporter
from mpi4py import MPI
class Importer:
def __init__(self, path=None):
if path is not None and not os.path.isdir(path):
raise ImportError
self.path = path
def find_module(self, fullname, path=None):
rank = MPI.COMM_WORLD.Get_rank()
# print "[%d] find_module %s %s" % (rank, fullname, path)
subname = fullname.split(".")[-1]
if subname != fullname and path is None:
return None
try:
file, filename, stuff = mpiimporter.find_module(subname, path)
except ImportError:
# print ImportError
return None
# print "[%d] find_module found: %s %s" % (rank, file, filename)
ignore, ext = os.path.splitext(filename)
if ext == '.so':
file, filename, stuff = imp.find_module(subname, path)
return ImpLoader(file, filename, stuff)
return Loader(file, filename, stuff)
class Loader:
def __init__(self, file, filename, stuff):
self.file = file
self.filename = filename
self.stuff = stuff
def load_module(self, fullname):
rank = MPI.COMM_WORLD.Get_rank()
# print "[%d] load_module: %s %s %s %s" % (rank, fullname, self.file, self.filename, self.stuff)
mod = mpiimporter.load_module(fullname, self.file, self.filename, self.stuff)
if self.file:
self.file.close()
mod.__loader__ = self # for introspection
# print "[%d] load_module loaded: %s" % (rank, mod)
return mod
class ImpLoader:
def __init__(self, file, filename, stuff):
self.file = file
self.filename = filename
self.stuff = stuff
def load_module(self, fullname):
mod = imp.load_module(fullname, self.file, self.filename, self.stuff)
if self.file:
self.file.close()
mod.__loader__ = self # for introspection
return mod