JupyterLab extension (#3954)

dependencies/develop.txt

@@ -10,3 +10,4 @@ pytest-cov
 pytest-azurepipelines
 coverage
 ipython
+jupyterlab

nni/tools/jupyter_extension/__init__.py

@@ -0,0 +1,4 @@
+from . import proxy
+
+load_jupyter_server_extension = proxy.setup
+_load_jupyter_server_extension = proxy.setup

nni/tools/jupyter_extension/management.py

@@ -0,0 +1,30 @@
+import json
+from pathlib import Path
+import shutil
+
+from jupyter_core.paths import jupyter_config_dir, jupyter_data_dir
+
+import nni_node
+
+_backend_config_file = Path(jupyter_config_dir(), 'jupyter_server_config.d', 'nni.json')
+_backend_config_content = {
+    'ServerApp': {
+        'jpserver_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')
+
+def install():
+    _backend_config_file.parent.mkdir(parents=True, exist_ok=True)
+    _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)
+
+def uninstall():
+    _backend_config_file.unlink()
+    shutil.rmtree(_frontend_dst)

nni/tools/jupyter_extension/proxy.py

@@ -0,0 +1,43 @@
+import json
+from pathlib import Path
+
+import requests
+from tornado.web import RequestHandler
+
+def setup(server):
+    base_url = server.web_app.settings['base_url']
+    url_pattern = base_url.rstrip('/') + '/nni/(.*)'
+    server.web_app.add_handlers('.*$', [(url_pattern, NniProxyHandler)])
+
+class NniProxyHandler(RequestHandler):
+    def get(self, path):
+        ports = _get_experiment_ports()
+        if not ports:
+            self.set_status(404)
+            return
+
+        if path == 'index':
+            if len(ports) > 1:  # if there is more than one running experiments, show experiment list
+                self.redirect('experiment')
+            else:  # if there is only one running experiment, show that experiment
+                self.redirect('oview')
+            return
+
+        r = requests.get(f'http://localhost:{ports[0]}/{path}')
+        self.set_status(r.status_code)
+        for key, value in r.headers.items():
+            self.add_header(key, value)
+        self.finish(r.content)
+
+    # TODO: post, put, etc
+
+    def set_default_headers(self):
+        self.clear_header('Content-Type')
+        self.clear_header('Date')
+
+def _get_experiment_ports():
+    experiment_list_path = Path.home() / 'nni-experiments/.experiment'
+    if not experiment_list_path.exists():
+        return None
+    experiments = json.load(open(experiment_list_path))
+    return [exp['port'] for exp in experiments.values() if exp['status'] != 'STOPPED']

nni/tools/nnictl/nnictl.py

@@ -6,6 +6,7 @@ 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
@@ -278,6 +279,14 @@ def parse_args():
     'the unit is second')
     parser_top.set_defaults(func=monitor_experiment)
 
+    # jupyter-extension command
+    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_uninstall_parser = jupyter_subparsers.add_parser('uninstall', help='uninstall JupyterLab extension')
+    jupyter_uninstall_parser.set_defaults(func=lambda _args: jupyter_management.uninstall())
+
     args = parser.parse_args()
     args.func(args)
 

pylintrc

@@ -47,4 +47,4 @@ ignore-patterns=test*
 # List of members which are set dynamically and missed by pylint inference
 generated-members=numpy.*,torch.*,tensorflow.*,pycuda.*,tensorrt.*
 
-ignored-modules=tensorflow,_winapi,msvcrt,tensorrt,pycuda
+ignored-modules=tensorflow,_winapi,msvcrt,tensorrt,pycuda,nni_node

setup_ts.py

@@ -153,6 +153,10 @@ def compile_ts():
     _yarn('ts/nasui')
     _yarn('ts/nasui', 'build')
 
+    _print('Building JupyterLab extension')
+    _yarn('ts/jupyter_extension')
+    _yarn('ts/jupyter_extension', 'build')
+
 
 def symlink_nni_node():
     """
@@ -172,6 +176,8 @@ def symlink_nni_node():
     _symlink('ts/nasui/build', 'nni_node/nasui/build')
     _symlink('ts/nasui/server.js', 'nni_node/nasui/server.js')
 
+    _symlink('ts/jupyter_extension/dist', 'nni_node/jupyter-extension')
+
 
 def copy_nni_node(version):
     """
@@ -205,6 +211,8 @@ def copy_nni_node(version):
     shutil.copytree('ts/nasui/build', 'nni_node/nasui/build')
     shutil.copyfile('ts/nasui/server.js', 'nni_node/nasui/server.js')
 
+    shutil.copytree('ts/jupyter_extension/dist', 'nni_node/jupyter-extension')
+
 
 _yarn_env = dict(os.environ)
 # `Path('nni_node').resolve()` does not work on Windows if the directory not exists

ts/jupyter_extension/.gitignore

@@ -0,0 +1,2 @@
+/build
+/nni

ts/jupyter_extension/package.json

@@ -0,0 +1,20 @@
+{
+    "name": "nni-jupyter-extension",
+    "version": "999.0.0-developing",
+    "license": "MIT",
+    "scripts": {
+        "build": "tsc && jupyter labextension build ."
+    },
+    "dependencies": {
+        "@jupyterlab/application": "^3.0.11",
+        "@jupyterlab/launcher": "^3.0.9"
+    },
+    "devDependencies": {
+        "@jupyterlab/builder": "^3.0.9"
+    },
+    "jupyterlab": {
+        "extension": true,
+        "outputDir": "dist"
+    },
+    "main": "build/index.js"
+}

ts/jupyter_extension/src/index.ts

@@ -0,0 +1,55 @@
+import { JupyterFrontEnd, JupyterFrontEndPlugin } from '@jupyterlab/application';
+import { ICommandPalette, IFrame } from '@jupyterlab/apputils';
+import { PageConfig } from '@jupyterlab/coreutils';
+import { ILauncher } from '@jupyterlab/launcher';
+
+class NniWidget extends IFrame {
+    constructor() {
+        super({
+            sandbox: [
+                'allow-same-origin',
+                'allow-scripts',
+            ]
+        });
+        this.url = PageConfig.getBaseUrl() + 'nni/index';
+        this.id = 'nni';
+        this.title.label = 'NNI';
+        this.title.closable = true;
+    }
+}
+
+async function activate(app: JupyterFrontEnd, palette: ICommandPalette, launcher: ILauncher | null) {
+    console.log('nni extension is activated');
+    const { commands, shell } = app;
+    const command = 'nni';
+    const category = 'Other';
+
+    commands.addCommand(command, {
+        label: 'NNI',
+        caption: 'NNI',
+        iconClass: (args) => (args.isPalette ? null : 'jp-Launcher-kernelIcon'),
+        execute: () => {
+            shell.add(new NniWidget(), 'main');
+        }
+    });
+
+    palette.addItem({ command, category });
+
+    if (launcher) {
+        launcher.add({
+            command,
+            category,
+            kernelIconUrl: '/nni/icon.png'  // FIXME: this field only works for "Notebook" category
+        });
+    }
+}
+
+const extension: JupyterFrontEndPlugin<void> = {
+    id: 'nni',
+    autoStart: true,
+    optional: [ILauncher],
+    requires: [ICommandPalette],
+    activate,
+};
+
+export default extension;

ts/jupyter_extension/tsconfig.json

@@ -0,0 +1,14 @@
+{
+    "compilerOptions": {
+        "allowSyntheticDefaultImports": true,
+        "jsx": "react",
+        "module": "esnext",
+        "moduleResolution": "node",
+        "outDir": "build",
+        "rootDir": "src",
+        "target": "es2017"
+    },
+    "include": [
+        "src/*"
+    ]
+}