Tensorflow网络传输性能分析

TensorFlow   网络传输

0. 写在前面

tensorflow分布式训练时，grpc的 慢 一直都被很多人所诟病。在早期的版本中，由于实现的一些原因，的确存在一些性能问题(可以参见这个 issue )。

但随着项目的迭代，现在性能如何，就有些莫衷一是了。这里通过对两个项目master分支代码的一些测试，希望能探讨下这些问题。

1. 直观的看传输速率

这里先用 一个测试程序 测试下tensor在两个机器中的传输速率。测试使用的两台机器配置的都是万兆以太网的网卡：

[work@host benchtools]$ ethtool eth0  
Settings for eth0:
...
        Speed: 10000Mb/s
...
在两台机器上分别跑测试程序的worker和ps：

[host1] python tensor_transfer_throughput.py --ps_hosts=host1:12222 --worker_hosts=host2:12222 --job=ps --task=0
[host2] python tensor_transfer_throughput.py --ps_hosts=host1:12222 --worker_hosts=host2:12222 --data_mb=100 --job=worker --task=0 --iters=100
测试程序干的事情很简单：在ps和worker上各创建一个相同大小的variable, 然后worker反复将自己的variable assign给ps。在上面的测试中，我们将variable的大小设置为100M，传输次数为100。

测试结果在worker运行结束后可以看到：

[host2] python tensor_transfer_throughput.py --ps_hosts=host1:12222 --worker_hosts=host2:12222 --data_mb=100 --job=worker --task=0 --iters=100
....
transfer rate: 173.488801 MB/s
利用ifstat工具也可以看到网络的传输性能：

[hosts1]$ ./ifstat
      eth0                eth1      
KB/s in  KB/s out   KB/s in  KB/s out
  191.95  176435.6      0.00      0.00
  206.18  170675.3      0.00      0.00
  222.45  220156.5      0.00      0.00
  162.84  169024.8      0.00      0.00
  224.44  211070.7      0.00      0.00
可以看到两种测试的througput效果差不多。理论上来说ifstat可能会比worker的输出稍微大一点，因为grpc要为每次传输额外添加一些header信息。但和100MB的数据相比，应该可以忽略不计。

但无论是哪个结果，离理论值的1.25GBps(10Gbps)差距仍旧非常大。所以初步来看，网卡的利用率是比较低的。

2. 单独测试grpc

为了验证问题是不是出在grpc这里，我利用 另一个测试程序 ，来测试grpc本身的传输效率。

程序不太复杂，要点包括：

client和server端的功能要简单，尽量减少额外操作所带来的时间开销：client只负责无脑发送，server端也要直接丢弃收到的数据。

直接利用grpc的ByteBuffer，从而避免掉在发送和接收时的memcpy。这点和tensorflow发送tensor的流程也是一致的。

server端可以创建多个completion queue, 从而可以指定多个worker线程。

client利用异步接口。可以指定传输并发度，也可以允许grpc创建多个channel。

可以指定发送数据和响应数据块的大小。

然后将程序部署到两台机器上开始测试。client每次向server发送100M数据，共发送1000条：

[host1] ./grpc_raw --job_type=server --server_threads=1 --message_size=10
[host2] ./grpc_raw --job_type=client --job_type=client --target_ip=host1 --total_message=1000 --message_size=104857600
利用ifstat看结果：

[work@host2 benchtools]$ ./ifstat
      eth0                 eth1      
KB/s in  KB/s out   KB/s in  KB/s out
  162.05  198529.9      0.00      0.00
  128.67  150799.5      0.00      0.00
  196.09  203136.0      0.00      0.00
  169.20  192864.8      0.00      0.00
  130.67  146532.7      0.00      0.00
可以看到和测tensor传输时类似，也是170MBps左右，离1.25GBps的理论值也差距较大。

3. 为什么慢

为了进一步确定问题，我用 iperf 工具对网络的throughput做了单独的测试：

[host1] ./iperf3 -s -i 5
[host2] ./iperf3 -c host1 -i 5 -t 1000
测试结果如下：

[host2]$ ./iperf3 -c host1 -i 5 -t 1000
...
[  5]   0.00-5.00   sec   983 MBytes  1.65 Gbits/sec  31545   2.49 MBytes      
[  5]   5.00-10.00  sec   839 MBytes  1.41 Gbits/sec  35645    889 KBytes      
[  5]  10.00-15.00  sec   830 MBytes  1.39 Gbits/sec  35863    954 KBytes
...
可以看到大概也就是1.4Gbps(175MBps)左右，和 grpc的测试结果差不多 。

为什么会这样呢？事实上，当提高socket数后，结果就会大大改观，总的传输速率会达到9.3 Gbps左右，从而和理论值接近：

[host2]$ ./iperf3 -c host1 -i 5 -t 1000 -P 8
...
[  5]  40.00-45.00  sec   621 MBytes  1.04 Gbits/sec  9936   2.06 MBytes      
....
[ 19]  40.00-45.00  sec   206 MBytes   346 Mbits/sec  922   90.5 KBytes      
[SUM]  40.00-45.00  sec  5.43 GBytes  9.33 Gbits/sec  33646
这里我们可以看到的一个结论是： 单个socket可能(远远)无法用满网卡的带宽 。

那么如果把grpc的socket数增加如何？遗憾的是，目前grpc还不支持这样的特性。在grpc里，通信是用 channel 来进行抽象的。 哪怕你在两个机器间创建多个channel, 他们在底层也是会共享socket的 。

4. 单个socket用不满网卡？

当我通过测试得出这个结论时，我内心也是无法接受的。我尝试了

手动调整拥塞窗口(事实上也没有必要，因为TCP会自发的增大它；稳定后的拥塞窗口大小，也没有达到Linux的上限)。

关闭Nagel算法
传输速率仍然没有变化。

后来在组里boss的建议下，我换了两台机器做测试。发现对于不同的机器组合，单socket的传输性能是不同的。 也存在一些机器，他们的单socket性能是可以达到网卡理论上限的 。

对于这一问题，现在怀疑可能和网络布局以及中间的交换机有关系。但具体的根源究竟是什么，还无从得知。

5. 继续测试

在我换了 单socket可以打满带宽 的两台机器后，我把1和2中的实验使用相同的参数重新做了一遍。结论如下：

grpc在单server单client的前提下，网卡传输的利用率还是非常高的。在我的实验中大概能到9Gbps左右，比iperf的结果稍逊一点，目测也就是5%左右。这可能和grpc在数据传输时的一些数据结构的分配、处理有关，但整理来说grpc性能已经比较可观了。

对于传输tensor的测试而言，传输速率大概能到5Gbps左右，是裸grpc的一多半。

这里有两个问题：

1. 为什么 传输tensor 的吞吐要低于 裸的grpc传输 ，问题在哪里？

2. 在我们最开始的两个实验中，由于单socket极限带宽较低，这二者的传输效率类似。为什么提高单socket的极限带宽后，二者开始体现出差别来？

其实这两个问题并不难解释：

在传输tensor时，除了有效的传输数据外，还有master驱动worker运行、序列化、反序列化、数据assign等其他操作。而我们测试看到的throughput，是把这些操作都当成 有效传输 而平均化后的一个结果。

两个机器间带宽越高，额外操作的占比就越大，对总throughput的影响就越大。

6. 验证假设

为了验证我们的假设，我们需要知道 tensorflow在传输tensor时，真正用于数据传输的时间是多少 ，从而可以根据数据量大致推算一下传输时的网络带宽。

可以先用timeline看一下每一步所有op的耗时，以及RecvTensor这个op的耗时。

run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
sess.run(add_op.op, options=run_options, run_metadata=run_metadata)

trace = timeline.Timeline(step_stats=run_metadata.step_stats)
trace_file = open('timeline.ctf.json', 'w')
trace_file.write(trace.generate_chrome_trace_format())
结果(dur表示op的耗时，单位为us)：

{
    "name": "RecvTensor",
...
    "dur": 183311
},
....
{
    "name": "Assign",
...
    "dur": 19925
}
耗时主要在RecvTensor和Assign上，总耗时有200ms左右。对于100M数据而言，这个耗时也和观察到的5Gbps的吞吐大致吻合。

但我们仍旧不能知道真正在 传输 的时候带宽能不能有效的利用。timeline所能给出的最小粒度就是op，而"RecvTensor"这个op，我们可以看到耗时是180ms左右。这比grpc的传输吞吐还是要低出不少来的。

我们知道，在Tensorflow中，一个RecvTensor是要分成如下几个步骤的：

1. RecvOp的AsyncCompute，通过rendezvous接口，最终调用到grpc这一层。

2. 发起RecvTensor的请求，包括获取一个grpc_remote_worker的handle，以及准备RecvTensorRequest的protobuf，然后创建和rpc call相关的数据结构

3. 调用grpc的API，将数据推到网络引擎，发送数据。

4. server端从rendezvous_manager中获取tensor, 并且和其他的meta信息包装成ByteBuffer返回给客户端。

5. 客户端将收到的ByteBuffer反序列化成Tensor。

所以整个传输过程的慢，可能会慢在以下几个地方：

1. 做准备工作时，一些线程调度或者加锁操作带来开销。

2. server的序列化费时间。

3. grpc的网络引擎就是慢，比如说引入额外的数据拷贝之类的，导致ByteBuffer传输很慢。

4. client的反序列化费时间。

第三点其实不太可能，因为我们已经拿裸的grpc+ByteBuffer做过测试，其带宽利用率是比较高的。当然，我们也可以在Tensorflow中通过更细致的metrics来验证下这一点。

因为没法用timeline，只能通过改tensorflow代码来测试。为此，我简单修改了 tensorflow的代码 ，来观察传输和客户端处理的耗时。测试的结论如下：

对于100M的tensor，grpc的传输的时间大概在100ms左右。大概的数据传输率应该有9Gbps左右，比较高效。

server数据序列化的时间占比很小。这点tensorflow的确做过专门处理：tensor的内存是作为ByteBuffer直接传输的，很大程度避免了内存拷贝。

客户端的消息反序列化会占用一定时间，大概占到了RecvTensor的1/4多一些。主要原因是grpc ByteBuffer中的Tensor数据不满足Tensor的内存布局要求，所以必须得通过内存拷贝来一次重新整理。

7. 扩展性

前面分析了grpc在传输效率方面的性能，接下来看下有关扩展性方面的问题。

首先明确下，当我们讨论扩展性时，应该从如下两个角度来衡量：

server端未到网卡的瓶颈时，通过增加client，server端的throughput能随着client的个数线性增加。
server端达到网卡瓶颈后，随着client个数的增加， server端的吞吐最好基本不会下降，而client端的latency则会线性的增加。
这里的测试细节就不再展开了。通过对这两个方面的测试，我发现grpc在这两个层面基本表现也比较良好。

8. 总结

测试的结论大致有如下几个：

在开发分布式程序时，机房间机器的拓扑结构需要注意下，可能会影响单socket的极限带宽。如果存在此类问题，多socket的rpc是一个可能可行的方案。

grpc在大数据包的传输上，带宽利用率和扩展性都还不错。

对于tensorflow的RecvTensor，收到数据后的后续处理，会占据一部分计算资源，对总体的网卡带宽会存在影响。

几个需要继续调研的方面有：

grpc在高并发处理小数据包上latency表现如何，可以调研一下。对与tensorflow而言，这其实不太重要。但对于latency敏感的在线服务而言，还是非常重要的。

在tensor的send方这边， tensor table 是用一个非常粗粒度的互斥锁保护的，在RecvTensor请求较多时候怀疑可能会成为瓶颈(比如很多个worker的分布式训练)。这点需要拿大的训练场景测试一下。