lazy import jupyter-related package (#4045)

nni/tools/nnictl/nnictl.py

@@ -6,7 +6,6 @@ import logging
 import os
 import pkg_resources
 from colorama import init
-import nni.tools.jupyter_extension.management as jupyter_management
 from .common_utils import print_error
 from .launcher import create_experiment, resume_experiment, view_experiment
 from .updater import update_searchspace, update_concurrency, update_duration, update_trialnum, import_data
@@ -283,12 +282,22 @@ def parse_args():
     jupyter_parser = subparsers.add_parser('jupyter-extension', help='install or uninstall JupyterLab extension (internal preview)')
     jupyter_subparsers = jupyter_parser.add_subparsers()
     jupyter_install_parser = jupyter_subparsers.add_parser('install', help='install JupyterLab extension')
-    jupyter_install_parser.set_defaults(func=lambda _args: jupyter_management.install())  # TODO: prompt message
+    jupyter_install_parser.set_defaults(func=_jupyter_install)
     jupyter_uninstall_parser = jupyter_subparsers.add_parser('uninstall', help='uninstall JupyterLab extension')
-    jupyter_uninstall_parser.set_defaults(func=lambda _args: jupyter_management.uninstall())
+    jupyter_uninstall_parser.set_defaults(func=_jupyter_uninstall)
 
     args = parser.parse_args()
     args.func(args)
 
+def _jupyter_install(_args):
+    import nni.tools.jupyter_extension.management as jupyter_management
+    jupyter_management.install()
+    print('Successfully installed JupyterLab extension')
+
+def _jupyter_uninstall(_args):
+    import nni.tools.jupyter_extension.management as jupyter_management
+    jupyter_management.uninstall()
+    print('Successfully uninstalled JupyterLab extension')
+
 if __name__ == '__main__':
     parse_args()