File tree 1 file changed +12
-2
lines changed
1 file changed +12
-2
lines changed Original file line number Diff line number Diff line change 7
7
import fsspec
8
8
import psutil
9
9
from dask .distributed import Client as DaskClient
10
+ from distributed import LocalCluster
10
11
from jupyter_core .paths import jupyter_data_dir
11
12
from jupyter_server .transutils import _i18n
12
13
from jupyter_server .utils import to_os_path
@@ -402,6 +403,12 @@ class Scheduler(BaseScheduler):
402
403
),
403
404
)
404
405
406
+ dask_cluster_url = Unicode (
407
+ allow_none = True ,
408
+ config = True ,
409
+ help = "URL of the Dask cluster to connect to." ,
410
+ )
411
+
405
412
db_url = Unicode (help = _i18n ("Scheduler database url" ))
406
413
407
414
task_runner = Instance (allow_none = True , klass = "jupyter_scheduler.task_runner.BaseTaskRunner" )
@@ -425,7 +432,10 @@ def __init__(
425
432
426
433
def _get_dask_client (self ):
427
434
"""Creates and configures a Dask client."""
428
- return DaskClient ()
435
+ if self .dask_cluster_url :
436
+ return DaskClient (self .dask_cluster_url )
437
+ cluster = LocalCluster (processes = True )
438
+ return DaskClient (cluster )
429
439
430
440
@property
431
441
def db_session (self ):
@@ -786,7 +796,7 @@ async def stop_extension(self):
786
796
Cleanup code to run when the server is stopping.
787
797
"""
788
798
if self .dask_client :
789
- self .dask_client .close ()
799
+ await self .dask_client .close ()
790
800
791
801
792
802
class ArchivingScheduler (Scheduler ):
You can’t perform that action at this time.
0 commit comments