diff --git a/conftest.py b/conftest.py index 5c584923..8de56f21 100644 --- a/conftest.py +++ b/conftest.py @@ -59,6 +59,7 @@ def jp_scheduler(jp_scheduler_db_url, jp_scheduler_root_dir, jp_scheduler_db): db_url=jp_scheduler_db_url, root_dir=str(jp_scheduler_root_dir), environments_manager=MockEnvironmentManager(), + update_last_activity=lambda: True, ) diff --git a/jupyter_scheduler/extension.py b/jupyter_scheduler/extension.py index 1a4ba373..4a36f792 100644 --- a/jupyter_scheduler/extension.py +++ b/jupyter_scheduler/extension.py @@ -1,4 +1,5 @@ import asyncio +from datetime import datetime, timezone from jupyter_core.paths import jupyter_data_dir from jupyter_server.extension.application import ExtensionApp @@ -78,6 +79,7 @@ def initialize_settings(self): environments_manager=environments_manager, db_url=self.db_url, config=self.config, + update_last_activity=self._update_last_activity, ) job_files_manager = self.job_files_manager_class(scheduler=scheduler) @@ -91,3 +93,15 @@ def initialize_settings(self): if scheduler.task_runner: loop = asyncio.get_event_loop() loop.create_task(scheduler.task_runner.start()) + + self._update_last_activity() + + def _update_last_activity(self): + """Update last activity in the web application + + Updates the `last_acitivity_times` dict in the web application. When running with jupyterhub, this prevents + the culler from thinking the server is idle. + """ + self.serverapp.web_app.settings["last_activity_times"]["jupyter_scheduler"] = datetime.now( + timezone.utc + ) diff --git a/jupyter_scheduler/scheduler.py b/jupyter_scheduler/scheduler.py index 867034c6..6b90f29d 100644 --- a/jupyter_scheduler/scheduler.py +++ b/jupyter_scheduler/scheduler.py @@ -96,11 +96,17 @@ def _default_staging_path(self): ) def __init__( - self, root_dir: str, environments_manager: Type[EnvironmentManager], config=None, **kwargs + self, + root_dir: str, + environments_manager: Type[EnvironmentManager], + config=None, + update_last_activity=None, + **kwargs, ): super().__init__(config=config, **kwargs) self.root_dir = root_dir self.environments_manager = environments_manager + self.update_last_activity = update_last_activity def create_job(self, model: CreateJob) -> str: """Creates a new job record, may trigger execution of the job. @@ -405,10 +411,15 @@ def __init__( environments_manager: Type[EnvironmentManager], db_url: str, config=None, + update_last_activity=None, **kwargs, ): super().__init__( - root_dir=root_dir, environments_manager=environments_manager, config=config, **kwargs + root_dir=root_dir, + environments_manager=environments_manager, + config=config, + update_last_activity=update_last_activity, + **kwargs, ) self.db_url = db_url if self.task_runner_class: @@ -438,6 +449,7 @@ def copy_input_folder(self, input_uri: str, nb_copy_to_path: str) -> List[str]: ) def create_job(self, model: CreateJob) -> str: + self.update_last_activity() if not model.job_definition_id and not self.file_exists(model.input_uri): raise InputUriError(model.input_uri)