Skip to content

Commit

Permalink
add scheduler.dask_cluster_url traitlet, create dask cluster based on it
Browse files Browse the repository at this point in the history
  • Loading branch information
andrii-i committed Aug 7, 2024
1 parent f4d8f8a commit 74aaf9a
Showing 1 changed file with 12 additions and 2 deletions.
14 changes: 12 additions & 2 deletions jupyter_scheduler/scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
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
Expand Down Expand Up @@ -402,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")
Expand All @@ -425,7 +432,10 @@ def __init__(

def _get_dask_client(self):
"""Creates and configures a Dask client."""
return DaskClient()
if self.dask_cluster_url:
return DaskClient(self.dask_cluster_url)
cluster = LocalCluster(processes=True)
return DaskClient(cluster)

@property
def db_session(self):
Expand Down Expand Up @@ -786,7 +796,7 @@ async def stop_extension(self):
Cleanup code to run when the server is stopping.
"""
if self.dask_client:
self.dask_client.close()
await self.dask_client.close()


class ArchivingScheduler(Scheduler):
Expand Down

0 comments on commit 74aaf9a

Please sign in to comment.