一种vGPU方案：控制GPU内存使用上限

发表于 | 更新于

在已经实现了k8s层面的GPU按小数申请支持的情况下(阿里是按内存数量支持），我们需要有个实际的资源限制的方案。
目前不太有好的办法控制gpu的算力（其实有，局限性比较大）。我们先控制下gpu显存的使用，具体实现就不方便贴了，讲下思路和效果

  1. 首先需要了解cuda api，所有框架如果细看代码都会发现底层是使用的cuda api去操作nvidia 的gpu，python只是个binding
  2. 需要了解hook 系统api的方法（ 搜索关键词ring3 hook)

我们从tensorflow的运行日志去翻代码 找设备初始化信息的地方 Found device ，跟进去，可以发现实际是调用了cuMemGetInfo这个函数获取的信息。 （这个页面中mem相关的函数都建议看下）

那么我们只需要重新写个so，实现这个函数就可以了，这个函数根据我们自己的逻辑设置最大的内存上限，与此同时由于cuMemGetInfo ( size_t free, size_t total ) 这个函数还返回了free的情况，所以我们也需要记录free的值，可以hook掉cuMemAlloc ，申请的时候记录下各个容器的使用情况，也可以利用外围的监控程序获取信息

最终实现的效果,设置环境变量后tensorflow看到的内存值是我们设置的值

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
root@cnn2-1553057347-worker-0:~# export GPU_MEMORY=2147483648
root@cnn2-1553057347-worker-0:~# python
Python 3.5.2 (default, Nov 23 2017, 16:37:01)
[GCC 5.4.0 20160609] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
>>> sess =tf.Session()
2019-03-22 14:57:33.394861: I tensorflow/core/platform/cpu_feature_guard.cc:140] Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.1 SSE4.2 AVX AVX2 FMA
hacking gpu memory from GPU_MEMORY env by lovejoy. mem=<2147483648>
2019-03-22 14:57:33.554589: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1344] Found device 0 with properties:
name: Tesla K20c major: 3 minor: 5 memoryClockRate(GHz): 0.7055
pciBusID: 0000:02:00.0
totalMemory: 2.00GiB freeMemory: 1.00GiB
2019-03-22 14:57:33.554637: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1423] Adding visible gpu devices: 0
2019-03-22 14:57:33.843047: I tensorflow/core/common_runtime/gpu/gpu_device.cc:911] Device interconnect StreamExecutor with strength 1 edge matrix:
2019-03-22 14:57:33.843094: I tensorflow/core/common_runtime/gpu/gpu_device.cc:917]      0
2019-03-22 14:57:33.843106: I tensorflow/core/common_runtime/gpu/gpu_device.cc:930] 0:   N