class SharedIO(object):
"""
A version of BytesIO which is shared between multiple processes, suitable
for IPC.
"""
CACHE = {}
def __init__(self, name, size):
self.shmem = None
self.shmem = SharedMemory(name, flags=O_CREAT, mode=0o666, size=size)
self.file = mmap(self.shmem.fd, size)
self.__init_file_methods()
self.__shmem_refs = [1]
def __del__(self):
if self.shmem is None:
return
self.__shmem_refs[0] -= 1
if self.__shmem_refs[0] == 0:
self.shmem.close_fd()
try:
self.shmem.unlink()
except ExistentialError:
pass
def __getstate__(self):
return {"name": self.name, "size": self.size, "pos": self.tell()}
def __setstate__(self, state):
name = state["name"]
size = state["size"]
cached = SharedIO.CACHE.get("%s:%d" % (name, size))
if cached is not None:
assert cached.size == size
self.shmem = cached.shmem
self.file = cached.file
self.__shmem_refs = cached.__shmem_refs
self.__shmem_refs[0] += 1
self.__init_file_methods()
else:
self.__init__(name, size)
SharedIO.CACHE["%s:%d" % (name, size)] = self
self.seek(state["pos"])
@property
def name(self):
return self.shmem.name
@property
def size(self):
return self.shmem.size
@property
def refs(self):
return self.__shmem_refs[0]
def __init_file_methods(self):
for name in ("read", "readline", "write", "tell", "close", "seek"):
setattr(self, name, getattr(self.file, name))
def clean_shared_memory(pid=None):
'''
This function cleans up shared memory segments created by the GUI
or a buggy Open MPI or CUDA.
'''
# this only works for linux that has /dev/shm
if not sys.platform.startswith('linux'):
print("This function works only under Linux. Stop.", file=sys.stderr)
return
assert os.path.isdir('/dev/shm/')
from posix_ipc import SharedMemory
shm_list = os.listdir('/dev/shm/')
user = getpass.getuser()
for shm in shm_list:
if (shm.startswith('ptycho') \
or shm.startswith('vader') \
or shm.startswith('cuda')) \
and user == find_owner('/dev/shm/'+shm):
if (pid is None) or (pid is not None and pid in shm):
s = SharedMemory("/" + shm)
s.close_fd()
s.unlink()
print("Done.")
def pi(n):
pids = []
unit = n / 10
sem_lock = Semaphore('/pi_sem_lock', flags=posix_ipc.O_CREX, initial_value=1)
memory = Memory('/pi_rw', size=8, flags=posix_ipc.O_CREX)
os.lseek(memory.fd, 0, os.SEEK_SET)
os.write(memory.fd, struct.pack('d', 0.0))
for i in range(10):
mink = unit * i
maxk = mink + unit
pid = os.fork()
if pid > 0:
pids.append(pid)
else:
s = calc_slice(mink, maxk)
sem_lock.acquire()
try:
os.lseek(memory.fd, 0, os.SEEK_SET)
bs = os.read(memory.fd, 8)
cur_val, = struct.unpack('d', bs)
cur_val += s
bs = struct.pack('d', cur_val) # 序列化
os.lseek(memory.fd, 0, os.SEEK_SET)
os.write(memory.fd, bs)
memory.close_fd()
finally:
sem_lock.release()
sys.exit(0)
sums = []
for pid in pids:
os.waitpid(pid, 0)
os.lseek(memory.fd, 0, os.SEEK_SET)
bs = os.read(memory.fd, 8)
sums, = struct.unpack('d', bs)
memory.close_fd()
memory.unlink()
sem_lock.unlink()
return math.sqrt(sums * 8)
def _create(self):
'''
attempt to create and initialize the shared memory
'''
LOG.info('shmem %s: attempting create shmem', self._name)
shmem = SharedMemory(self._name, size=SHMEM_SIZE, flags=O_CREX)
LOG.info('shmem %s: attempting create mmap', self._name)
self._mmap = mmap.mmap(shmem.fd, shmem.size)
shmem.close_fd()
LOG.info('shmem %s: attempting create semaphore', self._name)
self._lock = Semaphore('/pyspace_%s_lock' % self._name, flags=O_CREX)
LOG.info('shmem %s: create succeeded', self._name)
try:
self._initialize()
self._lock.release()
except:
LOG.exception('shmem %s: initialize failed; attempting unlink', self._name)
shmem.unlink()
self._lock.unlink()
raise
class SharedIO(object):
"""
A version of BytesIO which is shared between multiple processes, suitable
for IPC.
"""
CACHE = {}
def __init__(self, name, size):
self.shmem = None
self.shmem = SharedMemory(name, flags=O_CREAT, mode=0o666, size=size)
self.file = mmap(self.shmem.fd, size)
self.__init_file_methods()
self.__shmem_refs = [1]
def __del__(self):
if self.shmem is None:
return
self.__shmem_refs[0] -= 1
if self.__shmem_refs[0] == 0:
self.shmem.close_fd()
try:
self.shmem.unlink()
except ExistentialError:
pass
def __getstate__(self):
return {"name": self.name,
"size": self.size,
"pos": self.tell()}
def __setstate__(self, state):
name = state["name"]
size = state["size"]
cached = SharedIO.CACHE.get("%s:%d" % (name, size))
if cached is not None:
assert cached.size == size
self.shmem = cached.shmem
self.file = cached.file
self.__shmem_refs = cached.__shmem_refs
self.__shmem_refs[0] += 1
self.__init_file_methods()
else:
self.__init__(name, size)
SharedIO.CACHE["%s:%d" % (name, size)] = self
self.seek(state["pos"])
@property
def name(self):
return self.shmem.name
@property
def size(self):
return self.shmem.size
@property
def refs(self):
return self.__shmem_refs[0]
def __init_file_methods(self):
for name in ("read", "readline", "write", "tell", "close", "seek"):
setattr(self, name, getattr(self.file, name))
