From da2f0c935ff53d44033888c0edf9a26b004cefb1 Mon Sep 17 00:00:00 2001
From: Kellen Sunderland <kellen.sunderland@gmail.com>
Date: Fri, 24 Aug 2018 14:13:58 +0200
Subject: [PATCH] [MXNET-857] Add integration test.

---
 tests/python/profiling/simple_forward.py | 42 +++++++++++++++++++
 tests/python/profiling/test_nvtx.py      | 52 ++++++++++++++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 tests/python/profiling/simple_forward.py
 create mode 100644 tests/python/profiling/test_nvtx.py

diff --git a/tests/python/profiling/simple_forward.py b/tests/python/profiling/simple_forward.py
new file mode 100644
index 000000000000..0ad43c89a6f5
--- /dev/null
+++ b/tests/python/profiling/simple_forward.py
@@ -0,0 +1,42 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import mxnet as mx
+from mxnet.gluon import nn
+
+
+def simple_forward():
+    ctx = mx.gpu()
+    mx.profiler.set_config(profile_all=True)
+    mx.profiler.set_state('run')
+
+    # define simple gluon network with random weights
+    net = nn.Sequential()
+    with net.name_scope():
+        net.add(nn.Dense(128, activation='relu'))
+        net.add(nn.Dense(64, activation='relu'))
+        net.add(nn.Dense(10))
+    net.initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)
+
+    input = mx.nd.zeros((128,), ctx=ctx)
+    predictions = net(input)
+    print('Ran simple NN forward, results:')
+    print(predictions.asnumpy())
+
+
+if __name__ == '__main__':
+    simple_forward()
diff --git a/tests/python/profiling/test_nvtx.py b/tests/python/profiling/test_nvtx.py
new file mode 100644
index 000000000000..35b209ebb6eb
--- /dev/null
+++ b/tests/python/profiling/test_nvtx.py
@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import os
+import unittest
+
+import mxnet as mx
+import sys
+
+from subprocess import Popen, PIPE
+
+
+def test_nvtx_ranges_present_in_profile():
+
+    if not mx.test_utils.list_gpus():
+        unittest.skip('Test only applicable to machines with GPUs')
+
+    # Build a system independent wrapper to execute simple_forward with nvprof
+    # This requires nvprof to be on your path (which should be the case for most GPU workstations with cuda installed).
+    simple_forward_path = os.path.realpath(__file__)
+    simple_forward_path = simple_forward_path.replace('test_nvtx', 'simple_forward')
+
+    process = Popen(["nvprof", sys.executable, simple_forward_path], stdout=PIPE, stderr=PIPE)
+    (output, profiler_output) = process.communicate()
+    process.wait()
+    profiler_output = profiler_output.decode('ascii')
+
+    # Verify that some of the NVTX ranges we should have created are present
+    # Verify that we have NVTX ranges for our simple operators.
+    assert "Range \"FullyConnected\"" in profiler_output
+    assert "Range \"_zeros\"" in profiler_output
+
+    # Verify that we have some expected output from the engine.
+    assert "Range \"WaitForVar\"" in profiler_output
+
+
+if __name__ == '__main__':
+    import nose
+    nose.runmodule()