feat: Add the ability to override resources on a per-user basis

Closes #1951

backend/capellacollab/alembic/versions/4d42177579a2_add_resource_override.py

@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: Copyright DB InfraGO AG and contributors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Add resource override
+
+Revision ID: 4d42177579a2
+Revises: 320c5b39c509
+Create Date: 2024-11-12 17:43:23.486104
+
+"""
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "4d42177579a2"
+down_revision = "320c5b39c509"
+branch_labels = None
+depends_on = None
+
+
+t_tools = sa.Table(
+    "tools",
+    sa.MetaData(),
+    sa.Column("id", sa.Integer()),
+    sa.Column("integrations", postgresql.JSONB(astext_type=sa.Text())),
+    sa.Column("config", postgresql.JSONB(astext_type=sa.Text())),
+)
+
+
+def upgrade():
+    connection = op.get_bind()
+    results = connection.execute(sa.select(t_tools)).mappings().all()
+
+    for row in results:
+        config = row["config"]
+        config["resources"] = {
+            "profiles": {
+                "default": {
+                    "cpu": config["resources"]["cpu"],
+                    "memory": config["resources"]["memory"],
+                    "users": None,
+                }
+            }
+        }
+
+        connection.execute(
+            sa.update(t_tools)
+            .where(t_tools.c.id == row["id"])
+            .values(config=config)
+        )

backend/capellacollab/core/database/migration.py

@@ -132,10 +132,14 @@ def get_eclipse_session_configuration() -> (
     """
     return tools_models.ToolSessionConfiguration(
         resources=tools_models.Resources(
-            cpu=tools_models.CPUResources(requests=0.4, limits=2),
-            memory=tools_models.MemoryResources(
-                requests="1.6Gi", limits="6Gi"
-            ),
+            profiles={
+                "default": tools_models.ResourceProfile(
+                    cpu=tools_models.CPUResources(requests=0.4, limits=2),
+                    memory=tools_models.MemoryResources(
+                        requests="1.6Gi", limits="6Gi"
+                    ),
+                )
+            }
         ),
         environment={
             "RMT_PASSWORD": "{CAPELLACOLLAB_SESSION_TOKEN}",
@@ -279,10 +283,14 @@ def create_jupyter_tool(db: orm.Session) -> tools_models.DatabaseTool:
         integrations=tools_models.ToolIntegrations(jupyter=True),
         config=tools_models.ToolSessionConfiguration(
             resources=tools_models.Resources(
-                cpu=tools_models.CPUResources(requests=1, limits=2),
-                memory=tools_models.MemoryResources(
-                    requests="500Mi", limits="3Gi"
-                ),
+                profiles={
+                    "default": tools_models.ResourceProfile(
+                        cpu=tools_models.CPUResources(requests=1, limits=2),
+                        memory=tools_models.MemoryResources(
+                            requests="500Mi", limits="3Gi"
+                        ),
+                    )
+                }
             ),
             environment={
                 "JUPYTER_PORT": "8888",

backend/capellacollab/projects/toolmodels/backups/routes.py

@@ -118,7 +118,7 @@ def create_backup(
                 body.include_commit_history,
             ),
             labels=core.get_pipeline_labels(toolmodel),
-            tool_resources=toolmodel.tool.config.resources,
+            tool_resources=toolmodel.tool.config.resources.get_profile(None),
             command="backup",
             schedule=pipeline_config.cron,
             timezone=pipeline_config.timezone,

backend/capellacollab/projects/toolmodels/backups/runs/interface.py

@@ -80,7 +80,9 @@ def _schedule_pending_jobs():
                         pending_run.pipeline.t4c_password,
                         pending_run.pipeline.include_commit_history,
                     ),
-                    tool_resources=pending_run.pipeline.model.tool.config.resources,
+                    tool_resources=pending_run.pipeline.model.tool.config.resources.get_profile(
+                        None
+                    ),
                 )
                 pending_run.reference_id = job_name
                 pending_run.status = models.PipelineRunStatus.SCHEDULED

backend/capellacollab/sessions/operators/k8s.py

@@ -124,7 +124,7 @@ def start_session(
             ports=ports,
             volumes=volumes,
             init_volumes=init_volumes,
-            tool_resources=tool.config.resources,
+            tool_resources_profile=tool.config.resources.get_profile(username),
             annotations=annotations,
             labels=labels,
         )
@@ -284,7 +284,7 @@ def create_cronjob(
         image: str,
         command: str,
         labels: dict[str, str],
-        tool_resources: tools_models.Resources,
+        tool_resources: tools_models.ResourceProfile,
         environment: dict[str, str | None],
         schedule="* * * * *",
         timezone="UTC",
@@ -326,7 +326,7 @@ def create_job(
         command: str,
         labels: dict[str, str],
         environment: dict[str, str | None],
-        tool_resources: tools_models.Resources,
+        tool_resources: tools_models.ResourceProfile,
         timeout: int = 18000,
     ) -> str:
         _id = self._generate_id()
@@ -521,7 +521,7 @@ def _create_pod(
         ports: dict[str, int],
         volumes: list[models.Volume],
         init_volumes: list[models.Volume],
-        tool_resources: tools_models.Resources,
+        tool_resources_profile: tools_models.ResourceProfile,
         annotations: dict[str, str],
         labels: dict[str, str],
     ) -> client.V1Pod:
@@ -557,12 +557,12 @@ def _create_pod(
 
         resources = client.V1ResourceRequirements(
             limits={
-                "cpu": tool_resources.cpu.limits,
-                "memory": tool_resources.memory.limits,
+                "cpu": tool_resources_profile.cpu.limits,
+                "memory": tool_resources_profile.memory.limits,
             },
             requests={
-                "cpu": tool_resources.cpu.requests,
-                "memory": tool_resources.memory.requests,
+                "cpu": tool_resources_profile.cpu.requests,
+                "memory": tool_resources_profile.memory.requests,
             },
         )
 
@@ -786,7 +786,7 @@ def _create_job_spec(
         image: str,
         job_labels: dict[str, str],
         environment: dict[str, str | None],
-        tool_resources: tools_models.Resources,
+        tool_resources: tools_models.ResourceProfile,
         args: list[str] | None = None,
         timeout: int = 18000,
     ) -> client.V1JobSpec:

backend/capellacollab/tools/crud.py

@@ -67,6 +67,24 @@ def update_tool(
     return tool
 
 
+def update_tools_username(
+    db: orm.Session, old_username: str, new_username: str
+):
+    tools = get_tools(db)
+    for tool in tools:
+        updated = False
+        for profile in tool.config.resources.profiles.values():
+            if profile.usernames and old_username in profile.usernames:
+                profile.usernames = [
+                    new_username if username == old_username else username
+                    for username in profile.usernames
+                ]
+                updated = True
+        if updated:
+            orm.attributes.flag_modified(tool, "config")
+    db.commit()
+
+
 def delete_tool(db: orm.Session, tool: models.DatabaseTool) -> None:
     db.delete(tool)
     db.commit()

backend/capellacollab/tools/models.py

@@ -231,7 +231,7 @@ class MemoryResources(core_pydantic.BaseModel):
     )
 
 
-class Resources(core_pydantic.BaseModelStrict):
+class ResourceProfile(core_pydantic.BaseModel):
     cpu: CPUResources = pydantic.Field(
         default=CPUResources(),
         description="Configuration about the number of CPU cores that sessions can use.",
@@ -240,6 +240,70 @@ class Resources(core_pydantic.BaseModelStrict):
         default=MemoryResources(),
         description="Configuration about the amount of memory that sessions can use.",
     )
+    usernames: list[str] | None = pydantic.Field(
+        default=None,
+        description="List of usernames, which are allowed to use this resource profile.",
+    )
+
+
+class Resources(core_pydantic.BaseModelStrict):
+    profiles: dict[str, ResourceProfile] = pydantic.Field(
+        default={"default": ResourceProfile()},
+        description="Resource profiles, which can be used to limit the resource usage of sessions.",
+    )
+
+    def get_profile(self, username: str | None) -> ResourceProfile:
+        default_profile = self.profiles.get("default")
+        assert default_profile is not None
+        if username is None:
+            return default_profile
+
+        for profile_name, profile in self.profiles.items():
+            if profile_name == "default":
+                continue
+            if profile.usernames is not None and username in profile.usernames:
+                return profile
+        return default_profile
+
+    @pydantic.field_validator("profiles")
+    @classmethod
+    def check_default_profile(
+        cls, value: dict[str, ResourceProfile]
+    ) -> dict[str, ResourceProfile]:
+        if "default" not in value:
+            raise ValueError("A default profile must be defined.")
+
+        for profile_name, profile in value.items():
+            if profile_name != "default":
+                if profile.usernames is None:
+                    raise ValueError(
+                        f"Profile '{profile_name}' must have a list of usernames."
+                    )
+                if len(profile.usernames) != len(set(profile.usernames)):
+                    raise ValueError(
+                        f"Usernames in profile '{profile_name}' must be unique."
+                    )
+
+            if profile_name == "default" and profile.usernames is not None:
+                raise ValueError(
+                    "The default profile must not have a list of usernames."
+                )
+
+        usernames = [
+            set(profile.usernames)
+            for profile in value.values()
+            if profile.usernames is not None
+        ]
+
+        all_usernames = [
+            username for usernames in usernames for username in usernames
+        ]
+
+        # Check that usernames aren't in multiple profiles
+        if len(all_usernames) != len(set(all_usernames)):
+            raise ValueError("Usernames must be unique across all profiles.")
+
+        return value
 
 
 class PrometheusConfiguration(core_pydantic.BaseModel):

backend/capellacollab/users/crud.py

@@ -8,6 +8,7 @@
 from sqlalchemy import orm
 
 from capellacollab.core import database
+from capellacollab.tools import crud as tools_crud
 from capellacollab.users import models
 
 
@@ -80,6 +81,8 @@ def create_user(
 def update_user(
     db: orm.Session, user: models.DatabaseUser, patch_user: models.PatchUser
 ) -> models.DatabaseUser:
+    if patch_user.name:
+        tools_crud.update_tools_username(db, user.name, patch_user.name)
     database.patch_database_with_pydantic_object(user, patch_user)
     db.commit()
     return user

backend/tests/sessions/k8s_operator/test_session_k8s_operator.py

@@ -121,7 +121,7 @@ def test_create_job(monkeypatch: pytest.MonkeyPatch):
         command="fakecmd",
         labels={"key": "value"},
         environment={"ENVVAR": "value"},
-        tool_resources=tools_models.Resources(),
+        tool_resources=tools_models.ResourceProfile(),
     )
 
     assert result
@@ -139,7 +139,7 @@ def test_create_cronjob(monkeypatch: pytest.MonkeyPatch):
         command="fakecmd",
         environment={"ENVVAR": "value"},
         labels={},
-        tool_resources=tools_models.Resources(),
+        tool_resources=tools_models.ResourceProfile(),
     )
 
     assert result

backend/tests/tools/test_tools_resources.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright DB InfraGO AG and contributors
+# SPDX-License-Identifier: Apache-2.0
+
+import pydantic
+import pytest
+
+from capellacollab.tools import models as tools_models
+
+
+def test_validate_tools():
+    with pytest.raises(pydantic.ValidationError):
+        tools_models.Resources(
+            profiles={
+                "test1": tools_models.ResourceProfile(usernames=["test"])
+            }
+        )
+
+    with pytest.raises(pydantic.ValidationError):
+        tools_models.Resources(
+            profiles={
+                "default": tools_models.ResourceProfile(),
+                "test1": tools_models.ResourceProfile(),
+            }
+        )
+
+    with pytest.raises(pydantic.ValidationError):
+        tools_models.Resources(
+            profiles={
+                "default": tools_models.ResourceProfile(usernames=["test"])
+            }
+        )
+
+    with pytest.raises(pydantic.ValidationError):
+        tools_models.Resources(
+            profiles={
+                "default": tools_models.ResourceProfile(),
+                "test1": tools_models.ResourceProfile(usernames=["test"]),
+                "test2": tools_models.ResourceProfile(usernames=["test"]),
+            }
+        )
+
+
+def test_get_profile():
+    default_profile = tools_models.ResourceProfile(
+        memory=tools_models.MemoryResources(requests="1Gi", limits="2Gi"),
+        cpu=tools_models.CPUResources(requests=0.4, limits=2),
+    )
+    different_profile = tools_models.ResourceProfile(
+        usernames=["testuser"],
+        memory=tools_models.MemoryResources(requests="1Gi", limits="2Gi"),
+        cpu=tools_models.CPUResources(requests=0.4, limits=2),
+    )
+
+    resources = tools_models.Resources(
+        profiles={
+            "default": default_profile,
+            "test": different_profile,
+        }
+    )
+
+    resource_profile = resources.get_profile(None)
+    assert resource_profile == default_profile
+
+    resource_profile = resources.get_profile("fakeuser")
+    assert resource_profile == default_profile
+
+    resource_profile = resources.get_profile("testuser")
+    assert resource_profile == different_profile