Deep learning frameworks are core libraries that are increasingly used in a number of important real-world scenarios.  MXNet is no exception, being used  in range of production environments from embedded hardware, with very little RAM, to multi-million dollar web services running tens of thousands of requests per second.  MXNet is written in C++ and has a lot of raw pointer operations due to its high-performance, mathematical nature.  This introduces the potential for serious native coding errors to negatively affect important services and devices.  This document will describe how we're attempting to avoid these errors by introducing a heavily instrumented build (an ASAN build) that's designed to catch overflows and leaks in our automatic testing process.  The document will also describe how a developer can create an ASAN build and test for memory leaks locally when they're reported by users.

ASAN

ASAN, or the address sanitizer, is one of many C++ sanitizers developed by Google with the primary initial goal of securing Chrome from use-after-free and buffer-overflow errors.  It was originally launched as a feature for clang, but is now available in recent versions of GCC.  ASAN on it's own will detect:

The advantage of ASAN compared to other similar tools such as Valgrind is that it's fast, and well supported.  This is also why it's the primary mechanism of detecting leaks for the communities that manage both Chrome and Firefox.   For more information on ASAN (and other sanitizer) basics and motivations check out this talk.

GCC versus Clang compatibility

The capabilities of ASAN are similar whether the compiler used is clang or g++, as long as you're using very recent version of the compilers.  There's a good comparison of clang versus gcc7, in terms of ASAN capabilities, here.  With MXNet we've tested various versions of compilers and found that clang does not work well with our library.  We haven't completely tracked down the issue, but trying various methods of enabling ASAN has not worked for us when using clang, including forcing clang to dynamically link the ASAN library.  Luckily GCC ASAN seems to work correctly, and GCC 8 has the ASAN capabilities that we'd like to use.  Because of this we recommend using GCC 8 w/ ASAN when attempting to detect leaks or buffer overflows in MXNet, and we use GCC8 in CI and the Dockerfiles referred to below.

ASAN in CI

We're currently experimenting with a variety of methods for incorporating automatic ASAN checks in our CI system.  We will likely run these checks per-PR request while the detection is under active development.  We'll likely migrate the checks to nightly builds to reduce cost once the development is stable.  Thus far we've only enabled one variant of MXNet's build type, a basic CPU build.  We hope this will serve a basis for other developers to integrate ASAN with other MXNet build flavours.

C++ Training Test With Leak Detection

Example Output

Python Tests Without Leak Detection

Example Output

Using ASAN builds with MXNet

CI reports are nice, but it's sometimes more useful to build an ASAN build locally and to run some specific sections of code you're afraid may be leaking.  This is easy to do with MXNet.  We've installed all the prerequisites in or CI build Dockerfiles, so we can use docker to build ASAN builds without having to install or configure dependencies.  To build a CPU build with ASAN run the following commands in a new folder:


git clone --recurse https://github.com/apache/incubator-mxnet.git
cd incubator-mxnet/ci
# Build our dockerfile with all required deps for ASAN
docker build -f docker/Dockerfile.build.ubuntu_cpu -t mxnetci/build.ubuntu_cpu docker
cd ..
mkdir -p build
# Build an ASAN instrumented MXNet library
# Privileged probably not required for all steps, but in general ASAN requires some capabilities in order to inspect process memory.
docker run --privileged -v `pwd`:/work/mxnet -v `pwd`/build:/work/build  -ti mxnetci/build.ubuntu_cpu /work/runtime_functions.sh build_ubuntu_cpu_cmake_asan
# Choose an example of something you'd like to test with ASAN.  This could be a specific python test we want to run in a loop, it could be a C++ unit test written to expose leaks, etc.

# In our case we use a small Gluon python tests as an example.
# First we'll enter our container, and then we'll run tests within the container.
docker run --privileged -v `pwd`:/work/mxnet -v `pwd`/build:/work/build  -ti mxnetci/build.ubuntu_cpu bash

# Now within the container:
export PYTHONPATH=./python/
export MXNET_MKLDNN_DEBUG=1  # Ignored if not present
export MXNET_STORAGE_FALLBACK_LOG_VERBOSE=0
# Feel free to export any ASAN options via export ASAN_OPTIONS=...

# Importantly we need to make sure ASAN is the first library loaded (before our other libraries have a chance to allocate memory)
# To do this we add the library to the library preload list
export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libasan.so.5
nosetests-3.4 --verbose tests/python/unittest/test_rnn.py


The output should look similar to:


root@e77a083c4d00:/work/mxnet# nosetests-3.4 --verbose tests/python/unittest/test_rnn.py
test_rnn.test_deprecated ... ok
test_rnn.test_rnn ... ok
test_rnn.test_lstm ... ok
test_rnn.test_lstm_forget_bias ... ok
test_rnn.test_gru ... ok
test_rnn.test_residual ... ok
test_rnn.test_residual_bidirectional ... ok
test_rnn.test_stack ... ok
test_rnn.test_bidirectional ... ok
test_rnn.test_zoneout ... ok
test_rnn.test_unfuse ... ok
test_rnn.test_convrnn ... ok
test_rnn.test_convlstm ... ok
test_rnn.test_convgru ... ok
test_rnn.test_encode_sentences ... ok

----------------------------------------------------------------------
Ran 15 tests in 0.253s

OK

=================================================================
==93==ERROR: LeakSanitizer: detected memory leaks

Direct leak of 2414848 byte(s) in 572 object(s) allocated from:
    #0 0x7f7bacb38b60 in malloc (/usr/lib/x86_64-linux-gnu/libasan.so.5+0xedb60)
    #1 0x59c889  (/usr/bin/python3.5+0x59c889)

...

Direct leak of 1640 byte(s) in 1 object(s) allocated from:
    #0 0x7f7bacb3a970 in operator new[](unsigned long) (/usr/lib/x86_64-linux-gnu/libasan.so.5+0xef970)
    #1 0x7f7b53007bf0 in mxnet::profiler::Profiler::Profiler() /work/mxnet/src/profiler/profiler.cc:70
    #2 0x7f7b53028372 in void __gnu_cxx::new_allocator<mxnet::profiler::Profiler>::construct<mxnet::profiler::Profiler>(mxnet::profiler::Profiler*) (/work/mxnet/python/mxnet/../../build/libmxnet.so+0x68db372)
    #3 0x7f7b53027316 in void std::allocator_traits<std::allocator<mxnet::profiler::Profiler> >::construct<mxnet::profiler::Profiler>(std::allocator<mxnet::profiler::Profiler>&, mxnet::profiler::Profiler*) /usr/include/c++/8/bits/alloc_traits.h:475
    #4 0x7f7b53025b9d in std::_Sp_counted_ptr_inplace<mxnet::profiler::Profiler, std::allocator<mxnet::profiler::Profiler>, (__gnu_cxx::_Lock_policy)2>::_Sp_counted_ptr_inplace<>(std::allocator<mxnet::profiler::Profiler>) /usr/include/c++/8/bits/shared_ptr_base.h:549
    #5 0x7f7b530232fe in std::__shared_count<(__gnu_cxx::_Lock_policy)2>::__shared_count<mxnet::profiler::Profiler, std::allocator<mxnet::profiler::Profiler>>(std::_Sp_make_shared_tag, mxnet::profiler::Profiler*, std::allocator<mxnet::profiler::Profiler> const&) /usr/include/c++/8/bits/shared_ptr_base.h:662
    #6 0x7f7b5302053f in std::__shared_ptr<mxnet::profiler::Profiler, (__gnu_cxx::_Lock_policy)2>::__shared_ptr<std::allocator<mxnet::profiler::Profiler>>(std::_Sp_make_shared_tag, std::allocator<mxnet::profiler::Profiler> const&) /usr/include/c++/8/bits/shared_ptr_base.h:1328
    #7 0x7f7b5301cc03 in std::shared_ptr<mxnet::profiler::Profiler>::shared_ptr<std::allocator<mxnet::profiler::Profiler>>(std::_Sp_make_shared_tag, std::allocator<mxnet::profiler::Profiler> const&) /usr/include/c++/8/bits/shared_ptr.h:360
    #8 0x7f7b53019822 in std::shared_ptr<mxnet::profiler::Profiler> std::allocate_shared<mxnet::profiler::Profiler, std::allocator<mxnet::profiler::Profiler>>(std::allocator<mxnet::profiler::Profiler> const&) /usr/include/c++/8/bits/shared_ptr.h:707
    #9 0x7f7b53015d3b in std::shared_ptr<mxnet::profiler::Profiler> std::make_shared<mxnet::profiler::Profiler>() /usr/include/c++/8/bits/shared_ptr.h:723
    #10 0x7f7b530088e1 in mxnet::profiler::Profiler::Get(std::shared_ptr<mxnet::profiler::Profiler>*) /work/mxnet/src/profiler/profiler.cc:106
    #11 0x7f7b531d9778 in mxnet::engine::ThreadedEngine::ThreadedEngine() /work/mxnet/src/engine/./threaded_engine.h:310
    #12 0x7f7b531dc9ac in mxnet::engine::ThreadedEnginePerDevice::ThreadedEnginePerDevice() /work/mxnet/src/engine/threaded_engine_perdevice.cc:54
    #13 0x7f7b531d73a9 in mxnet::engine::CreateThreadedEnginePerDevice() /work/mxnet/src/engine/threaded_engine_perdevice.cc:342
    #14 0x7f7b5320633e in mxnet::engine::CreateEngine() /work/mxnet/src/engine/engine.cc:45
    #15 0x7f7b53205e4c in mxnet::Engine::_GetSharedRef() /work/mxnet/src/engine/engine.cc:62
    #16 0x7f7b53205fab in mxnet::Engine::Get() /work/mxnet/src/engine/engine.cc:67
    #17 0x7f7b5308422c in mxnet::LibraryInitializer::LibraryInitializer()::{lambda()#1}::operator()() const /work/mxnet/src/initialize.cc:54
    #18 0x7f7b5308428d in mxnet::LibraryInitializer::LibraryInitializer()::{lambda()#1}::_FUN() /work/mxnet/src/initialize.cc:55
    #19 0x7f7bac5303a4 in __fork (/lib/x86_64-linux-gnu/libc.so.6+0xcc3a4)
    #20 0x5e9abc  (/usr/bin/python3.5+0x5e9abc)
...

SUMMARY: AddressSanitizer: 7521184 byte(s) leaked in 3071 allocation(s).


Once the test is run, a developer will have to look at the stacks of memory allocations and determine which are important, and which are running as designed.  All of the reports in the sample output are displaying memory leak information, but other bugs will also be output and included in the summary if they are found.  If you are trying to reproduce a memory leak reported by users the recommend approach would be to reproduce the error by emulating the user's use case and running it in a loop with ASAN enabled.  You can then reduce the scope (for example to a C++ unit test) while continuing to run in a loop.  Eventually the stack traces should make it clear why the leak occurs.

Running in CLion

If you prefer debugging in an IDE, or you wish to break on ASAN errors you can use CLion.  You'll need to install all the required dependencies manually (but you can reference the CI Dockerfiles for help).  On Ubunut 16.04 this is as simple as installing gcc-8, and setting it as the project compiler in CLion.  You can then add the -DUSE_ASAN build flag for your project to enable ASAN support.  Finally you must add LD_PRELOAD to your run environment variables for your launch target and point it at the version of ASAN you have installed (libasan.so.5 for gcc8).

Disabling Memory Leak Detection

When ASAN builds are enabled we have leaks that are reported when running almost any MXNet test.  If you want to focus on the possibly more important memory errors such as buffer overflows, you can turn off leak detection by setting ASAN_OPTIONS=detect_leaks=0.

Other Sanitizers

After enabling and addressing issues reported by ASAN we can enable other sanitizers following the same template.  The two most applicable sanitizers are described below.

TSAN

TSAN is a sanitizer that detects data races and other thread-saftey errors in native libraries. TSAN works in a similar fashion to ASAN.  It instruments builds and surrounds memory with protect access buffers.  It then uses this instrumented code and specially protected buffers to ensures that each thread accesses memory in a threadsafe way.  TSAN supports C++11 atomics and other modern C++ features.  TSAN has more overhead (especially in memory usage) than ASAN.

MSAN

MSAN detects uninitialized memory accesses.  This could help us reduce errors in MXNet, especially difficult to reproduce, non-deterministic errors.  MSAN has a slowdown of roughly 3x when it instruments MXNet.