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

zenml/orchestrators/step_run_utils.py

@@ -14,7 +14,7 @@
 """Utilities for creating step runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
+from typing import Dict, List, Optional, Set, Tuple
 
 from zenml.client import Client
 from zenml.config.step_configurations import Step
@@ -24,21 +24,13 @@ from zenml.logger import get_logger
 from zenml.model.utils import link_artifact_version_to_model_version
 from zenml.models import (
     ArtifactVersionResponse,
-    ModelVersionPipelineRunRequest,
     ModelVersionResponse,
     PipelineDeploymentResponse,
     PipelineRunResponse,
-    PipelineRunUpdate,
     StepRunRequest,
-    StepRunResponse,
-    StepRunUpdate,
 )
 from zenml.orchestrators import cache_utils, input_utils, utils
 from zenml.stack import Stack
-from zenml.utils import pagination_utils, string_utils
-
-if TYPE_CHECKING:
-    from zenml.model.model import Model
 
 logger = get_logger(__name__)
 
@@ -293,10 +285,6 @@ def create_cached_step_runs(
         deployment=deployment, pipeline_run=pipeline_run, stack=stack
     )
 
-    pipeline_model_version, pipeline_run = prepare_pipeline_run_model_version(
-        pipeline_run=pipeline_run
-    )
-
     while (
         cache_candidates := find_cacheable_invocation_candidates(
             deployment=deployment,
@@ -311,7 +299,9 @@ def create_cached_step_runs(
 
             # Make sure the request factory has the most up to date pipeline
             # run to avoid hydration calls
-            request_factory.pipeline_run = pipeline_run
+            request_factory.pipeline_run = Client().get_pipeline_run(
+                pipeline_run.id
+            )
             try:
                 step_run_request = request_factory.create_request(
                     invocation_id
@@ -336,15 +326,10 @@ def create_cached_step_runs(
 
             step_run = Client().zen_store.create_run_step(step_run_request)
 
-            # Refresh the pipeline run here to make sure we have the latest
-            # state
-            pipeline_run = Client().get_pipeline_run(pipeline_run.id)
-
-            step_model_version, step_run = prepare_step_run_model_version(
-                step_run=step_run, pipeline_run=pipeline_run
-            )
-
-            if model_version := step_model_version or pipeline_model_version:
+            if (
+                model_version := step_run.model_version
+                or pipeline_run.model_version
+            ):
                 link_output_artifacts_to_model_version(
                     artifacts=step_run.outputs,
                     model_version=model_version,
@@ -356,169 +341,6 @@ def create_cached_step_runs(
     return cached_invocations
 
 
-def get_or_create_model_version_for_pipeline_run(
-    model: "Model",
-    pipeline_run: PipelineRunResponse,
-    substitutions: Dict[str, str],
-) -> Tuple[ModelVersionResponse, bool]:
-    """Get or create a model version as part of a pipeline run.
-
-    Args:
-        model: The model to get or create.
-        pipeline_run: The pipeline run for which the model should be created.
-        substitutions: Substitutions to apply to the model version name.
-
-    Returns:
-        The model version and a boolean indicating whether it was newly created
-        or not.
-    """
-    # Copy the model before modifying it so we don't accidently modify
-    # configurations in which the model object is potentially referenced
-    model = model.model_copy()
-
-    if model.model_version_id:
-        return model._get_model_version(), False
-    elif model.version:
-        if isinstance(model.version, str):
-            model.version = string_utils.format_name_template(
-                model.version,
-                substitutions=substitutions,
-            )
-        model.name = string_utils.format_name_template(
-            model.name,
-            substitutions=substitutions,
-        )
-
-        return (
-            model._get_or_create_model_version(),
-            model._created_model_version,
-        )
-
-    # The model version should be created as part of this run
-    # -> We first check if it was already created as part of this run, and if
-    # not we do create it. If this is running in two parallel steps, we might
-    # run into issues that this will create two versions. Ideally, all model
-    # versions required for a pipeline run and its steps could be created
-    # server-side at run creation time before the first step starts.
-    if model_version := get_model_version_created_by_pipeline_run(
-        model_name=model.name, pipeline_run=pipeline_run
-    ):
-        return model_version, False
-    else:
-        return model._get_or_create_model_version(), True
-
-
-def get_model_version_created_by_pipeline_run(
-    model_name: str, pipeline_run: PipelineRunResponse
-) -> Optional[ModelVersionResponse]:
-    """Get a model version that was created by a specific pipeline run.
-
-    This function does not refresh the pipeline run, so it will only try to
-    fetch the model version from existing steps if they're already part of the
-    response.
-
-    Args:
-        model_name: The model name for which to get the version.
-        pipeline_run: The pipeline run for which to get the version.
-
-    Returns:
-        A model version with the given name created by the run, or None if such
-        a model version does not exist.
-    """
-    if pipeline_run.config.model and pipeline_run.model_version:
-        if (
-            pipeline_run.config.model.name == model_name
-            and pipeline_run.config.model.version is None
-        ):
-            return pipeline_run.model_version
-
-    # We fetch a list of hydrated step runs here in order to avoid hydration
-    # calls for each step separately.
-    candidate_step_runs = pagination_utils.depaginate(
-        Client().list_run_steps,
-        pipeline_run_id=pipeline_run.id,
-        model=model_name,
-        hydrate=True,
-    )
-    for step_run in candidate_step_runs:
-        if step_run.config.model and step_run.model_version:
-            if (
-                step_run.config.model.name == model_name
-                and step_run.config.model.version is None
-            ):
-                return step_run.model_version
-
-    return None
-
-
-def prepare_pipeline_run_model_version(
-    pipeline_run: PipelineRunResponse,
-) -> Tuple[Optional[ModelVersionResponse], PipelineRunResponse]:
-    """Prepare the model version for a pipeline run.
-
-    Args:
-        pipeline_run: The pipeline run for which to prepare the model version.
-
-    Returns:
-        The prepared model version and the updated pipeline run.
-    """
-    model_version = None
-
-    if pipeline_run.model_version:
-        model_version = pipeline_run.model_version
-    elif config_model := pipeline_run.config.model:
-        model_version, _ = get_or_create_model_version_for_pipeline_run(
-            model=config_model,
-            pipeline_run=pipeline_run,
-            substitutions=pipeline_run.config.substitutions,
-        )
-        pipeline_run = Client().zen_store.update_run(
-            run_id=pipeline_run.id,
-            run_update=PipelineRunUpdate(model_version_id=model_version.id),
-        )
-        link_pipeline_run_to_model_version(
-            pipeline_run=pipeline_run, model_version=model_version
-        )
-        log_model_version_dashboard_url(model_version)
-
-    return model_version, pipeline_run
-
-
-def prepare_step_run_model_version(
-    step_run: StepRunResponse, pipeline_run: PipelineRunResponse
-) -> Tuple[Optional[ModelVersionResponse], StepRunResponse]:
-    """Prepare the model version for a step run.
-
-    Args:
-        step_run: The step run for which to prepare the model version.
-        pipeline_run: The pipeline run of the step.
-
-    Returns:
-        The prepared model version and the updated step run.
-    """
-    model_version = None
-
-    if step_run.model_version:
-        model_version = step_run.model_version
-    elif config_model := step_run.config.model:
-        model_version, created = get_or_create_model_version_for_pipeline_run(
-            model=config_model,
-            pipeline_run=pipeline_run,
-            substitutions=step_run.config.substitutions,
-        )
-        step_run = Client().zen_store.update_run_step(
-            step_run_id=step_run.id,
-            step_run_update=StepRunUpdate(model_version_id=model_version.id),
-        )
-        link_pipeline_run_to_model_version(
-            pipeline_run=pipeline_run, model_version=model_version
-        )
-        if created:
-            log_model_version_dashboard_url(model_version)
-
-    return model_version, step_run
-
-
 def log_model_version_dashboard_url(
     model_version: ModelVersionResponse,
 ) -> None:
@@ -546,24 +368,6 @@ def log_model_version_dashboard_url(
         )
 
 
-def link_pipeline_run_to_model_version(
-    pipeline_run: PipelineRunResponse, model_version: ModelVersionResponse
-) -> None:
-    """Link a pipeline run to a model version.
-
-    Args:
-        pipeline_run: The pipeline run to link.
-        model_version: The model version to link.
-    """
-    client = Client()
-    client.zen_store.create_model_version_pipeline_run_link(
-        ModelVersionPipelineRunRequest(
-            pipeline_run=pipeline_run.id,
-            model_version=model_version.id,
-        )
-    )
-
-
 def link_output_artifacts_to_model_version(
     artifacts: Dict[str, List[ArtifactVersionResponse]],
     model_version: ModelVersionResponse,

zenml/pipelines/build_utils.py

@@ -249,6 +249,11 @@ def find_existing_build(
     client = Client()
     stack = client.active_stack
 
+    if not stack.container_registry:
+        # There can be no non-local builds that we can reuse if there is no
+        # container registry in the stack.
+        return None
+
     python_version_prefix = ".".join(platform.python_version_tuple()[:2])
     required_builds = stack.get_docker_builds(deployment=deployment)
 
@@ -263,6 +268,13 @@ def find_existing_build(
         sort_by="desc:created",
         size=1,
         stack_id=stack.id,
+        # Until we implement stack versioning, users can still update their
+        # stack to update/remove the container registry. In that case, we might
+        # try to pull an image from a container registry that we don't have
+        # access to. This is why we add an additional check for the container
+        # registry ID here. (This is still not perfect as users can update the
+        # container registry URI or config, but the best we can do)
+        container_registry_id=stack.container_registry.id,
         # The build is local and it's not clear whether the images
         # exist on the current machine or if they've been overwritten.
         # TODO: Should we support this by storing the unique Docker ID for

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/26351d482b9e_add_step_run_unique_constraint.py

@@ -0,0 +1,37 @@
+"""Add step run unique constraint [26351d482b9e].
+
+Revision ID: 26351d482b9e
+Revises: 0.71.0
+Create Date: 2024-12-03 11:46:57.541578
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "26351d482b9e"
+down_revision = "0.71.0"
+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("step_run", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_step_name_for_pipeline_run", ["name", "pipeline_run_id"]
+        )
+
+    # ### 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("step_run", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_step_name_for_pipeline_run", type_="unique"
+        )
+
+    # ### end Alembic commands ###

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 ###