-
Notifications
You must be signed in to change notification settings - Fork 23
/
extension.py
93 lines (74 loc) · 3.15 KB
/
extension.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import asyncio
from jupyter_core.paths import jupyter_data_dir
from jupyter_server.extension.application import ExtensionApp
from jupyter_server.transutils import _i18n
from traitlets import Bool, Type, Unicode, default
from jupyter_scheduler.orm import create_tables
from .handlers import (
BatchJobHandler,
ConfigHandler,
FilesDownloadHandler,
JobDefinitionHandler,
JobFromDefinitionHandler,
JobHandler,
JobsCountHandler,
RuntimeEnvironmentsHandler,
)
JOB_DEFINITION_ID_REGEX = r"(?P<job_definition_id>\w+(?:-\w+)+)"
JOB_ID_REGEX = r"(?P<job_id>\w+(?:-\w+)+)"
class SchedulerApp(ExtensionApp):
name = "jupyter_scheduler"
handlers = [
(r"scheduler/jobs", JobHandler),
(r"scheduler/jobs/count", JobsCountHandler),
(r"scheduler/jobs/%s" % JOB_ID_REGEX, JobHandler),
(r"scheduler/jobs/%s/download_files" % JOB_ID_REGEX, FilesDownloadHandler),
(r"scheduler/batch/jobs", BatchJobHandler),
(r"scheduler/job_definitions", JobDefinitionHandler),
(r"scheduler/job_definitions/%s" % JOB_DEFINITION_ID_REGEX, JobDefinitionHandler),
(r"scheduler/job_definitions/%s/jobs" % JOB_DEFINITION_ID_REGEX, JobFromDefinitionHandler),
(r"scheduler/runtime_environments", RuntimeEnvironmentsHandler),
(r"scheduler/config", ConfigHandler),
]
drop_tables = Bool(False, config=True, help="Drop the database tables before starting.")
db_url = Unicode(config=True, help="URI for the scheduler database")
@default("db_url")
def _db_url_default(self):
return f"sqlite:///{jupyter_data_dir()}/scheduler.sqlite"
environment_manager_class = Type(
default_value="jupyter_scheduler.environments.CondaEnvironmentManager",
klass="jupyter_scheduler.environments.EnvironmentManager",
config=True,
help=_i18n("The runtime environment manager class to use."),
)
scheduler_class = Type(
default_value="jupyter_scheduler.scheduler.Scheduler",
klass="jupyter_scheduler.scheduler.BaseScheduler",
config=True,
help=_i18n("The scheduler class to use."),
)
job_files_manager_class = Type(
default_value="jupyter_scheduler.job_files_manager.JobFilesManager",
klass="jupyter_scheduler.job_files_manager.JobFilesManager",
config=True,
help=_i18n("The job files manager class to use."),
)
def initialize_settings(self):
super().initialize_settings()
create_tables(self.db_url, self.drop_tables)
environments_manager = self.environment_manager_class()
scheduler = self.scheduler_class(
root_dir=self.serverapp.root_dir,
environments_manager=environments_manager,
db_url=self.db_url,
config=self.config,
)
job_files_manager = self.job_files_manager_class(scheduler=scheduler)
self.settings.update(
environments_manager=environments_manager,
scheduler=scheduler,
job_files_manager=job_files_manager,
)
if scheduler.task_runner:
loop = asyncio.get_event_loop()
loop.create_task(scheduler.task_runner.start())