zenml-nightly 0.70.0.dev20241120__py3-none-any.whl → 0.70.0.dev20241121__py3-none-any.whl

zenml/models/v2/core/model_version_pipeline_run.py

@@ -23,23 +23,19 @@ from zenml.enums import GenericFilterOps
 from zenml.models.v2.base.base import (
     BaseDatedResponseBody,
     BaseIdentifiedResponse,
+    BaseRequest,
     BaseResponseMetadata,
     BaseResponseResources,
 )
-from zenml.models.v2.base.filter import StrFilter
-from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
-    WorkspaceScopedRequest,
-)
+from zenml.models.v2.base.filter import BaseFilter, StrFilter
 from zenml.models.v2.core.pipeline_run import PipelineRunResponse
 
 # ------------------ Request Model ------------------
 
 
-class ModelVersionPipelineRunRequest(WorkspaceScopedRequest):
+class ModelVersionPipelineRunRequest(BaseRequest):
     """Request model for links between model versions and pipeline runs."""
 
-    model: UUID
     model_version: UUID
     pipeline_run: UUID
 
@@ -62,7 +58,6 @@ class ModelVersionPipelineRunRequest(WorkspaceScopedRequest):
 class ModelVersionPipelineRunResponseBody(BaseDatedResponseBody):
     """Response body for links between model versions and pipeline runs."""
 
-    model: UUID
     model_version: UUID
     pipeline_run: PipelineRunResponse
 
@@ -88,16 +83,6 @@ class ModelVersionPipelineRunResponse(
 ):
     """Response model for links between model versions and pipeline runs."""
 
-    # Body and metadata properties
-    @property
-    def model(self) -> UUID:
-        """The `model` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().model
-
     @property
     def model_version(self) -> UUID:
         """The `model_version` property.
@@ -120,39 +105,21 @@ class ModelVersionPipelineRunResponse(
 # ------------------ Filter Model ------------------
 
 
-class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
+class ModelVersionPipelineRunFilter(BaseFilter):
     """Model version pipeline run links filter model."""
 
     FILTER_EXCLUDE_FIELDS = [
-        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        *BaseFilter.FILTER_EXCLUDE_FIELDS,
         "pipeline_run_name",
         "user",
     ]
     CLI_EXCLUDE_FIELDS = [
-        *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
-        "model_id",
+        *BaseFilter.CLI_EXCLUDE_FIELDS,
         "model_version_id",
-        "user_id",
-        "workspace_id",
         "updated",
         "id",
     ]
 
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="The workspace of the Model Version",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="The user of the Model Version",
-        union_mode="left_to_right",
-    )
-    model_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Filter by model ID",
-        union_mode="left_to_right",
-    )
     model_version_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Filter by model version ID",

zenml/orchestrators/step_launcher.py

@@ -294,7 +294,6 @@ class StepLauncher:
                     ):
                         step_run_utils.link_output_artifacts_to_model_version(
                             artifacts=step_run.outputs,
-                            output_configurations=step_run.config.outputs,
                             model_version=model_version,
                         )
 

zenml/orchestrators/step_run_utils.py

@@ -14,12 +14,12 @@
 """Utilities for creating step runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
 
 from zenml.client import Client
-from zenml.config.step_configurations import ArtifactConfiguration, Step
+from zenml.config.step_configurations import Step
 from zenml.constants import CODE_HASH_PARAMETER_NAME, TEXT_FIELD_MAX_LENGTH
-from zenml.enums import ArtifactSaveType, ExecutionStatus
+from zenml.enums import ExecutionStatus
 from zenml.logger import get_logger
 from zenml.model.utils import link_artifact_version_to_model_version
 from zenml.models import (
@@ -344,7 +344,6 @@ def create_cached_step_runs(
             if model_version := step_model_version or pipeline_model_version:
                 link_output_artifacts_to_model_version(
                     artifacts=step_run.outputs,
-                    output_configurations=step_run.config.outputs,
                     model_version=model_version,
                 )
 
@@ -542,10 +541,7 @@ def link_pipeline_run_to_model_version(
     client = Client()
     client.zen_store.create_model_version_pipeline_run_link(
         ModelVersionPipelineRunRequest(
-            user=client.active_user.id,
-            workspace=client.active_workspace.id,
             pipeline_run=pipeline_run.id,
-            model=model_version.model.id,
             model_version=model_version.id,
         )
     )
@@ -553,26 +549,17 @@ def link_pipeline_run_to_model_version(
 
 def link_output_artifacts_to_model_version(
     artifacts: Dict[str, List[ArtifactVersionResponse]],
-    output_configurations: Mapping[str, ArtifactConfiguration],
     model_version: ModelVersionResponse,
 ) -> None:
     """Link the outputs of a step run to a model version.
 
     Args:
         artifacts: The step output artifacts.
-        output_configurations: The output configurations for the step.
         model_version: The model version to link.
     """
-    for output_name, output_artifacts in artifacts.items():
+    for output_artifacts in artifacts.values():
         for output_artifact in output_artifacts:
-            artifact_config = None
-            if output_artifact.save_type == ArtifactSaveType.STEP_OUTPUT and (
-                output_config := output_configurations.get(output_name, None)
-            ):
-                artifact_config = output_config.artifact_config
-
             link_artifact_version_to_model_version(
                 artifact_version=output_artifact,
                 model_version=model_version,
-                artifact_config=artifact_config,
             )

zenml/orchestrators/step_runner.py

@@ -247,7 +247,6 @@ class StepRunner:
                                 artifacts={
                                     k: [v] for k, v in output_artifacts.items()
                                 },
-                                output_configurations=step_run.config.outputs,
                                 model_version=model_version,
                             )
                 finally:
@@ -577,9 +576,11 @@ class StepRunner:
             uri = output_artifact_uris[output_name]
             artifact_config = output_annotations[output_name].artifact_config
 
+            artifact_type = None
             if artifact_config is not None:
                 has_custom_name = bool(artifact_config.name)
                 version = artifact_config.version
+                artifact_type = artifact_config.artifact_type
             else:
                 has_custom_name, version = False, None
 
@@ -605,6 +606,7 @@ class StepRunner:
                 data=return_value,
                 materializer_class=materializer_class,
                 uri=uri,
+                artifact_type=artifact_type,
                 store_metadata=artifact_metadata_enabled,
                 store_visualizations=artifact_visualization_enabled,
                 has_custom_name=has_custom_name,

zenml/zen_server/routers/model_versions_endpoints.py

@@ -29,9 +29,11 @@ from zenml.constants import (
 )
 from zenml.models import (
     ModelVersionArtifactFilter,
+    ModelVersionArtifactRequest,
     ModelVersionArtifactResponse,
     ModelVersionFilter,
     ModelVersionPipelineRunFilter,
+    ModelVersionPipelineRunRequest,
     ModelVersionPipelineRunResponse,
     ModelVersionResponse,
     ModelVersionUpdate,
@@ -198,6 +200,34 @@ model_version_artifacts_router = APIRouter(
 )
 
 
+@model_version_artifacts_router.post(
+    "",
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_model_version_artifact_link(
+    model_version_artifact_link: ModelVersionArtifactRequest,
+    _: AuthContext = Security(authorize),
+) -> ModelVersionArtifactResponse:
+    """Create a new model version to artifact link.
+
+    Args:
+        model_version_artifact_link: The model version to artifact link to create.
+
+    Returns:
+        The created model version to artifact link.
+    """
+    model_version = zen_store().get_model_version(
+        model_version_artifact_link.model_version
+    )
+    verify_permission_for_model(model_version, action=Action.UPDATE)
+
+    mv = zen_store().create_model_version_artifact_link(
+        model_version_artifact_link
+    )
+    return mv
+
+
 @model_version_artifacts_router.get(
     "",
     response_model=Page[ModelVersionArtifactResponse],
@@ -291,6 +321,35 @@ model_version_pipeline_runs_router = APIRouter(
 )
 
 
+@model_version_pipeline_runs_router.post(
+    "",
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_model_version_pipeline_run_link(
+    model_version_pipeline_run_link: ModelVersionPipelineRunRequest,
+    _: AuthContext = Security(authorize),
+) -> ModelVersionPipelineRunResponse:
+    """Create a new model version to pipeline run link.
+
+    Args:
+        model_version_pipeline_run_link: The model version to pipeline run link to create.
+
+    Returns:
+        - If Model Version to Pipeline Run Link already exists - returns the existing link.
+        - Otherwise, returns the newly created model version to pipeline run link.
+    """
+    model_version = zen_store().get_model_version(
+        model_version_pipeline_run_link.model_version, hydrate=False
+    )
+    verify_permission_for_model(model_version, action=Action.UPDATE)
+
+    mv = zen_store().create_model_version_pipeline_run_link(
+        model_version_pipeline_run_link
+    )
+    return mv
+
+
 @model_version_pipeline_runs_router.get(
     "",
     response_model=Page[ModelVersionPipelineRunResponse],

zenml/zen_server/routers/workspaces_endpoints.py

@@ -20,7 +20,6 @@ from fastapi import APIRouter, Depends, Security
 
 from zenml.constants import (
     API,
-    ARTIFACTS,
     CODE_REPOSITORIES,
     GET_OR_CREATE,
     MODEL_VERSIONS,
@@ -54,10 +53,6 @@ from zenml.models import (
     ComponentResponse,
     ModelRequest,
     ModelResponse,
-    ModelVersionArtifactRequest,
-    ModelVersionArtifactResponse,
-    ModelVersionPipelineRunRequest,
-    ModelVersionPipelineRunResponse,
     ModelVersionRequest,
     ModelVersionResponse,
     Page,
@@ -1442,131 +1437,6 @@ def create_model_version(
     )
 
 
-@router.post(
-    WORKSPACES
-    + "/{workspace_name_or_id}"
-    + MODEL_VERSIONS
-    + "/{model_version_id}"
-    + ARTIFACTS,
-    response_model=ModelVersionArtifactResponse,
-    responses={401: error_response, 409: error_response, 422: error_response},
-)
-@handle_exceptions
-def create_model_version_artifact_link(
-    workspace_name_or_id: Union[str, UUID],
-    model_version_id: UUID,
-    model_version_artifact_link: ModelVersionArtifactRequest,
-    auth_context: AuthContext = Security(authorize),
-) -> ModelVersionArtifactResponse:
-    """Create a new model version to artifact link.
-
-    Args:
-        workspace_name_or_id: Name or ID of the workspace.
-        model_version_id: ID of the model version.
-        model_version_artifact_link: The model version to artifact link to create.
-        auth_context: Authentication context.
-
-    Returns:
-        The created model version to artifact link.
-
-    Raises:
-        IllegalOperationError: If the workspace or user specified in the
-            model version does not match the current workspace or authenticated
-            user.
-    """
-    workspace = zen_store().get_workspace(workspace_name_or_id)
-    if str(model_version_id) != str(model_version_artifact_link.model_version):
-        raise IllegalOperationError(
-            f"The model version id in your path `{model_version_id}` does not "
-            f"match the model version specified in the request model "
-            f"`{model_version_artifact_link.model_version}`"
-        )
-
-    if model_version_artifact_link.workspace != workspace.id:
-        raise IllegalOperationError(
-            "Creating model version to artifact links outside of the workspace scope "
-            f"of this endpoint `{workspace_name_or_id}` is "
-            f"not supported."
-        )
-    if model_version_artifact_link.user != auth_context.user.id:
-        raise IllegalOperationError(
-            "Creating model to artifact links for a user other than yourself "
-            "is not supported."
-        )
-
-    model_version = zen_store().get_model_version(model_version_id)
-    verify_permission_for_model(model_version, action=Action.UPDATE)
-
-    mv = zen_store().create_model_version_artifact_link(
-        model_version_artifact_link
-    )
-    return mv
-
-
-@router.post(
-    WORKSPACES
-    + "/{workspace_name_or_id}"
-    + MODEL_VERSIONS
-    + "/{model_version_id}"
-    + RUNS,
-    response_model=ModelVersionPipelineRunResponse,
-    responses={401: error_response, 409: error_response, 422: error_response},
-)
-@handle_exceptions
-def create_model_version_pipeline_run_link(
-    workspace_name_or_id: Union[str, UUID],
-    model_version_id: UUID,
-    model_version_pipeline_run_link: ModelVersionPipelineRunRequest,
-    auth_context: AuthContext = Security(authorize),
-) -> ModelVersionPipelineRunResponse:
-    """Create a new model version to pipeline run link.
-
-    Args:
-        workspace_name_or_id: Name or ID of the workspace.
-        model_version_id: ID of the model version.
-        model_version_pipeline_run_link: The model version to pipeline run link to create.
-        auth_context: Authentication context.
-
-    Returns:
-        - If Model Version to Pipeline Run Link already exists - returns the existing link.
-        - Otherwise, returns the newly created model version to pipeline run link.
-
-    Raises:
-        IllegalOperationError: If the workspace or user specified in the
-            model version does not match the current workspace or authenticated
-            user.
-    """
-    workspace = zen_store().get_workspace(workspace_name_or_id)
-    if str(model_version_id) != str(
-        model_version_pipeline_run_link.model_version
-    ):
-        raise IllegalOperationError(
-            f"The model version id in your path `{model_version_id}` does not "
-            f"match the model version specified in the request model "
-            f"`{model_version_pipeline_run_link.model_version}`"
-        )
-
-    if model_version_pipeline_run_link.workspace != workspace.id:
-        raise IllegalOperationError(
-            "Creating model versions outside of the workspace scope "
-            f"of this endpoint `{workspace_name_or_id}` is "
-            f"not supported."
-        )
-    if model_version_pipeline_run_link.user != auth_context.user.id:
-        raise IllegalOperationError(
-            "Creating models for a user other than yourself "
-            "is not supported."
-        )
-
-    model_version = zen_store().get_model_version(model_version_id)
-    verify_permission_for_model(model_version, action=Action.UPDATE)
-
-    mv = zen_store().create_model_version_pipeline_run_link(
-        model_version_pipeline_run_link
-    )
-    return mv
-
-
 @router.post(
     WORKSPACES + "/{workspace_name_or_id}" + SERVICES,
     response_model=ServiceResponse,

zenml/zen_stores/base_zen_store.py

@@ -42,6 +42,7 @@ from zenml.enums import (
     SecretsStoreType,
     StoreType,
 )
+from zenml.exceptions import IllegalOperationError
 from zenml.logger import get_logger
 from zenml.models import (
     ServerDatabaseType,
@@ -335,7 +336,7 @@ class BaseZenStore(
             # Ensure that the active stack is still valid
             try:
                 active_stack = self.get_stack(stack_id=active_stack_id)
-            except KeyError:
+            except (KeyError, IllegalOperationError):
                 logger.warning(
                     "The current %s active stack is no longer available. "
                     "Resetting the active stack to default.",

zenml/zen_stores/migrations/versions/ec6307720f92_simplify_model_version_links.py

@@ -0,0 +1,118 @@
+"""Simplify model version links [ec6307720f92].
+
+Revision ID: ec6307720f92
+Revises: 0.70.0
+Create Date: 2024-11-06 16:16:43.344569
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "ec6307720f92"
+down_revision = "0.70.0"
+branch_labels = None
+depends_on = None
+
+
+def _migrate_artifact_type() -> None:
+    """Migrate the artifact type."""
+    meta = sa.MetaData()
+    meta.reflect(
+        bind=op.get_bind(),
+        only=("artifact_version", "model_versions_artifacts"),
+    )
+    artifact_version_table = sa.Table("artifact_version", meta)
+    model_versions_artifacts = sa.Table("model_versions_artifacts", meta)
+
+    connection = op.get_bind()
+
+    query = sa.select(
+        model_versions_artifacts.c.artifact_version_id,
+        model_versions_artifacts.c.is_model_artifact,
+        model_versions_artifacts.c.is_deployment_artifact,
+    )
+    result = connection.execute(query)
+
+    updated = set()
+    updates = []
+    for (
+        artifact_version_id,
+        is_model_artifact,
+        is_deployment_artifact,
+    ) in result:
+        if artifact_version_id in updated:
+            # If an artifact was a model artifact in one model version, and a
+            # deployment artifact in another model version, we only update it
+            # once.
+            continue
+
+        if is_model_artifact:
+            updated.add(artifact_version_id)
+            updates.append(
+                {"id_": artifact_version_id, "type": "ModelArtifact"}
+            )
+        elif is_deployment_artifact:
+            updated.add(artifact_version_id)
+            updates.append(
+                {"id_": artifact_version_id, "type": "ServiceArtifact"}
+            )
+
+    connection.execute(
+        sa.update(artifact_version_table).where(
+            artifact_version_table.c.id == sa.bindparam("id_")
+        ),
+        updates,
+    )
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    _migrate_artifact_type()
+
+    with op.batch_alter_table(
+        "model_versions_artifacts", schema=None
+    ) as batch_op:
+        batch_op.drop_constraint(
+            "fk_model_versions_artifacts_workspace_id_workspace",
+            type_="foreignkey",
+        )
+        batch_op.drop_constraint(
+            "fk_model_versions_artifacts_user_id_user", type_="foreignkey"
+        )
+        batch_op.drop_constraint(
+            "fk_model_versions_artifacts_model_id_model", type_="foreignkey"
+        )
+        batch_op.drop_column("user_id")
+        batch_op.drop_column("model_id")
+        batch_op.drop_column("is_deployment_artifact")
+        batch_op.drop_column("workspace_id")
+        batch_op.drop_column("is_model_artifact")
+
+    with op.batch_alter_table("model_versions_runs", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "fk_model_versions_runs_model_id_model", type_="foreignkey"
+        )
+        batch_op.drop_constraint(
+            "fk_model_versions_runs_workspace_id_workspace", type_="foreignkey"
+        )
+        batch_op.drop_constraint(
+            "fk_model_versions_runs_user_id_user", type_="foreignkey"
+        )
+        batch_op.drop_column("model_id")
+        batch_op.drop_column("workspace_id")
+        batch_op.drop_column("user_id")
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision.
+
+    Raises:
+        NotImplementedError: Downgrade is not supported for this migration.
+    """
+    raise NotImplementedError("Downgrade is not supported for this migration.")

zenml/zen_stores/rest_zen_store.py

@@ -3672,10 +3672,10 @@ class RestZenStore(BaseZenStore):
         Returns:
             The newly created model version to artifact link.
         """
-        return self._create_workspace_scoped_resource(
+        return self._create_resource(
             resource=model_version_artifact_link,
             response_model=ModelVersionArtifactResponse,
-            route=f"{MODEL_VERSIONS}/{model_version_artifact_link.model_version}{ARTIFACTS}",
+            route=MODEL_VERSION_ARTIFACTS,
         )
 
     def list_model_version_artifact_links(
@@ -3752,10 +3752,10 @@ class RestZenStore(BaseZenStore):
             - Otherwise, returns the newly created model version to pipeline
                 run link.
         """
-        return self._create_workspace_scoped_resource(
+        return self._create_resource(
             resource=model_version_pipeline_run_link,
             response_model=ModelVersionPipelineRunResponse,
-            route=f"{MODEL_VERSIONS}/{model_version_pipeline_run_link.model_version}{RUNS}",
+            route=MODEL_VERSION_PIPELINE_RUNS,
         )
 
     def list_model_version_pipeline_run_links(