jupyter_ext_dev: pull request for jupyter extension 2.x and 3.x recap #4141 (#4167)

nni/tools/jupyter_extension/management.py

 import json
 from pathlib import Path
 import shutil
+import os
 
 from jupyter_core.paths import jupyter_config_dir, jupyter_data_dir
 
 import nni_node
 
+def _get_jupyter_lab_version():
+    try:
+        import jupyterlab
+        return jupyterlab.__version__
+    except ImportError:
+        return '3.x'
+
+jupyter_lab_major_version = _get_jupyter_lab_version().split('.')[0]
+
 _backend_config_file = Path(jupyter_config_dir(), 'jupyter_server_config.d', 'nni.json')
 _backend_config_content = {
     'ServerApp': {
@@ -14,6 +24,14 @@ _backend_config_content = {
         }
     }
 }
+_v2_backend_config_file = Path(jupyter_config_dir(), 'jupyter_notebook_config.d', 'nni.json')
+_v2_backend_config_content = {
+    "NotebookApp": {
+        "nbserver_extensions": {
+            "nni.tools.jupyter_extension": True
+        }
+    }
+}
 
 _frontend_src = Path(nni_node.__path__[0], 'jupyter-extension')
 _frontend_dst = Path(jupyter_data_dir(), 'labextensions', 'nni-jupyter-extension')
@@ -23,8 +41,21 @@ def install():
     _backend_config_file.write_text(json.dumps(_backend_config_content))
 
     _frontend_dst.parent.mkdir(parents=True, exist_ok=True)
-    shutil.copytree(_frontend_src, _frontend_dst)
+
+    if jupyter_lab_major_version == '2':
+        _v2_backend_config_file.parent.mkdir(parents=True, exist_ok=True)
+        _v2_backend_config_file.write_text(json.dumps(_v2_backend_config_content))
+
+        if (_frontend_src.is_symlink()):
+            linkto = os.path.realpath(_frontend_src)
+            os.symlink(linkto, _frontend_dst)
+        else:
+            shutil.copytree(_frontend_src, _frontend_dst)
+    else:
+        shutil.copytree(_frontend_src, _frontend_dst)
 
 def uninstall():
     _backend_config_file.unlink()
+    if jupyter_lab_major_version == '2':
+        _v2_backend_config_file.unlink()
     shutil.rmtree(_frontend_dst)

setup.py

@@ -43,6 +43,10 @@ Build wheel package:
   $ NNI_RELEASE=2.0 python setup.py build_ts
   $ NNI_RELEASE=2.0 python setup.py bdist_wheel -p manylinux1_x86_64
 
+for jupyterlab 2.x package:
+  $ JUPYTER_LAB_VERSION=2.3.1 NNI_RELEASE=2.0 python setup.py build_ts
+  $ JUPYTER_LAB_VERSION=2.3.1 NNI_RELEASE=2.0 python setup.py bdist_wheel -p manylinux1_x86_64
+
 Where "2.0" is version string and "manylinux1_x86_64" is platform.
 The platform may also be "macosx_10_9_x86_64" or "win_amd64".
 
@@ -65,6 +69,26 @@ import setup_ts
 
 release = os.environ.get('NNI_RELEASE')
 
+def _get_jupyter_lab_version():
+    try:
+        import jupyterlab
+        return jupyterlab.__version__
+    except ImportError:
+        return '3.x'
+
+jupyter_lab_major_version = _get_jupyter_lab_version().split('.')[0]
+
+def check_jupyter_lab_version():
+    environ_version = os.environ.get('JUPYTER_LAB_VERSION')
+
+    jupyter_lab_version = _get_jupyter_lab_version()
+
+    if environ_version:
+        if jupyter_lab_version.split('.')[0] != environ_version.split('.')[0]:
+            sys.exit(f'ERROR: To build a jupyter lab extension, run "JUPYTER_LAB_VERSION={jupyter_lab_version}", current: {environ_version} ')
+    elif jupyter_lab_version.split('.')[0] != '3':
+        sys.exit(f'ERROR: To build a jupyter lab extension, run "JUPYTER_LAB_VERSION={jupyter_lab_version}" first for nondefault version(3.x)')
+
 def _setup():
     setuptools.setup(
         name = 'nni',
@@ -91,6 +115,8 @@ def _setup():
             'nni_node': _find_node_files()  # note: this does not work before building
         },
 
+        data_files = _get_data_files(),
+
         python_requires = '>=3.6',
         install_requires = _read_requirements_txt('dependencies/required.txt'),
         extras_require = {
@@ -115,6 +141,12 @@ def _setup():
         }
     )
 
+def _get_data_files():
+    data_files = []
+    if jupyter_lab_major_version == '2':
+        extension_file = glob.glob("nni_node/jupyter-extension/extensions/nni-jupyter-extension*.tgz")
+        data_files = [('share/jupyter/lab/extensions', extension_file)]
+    return data_files
 
 def _find_python_packages():
     packages = []
@@ -178,12 +210,16 @@ class BuildTs(Command):
         pass
 
     def run(self):
+        check_jupyter_lab_version()
         setup_ts.build(release)
 
 class Build(build):
     def run(self):
         if not release:
             sys.exit('Please set environment variable "NNI_RELEASE=<release_version>"')
+
+        check_jupyter_lab_version()
+
         if os.path.islink('nni_node/main.js'):
             sys.exit('A development build already exists. Please uninstall NNI and run "python3 setup.py clean --all".')
         open('nni/version.py', 'w').write(f"__version__ = '{release}'")

setup_ts.py

@@ -26,6 +26,14 @@ from zipfile import ZipFile
 node_version = 'v16.3.0'
 yarn_version = 'v1.22.10'
 
+def _get_jupyter_lab_version():
+    try:
+        import jupyterlab
+        return jupyterlab.__version__
+    except ImportError:
+        return '3.x'
+
+jupyter_lab_major_version = _get_jupyter_lab_version().split('.')[0]
 
 def build(release):
     """
@@ -40,11 +48,13 @@ def build(release):
     if release or not os.environ.get('GLOBAL_TOOLCHAIN'):
         download_toolchain()
     prepare_nni_node()
+    update_package()
     compile_ts(release)
     if release or sys.platform == 'win32':
         copy_nni_node(release)
     else:
         symlink_nni_node()
+    restore_package()
 
 def clean(clean_all=False):
     """
@@ -120,6 +130,25 @@ def download_toolchain():
     shutil.rmtree('toolchain/yarn', ignore_errors=True)
     Path(f'toolchain/yarn-{yarn_version}').rename('toolchain/yarn')
 
+def update_package():
+    if jupyter_lab_major_version == '2':
+        package_json = json.load(open('ts/jupyter_extension/package.json'))
+        json.dump(package_json, open('ts/jupyter_extension/.package_default.json', 'w'), indent=2)
+
+        package_json['scripts']['build'] = 'tsc && jupyter labextension link .'
+        package_json['dependencies']['@jupyterlab/application'] = '^2.3.0'
+        package_json['dependencies']['@jupyterlab/launcher'] = '^2.3.0'
+
+        package_json['jupyterlab']['outputDir'] = 'build'
+        json.dump(package_json, open('ts/jupyter_extension/package.json', 'w'), indent=2)
+        print(f'updated package.json with {json.dumps(package_json, indent=2)}')
+
+def restore_package():
+    if jupyter_lab_major_version == '2':
+        package_json = json.load(open('ts/jupyter_extension/.package_default.json'))
+        print(f'stored package.json with {json.dumps(package_json, indent=2)}')
+        json.dump(package_json, open('ts/jupyter_extension/package.json', 'w'), indent=2)
+        os.remove('ts/jupyter_extension/.package_default.json')
 
 def prepare_nni_node():
     """
@@ -177,7 +206,10 @@ def symlink_nni_node():
 
     _symlink('ts/webui/build', 'nni_node/static')
 
-    if Path('ts/jupyter_extension/dist').exists():
+    if jupyter_lab_major_version == '2':
+        _symlink('ts/jupyter_extension/build', 'nni_node/jupyter-extension')
+        _symlink(os.path.join(sys.exec_prefix, 'share/jupyter/lab/extensions'), 'nni_node/jupyter-extension/extensions')
+    elif Path('ts/jupyter_extension/dist').exists():
         _symlink('ts/jupyter_extension/dist', 'nni_node/jupyter-extension')
 
 
@@ -209,7 +241,10 @@ def copy_nni_node(version):
 
     shutil.copytree('ts/webui/build', 'nni_node/static')
 
-    if version or Path('ts/jupyter_extension/dist').exists():
+    if jupyter_lab_major_version == '2':
+        shutil.copytree('ts/jupyter_extension/build', 'nni_node/jupyter-extension/build')
+        shutil.copytree(os.path.join(sys.exec_prefix, 'share/jupyter/lab/extensions'), 'nni_node/jupyter-extension/extensions')
+    elif version or Path('ts/jupyter_extension/dist').exists():
         shutil.copytree('ts/jupyter_extension/dist', 'nni_node/jupyter-extension')