Improve data loading speed for tensorflow_stub

Overview

We're using the tensorflow_stub library for data loading (due to not installing tensorflow), on top of an internal cloud Filesystem with non negligible latency(SLA for 1MB read p50 100ms, p90 150ms). We saw data loading speed can be very slow and have big room for improvement.

On the high level here're some of the things we're thinking about:

• Cache listdir results in directory_watcher https://github.com/tensorflow/tensorboard/blob/74214c309f72c4ecab000a983efb224679dfff73/tensorboard/backend/event_processing/directory_watcher.py#L212
• Bypass or parallelize exists check in PyRecordReader_New https://github.com/tensorflow/tensorboard/blob/b871e24dad00ba4e50626346d2cbf5258c016917/tensorboard/compat/tensorflow_stub/pywrap_tensorflow.py#L179
• Reduce read amplification in Gfile https://github.com/tensorflow/tensorboard/blob/14bdef2e8f4e97b0aee40b95e5d31e734786f542/tensorboard/compat/tensorflow_stub/io/gfile.py#L493
• Option to knob to toggle off crc32c https://github.com/tensorflow/tensorboard/blob/b871e24dad00ba4e50626346d2cbf5258c016917/tensorboard/compat/tensorflow_stub/pywrap_tensorflow.py#L224
• Make _DEFAULT_BLOCK_SIZE tunable https://github.com/tensorflow/tensorboard/blob/14bdef2e8f4e97b0aee40b95e5d31e734786f542/tensorboard/compat/tensorflow_stub/io/gfile.py#L52
• Explore parallelizing file load / asyncio

Below are the profiling results that explains why. The profiling was done on TensorBoard 3.0, seems a lot of changes has happened in https://github.com/tensorflow/tensorboard/issues/1711, but some of our preliminary results shows the same bottleneck still exists with the most recent version.

Profiling code

mul = event_multiplexer.EventMultiplexer()
mul.AddRunsFromDirectory("manifold://xxx", None)
cProfile.runctx("mul.Reload()", globals(), locals(), sort="tottime")

Typical use case 1: Lots of small files

• Single dir with 546 file * 42k each = 22M total
• Total time is 413.518s (6.8min)

Profile results

14259238 function calls (14258115 primitive calls) in *413.518* seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.002    0.002  413.534  413.534 {built-in method builtins.exec}
        1    0.000    0.000  413.533  413.533 <string>:1(<module>)
        1    0.000    0.000  413.533  413.533 plugin_event_multiplexer.py:214(Reload)
        1    0.000    0.000  413.533  413.533 plugin_event_multiplexer.py:232(Worker)
        1    0.082    0.082  413.533  413.533 plugin_event_accumulator.py:168(Reload)
    35635    0.011    0.000  409.643    0.011 directory_watcher.py:71(Load)
    35635    0.027    0.000  409.632    0.011 directory_watcher.py:96(_LoadInternal)
     5625    0.022    0.000  398.274    0.071 common.py:77(runWithRetries)
     5706    0.026    0.000  393.933    0.069 THeaderProtocol.py:152(readMessageBegin)
     5625    0.015    0.000  393.897    0.070 srproxy.py:625(readMessageBegin)
     5706    0.021    0.000  393.774    0.069 THeaderTransport.py:228(_reset_protocol)
     5706    0.154    0.000  393.753    0.069 THeaderTransport.py:263(readFrame)
    22824    0.068    0.000  393.183    0.017 TTransport.py:66(readAll)
    22583    0.052    0.000  393.065    0.017 TSocket.py:297(read)
    22502    0.029    0.000  393.052    0.017 srproxy_base.py:179(read)
    22583  393.003    0.017  393.003    0.017 {method 'recv' of '_socket.socket' objects}
     2810    0.016    0.000  268.089    0.095 manifold_client_deprecated.py:79(wrapper)
    36757    0.043    0.000  214.215    0.006 event_file_loader.py:85(Load)
    36757    0.070    0.000  213.981    0.006 event_file_loader.py:49(Load)
    36757    0.154    0.000  213.779    0.006 pywrap_tensorflow.py:190(GetNext)
*  143659    0.150    0.000  208.132    0.001 gfile.py:441(read)
*  1685    0.006    0.000  207.772    0.123 manifoldio.py:54(wrap)
     1685    0.027    0.000  207.766    0.123 manifoldio.py:98(read)
     1685    0.009    0.000  207.627    0.123 manifold_client_deprecated.py:144(get)
     1685    0.020    0.000  207.606    0.123 default.py:300(get)
     1685    0.040    0.000  207.586    0.123 default.py:478(getSerial)
     1685    0.008    0.000  206.252    0.122 BlobStore.py:8251(rangeReadObj)
     1685    0.021    0.000  205.967    0.122 BlobStore.py:8273(recv_rangeReadObj)
      563    0.020    0.000  134.943    0.240 directory_watcher.py:199(_GetNextPath)
      563    0.004    0.000  133.164    0.237 io_wrapper.py:63(ListDirectoryAbsolute)
*    563    0.003    0.000  133.160    0.237 gfile.py:592(listdir)
*    563    0.006    0.000  133.151    0.237 manifoldio.py:199(listdir)
      563    0.088    0.000  133.122    0.236 manifoldio.py:204(<listcomp>)
   316969    0.517    0.000  133.034    0.000 default.py:1024(ls)
     2815    0.009    0.000   83.334    0.030 BlobStore.py:8487(continueBlockEnumeration)
     2815    0.029    0.000   82.951    0.029 BlobStore.py:8505(recv_continueBlockEnumeration)
      562    0.005    0.000   60.445    0.108 directory_watcher.py:177(_SetPath)
      562    0.011    0.000   60.438    0.108 event_file_loader.py:36(__init__)
      562    0.006    0.000   60.423    0.108 pywrap_tensorflow.py:167(__init__)
*    562    0.003    0.000   60.401    0.107 gfile.py:557(exists)
*    562    0.005    0.000   60.391    0.107 manifoldio.py:88(exists)
      562    0.003    0.000   60.354    0.107 manifold_client_deprecated.py:481(exists)
      562    0.002    0.000   60.346    0.107 default.py:1265(exists)
      562    0.005    0.000   60.345    0.107 default.py:541(getMeta)
      562    0.009    0.000   60.339    0.107 default.py:392(getStoredObject)
      562    0.003    0.000   59.918    0.107 BlobStore.py:8215(readObj)
      562    0.007    0.000   59.824    0.106 BlobStore.py:8235(recv_readObj)
      563    0.003    0.000   48.749    0.087 BlobStore.py:8663(startOrderedNodeEnumeration)
      563    0.007    0.000   48.653    0.086 BlobStore.py:8685(recv_startOrderedNodeEnumeration)
    71268    0.086    0.000    5.437    0.000 pywrap_tensorflow.py:42(masked_crc32c)
    71268    0.045    0.000    5.325    0.000 pywrap_tensorflow.py:156(crc32c)
    71268    5.268    0.000    5.268    0.000 pywrap_tensorflow.py:124(crc_update)
    35634    0.259    0.000    3.807    0.000 plugin_event_accumulator.py:268(_ProcessEvent)
     5706    2.243    0.000    3.334    0.001 {built-in method thrift.protocol.fastproto.decode}
    35071    0.091    0.000    3.054    0.000 data_compat.py:32(migrate_value)
     2815    0.023    0.000    2.617    0.001 BlobStore.py:4359(read)
    35028    0.281    0.000    1.964    0.000 data_compat.py:62(make_summary)
     2810    0.015    0.000    1.890    0.001 default.py:256(client)

Composition:

• 50% on file reading
• 32% on listdir
• 15% on exists

Deep dive

File reading:

• Event_file_loader will continue read till OutOfRangeError thrown
  https://github.com/tensorflow/tensorboard/blob/74214c309f72c4ecab000a983efb224679dfff73/tensorboard/backend/event_processing/event_file_loader.py#L76

• When entire file was read, the next read call will still trigger a filesystem read, causing the file to be at least read twice
  https://github.com/tensorflow/tensorboard/blob/14bdef2e8f4e97b0aee40b95e5d31e734786f542/tensorboard/compat/tensorflow_stub/io/gfile.py#L493

• But I don’t know why in the above profiling, each file was read 3 times

• On the other side we can try to log to single event file in client side

• parallelizing file load / asyncio?

Listdir:

• directory_watcher will serially load each file https://github.com/tensorflow/tensorboard/blob/74214c309f72c4ecab000a983efb224679dfff73/tensorboard/backend/event_processing/directory_watcher.py#L123

• And the way to find the next file is to call listdir, which is unnecessary https://github.com/tensorflow/tensorboard/blob/74214c309f72c4ecab000a983efb224679dfff73/tensorboard/backend/event_processing/directory_watcher.py#L210-L213

exists:

• PyRecordReader_New will call exists on construction, we don’t need this
  https://github.com/tensorflow/tensorboard/blob/b871e24dad00ba4e50626346d2cbf5258c016917/tensorboard/compat/tensorflow_stub/pywrap_tensorflow.py#L179

How much can we save

Read all files in take 77s, means upper bound saving is 82%. If we read files in parallel it’ll be faster

Typical use case 1: Single big file

• Two dirs with 131M + 40M = 170M
• Total time is 55s

Profile results

3881130 function calls (3881126 primitive calls) in 55.790 seconds

   Ordered by: cumulative time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000   55.803   55.803 plugin_event_multiplexer.py:225(Reload)
        1    0.000    0.000   55.803   55.803 plugin_event_multiplexer.py:243(Worker)
        2    0.337    0.169   55.803   27.902 plugin_event_accumulator.py:168(Reload)
    32482    0.010    0.000   44.229    0.001 directory_watcher.py:71(Load)
    32482    0.011    0.000   44.219    0.001 directory_watcher.py:96(_LoadInternal)
    32482    0.092    0.000   43.669    0.001 event_file_loader.py:85(Load)
    32482    0.048    0.000   42.391    0.001 event_file_loader.py:49(Load)
    32482    0.148    0.000   42.338    0.001 pywrap_tensorflow.py:190(GetNext)
*    64960    0.091    0.000   36.777    0.001 pywrap_tensorflow.py:42(masked_crc32c)
*    64960    0.088    0.000   36.661    0.001 pywrap_tensorflow.py:156(crc32c)
    64960   36.562    0.001   36.562    0.001 pywrap_tensorflow.py:124(crc_update)
    32480    0.350    0.000   11.237    0.000 plugin_event_accumulator.py:268(_ProcessEvent)
*    32478    0.374    0.000   10.335    0.000 data_compat.py:32(migrate_value)
*    28078    2.672    0.000    9.493    0.000 data_compat.py:69(_migrate_histogram_value)
   129922    0.193    0.000    5.362    0.000 gfile.py:441(read)
       19    0.000    0.000    5.051    0.266 common.py:77(runWithRetries)
       19    0.000    0.000    4.935    0.260 manifold_client_deprecated.py:79(wrapper)
       13    0.000    0.000    4.861    0.374 manifoldio.py:54(wrap)
*       13    0.000    0.000    4.861    0.374 manifoldio.py:98(read)
*       28    0.000    0.000    4.792    0.171 THeaderProtocol.py:152(readMessageBegin)
       28    0.037    0.001    4.791    0.171 THeaderTransport.py:228(_reset_protocol)
       19    0.000    0.000    4.787    0.252 srproxy.py:625(readMessageBegin)
       28    0.080    0.003    4.754    0.170 THeaderTransport.py:263(readFrame)
       13    0.000    0.000    4.728    0.364 manifold_client_deprecated.py:144(get)
       13    0.078    0.006    4.728    0.364 default.py:300(get)
       13    0.000    0.000    4.650    0.358 default.py:478(getSerial)
       13    0.000    0.000    4.519    0.348 BlobStore.py:8251(rangeReadObj)
       13    0.000    0.000    4.515    0.347 BlobStore.py:8273(recv_rangeReadObj)
      112    0.001    0.000    4.459    0.040 TTransport.py:66(readAll)
      171    0.001    0.000    4.347    0.025 TSocket.py:297(read)

Composition:

• 65% on crc32c
• 18% on migrate Histogram
• 9% on file reading

Deep dive

crc32c:

• Need a way to turn crc32c off
  https://github.com/tensorflow/tensorboard/blob/b871e24dad00ba4e50626346d2cbf5258c016917/tensorboard/compat/tensorflow_stub/pywrap_tensorflow.py#L224

Histogram:

• We should migrate write path ourselves

file reading:

• Current buffer size is 16M, so for a 70M file we need 3 read calls; we should support increase buffer size

https://github.com/tensorflow/tensorboard/blob/14bdef2e8f4e97b0aee40b95e5d31e734786f542/tensorboard/compat/tensorflow_stub/io/gfile.py#L52

• Asyncio?

How much can we save

Reading entire file in take 1.6s, upper bound saving 98%

What we need

We want to get the high level opinion from TensorBoard team on whether there're existing efforts that can already cover these bottlenecks, or should we spend engineering efforts to make some of these happen.

As always, we deeply appreciate all the help and support we received from TensorBoard team & community, and look forward to this discussion!