feat: add disk storage to session launchers (#590)

Closes #589.

Add a new `disk_storage` field to session launchers so that a custom disk size can be persisted.

Note: the disk storage size is not validated when saved.

components/renku_data_services/migrations/versions/939c7c649bef_add_disk_storage_to_session_launchers.py

@@ -0,0 +1,28 @@
+"""Add disk storage to session launchers
+
+Revision ID: 939c7c649bef
+Revises: d1cdcbb2adc3
+Create Date: 2024-12-20 15:06:01.937878
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "939c7c649bef"
+down_revision = "d1cdcbb2adc3"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("launchers", sa.Column("disk_storage", sa.BigInteger(), nullable=True), schema="sessions")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("launchers", "disk_storage", schema="sessions")
+    # ### end Alembic commands ###

components/renku_data_services/session/api.spec.yaml

@@ -419,32 +419,15 @@ components:
           $ref: "#/components/schemas/EnvironmentGetInLauncher"
         resource_class_id:
           $ref: "#/components/schemas/ResourceClassId"
+        disk_storage:
+          $ref: "#/components/schemas/DiskStorage"
       required:
         - id
         - project_id
         - name
         - creation_date
         - environment
         - resource_class_id
-      example:
-        id: 01AN4Z79ZS5XN0F25N3DB94T4R
-        project_id: 01AN4Z79ZS5XN0F25N3DB94T4R
-        name: Renku R Session
-        creation_date: "2023-11-01T17:32:28Z"
-        description: R compute session
-        environment:
-          id: 01AN4Z79ZS6XX96588FDX0H099
-          name: Rstudio
-          creation_date: "2023-11-01T17:32:28Z"
-          description: JupyterLab session environment
-          environment_kind: GLOBAL
-          container_image: rocker/rstudio
-          default_url: "/rstudio"
-          port: 8080
-          working_directory: /home/rstudio/work
-          mount_directory: /home/rstudio/work
-          uid: 1000
-          gid: 1000
     SessionLauncherPost:
       description: Data required to create a session launcher
       type: object
@@ -458,6 +441,8 @@ components:
           $ref: "#/components/schemas/Description"
         resource_class_id:
           $ref: "#/components/schemas/ResourceClassId"
+        disk_storage:
+          $ref: "#/components/schemas/DiskStorage"
         environment:
           oneOf:
             - $ref: "#/components/schemas/EnvironmentPostInLauncher"
@@ -482,6 +467,8 @@ components:
           $ref: "#/components/schemas/Description"
         resource_class_id:
           $ref: "#/components/schemas/ResourceClassId"
+        disk_storage:
+          $ref: "#/components/schemas/DiskStoragePatch"
         environment:
           oneOf:
             - $ref: "#/components/schemas/EnvironmentPatchInLauncher"
@@ -549,6 +536,15 @@ components:
       type: integer
       default: null
       nullable: true
+    DiskStorage:
+      description: The size of disk storage for the session, in gigabytes
+      type: integer
+      minimum: 1
+      example: 8
+    DiskStoragePatch:
+      type: integer
+      minimum: 1
+      nullable: true
     EnvironmentPort:
       type: integer
       minimum: 0

components/renku_data_services/session/apispec.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api.spec.yaml
-#   timestamp: 2024-12-19T08:38:19+00:00
+#   timestamp: 2024-12-23T08:57:28+00:00
 
 from __future__ import annotations
 
@@ -252,6 +252,12 @@ class SessionLauncher(BaseAPISpec):
     resource_class_id: Optional[int] = Field(
         ..., description="The identifier of a resource class"
     )
+    disk_storage: Optional[int] = Field(
+        None,
+        description="The size of disk storage for the session, in gigabytes",
+        example=8,
+        ge=1,
+    )
 
 
 class EnvironmentIdOnlyPatch(BaseAPISpec):
@@ -314,6 +320,12 @@ class SessionLauncherPost(BaseAPISpec):
     resource_class_id: Optional[int] = Field(
         None, description="The identifier of a resource class"
     )
+    disk_storage: Optional[int] = Field(
+        None,
+        description="The size of disk storage for the session, in gigabytes",
+        example=8,
+        ge=1,
+    )
     environment: Union[EnvironmentPostInLauncher, EnvironmentIdOnlyPost]
 
 
@@ -334,6 +346,7 @@ class SessionLauncherPatch(BaseAPISpec):
     resource_class_id: Optional[int] = Field(
         None, description="The identifier of a resource class"
     )
+    disk_storage: Optional[int] = Field(None, ge=1)
     environment: Optional[Union[EnvironmentPatchInLauncher, EnvironmentIdOnlyPatch]] = (
         None
     )

components/renku_data_services/session/core.py

@@ -69,6 +69,7 @@ def validate_unsaved_session_launcher(launcher: apispec.SessionLauncherPost) ->
         name=launcher.name,
         description=launcher.description,
         resource_class_id=launcher.resource_class_id,
+        disk_storage=launcher.disk_storage,
         # NOTE: When you create an environment with a launcher the environment can only be custom
         environment=validate_unsaved_environment(launcher.environment, models.EnvironmentKind.CUSTOM)
         if isinstance(launcher.environment, apispec.EnvironmentPostInLauncher)
@@ -118,9 +119,11 @@ def validate_session_launcher_patch(
         resource_class_id = RESET
     else:
         resource_class_id = patch.resource_class_id
+    disk_storage = RESET if "disk_storage" in data_dict and data_dict["disk_storage"] is None else patch.disk_storage
     return models.SessionLauncherPatch(
         name=patch.name,
         description=patch.description,
         environment=environment,
         resource_class_id=resource_class_id,
+        disk_storage=disk_storage,
     )

components/renku_data_services/session/db.py

@@ -316,6 +316,7 @@ async def insert_launcher(
                 description=launcher.description if launcher.description else None,
                 environment_id=environment_id,
                 resource_class_id=launcher.resource_class_id,
+                disk_storage=launcher.disk_storage,
                 created_by_id=user.id,
                 creation_date=datetime.now(UTC).replace(microsecond=0),
             )
@@ -351,6 +352,7 @@ async def copy_launcher(
                 description=launcher.description,
                 environment_id=launcher.environment.id,
                 resource_class_id=launcher.resource_class_id,
+                disk_storage=launcher.disk_storage,
                 created_by_id=user.id,
                 creation_date=datetime.now(UTC).replace(microsecond=0),
             )
@@ -425,6 +427,10 @@ async def update_launcher(
                 launcher.resource_class_id = patch.resource_class_id
             elif patch.resource_class_id is RESET:
                 launcher.resource_class_id = None
+            if isinstance(patch.disk_storage, int):
+                launcher.disk_storage = patch.disk_storage
+            elif patch.disk_storage is RESET:
+                launcher.disk_storage = None
 
             if patch.environment is None:
                 return launcher.dump()

components/renku_data_services/session/models.py

@@ -98,6 +98,7 @@ class UnsavedSessionLauncher:
     name: str
     description: str | None
     resource_class_id: int | None
+    disk_storage: int | None
     environment: str | UnsavedEnvironment
     """When a string is passed for the environment it should be the ID of an existing environment."""
 
@@ -122,3 +123,4 @@ class SessionLauncherPatch:
     # launcher with the update of the launcher.
     environment: str | EnvironmentPatch | UnsavedEnvironment | None = None
     resource_class_id: int | None | ResetType = None
+    disk_storage: int | None | ResetType = None

components/renku_data_services/session/orm.py

@@ -3,7 +3,7 @@
 from datetime import datetime
 from pathlib import PurePosixPath
 
-from sqlalchemy import JSON, DateTime, MetaData, String, func
+from sqlalchemy import JSON, BigInteger, DateTime, MetaData, String, func
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, mapped_column, relationship
 from sqlalchemy.schema import ForeignKey
@@ -128,6 +128,9 @@ class SessionLauncherORM(BaseORM):
     )
     """Id of the resource class."""
 
+    disk_storage: Mapped[int | None] = mapped_column("disk_storage", BigInteger, default=None, nullable=True)
+    """Default value for requested disk storage."""
+
     @classmethod
     def load(cls, launcher: models.SessionLauncher) -> "SessionLauncherORM":
         """Create SessionLauncherORM from the session launcher model."""
@@ -139,6 +142,7 @@ def load(cls, launcher: models.SessionLauncher) -> "SessionLauncherORM":
             project_id=launcher.project_id,
             environment_id=launcher.environment.id,
             resource_class_id=launcher.resource_class_id,
+            disk_storage=launcher.disk_storage,
         )
 
     def dump(self) -> models.SessionLauncher:
@@ -151,5 +155,6 @@ def dump(self) -> models.SessionLauncher:
             creation_date=self.creation_date,
             description=self.description,
             resource_class_id=self.resource_class_id,
+            disk_storage=self.disk_storage,
             environment=self.environment.dump(),
         )

test/bases/renku_data_services/data_api/conftest.py

@@ -398,11 +398,11 @@ async def create_data_connector_and_link_project_helper(
 
 
 @pytest_asyncio.fixture
-async def create_resource_pool(sanic_client, user_headers, admin_headers):
+async def create_resource_pool(sanic_client, user_headers, admin_headers, valid_resource_pool_payload):
     async def create_resource_pool_helper(admin: bool = False, **payload) -> dict[str, Any]:
         headers = admin_headers if admin else user_headers
-        payload = payload.copy()
-        _, res = await sanic_client.post("/api/data/resource_pools", headers=headers, json=payload)
+        valid_resource_pool_payload.update(payload)
+        _, res = await sanic_client.post("/api/data/resource_pools", headers=headers, json=valid_resource_pool_payload)
         assert res.status_code == 201, res.text
         assert res.json is not None
         return res.json

test/bases/renku_data_services/data_api/test_sessions.py

@@ -334,6 +334,7 @@ async def test_post_session_launcher(
         "project_id": project["id"],
         "description": "A session launcher.",
         "resource_class_id": resource_pool["classes"][0]["id"],
+        "disk_storage": 2,
         "environment": {
             "container_image": "some_image:some_tag",
             "name": "custom_name",
@@ -353,6 +354,7 @@ async def test_post_session_launcher(
     assert environment.get("container_image") == "some_image:some_tag"
     assert environment.get("id") is not None
     assert res.json.get("resource_class_id") == resource_pool["classes"][0]["id"]
+    assert res.json.get("disk_storage") == 2
 
 
 @pytest.mark.asyncio
@@ -437,11 +439,13 @@ async def test_patch_session_launcher(
     assert environment.get("container_image") == "some_image:some_tag"
     assert environment.get("id") is not None
     assert res.json.get("resource_class_id") == resource_pool["classes"][0]["id"]
+    assert res.json.get("disk_storage") is None
 
     patch_payload = {
         "name": "New Name",
         "description": "An updated session launcher.",
         "resource_class_id": resource_pool["classes"][1]["id"],
+        "disk_storage": 3,
     }
     _, res = await sanic_client.patch(
         f"/api/data/session_launchers/{res.json['id']}", headers=user_headers, json=patch_payload
@@ -451,6 +455,7 @@ async def test_patch_session_launcher(
     assert res.json.get("name") == patch_payload["name"]
     assert res.json.get("description") == patch_payload["description"]
     assert res.json.get("resource_class_id") == patch_payload["resource_class_id"]
+    assert res.json.get("disk_storage") == 3
 
 
 @pytest.mark.asyncio
@@ -541,13 +546,98 @@ async def test_patch_session_launcher_environment(
     assert res.json["environment"].get("command") is None
 
 
+@pytest.mark.asyncio
+async def test_patch_session_launcher_reset_fields(
+    sanic_client: SanicASGITestClient,
+    valid_resource_pool_payload: dict[str, Any],
+    user_headers,
+    create_project,
+    create_resource_pool,
+) -> None:
+    project = await create_project("Some project 1")
+    resource_pool_data = valid_resource_pool_payload
+    resource_pool = await create_resource_pool(admin=True, **resource_pool_data)
+
+    payload = {
+        "name": "Launcher 1",
+        "project_id": project["id"],
+        "description": "A session launcher.",
+        "resource_class_id": resource_pool["classes"][0]["id"],
+        "disk_storage": 2,
+        "environment": {
+            "container_image": "some_image:some_tag",
+            "name": "custom_name",
+            "environment_kind": "CUSTOM",
+        },
+    }
+
+    _, res = await sanic_client.post("/api/data/session_launchers", headers=user_headers, json=payload)
+
+    assert res.status_code == 201, res.text
+    assert res.json is not None
+    assert res.json.get("name") == "Launcher 1"
+    assert res.json.get("description") == "A session launcher."
+    environment = res.json.get("environment", {})
+    assert environment.get("environment_kind") == "CUSTOM"
+    assert environment.get("container_image") == "some_image:some_tag"
+    assert environment.get("id") is not None
+    assert res.json.get("resource_class_id") == resource_pool["classes"][0]["id"]
+    assert res.json.get("disk_storage") == 2
+
+    patch_payload = {"resource_class_id": None, "disk_storage": None}
+    _, res = await sanic_client.patch(
+        f"/api/data/session_launchers/{res.json['id']}", headers=user_headers, json=patch_payload
+    )
+    assert res.status_code == 200, res.text
+    assert res.json is not None
+    assert res.json.get("resource_class_id") is None
+    assert res.json.get("disk_storage") is None
+
+
+@pytest.mark.asyncio
+async def test_patch_session_launcher_keeps_unset_values(
+    sanic_client, user_headers, create_project, create_resource_pool, create_session_launcher
+) -> None:
+    project = await create_project("Some project")
+    resource_pool = await create_resource_pool(admin=True)
+    session_launcher = await create_session_launcher(
+        name="Session Launcher",
+        project_id=project["id"],
+        description="A session launcher.",
+        resource_class_id=resource_pool["classes"][0]["id"],
+        disk_storage=42,
+        environment={
+            "container_image": "some_image:some_tag",
+            "environment_kind": "CUSTOM",
+            "name": "custom_name",
+        },
+    )
+
+    _, response = await sanic_client.patch(
+        f"/api/data/session_launchers/{session_launcher['id']}", headers=user_headers, json={}
+    )
+
+    assert response.status_code == 200, response.text
+    assert response.json is not None
+    assert response.json.get("name") == "Session Launcher"
+    assert response.json.get("project_id") == project["id"]
+    assert response.json.get("description") == "A session launcher."
+    assert response.json.get("resource_class_id") == resource_pool["classes"][0]["id"]
+    assert response.json.get("disk_storage") == 42
+    environment = response.json.get("environment", {})
+    assert environment.get("container_image") == "some_image:some_tag"
+    assert environment.get("environment_kind") == "CUSTOM"
+    assert environment.get("name") == "custom_name"
+    assert environment.get("id") is not None
+
+
 @pytest.fixture
 def anonymous_user_headers() -> dict[str, str]:
     return {"Renku-Auth-Anon-Id": "some-random-value-1234"}
 
 
 @pytest.mark.asyncio
-@pytest.mark.skip(reason="Setup for testing sessions is not done yet.")  # TODO: enable in follwup PR
+@pytest.mark.skip(reason="Setup for testing sessions is not done yet.")  # TODO: enable in followup PR
 async def test_starting_session_anonymous(
     sanic_client: SanicASGITestClient,
     create_project,