Merge pull request #93 from fcollonval/fcollonval/issue92

Clarify notebook/jupyter_server usage

Author: jtpio

jupyter_resource_usage/__init__.py

@@ -1,7 +1,7 @@
 import json
 import os.path as osp
 
-from notebook.utils import url_path_join
+from jupyter_server.utils import url_path_join
 from tornado import ioloop
 
 from jupyter_resource_usage.api import ApiHandler
@@ -19,7 +19,7 @@ def _jupyter_labextension_paths():
     return [{"src": "labextension", "dest": data["name"]}]
 
 
-def _jupyter_server_extension_paths():
+def _jupyter_server_extension_points():
     """
     Set up the server extension for collecting metrics
     """
@@ -40,19 +40,22 @@ def _jupyter_nbextension_paths():
     ]
 
 
-def load_jupyter_server_extension(nbapp):
+def _load_jupyter_server_extension(server_app):
     """
     Called during notebook start
     """
-    resuseconfig = ResourceUseDisplay(parent=nbapp)
-    nbapp.web_app.settings["jupyter_resource_usage_display_config"] = resuseconfig
-    base_url = nbapp.web_app.settings["base_url"]
+    resuseconfig = ResourceUseDisplay(parent=server_app)
+    server_app.web_app.settings["jupyter_resource_usage_display_config"] = resuseconfig
+    base_url = server_app.web_app.settings["base_url"]
 
-    nbapp.web_app.add_handlers(
+    server_app.web_app.add_handlers(
         ".*", [(url_path_join(base_url, "/api/metrics/v1"), ApiHandler)]
     )
 
     callback = ioloop.PeriodicCallback(
-        PrometheusHandler(PSUtilMetricsLoader(nbapp)), 1000
+        PrometheusHandler(PSUtilMetricsLoader(server_app)), 1000
     )
     callback.start()
+
+
+load_jupyter_server_extension = _load_jupyter_server_extension

jupyter_resource_usage/api.py

@@ -2,7 +2,7 @@
 from concurrent.futures import ThreadPoolExecutor
 
 import psutil
-from notebook.base.handlers import IPythonHandler
+from jupyter_server.base.handlers import APIHandler
 from tornado import web
 from tornado.concurrent import run_on_executor
 
@@ -13,9 +13,14 @@
     from .utils import Callable
 
 
-class ApiHandler(IPythonHandler):
+class ApiHandler(APIHandler):
+    def initialize(self):
+        self.executor = ThreadPoolExecutor(max_workers=5)
 
-    executor = ThreadPoolExecutor(max_workers=5)
+    def __del__(self):
+        # Ensure the executor is properly shutdown
+        self.executor.shutdown(wait=False)
+        super().__del__()
 
     @web.authenticated
     async def get(self):

jupyter_resource_usage/metrics.py

@@ -3,13 +3,15 @@
 except ImportError:
     psutil = None
 
-from notebook.notebookapp import NotebookApp
+from jupyter_server.serverapp import ServerApp
 
 
 class PSUtilMetricsLoader:
-    def __init__(self, nbapp: NotebookApp):
-        self.config = nbapp.web_app.settings["jupyter_resource_usage_display_config"]
-        self.nbapp = nbapp
+    def __init__(self, server_app: ServerApp):
+        self.config = server_app.web_app.settings[
+            "jupyter_resource_usage_display_config"
+        ]
+        self.server_app = server_app
 
     def get_process_metric_value(self, process, name, kwargs, attribute=None):
         try:

jupyter_resource_usage/prometheus.py

@@ -1,9 +1,9 @@
 from typing import Optional
 
-from jupyter_resource_usage.metrics import PSUtilMetricsLoader
-from notebook.notebookapp import NotebookApp
 from prometheus_client import Gauge
 
+from jupyter_resource_usage.metrics import PSUtilMetricsLoader
+
 try:
     # Traitlets >= 4.3.3
     from traitlets import Callable
@@ -16,7 +16,7 @@ def __init__(self, metricsloader: PSUtilMetricsLoader):
         super().__init__()
         self.metricsloader = metricsloader
         self.config = metricsloader.config
-        self.session_manager = metricsloader.nbapp.session_manager
+        self.session_manager = metricsloader.server_app.session_manager
 
         gauge_names = ["total_memory", "max_memory", "total_cpu", "max_cpu"]
         for name in gauge_names:

jupyter_resource_usage/tests/test_basic.py

@@ -8,12 +8,12 @@ class TestBasic:
     def test_import_serverextension(self):
         """Check that serverextension hooks are available"""
         from jupyter_resource_usage import (
-            _jupyter_server_extension_paths,
+            _jupyter_server_extension_points,
             _jupyter_nbextension_paths,
-            load_jupyter_server_extension,
+            _load_jupyter_server_extension,
         )
 
-        assert _jupyter_server_extension_paths() == [
+        assert _jupyter_server_extension_points() == [
             {"module": "jupyter_resource_usage"}
         ]
         assert _jupyter_nbextension_paths() == [
@@ -39,7 +39,7 @@ def test_import_serverextension(self):
         ) as psutil_metrics_loader:
 
             # load up with mock
-            load_jupyter_server_extension(nbapp_mock)
+            _load_jupyter_server_extension(nbapp_mock)
 
             # assert that we installed the application in settings
             print(nbapp_mock.web_app.settings)

setup.py

@@ -75,7 +75,7 @@
     long_description_content_type="text/markdown",
     cmdclass=cmdclass,
     packages=setuptools.find_packages(),
-    install_requires=["notebook>=5.6.0", "prometheus_client", "psutil>=5.6.0"],
+    install_requires=["jupyter_server>=1.0.0", "prometheus_client", "psutil>=5.6.0"],
     extras_require={
         "dev": ["autopep8", "black", "pytest", "flake8", "pytest-cov>=2.6.1", "mock"]
     },