feat: custom culling times (#1852)

SwissDataScienceCenter · May 22, 2024 · 4303bb9 · 4303bb9
1 parent e9b5267
commit 4303bb9
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 10 deletions.
diff --git a/renku_notebooks/api/classes/data_service.py b/renku_notebooks/api/classes/data_service.py
@@ -131,6 +131,8 @@ def validate_class_storage(
         if storage > res_class.get("max_storage"):
             raise InvalidComputeResourceError(message="The requested storage surpasses the maximum value allowed.")
         options = ServerOptions.from_resource_class(res_class)
+        options.idle_threshold_seconds = pool.get("idle_threshold")
+        options.hibernation_threshold_seconds = pool.get("hibernation_threshold")
         options.set_storage(storage, gigabytes=True)
         quota = pool.get("quota")
         if quota is not None and isinstance(quota, dict):

diff --git a/renku_notebooks/api/classes/server.py b/renku_notebooks/api/classes/server.py
@@ -59,16 +59,24 @@ def __init__(
         self.work_dir = work_dir
         self.cloudstorage: list[ICloudStorageRequest] | None = cloudstorage
         self.is_image_private = is_image_private
-        self.idle_seconds_threshold: int = (
-            config.sessions.culling.registered.idle_seconds
-            if isinstance(self._user, RegisteredUser)
-            else config.sessions.culling.anonymous.idle_seconds
-        )
-        self.hibernated_seconds_threshold: int = (
-            config.sessions.culling.registered.hibernated_seconds
-            if isinstance(user, RegisteredUser)
-            else config.sessions.culling.anonymous.hibernated_seconds
-        )
+
+        if self.server_options.idle_threshold_seconds is not None:
+            self.idle_seconds_threshold = self.server_options.idle_threshold_seconds
+        else:
+            self.idle_seconds_threshold: int = (
+                config.sessions.culling.registered.idle_seconds
+                if isinstance(self._user, RegisteredUser)
+                else config.sessions.culling.anonymous.idle_seconds
+            )
+
+        if self.server_options.hibernation_threshold_seconds is not None:
+            self.hibernated_seconds_threshold: int = self.server_options.hibernation_threshold_seconds
+        else:
+            self.hibernated_seconds_threshold: int = (
+                config.sessions.culling.registered.hibernated_seconds
+                if isinstance(user, RegisteredUser)
+                else config.sessions.culling.anonymous.hibernated_seconds
+            )
         self._repositories: list[Repository] = repositories
         self._git_providers: list[GitProvider] | None = None
         self._has_configured_git_providers = False

diff --git a/renku_notebooks/api/notebooks.py b/renku_notebooks/api/notebooks.py
@@ -76,6 +76,7 @@ def version():
             application/json:
               schema: VersionResponse
     """
+    culling = config.sessions.culling
     info = {
         "name": "renku-notebooks",
         "versions": [
@@ -86,6 +87,16 @@ def version():
                     "cloudstorageEnabled": config.cloud_storage.enabled,
                     "cloudstorageClass": config.cloud_storage.storage_class,
                     "sshEnabled": config.ssh_enabled,
+                    "defaultCullingThresholds": {
+                        "registered": {
+                            "idle": culling.registered.idle_seconds,
+                            "hibernation": culling.registered.hibernated_seconds,
+                        },
+                        "anonymous": {
+                            "idle": culling.anonymous.idle_seconds,
+                            "hibernation": culling.anonymous.hibernated_seconds,
+                        },
+                    },
                 },
             }
         ],

diff --git a/renku_notebooks/api/schemas/server_options.py b/renku_notebooks/api/schemas/server_options.py
@@ -56,6 +56,8 @@ class ServerOptions:
     node_affinities: list[NodeAffinity] = field(default_factory=list)
     tolerations: list[Toleration] = field(default_factory=list)
     resource_class_id: Optional[int] = None
+    idle_threshold_seconds: Optional[int] = None
+    hibernation_threshold_seconds: Optional[int] = None
 
     def __post_init__(self):
         if self.default_url is None:

diff --git a/renku_notebooks/api/schemas/version.py b/renku_notebooks/api/schemas/version.py
@@ -3,12 +3,27 @@
 from marshmallow import Schema, fields
 
 
+class CullingThreshold(Schema):
+    """Culling thresholds info."""
+
+    idle = fields.Int(required=True)
+    hibernation = fields.Int(required=True)
+
+
+class DefaultCullingThresholds(Schema):
+    """Culling thresholds for this deployment."""
+
+    registered = fields.Nested(CullingThreshold, required=True)
+    anonymous = fields.Nested(CullingThreshold, required=True)
+
+
 class NotebooksServiceInfo(Schema):
     """Various notebooks service info."""
 
     anonymousSessionsEnabled = fields.Boolean(required=True)
     cloudstorageEnabled = fields.Boolean(required=True)
     sshEnabled = fields.Boolean(required=True)
+    defaultCullingThresholds = fields.Nested(DefaultCullingThresholds, required=True)
 
 
 class NotebooksServiceVersions(Schema):

diff --git a/tests/integration/test_session_creation.py b/tests/integration/test_session_creation.py
@@ -49,6 +49,11 @@ def test_version_endpoint(base_url):
     storage = data.get("cloudstorageEnabled", {})
     assert isinstance(storage.get("s3"), bool)
 
+    assert isinstance(data.get("registeredUsersIdleThreshold"), int)
+    assert isinstance(data.get("registeredUsersHibernationThreshold"), int)
+    assert isinstance(data.get("anonymousUsersIdleThreshold"), int)
+    assert isinstance(data.get("anonymousUsersHibernationThreshold"), int)
+
 
 def test_getting_session_and_logs_after_creation(
     headers, start_session_and_wait_until_ready, base_url, valid_payload, gitlab_project