| 6 | |
| 7 | |
| 8 | class ShmArray: |
| 9 | def __init__(self, name, shape, dtype): |
| 10 | self.shm = None |
| 11 | self.arr = None |
| 12 | self.name = name |
| 13 | self.dtype_byte_num = np.array([1], dtype=dtype).dtype.itemsize |
| 14 | self.dest_size = np.prod(shape) * self.dtype_byte_num |
| 15 | self.shape = shape |
| 16 | self.dtype = dtype |
| 17 | |
| 18 | def create_shm(self): |
| 19 | try: |
| 20 | shm = shared_memory.SharedMemory(name=self.name, create=True, size=self.dest_size) |
| 21 | except: |
| 22 | shm = shared_memory.SharedMemory(name=self.name, create=False, size=self.dest_size) |
| 23 | |
| 24 | if shm.size != self.dest_size: |
| 25 | logger.warning(f"size not same, unlink shm {self.name} and create again") |
| 26 | shm.close() |
| 27 | shm.unlink() |
| 28 | try: |
| 29 | shm = shared_memory.SharedMemory(name=self.name, create=True, size=self.dest_size) |
| 30 | logger.info(f"create shm {self.name}") |
| 31 | except: |
| 32 | shm = shared_memory.SharedMemory(name=self.name, create=False, size=self.dest_size) |
| 33 | logger.info(f"link shm {self.name}") |
| 34 | |
| 35 | self.shm = shm # SharedMemory 对象一定要被持有，否则会被释放 |
| 36 | self.arr = np.ndarray(self.shape, dtype=self.dtype, buffer=self.shm.buf) |
| 37 | |
| 38 | def link_shm(self): |
| 39 | shm = shared_memory.SharedMemory(name=self.name, create=False, size=self.dest_size) |
| 40 | assert shm.size == self.dest_size |
| 41 | self.shm = shm |
| 42 | self.arr = np.ndarray(self.shape, dtype=self.dtype, buffer=self.shm.buf) |
| 43 | return |
| 44 | |
| 45 | def close_shm(self): |
| 46 | if self.shm is not None: |
| 47 | self.shm.close() |
| 48 | self.shm.unlink() |
| 49 | self.shm = None |
| 50 | self.arr = None |
no outgoing calls