经验教训写在最前面:以后用Linux接口要认真看官方的manual,网上的大部分blog都是复制来复制去,一个
man xxx
比十篇blog都管用
Warning:本文的方法
可能会带来内存泄漏,到底有没有泄露还没测,反正也不是写线上代码 : )
最近做一些分布式训练的实验时,需要让一个训练的进程与另外几个维护数据的进程交互得到数据,数据格式是numpy.ndarray,于是就考虑能不能开一个numpy.ndarray格式的buffer来放一整个batch的数据,然后把numpy的内存空间映射成shared memory,这样另一个进程就可以直接通过numpy来读数据
选用这种方案的时候查了若干资料,目前来说多进程内存数据共享问题,shared memory应该是效率最高的方案了,其原理是将不同进程内的逻辑地址通过页表映射到一片相同的物理地址,一旦shared memory构造完成,进程读取数据的速度又是甚至比普通的访问内存还要快(因为是页对齐的,后面再讲)
API介绍
# 查看系统中的共享内存
ipcs -m
# 通过shmid删除某片共享内存
ipcrm -m ${shmid}
首先是shared memory的创建,对应函数shmget
,在物理内存空间中创建共享内存,成功返回shmid,错误返回-1,shmid在Linux系统中唯一标定了一片shared memory,后续对shared memory的操作大多都要用到shmid
int shmget(key_t key, size_t size, int shmflg);
key
:大于零的整数即可,一般可以用ftok
来生成,ftok
中的参数可以随便定,系统会保证给你返回一个不冲突的shmid;但是考虑到从Python向C传字符串很恶心,本渣选择直接手动指定size
:需要注意,这个函数中会将size
变量自动round up到getpagesize()
返回值的整数倍大小,有些blog里给的代码会先把size手动round up,其实没有必要- 对于shmflg参数:
IPC_CREAT
代表创建,可以通过或操作符设置权限,一个比较常用的组合是IPC_CREATE|0666
,其中0666代表各用户的读写权限;其他一些常用的flag可以直接用man shmget
查看
shmat
,将进程内的一篇连续空间使用页表映射到物理内存中的共享内存,成功返回进程内被关联到的共享内存逻辑地址指针,错误返回-1
void *shmat(int shmid, const void *shmaddr, int shmflg);
shmid
:shmget
的返回值shmaddr
:进程空间内被映射的首地址,如果传NULL则系统会自动在进程空间里开一片新的内存来做shared memory的映射(但这显然不符合我们的需求,因为我们的目标是希望去映射一个已经存在的numpy.ndarray的内存空间)shmflg
:最后发现用SHM_RND|SHM_REMAP
组合是没有问题的
对于本文要实现的功能,主要的幺蛾子就出在这个函数身上,无论你去搜中文blog还是stackoverflow,大多数人都会告诉你第二个参数传NULL就可以了,这样系统就会自动在进程空间里开一片新的内存来做shared memory,但我们的目标是希望去映射一个已经存在的numpy.ndarray的内存空间
那么直接把numpy.ndarray的指针传给shmaddr
参数可不可以呢?
答案是不行,文档里表示如果用户想要自定义这个参数,那么shmaddr
指针必须得是页对齐的,我搜了半天文档也没找到怎么把numpy数组开成页对齐的方法
幸运的是官方文档里给了一个方案:对shmflg设置SHM_RND
,这样shmat
函数会自动将指针shmaddr
的地址round down到页对齐的地址,于是本渣就考虑在Python里面开numpy.ndarray的时候预先多留出一页的空间,这样round down的时候就页表映射不会映射到越界的内存了
结果:报错,errno==22
,22的意思是Invalid argument
联系shared memory的原理,想到问题应该是出在numpy开出来的内存空间已经在页表中被映射过了,回去看文档,发现文档里还给了一个shmflg
的参数叫做SHM_REMAP
,这个flag会强制操作系统对shmaddr
对应的内存进行重新映射,加上问题解决
关键代码
初始化shared memory代码如下
void init_shm(float* shmptr, int _shmsz, int _shmkey, int* additional) {
int shmid = -1;
shmid = shmget((key_t)_shmkey, _shmsz, 0666|IPC_CREAT);
if (shmid == -1) {
perror("create shm failed, maybe run `ipcs -m` to check...`n");
exit(-1);
}
void *shm_start_addr = (void*)((char*)shmptr + getpagesize());
void *shared_ptr = (void*)shmat(shmid, shm_start_addr, SHM_RND|SHM_REMAP);
if ((long long)shared_ptr == -1) {
perror("shmat failedn");
exit(-1);
}
additional[0] = shmid;
additional[1] = (char*)shared_ptr - (char*)shmptr;
}
销毁shared memory代码如下(变量addinfo
对应上面的参数additional
)
void close_shm(float* shmbuf, int* addinfo) {
int shmid = addinfo[0];
int offset = addinfo[1];
if (shmdt((char*)shmbuf + offset) == -1) {
perror("shmdt failedn");
exit(-1);
}
if (shmctl(shmid, IPC_RMID, 0) == -1) {
perror("remove shm errorn");
exit(-1);
}
}
numpy与C交互
这部分网上随便搜资料应该是很多的,简单总结,步骤可以分为三步:
- 在Python中开numpy.ndarray,一般要对array做连续化处理,然后转成C指针
- 写一个C或者C++的代码文件,编译成.so动态链接库文件,注意如果是C++文件的话,函数接口的声明一定要放进 extern "C" 里面
- 在Python中导入动态链接库,然后就可以直接调用C的接口了
需要注意的几个坑:首先,numpy.ndarray在内存空间上可能是不连续的,如果不加处理就扔进C里面,很可能会出现莫名其妙的segmentation fault,解决方案:
import numpy as np
import ctypes as ct
# ...
shmbuf = np.zeros((MAX_SHM_BYTES_ROUNDED), dtype=np.float32)
if not shmbuf.flags['C_CONTIGUOUS']:
shmbuf = np.ascontiguousarray(shmbuf, dtype=shmbuf.dtype)
# 转C指针
c_shmbuf = shmbuf.ctypes.data_as(ct.c_void_p)
c_max_shm_bytes_rounded = ct.c_int(MAX_SHM_BYTES_ROUNDED)
编译的时候输出一定要是fPIC的,比方说cpp代码文件叫做 shm.cpp:
g++ -c -fpic shm.cpp
g++ -shared -o libshm.so shm.o
最后把.so文件copy或者软链到Python的工作目录下,导入的时候就可以直接调用C里面定义的 call_method 方法
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
dll = np.ctypeslib.load_library(os.path.join(BASE_DIR, DLLNAME), '.')
dll.call_method(c_shmbuf, c_max_shm_bytes_rounded)