initialize dask cluster

Author: andrii-i

jupyter_scheduler/extension.py

@@ -91,3 +91,25 @@ def initialize_settings(self):
         if scheduler.task_runner:
             loop = asyncio.get_event_loop()
             loop.create_task(scheduler.task_runner.start())
+
+    async def stop_extension(self):
+        """
+        Public method called by Jupyter Server when the server is stopping.
+        This calls the cleanup code defined in `self._stop_exception()` inside
+        an exception handler, as the server halts if this method raises an
+        exception.
+        """
+        try:
+            await self._stop_extension()
+        except Exception as e:
+            self.log.error("Jupyter Scheduler raised an exception while stopping:")
+            self.log.exception(e)
+
+    async def _stop_extension(self):
+        """
+        Private method that defines the cleanup code to run when the server is
+        stopping.
+        """
+        if "scheduler" in self.settings:
+            scheduler: SchedulerApp = self.settings["scheduler"]
+            await scheduler.stop_extension()

jupyter_scheduler/scheduler.py

@@ -6,6 +6,8 @@
 
 import fsspec
 import psutil
+from dask.distributed import Client as DaskClient
+from distributed import LocalCluster
 from jupyter_core.paths import jupyter_data_dir
 from jupyter_server.transutils import _i18n
 from jupyter_server.utils import to_os_path
@@ -381,6 +383,12 @@ def get_local_output_path(
         else:
             return os.path.join(self.root_dir, self.output_directory, output_dir_name)
 
+    async def stop_extension(self):
+        """
+        Placeholder method for a cleanup code to run when the server is stopping.
+        """
+        pass
+
 
 class Scheduler(BaseScheduler):
     _db_session = None
@@ -395,6 +403,12 @@ class Scheduler(BaseScheduler):
         ),
     )
 
+    dask_cluster_url = Unicode(
+        allow_none=True,
+        config=True,
+        help="URL of the Dask cluster to connect to.",
+    )
+
     db_url = Unicode(help=_i18n("Scheduler database url"))
 
     task_runner = Instance(allow_none=True, klass="jupyter_scheduler.task_runner.BaseTaskRunner")
@@ -414,6 +428,15 @@ def __init__(
         if self.task_runner_class:
             self.task_runner = self.task_runner_class(scheduler=self, config=config)
 
+        self.dask_client: DaskClient = self._get_dask_client()
+
+    def _get_dask_client(self):
+        """Creates and configures a Dask client."""
+        if self.dask_cluster_url:
+            return DaskClient(self.dask_cluster_url)
+        cluster = LocalCluster(processes=True)
+        return DaskClient(cluster)
+
     @property
     def db_session(self):
         if not self._db_session:
@@ -777,6 +800,13 @@ def get_staging_paths(self, model: Union[DescribeJob, DescribeJobDefinition]) ->
 
         return staging_paths
 
+    async def stop_extension(self):
+        """
+        Cleanup code to run when the server is stopping.
+        """
+        if self.dask_client:
+            await self.dask_client.close()
+
 
 class ArchivingScheduler(Scheduler):
     """Scheduler that captures all files in output directory in an archive."""