zenml-nightly 0.71.0.dev20241213__py3-none-any.whl → 0.71.0.dev20241214__py3-none-any.whl

zenml/zen_server/rbac/rbac_sql_zen_store.py

@@ -0,0 +1,173 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""RBAC SQL Zen Store implementation."""
+
+from typing import (
+    Optional,
+    Tuple,
+)
+from uuid import UUID
+
+from zenml.logger import get_logger
+from zenml.models import (
+    ModelRequest,
+    ModelResponse,
+    ModelVersionRequest,
+    ModelVersionResponse,
+)
+from zenml.zen_server.feature_gate.endpoint_utils import (
+    check_entitlement,
+    report_usage,
+)
+from zenml.zen_server.rbac.models import Action, ResourceType
+from zenml.zen_server.rbac.utils import (
+    verify_permission,
+    verify_permission_for_model,
+)
+from zenml.zen_stores.sql_zen_store import SqlZenStore
+
+logger = get_logger(__name__)
+
+
+class RBACSqlZenStore(SqlZenStore):
+    """Wrapper around the SQLZenStore that implements RBAC functionality."""
+
+    def _get_or_create_model(
+        self, model_request: ModelRequest
+    ) -> Tuple[bool, ModelResponse]:
+        """Get or create a model.
+
+        Args:
+            model_request: The model request.
+
+        # noqa: DAR401
+        Raises:
+            Exception: If the user is not allowed to create a model.
+
+        Returns:
+            A boolean whether the model was created or not, and the model.
+        """
+        allow_model_creation = True
+        error = None
+
+        try:
+            verify_permission(
+                resource_type=ResourceType.MODEL, action=Action.CREATE
+            )
+            check_entitlement(resource_type=ResourceType.MODEL)
+        except Exception as e:
+            allow_model_creation = False
+            error = e
+
+        if allow_model_creation:
+            created, model_response = super()._get_or_create_model(
+                model_request
+            )
+        else:
+            try:
+                model_response = self.get_model(model_request.name)
+                created = False
+            except KeyError:
+                # The model does not exist. We now raise the error that
+                # explains why the model could not be created, instead of just
+                # the KeyError that it doesn't exist
+                assert error
+                raise error from None
+
+        if created:
+            report_usage(
+                resource_type=ResourceType.MODEL, resource_id=model_response.id
+            )
+        else:
+            verify_permission_for_model(model_response, action=Action.READ)
+
+        return created, model_response
+
+    def _get_model_version(
+        self,
+        model_id: UUID,
+        version_name: Optional[str] = None,
+        producer_run_id: Optional[UUID] = None,
+    ) -> ModelVersionResponse:
+        """Get a model version.
+
+        Args:
+            model_id: The ID of the model.
+            version_name: The name of the model version.
+            producer_run_id: The ID of the producer pipeline run. If this is
+                set, only numeric versions created as part of the pipeline run
+                will be returned.
+
+        Returns:
+            The model version.
+        """
+        model_version = super()._get_model_version(
+            model_id=model_id,
+            version_name=version_name,
+            producer_run_id=producer_run_id,
+        )
+        verify_permission_for_model(model_version, action=Action.READ)
+        return model_version
+
+    def _get_or_create_model_version(
+        self,
+        model_version_request: ModelVersionRequest,
+        producer_run_id: Optional[UUID] = None,
+    ) -> Tuple[bool, ModelVersionResponse]:
+        """Get or create a model version.
+
+        Args:
+            model_version_request: The model version request.
+            producer_run_id: ID of the producer pipeline run.
+
+        # noqa: DAR401
+        Raises:
+            Exception: If the authenticated user is not allowed to
+                create a model version.
+
+        Returns:
+            A boolean whether the model version was created or not, and the
+            model version.
+        """
+        allow_creation = True
+        error = None
+
+        try:
+            verify_permission(
+                resource_type=ResourceType.MODEL_VERSION, action=Action.CREATE
+            )
+        except Exception as e:
+            allow_creation = False
+            error = e
+
+        if allow_creation:
+            created, model_version_response = (
+                super()._get_or_create_model_version(model_version_request, producer_run_id=producer_run_id)
+            )
+        else:
+            try:
+                model_version_response = self._get_model_version(
+                    model_id=model_version_request.model,
+                    version_name=model_version_request.name,
+                    producer_run_id=producer_run_id,
+                )
+                created = False
+            except KeyError:
+                # The model version does not exist. We now raise the error that
+                # explains why the version could not be created, instead of just
+                # the KeyError that it doesn't exist
+                assert error
+                raise error from None
+
+        return created, model_version_response

zenml/zen_server/utils.py

@@ -421,6 +421,8 @@ def make_dependable(cls: Type[BaseModel]) -> Callable[..., Any]:
     """
     from fastapi import Query
 
+    from zenml.zen_server.exceptions import error_detail
+
     def init_cls_and_handle_errors(*args: Any, **kwargs: Any) -> BaseModel:
         from fastapi import HTTPException
 
@@ -428,9 +430,8 @@ def make_dependable(cls: Type[BaseModel]) -> Callable[..., Any]:
             inspect.signature(init_cls_and_handle_errors).bind(*args, **kwargs)
             return cls(*args, **kwargs)
         except ValidationError as e:
-            for error in e.errors():
-                error["loc"] = tuple(["query"] + list(error["loc"]))
-            raise HTTPException(422, detail=e.errors())
+            detail = error_detail(e, exception_type=ValueError)
+            raise HTTPException(422, detail=detail)
 
     params = {v.name: v for v in inspect.signature(cls).parameters.values()}
     query_params = getattr(cls, "API_MULTI_INPUT_PARAMS", [])

zenml/zen_stores/base_zen_store.py

@@ -36,6 +36,7 @@ from zenml.constants import (
     DEFAULT_STACK_AND_COMPONENT_NAME,
     DEFAULT_WORKSPACE_NAME,
     ENV_ZENML_DEFAULT_WORKSPACE_NAME,
+    ENV_ZENML_SERVER,
     IS_DEBUG_ENV,
 )
 from zenml.enums import (
@@ -155,9 +156,16 @@ class BaseZenStore(
             TypeError: If the store type is unsupported.
         """
         if store_type == StoreType.SQL:
-            from zenml.zen_stores.sql_zen_store import SqlZenStore
+            if os.environ.get(ENV_ZENML_SERVER):
+                from zenml.zen_server.rbac.rbac_sql_zen_store import (
+                    RBACSqlZenStore,
+                )
+
+                return RBACSqlZenStore
+            else:
+                from zenml.zen_stores.sql_zen_store import SqlZenStore
 
-            return SqlZenStore
+                return SqlZenStore
         elif store_type == StoreType.REST:
             from zenml.zen_stores.rest_zen_store import RestZenStore
 

zenml/zen_stores/migrations/versions/a1237ba94fd8_add_model_version_producer_run_unique_.py

@@ -0,0 +1,68 @@
+"""Add model version producer run unique constraint [a1237ba94fd8].
+
+Revision ID: a1237ba94fd8
+Revises: 26351d482b9e
+Create Date: 2024-12-13 10:28:55.432414
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "a1237ba94fd8"
+down_revision = "26351d482b9e"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("model_version", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "producer_run_id_if_numeric",
+                sqlmodel.sql.sqltypes.GUID(),
+                nullable=True,
+            )
+        )
+
+    # Set the producer_run_id_if_numeric column to the model version ID for
+    # existing rows
+    connection = op.get_bind()
+    metadata = sa.MetaData()
+    metadata.reflect(only=("model_version",), bind=connection)
+    model_version_table = sa.Table("model_version", metadata)
+
+    connection.execute(
+        model_version_table.update().values(
+            producer_run_id_if_numeric=model_version_table.c.id
+        )
+    )
+
+    with op.batch_alter_table("model_version", schema=None) as batch_op:
+        batch_op.alter_column(
+            "producer_run_id_if_numeric",
+            existing_type=sqlmodel.sql.sqltypes.GUID(),
+            nullable=False,
+        )
+        batch_op.create_unique_constraint(
+            "unique_numeric_version_for_pipeline_run",
+            ["model_id", "producer_run_id_if_numeric"],
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("model_version", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_numeric_version_for_pipeline_run", type_="unique"
+        )
+        batch_op.drop_column("producer_run_id_if_numeric")
+
+    # ### end Alembic commands ###

zenml/zen_stores/schemas/model_schemas.py

@@ -15,10 +15,16 @@
 
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
-from uuid import UUID
+from uuid import UUID, uuid4
 
 from pydantic import ConfigDict
-from sqlalchemy import BOOLEAN, INTEGER, TEXT, Column, UniqueConstraint
+from sqlalchemy import (
+    BOOLEAN,
+    INTEGER,
+    TEXT,
+    Column,
+    UniqueConstraint,
+)
 from sqlmodel import Field, Relationship
 
 from zenml.enums import (
@@ -228,11 +234,13 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
 
     __tablename__ = MODEL_VERSION_TABLENAME
     __table_args__ = (
-        # We need two unique constraints here:
+        # We need three unique constraints here:
         # - The first to ensure that each model version for a
         #   model has a unique version number
         # - The second one to ensure that explicit names given by
         #   users are unique
+        # - The third one to ensure that a pipeline run only produces a single
+        #   auto-incremented version per model
         UniqueConstraint(
             "number",
             "model_id",
@@ -243,6 +251,11 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
             "model_id",
             name="unique_version_for_model_id",
         ),
+        UniqueConstraint(
+            "model_id",
+            "producer_run_id_if_numeric",
+            name="unique_numeric_version_for_pipeline_run",
+        ),
     )
 
     workspace_id: UUID = build_foreign_key_field(
@@ -312,12 +325,23 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
         ),
     )
     pipeline_runs: List["PipelineRunSchema"] = Relationship(
-        back_populates="model_version"
+        back_populates="model_version",
     )
     step_runs: List["StepRunSchema"] = Relationship(
         back_populates="model_version"
     )
 
+    # We want to make sure each pipeline run only creates a single numeric
+    # version for each model. To solve this, we need to add a unique constraint.
+    # If a value of a unique constraint is NULL it is ignored and the
+    # remaining values in the unique constraint have to be unique. In
+    # our case however, we only want the unique constraint applied in
+    # case there is a producer run and only for numeric versions. To solve this,
+    # we fall back to the model version ID (which is the primary key and
+    # therefore unique) in case there is no producer run or the version is not
+    # numeric.
+    producer_run_id_if_numeric: UUID
+
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.
     #  It is possible to suppress this warning message with the following
@@ -328,24 +352,36 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
 
     @classmethod
     def from_request(
-        cls, model_version_request: ModelVersionRequest
+        cls,
+        model_version_request: ModelVersionRequest,
+        model_version_number: int,
+        producer_run_id: Optional[UUID] = None,
     ) -> "ModelVersionSchema":
         """Convert an `ModelVersionRequest` to an `ModelVersionSchema`.
 
         Args:
             model_version_request: The request model version to convert.
+            model_version_number: The model version number.
+            producer_run_id: The ID of the producer run.
 
         Returns:
             The converted schema.
         """
+        id_ = uuid4()
+        is_numeric = str(model_version_number) == model_version_request.name
+
         return cls(
+            id=id_,
             workspace_id=model_version_request.workspace,
             user_id=model_version_request.user,
             model_id=model_version_request.model,
             name=model_version_request.name,
-            number=model_version_request.number,
+            number=model_version_number,
             description=model_version_request.description,
             stage=model_version_request.stage,
+            producer_run_id_if_numeric=producer_run_id
+            if (producer_run_id and is_numeric)
+            else id_,
         )
 
     def to_model(

zenml/zen_stores/schemas/pipeline_deployment_schemas.py

@@ -228,13 +228,6 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
         Returns:
             The created `PipelineDeploymentResponse`.
         """
-        pipeline_configuration = PipelineConfiguration.model_validate_json(
-            self.pipeline_configuration
-        )
-        step_configurations = json.loads(self.step_configurations)
-        for s, c in step_configurations.items():
-            step_configurations[s] = Step.model_validate(c)
-
         body = PipelineDeploymentResponseBody(
             user=self.user.to_model() if self.user else None,
             created=self.created,
@@ -242,6 +235,13 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
         )
         metadata = None
         if include_metadata:
+            pipeline_configuration = PipelineConfiguration.model_validate_json(
+                self.pipeline_configuration
+            )
+            step_configurations = json.loads(self.step_configurations)
+            for s, c in step_configurations.items():
+                step_configurations[s] = Step.model_validate(c)
+
             metadata = PipelineDeploymentResponseMetadata(
                 workspace=self.workspace.to_model(),
                 run_name_template=self.run_name_template,

zenml/zen_stores/schemas/pipeline_schemas.py

@@ -156,7 +156,12 @@ class PipelineSchema(NamedSchema, table=True):
 
         resources = None
         if include_resources:
+            latest_run_user = self.runs[-1].user if self.runs else None
+
             resources = PipelineResponseResources(
+                latest_run_user=latest_run_user.to_model()
+                if latest_run_user
+                else None,
                 tags=[t.tag.to_model() for t in self.tags],
             )