From 13d8e1e30f29aefef940d9238c028fecc64771e6 Mon Sep 17 00:00:00 2001
From: Pedro Larroy <pedro.larroy.lists@gmail.com>
Date: Wed, 31 Jul 2019 18:57:38 -0700
Subject: [PATCH] Improve diagnose.py to display environment variables

---
 tools/diagnose.py | 53 ++++++++++++++++++++++++++++++++---------------
 1 file changed, 36 insertions(+), 17 deletions(-)

diff --git a/tools/diagnose.py b/tools/diagnose.py
index f3026e635844..9cd5792a2f1a 100755
--- a/tools/diagnose.py
+++ b/tools/diagnose.py
@@ -30,23 +30,6 @@
     from urllib2 import urlopen
 import argparse
 
-def parse_args():
-    """Parse arguments."""
-    parser = argparse.ArgumentParser(
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-        description='Diagnose script for checking the current system.')
-    choices = ['python', 'pip', 'mxnet', 'os', 'hardware', 'network']
-    for choice in choices:
-        parser.add_argument('--' + choice, default=1, type=int,
-                            help='Diagnose {}.'.format(choice))
-    parser.add_argument('--region', default='', type=str,
-                        help="Additional sites in which region(s) to test. \
-                        Specify 'cn' for example to test mirror sites in China.")
-    parser.add_argument('--timeout', default=10, type=int,
-                        help="Connection test timeout threshold, 0 to disable.")
-    args = parser.parse_args()
-    return args
-
 URLS = {
     'MXNet': 'https://github.com/apache/incubator-mxnet',
     'Gluon Tutorial(en)': 'http://gluon.mxnet.io',
@@ -55,6 +38,7 @@ def parse_args():
     'PYPI': 'https://pypi.python.org/pypi/pip',
     'Conda': 'https://repo.continuum.io/pkgs/free/',
 }
+
 REGIONAL_URLS = {
     'cn': {
         'PYPI(douban)': 'https://pypi.douban.com/',
@@ -81,6 +65,7 @@ def test_connection(name, url, timeout=10):
     load_elapsed = time.time() - start
     print("Timing for {}: {}, DNS: {:.4f} sec, LOAD: {:.4f} sec.".format(name, url, dns_elapsed, load_elapsed))
 
+
 def check_python():
     print('----------Python Info----------')
     print('Version      :', platform.python_version())
@@ -97,11 +82,13 @@ def check_pip():
     except ImportError:
         print('No corresponding pip install for current python.')
 
+
 def get_build_features_str():
     import mxnet.runtime
     features = mxnet.runtime.Features()
     return '\n'.join(map(str, list(features.values())))
 
+
 def check_mxnet():
     print('----------MXNet Info-----------')
     try:
@@ -131,6 +118,7 @@ def check_mxnet():
             print("This is very likely due to missing missing or incompatible library files.")
         print(traceback.format_exc())
 
+
 def check_os():
     print('----------System Info----------')
     print('Platform     :', platform.platform())
@@ -139,6 +127,7 @@ def check_os():
     print('release      :', platform.release())
     print('version      :', platform.version())
 
+
 def check_hardware():
     print('----------Hardware Info----------')
     print('machine      :', platform.machine())
@@ -154,6 +143,7 @@ def check_hardware():
     elif sys.platform.startswith('win32'):
         subprocess.call(['wmic', 'cpu', 'get', 'name'])
 
+
 def check_network(args):
     print('----------Network Test----------')
     if args.timeout > 0:
@@ -171,6 +161,32 @@ def check_network(args):
     for name, url in URLS.items():
         test_connection(name, url, args.timeout)
 
+
+def check_environment():
+    print('----------Environment----------')
+    for k,v in os.environ.items():
+        if k.startswith('MXNET_') or k.startswith('OMP_') or k.startswith('KMP_') or k == 'CC' or k == 'CXX':
+            print('{}="{}"'.format(k,v))
+
+
+def parse_args():
+    """Parse arguments."""
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description='Diagnose script for checking the current system.')
+    choices = ['python', 'pip', 'mxnet', 'os', 'hardware', 'network', 'environment']
+    for choice in choices:
+        parser.add_argument('--' + choice, default=1, type=int,
+                            help='Diagnose {}.'.format(choice))
+    parser.add_argument('--region', default='', type=str,
+                        help="Additional sites in which region(s) to test. \
+                        Specify 'cn' for example to test mirror sites in China.")
+    parser.add_argument('--timeout', default=10, type=int,
+                        help="Connection test timeout threshold, 0 to disable.")
+    args = parser.parse_args()
+    return args
+
+
 if __name__ == '__main__':
     args = parse_args()
     if args.python:
@@ -190,3 +206,6 @@ def check_network(args):
 
     if args.network:
         check_network(args)
+
+    if args.environment:
+        check_environment()