Add async start hook to ExtensionApp API (#1417)

Co-authored-by: Darshan808 <[email protected]>

docs/source/developers/extensions.rst

@@ -65,6 +65,33 @@ Then add this handler to Jupyter Server's Web Application through the ``_load_ju
         serverapp.web_app.add_handlers(".*$", handlers)
 
 
+Starting asynchronous tasks from an extension
+---------------------------------------------
+
+.. versionadded:: 2.15.0
+
+Jupyter Server offers a simple API for starting asynchronous tasks from a server extension. This is useful for calling
+async tasks after the event loop is running.
+
+The function should be named ``_start_jupyter_server_extension`` and found next to the ``_load_jupyter_server_extension`` function.
+
+Here is basic example:
+
+.. code-block:: python
+
+    import asyncio
+
+    async def _start_jupyter_server_extension(serverapp: jupyter_server.serverapp.ServerApp):
+        """
+        This function is called after the server's event loop is running.
+        """
+        await asyncio.sleep(.1)
+
+.. note:: The server startup banner (displaying server info and access URLs) is printed before starting asynchronous tasks, so those tasks might still be running even after the banner appears.
+
+..  WARNING: This note is also present in the "Starting asynchronous tasks from an ExtensionApp" section.
+   If you update it here, please update it there as well.
+
 Making an extension discoverable
 --------------------------------
 
@@ -117,6 +144,7 @@ An ExtensionApp:
     - has an entrypoint, ``jupyter <name>``.
     - can serve static content from the ``/static/<name>/`` endpoint.
     - can add new endpoints to the Jupyter Server.
+    - can start asynchronous tasks after the server has started.
 
 The basic structure of an ExtensionApp is shown below:
 
@@ -156,6 +184,11 @@ The basic structure of an ExtensionApp is shown below:
             ...
             # Change the jinja templating environment
 
+        async def _start_jupyter_server_extension(self):
+            ...
+            # Extend this method to start any (e.g. async) tasks
+            # after the main Server's Event Loop is running.
+
         async def stop_extension(self):
             ...
             # Perform any required shut down steps
@@ -171,6 +204,7 @@ Methods
 * ``initialize_settings()``: adds custom settings to the Tornado Web Application.
 * ``initialize_handlers()``: appends handlers to the Tornado Web Application.
 * ``initialize_templates()``: initialize the templating engine (e.g. jinja2) for your frontend.
+* ``_start_jupyter_server_extension()``: enables the extension to start (async) tasks _after_ the server's main Event Loop has started.
 * ``stop_extension()``: called on server shut down.
 
 Properties
@@ -320,6 +354,48 @@ pointing at the ``load_classic_server_extension`` method:
 If the extension is enabled, the extension will be loaded when the server starts.
 
 
+Starting asynchronous tasks from an ExtensionApp
+------------------------------------------------
+
+.. versionadded:: 2.15.0
+
+
+An ``ExtensionApp`` can start asynchronous tasks after Jupyter Server's event loop is started by overriding its
+``_start_jupyter_server_extension()`` method. This can be helpful for setting up e.g. background tasks.
+
+Here is a basic (pseudo) code example:
+
+.. code-block:: python
+
+    import asyncio
+    import time
+
+
+    async def log_time_periodically(log, dt=1):
+        """Log the current time from a periodic loop."""
+        while True:
+            current_time = time.time()
+            log.info(current_time)
+            await sleep(dt)
+
+
+    class MyExtension(ExtensionApp):
+        ...
+
+        async def _start_jupyter_server_extension(self):
+            self.my_background_task = asyncio.create_task(
+                log_time_periodically(self.log)
+            )
+
+        async def stop_extension(self):
+            self.my_background_task.cancel()
+
+.. note:: The server startup banner (displaying server info and access URLs) is printed before starting asynchronous tasks, so those tasks might still be running even after the banner appears.
+
+..  WARNING: This note is also present in the "Starting asynchronous tasks from an extension" section.
+   If you update it here, please update it there as well.
+
+
 Distributing a server extension
 ===============================
 

examples/simple/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "hatchling.build"
 name = "jupyter-server-example"
 description = "Jupyter Server Example"
 readme = "README.md"
-license = "MIT"
+license = "BSD-3-Clause"
 requires-python = ">=3.9"
 dependencies = [
     "jinja2",

jupyter_server/extension/application.py

@@ -475,6 +475,18 @@ def _load_jupyter_server_extension(cls, serverapp):
         extension.initialize()
         return extension
 
+    async def _start_jupyter_server_extension(self, serverapp):
+        """
+        An async hook to start e.g. tasks from the extension after
+        the server's event loop is running.
+
+        Override this method (no need to call `super()`) to
+        start (async) tasks from an extension.
+
+        This is useful for starting e.g. background tasks from
+        an extension.
+        """
+
     @classmethod
     def load_classic_server_extension(cls, serverapp):
         """Enables extension to be loaded as classic Notebook (jupyter/notebook) extension."""

jupyter_server/extension/manager.py

@@ -119,6 +119,24 @@ def _get_loader(self):
         loader = get_loader(loc)
         return loader
 
+    def _get_starter(self):
+        """Get a starter function."""
+        if self.app:
+            linker = self.app._start_jupyter_server_extension
+        else:
+
+            async def _noop_start(serverapp):
+                return
+
+            linker = getattr(
+                self.module,
+                # Search for a _start_jupyter_extension
+                "_start_jupyter_server_extension",
+                # Otherwise return a no-op function.
+                _noop_start,
+            )
+        return linker
+
     def validate(self):
         """Check that both a linker and loader exists."""
         try:
@@ -150,6 +168,13 @@ def load(self, serverapp):
         loader = self._get_loader()
         return loader(serverapp)
 
+    async def start(self, serverapp):
+        """Call's the extensions 'start' hook where it can
+        start (possibly async) tasks _after_ the event loop is running.
+        """
+        starter = self._get_starter()
+        return await starter(serverapp)
+
 
 class ExtensionPackage(LoggingConfigurable):
     """An API for interfacing with a Jupyter Server extension package.
@@ -222,6 +247,11 @@ def load_point(self, point_name, serverapp):
         point = self.extension_points[point_name]
         return point.load(serverapp)
 
+    async def start_point(self, point_name, serverapp):
+        """Load an extension point."""
+        point = self.extension_points[point_name]
+        return await point.start(serverapp)
+
     def link_all_points(self, serverapp):
         """Link all extension points."""
         for point_name in self.extension_points:
@@ -231,9 +261,14 @@ def load_all_points(self, serverapp):
         """Load all extension points."""
         return [self.load_point(point_name, serverapp) for point_name in self.extension_points]
 
+    async def start_all_points(self, serverapp):
+        """Load all extension points."""
+        for point_name in self.extension_points:
+            await self.start_point(point_name, serverapp)
+
 
 class ExtensionManager(LoggingConfigurable):
-    """High level interface for findind, validating,
+    """High level interface for finding, validating,
     linking, loading, and managing Jupyter Server extensions.
 
     Usage:
@@ -367,6 +402,22 @@ def load_extension(self, name):
             else:
                 self.log.info("%s | extension was successfully loaded.", name)
 
+    async def start_extension(self, name):
+        """Start an extension by name."""
+        extension = self.extensions.get(name)
+
+        if extension and extension.enabled:
+            try:
+                await extension.start_all_points(self.serverapp)
+            except Exception as e:
+                if self.serverapp and self.serverapp.reraise_server_extension_failures:
+                    raise
+                self.log.warning(
+                    "%s | extension failed starting with message: %r", name, e, exc_info=True
+                )
+            else:
+                self.log.debug("%s | extension was successfully started.", name)
+
     async def stop_extension(self, name, apps):
         """Call the shutdown hooks in the specified apps."""
         for app in apps:
@@ -392,6 +443,12 @@ def load_all_extensions(self):
         for name in self.sorted_extensions:
             self.load_extension(name)
 
+    async def start_all_extensions(self):
+        """Start all enabled extensions."""
+        # Sort the extension names to enforce deterministic loading
+        # order.
+        await multi([self.start_extension(name) for name in self.sorted_extensions])
+
     async def stop_all_extensions(self):
         """Call the shutdown hooks in all extensions."""
         await multi(list(starmap(self.stop_extension, sorted(dict(self.extension_apps).items()))))

jupyter_server/serverapp.py

@@ -3164,6 +3164,7 @@ def start_ioloop(self) -> None:
             pc = ioloop.PeriodicCallback(lambda: None, 5000)
             pc.start()
         try:
+            self.io_loop.add_callback(self._post_start)
             self.io_loop.start()
         except KeyboardInterrupt:
             self.log.info(_i18n("Interrupted..."))
@@ -3172,6 +3173,17 @@ def init_ioloop(self) -> None:
         """init self.io_loop so that an extension can use it by io_loop.call_later() to create background tasks"""
         self.io_loop = ioloop.IOLoop.current()
 
+    async def _post_start(self):
+        """Add an async hook to start tasks after the event loop is running.
+
+        This will also attempt to start all tasks found in
+        the `start_extension` method in Extension Apps.
+        """
+        try:
+            await self.extension_manager.start_all_extensions()
+        except Exception as err:
+            self.log.error(err)
+
     def start(self) -> None:
         """Start the Jupyter server app, after initialization
 

tests/extension/mockextensions/app.py

@@ -5,7 +5,7 @@
 from jupyter_events import EventLogger
 from jupyter_events.schema_registry import SchemaRegistryException
 from tornado import web
-from traitlets import List, Unicode
+from traitlets import Bool, List, Unicode
 
 from jupyter_server.base.handlers import JupyterHandler
 from jupyter_server.extension.application import ExtensionApp, ExtensionAppJinjaMixin
@@ -56,6 +56,7 @@ class MockExtensionApp(ExtensionAppJinjaMixin, ExtensionApp):
     static_paths = [STATIC_PATH]  # type:ignore[assignment]
     mock_trait = Unicode("mock trait", config=True)
     loaded = False
+    started = Bool(False)
 
     serverapp_config = {
         "jpserver_extensions": {
@@ -64,6 +65,9 @@ class MockExtensionApp(ExtensionAppJinjaMixin, ExtensionApp):
         }
     }
 
+    async def _start_jupyter_server_extension(self, serverapp):
+        self.started = True
+
     @staticmethod
     def get_extension_package():
         return "tests.extension.mockextensions"
@@ -96,6 +100,9 @@ def initialize_handlers(self):
         self.handlers.append(("/mock_error_notemplate", MockExtensionErrorHandler))
         self.loaded = True
 
+    async def _start_jupyter_server_extension(self, serverapp):
+        self.started = True
+
 
 if __name__ == "__main__":
     MockExtensionApp.launch_instance()

tests/extension/mockextensions/mock1.py

@@ -1,5 +1,7 @@
 """A mock extension named `mock1` for testing purposes."""
+
 # by the test functions.
+import asyncio
 
 
 def _jupyter_server_extension_paths():
@@ -9,3 +11,8 @@ def _jupyter_server_extension_paths():
 def _load_jupyter_server_extension(serverapp):
     serverapp.mockI = True
     serverapp.mock_shared = "I"
+
+
+async def _start_jupyter_server_extension(serverapp):
+    await asyncio.sleep(0.1)
+    serverapp.mock1_started = True

tests/extension/test_app.py

@@ -140,6 +140,15 @@ async def test_load_parallel_extensions(monkeypatch, jp_environ):
     assert exts["tests.extension.mockextensions"]
 
 
+async def test_start_extension(jp_serverapp, mock_extension):
+    await jp_serverapp._post_start()
+    assert mock_extension.started
+    assert hasattr(
+        jp_serverapp, "mock1_started"
+    ), "Failed because the `_start_jupyter_server_extension` function in 'mock1.py' was never called"
+    assert jp_serverapp.mock1_started
+
+
 async def test_stop_extension(jp_serverapp, caplog):
     """Test the stop_extension method.
 

tests/test_serverapp.py

@@ -607,8 +607,9 @@ def test_running_server_info(jp_serverapp):
 
 
 @pytest.mark.parametrize("should_exist", [True, False])
-def test_browser_open_files(jp_configurable_serverapp, should_exist, caplog):
+async def test_browser_open_files(jp_configurable_serverapp, should_exist, caplog):
     app = jp_configurable_serverapp(no_browser_open_file=not should_exist)
+    await app._post_start()
     assert os.path.exists(app.browser_open_file) == should_exist
     url = urljoin("file:", pathname2url(app.browser_open_file))
     url_messages = [rec.message for rec in caplog.records if url in rec.message]