Skip to content

Commit

Permalink
add idle timeout
Browse files Browse the repository at this point in the history
  • Loading branch information
Tianhao-Gu committed Nov 18, 2024
1 parent 3b2b59a commit 7ba01d2
Showing 1 changed file with 56 additions and 0 deletions.
56 changes: 56 additions & 0 deletions src/jupyterhub_config/custom_docker_spawner.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import os
import shutil
import venv
from datetime import timedelta, datetime
from pathlib import Path

import json5
Expand All @@ -10,11 +11,13 @@

class CustomDockerSpawner(DockerSpawner):
RW_MINIO_GROUP = 'minio_rw'
DEFAULT_IDLE_TIMEOUT_MINUTES = 180

def start(self):
username = self.user.name
global_home = Path(os.environ['JUPYTERHUB_USER_HOME'])
user_dir = global_home / username
self.idle_timeout = self._get_idle_timeout()

Check warning on line 20 in src/jupyterhub_config/custom_docker_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_docker_spawner.py#L20

Added line #L20 was not covered by tests

# Ensure the user directory exists
self._ensure_user_directory(user_dir, username)
Expand All @@ -41,6 +44,59 @@ def start(self):

return super().start()

def _get_idle_timeout(self):
"""
Retrieves the idle timeout from the environment variable `IDLE_TIMEOUT_MINUTES`.
If not set, defaults to 180 minutes.
Returns:
timedelta: Idle timeout duration.
"""
idle_timeout_minutes = int(os.getenv("IDLE_TIMEOUT_MINUTES", self.DEFAULT_IDLE_TIMEOUT_MINUTES))
self.log.info(f"Idle timeout set to {idle_timeout_minutes} minutes")
return timedelta(minutes=idle_timeout_minutes)

Check warning on line 57 in src/jupyterhub_config/custom_docker_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_docker_spawner.py#L55-L57

Added lines #L55 - L57 were not covered by tests

async def poll(self):
"""
Overrides the poll method to periodically checks the status of the user’s JupyterHub container.
ref:
https://github.com/jupyterhub/dockerspawner/blob/main/dockerspawner/dockerspawner.py#L1004
https://jupyterhub-dockerspawner.readthedocs.io/en/latest/api/index.html#dockerspawner.DockerSpawner.poll
- If the container is stopped, returns the status immediately.
- If the container is running, checks how long the user has been idle.
- If idle time exceeds the defined threshold, stops the container to save resources.
The poll method is invoked at regular intervals by the Spawner, with the frequency determined by the JupyterHub
server's configuration (default is 30 seconds).
Returns:
int or None: Returns an exit code (0) if the container has been stopped due
to inactivity. Returns None if the container is still active
and running.
"""
# Check if the container has already stopped
status = await super().poll()
if status is not None:

Check warning on line 81 in src/jupyterhub_config/custom_docker_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_docker_spawner.py#L80-L81

Added lines #L80 - L81 were not covered by tests
# Container has already stopped, return its status code immediately
return status

Check warning on line 83 in src/jupyterhub_config/custom_docker_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_docker_spawner.py#L83

Added line #L83 was not covered by tests

last_activity = self.user.last_activity
self.log.info(f"Last activity for {self.container_name}: {last_activity}")

Check warning on line 86 in src/jupyterhub_config/custom_docker_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_docker_spawner.py#L85-L86

Added lines #L85 - L86 were not covered by tests

if last_activity:
idle_time = datetime.utcnow() - last_activity
self.log.info(f"Idle time for {self.container_name}: {idle_time}")

Check warning on line 90 in src/jupyterhub_config/custom_docker_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_docker_spawner.py#L88-L90

Added lines #L88 - L90 were not covered by tests

if idle_time > self.idle_timeout:
self.log.warn(f"Container {self.container_name} has been idle for {idle_time}. Stopping...")
await self.stop()
return 0 # Return an exit code to indicate the container has stopped

Check warning on line 95 in src/jupyterhub_config/custom_docker_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_docker_spawner.py#L92-L95

Added lines #L92 - L95 were not covered by tests

# Return status (None) to indicate that the container is still running and active
return status

Check warning on line 98 in src/jupyterhub_config/custom_docker_spawner.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/custom_docker_spawner.py#L98

Added line #L98 was not covered by tests

def _ensure_user_directory(self, user_dir: Path, username: str):
"""
Ensure the user's home directory exists.
Expand Down

0 comments on commit 7ba01d2

Please sign in to comment.