zenml-nightly 0.68.1.dev20241106__py3-none-any.whl → 0.70.0.dev20241116__py3-none-any.whl

zenml/zen_stores/migrations/versions/1cb6477f72d6_move_artifact_save_type.py

@@ -0,0 +1,99 @@
+"""move artifact save type [1cb6477f72d6].
+
+Revision ID: 1cb6477f72d6
+Revises: c22561cbb3a9
+Create Date: 2024-10-10 15:44:09.465210
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "1cb6477f72d6"
+down_revision = "c22561cbb3a9"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # Step 1: Add nullable save_type column to artifact_version
+    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("save_type", sa.TEXT(), nullable=True))
+
+    # Step 2: Move data from step_run_output_artifact.type to artifact_version.save_type
+    op.execute("""
+        UPDATE artifact_version
+        SET save_type = (
+            SELECT max(step_run_output_artifact.type)
+            FROM step_run_output_artifact
+            WHERE step_run_output_artifact.artifact_id = artifact_version.id
+            GROUP BY artifact_id
+        )
+    """)
+    op.execute("""
+        UPDATE artifact_version
+        SET save_type = 'step_output'
+        WHERE artifact_version.save_type = 'default'
+    """)
+    op.execute("""
+        UPDATE artifact_version
+        SET save_type = 'external'
+        WHERE save_type is NULL
+    """)
+
+    # # Step 3: Set save_type to non-nullable
+    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
+        batch_op.alter_column(
+            "save_type",
+            existing_type=sa.TEXT(),
+            nullable=False,
+        )
+
+    # Step 4: Remove type column from step_run_output_artifact
+    with op.batch_alter_table(
+        "step_run_output_artifact", schema=None
+    ) as batch_op:
+        batch_op.drop_column("type")
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # Add type column back to step_run_output_artifact
+    with op.batch_alter_table(
+        "step_run_output_artifact", schema=None
+    ) as batch_op:
+        batch_op.add_column(
+            sa.Column("type", sa.TEXT(), nullable=True),
+        )
+
+    # Move data back from artifact_version.save_type to step_run_output_artifact.type
+    op.execute("""
+        UPDATE step_run_output_artifact
+        SET type = (
+            SELECT max(artifact_version.save_type)
+            FROM artifact_version
+            WHERE step_run_output_artifact.artifact_id = artifact_version.id
+            GROUP BY artifact_id
+        )
+    """)
+    op.execute("""
+        UPDATE step_run_output_artifact
+        SET type = 'default'
+        WHERE step_run_output_artifact.type = 'step_output'
+    """)
+
+    # Set type to non-nullable
+    with op.batch_alter_table(
+        "step_run_output_artifact", schema=None
+    ) as batch_op:
+        batch_op.alter_column(
+            "type",
+            existing_type=sa.TEXT(),
+            nullable=False,
+        )
+
+    # Remove save_type column from artifact_version
+    with op.batch_alter_table("artifact_version", schema=None) as batch_op:
+        batch_op.drop_column("save_type")

zenml/zen_stores/migrations/versions/904464ea4041_add_pipeline_model_run_unique_constraints.py

@@ -0,0 +1,192 @@
+"""Add pipeline, model and run unique constraints [904464ea4041].
+
+Revision ID: 904464ea4041
+Revises: b557b2871693
+Create Date: 2024-11-04 10:27:05.450092
+
+"""
+
+from collections import defaultdict
+from typing import Any, Dict, Set
+
+import sqlalchemy as sa
+from alembic import op
+
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
+
+# revision identifiers, used by Alembic.
+revision = "904464ea4041"
+down_revision = "b557b2871693"
+branch_labels = None
+depends_on = None
+
+
+def resolve_duplicate_entities() -> None:
+    """Resolve duplicate entities."""
+    connection = op.get_bind()
+    meta = sa.MetaData()
+    meta.reflect(
+        bind=connection,
+        only=("pipeline_run", "pipeline", "model", "model_version"),
+    )
+
+    # Remove duplicate names for runs, pipelines and models
+    for table_name in ["pipeline_run", "pipeline", "model"]:
+        table = sa.Table(table_name, meta)
+        result = connection.execute(
+            sa.select(table.c.id, table.c.name, table.c.workspace_id)
+        ).all()
+        existing: Dict[str, Set[str]] = defaultdict(set)
+
+        for id_, name, workspace_id in result:
+            names_in_workspace = existing[workspace_id]
+
+            if name in names_in_workspace:
+                new_name = f"{name}_{id_[:6]}"
+                logger.warning(
+                    "Migrating %s name from %s to %s to resolve duplicate name.",
+                    table_name,
+                    name,
+                    new_name,
+                )
+                connection.execute(
+                    sa.update(table)
+                    .where(table.c.id == id_)
+                    .values(name=new_name)
+                )
+                names_in_workspace.add(new_name)
+            else:
+                names_in_workspace.add(name)
+
+    # Remove duplicate names and version numbers for model versions
+    model_version_table = sa.Table("model_version", meta)
+    result = connection.execute(
+        sa.select(
+            model_version_table.c.id,
+            model_version_table.c.name,
+            model_version_table.c.number,
+            model_version_table.c.model_id,
+        )
+    ).all()
+
+    existing_names: Dict[str, Set[str]] = defaultdict(set)
+    existing_numbers: Dict[str, Set[int]] = defaultdict(set)
+
+    needs_update = []
+
+    for id_, name, number, model_id in result:
+        names_for_model = existing_names[model_id]
+        numbers_for_model = existing_numbers[model_id]
+
+        needs_new_name = name in names_for_model
+        needs_new_number = number in numbers_for_model
+
+        if needs_new_name or needs_new_number:
+            needs_update.append(
+                (id_, name, number, model_id, needs_new_name, needs_new_number)
+            )
+
+        names_for_model.add(name)
+        numbers_for_model.add(number)
+
+    for (
+        id_,
+        name,
+        number,
+        model_id,
+        needs_new_name,
+        needs_new_number,
+    ) in needs_update:
+        values: Dict[str, Any] = {}
+
+        is_numeric_version = str(number) == name
+        next_numeric_version = max(existing_numbers[model_id]) + 1
+
+        if is_numeric_version:
+            # No matter if the name or number clashes, we need to update both
+            values["number"] = next_numeric_version
+            values["name"] = str(next_numeric_version)
+            existing_numbers[model_id].add(next_numeric_version)
+            logger.warning(
+                "Migrating model version %s to %s to resolve duplicate name.",
+                name,
+                values["name"],
+            )
+        else:
+            if needs_new_name:
+                values["name"] = f"{name}_{id_[:6]}"
+                logger.warning(
+                    "Migrating model version %s to %s to resolve duplicate name.",
+                    name,
+                    values["name"],
+                )
+
+            if needs_new_number:
+                values["number"] = next_numeric_version
+                existing_numbers[model_id].add(next_numeric_version)
+
+        connection.execute(
+            sa.update(model_version_table)
+            .where(model_version_table.c.id == id_)
+            .values(**values)
+        )
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    resolve_duplicate_entities()
+
+    with op.batch_alter_table("pipeline", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_pipeline_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    with op.batch_alter_table("pipeline_run", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_run_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    with op.batch_alter_table("model", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_model_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    with op.batch_alter_table("model_version", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_version_for_model_id", ["name", "model_id"]
+        )
+        batch_op.create_unique_constraint(
+            "unique_version_number_for_model_id", ["number", "model_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("model_version", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_version_number_for_model_id", type_="unique"
+        )
+        batch_op.drop_constraint("unique_version_for_model_id", type_="unique")
+
+    with op.batch_alter_table("model", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_model_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("pipeline_run", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_run_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("pipeline", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_pipeline_name_in_workspace", type_="unique"
+        )
+
+    # ### end Alembic commands ###

zenml/zen_stores/migrations/versions/b557b2871693_update_step_run_input_types.py

@@ -0,0 +1,33 @@
+"""Update step run input types [b557b2871693].
+
+Revision ID: b557b2871693
+Revises: 1cb6477f72d6
+Create Date: 2024-10-30 13:06:55.147202
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b557b2871693"
+down_revision = "1cb6477f72d6"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    op.execute("""
+        UPDATE step_run_input_artifact
+        SET type = 'step_output'
+        WHERE type = 'default'
+    """)
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    op.execute("""
+        UPDATE step_run_input_artifact
+        SET type = 'default'
+        WHERE type = 'step_output'
+    """)

zenml/zen_stores/rest_zen_store.py

@@ -57,6 +57,7 @@ from zenml.constants import (
     ARTIFACT_VERSIONS,
     ARTIFACT_VISUALIZATIONS,
     ARTIFACTS,
+    BATCH,
     CODE_REFERENCES,
     CODE_REPOSITORIES,
     CONFIG,
@@ -109,6 +110,7 @@ from zenml.constants import (
     WORKSPACES,
 )
 from zenml.enums import (
+    APITokenType,
     OAuthGrantTypes,
     StackDeploymentProvider,
     StoreType,
@@ -200,9 +202,7 @@ from zenml.models import (
     PipelineRunResponse,
     PipelineRunUpdate,
     PipelineUpdate,
-    RunMetadataFilter,
     RunMetadataRequest,
-    RunMetadataResponse,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -991,6 +991,23 @@ class RestZenStore(BaseZenStore):
             route=ARTIFACT_VERSIONS,
         )
 
+    def batch_create_artifact_versions(
+        self, artifact_versions: List[ArtifactVersionRequest]
+    ) -> List[ArtifactVersionResponse]:
+        """Creates a batch of artifact versions.
+
+        Args:
+            artifact_versions: The artifact versions to create.
+
+        Returns:
+            The created artifact versions.
+        """
+        return self._batch_create_resources(
+            resources=artifact_versions,
+            response_model=ArtifactVersionResponse,
+            route=ARTIFACT_VERSIONS,
+        )
+
     def get_artifact_version(
         self, artifact_version_id: UUID, hydrate: bool = True
     ) -> ArtifactVersionResponse:
@@ -1996,9 +2013,7 @@ class RestZenStore(BaseZenStore):
 
     # ----------------------------- Run Metadata -----------------------------
 
-    def create_run_metadata(
-        self, run_metadata: RunMetadataRequest
-    ) -> List[RunMetadataResponse]:
+    def create_run_metadata(self, run_metadata: RunMetadataRequest) -> None:
         """Creates run metadata.
 
         Args:
@@ -2008,55 +2023,8 @@ class RestZenStore(BaseZenStore):
             The created run metadata.
         """
         route = f"{WORKSPACES}/{str(run_metadata.workspace)}{RUN_METADATA}"
-        response_body = self.post(f"{route}", body=run_metadata)
-        result: List[RunMetadataResponse] = []
-        if isinstance(response_body, list):
-            for metadata in response_body or []:
-                result.append(RunMetadataResponse.model_validate(metadata))
-        return result
-
-    def get_run_metadata(
-        self, run_metadata_id: UUID, hydrate: bool = True
-    ) -> RunMetadataResponse:
-        """Gets run metadata with the given ID.
-
-        Args:
-            run_metadata_id: The ID of the run metadata to get.
-            hydrate: Flag deciding whether to hydrate the output model(s)
-                by including metadata fields in the response.
-
-        Returns:
-            The run metadata.
-        """
-        return self._get_resource(
-            resource_id=run_metadata_id,
-            route=RUN_METADATA,
-            response_model=RunMetadataResponse,
-            params={"hydrate": hydrate},
-        )
-
-    def list_run_metadata(
-        self,
-        run_metadata_filter_model: RunMetadataFilter,
-        hydrate: bool = False,
-    ) -> Page[RunMetadataResponse]:
-        """List run metadata.
-
-        Args:
-            run_metadata_filter_model: All filter parameters including
-                pagination params.
-            hydrate: Flag deciding whether to hydrate the output model(s)
-                by including metadata fields in the response.
-
-        Returns:
-            The run metadata.
-        """
-        return self._list_paginated_resources(
-            route=RUN_METADATA,
-            response_model=RunMetadataResponse,
-            filter_model=run_metadata_filter_model,
-            params={"hydrate": hydrate},
-        )
+        self.post(f"{route}", body=run_metadata)
+        return None
 
     # ----------------------------- Schedules -----------------------------
 
@@ -3905,17 +3873,16 @@ class RestZenStore(BaseZenStore):
 
     def get_api_token(
         self,
-        pipeline_id: Optional[UUID] = None,
         schedule_id: Optional[UUID] = None,
-        expires_minutes: Optional[int] = None,
+        pipeline_run_id: Optional[UUID] = None,
+        step_run_id: Optional[UUID] = None,
     ) -> str:
         """Get an API token for a workload.
 
         Args:
-            pipeline_id: The ID of the pipeline to get a token for.
             schedule_id: The ID of the schedule to get a token for.
-            expires_minutes: The number of minutes for which the token should
-                be valid. If not provided, the token will be valid indefinitely.
+            pipeline_run_id: The ID of the pipeline run to get a token for.
+            step_run_id: The ID of the step run to get a token for.
 
         Returns:
             The API token.
@@ -3923,13 +3890,16 @@ class RestZenStore(BaseZenStore):
         Raises:
             ValueError: if the server response is not valid.
         """
-        params: Dict[str, Any] = {}
-        if pipeline_id:
-            params["pipeline_id"] = pipeline_id
+        params: Dict[str, Any] = {
+            # Python clients may only request workload tokens.
+            "token_type": APITokenType.WORKLOAD.value,
+        }
         if schedule_id:
             params["schedule_id"] = schedule_id
-        if expires_minutes:
-            params["expires_minutes"] = expires_minutes
+        if pipeline_run_id:
+            params["pipeline_run_id"] = pipeline_run_id
+        if step_run_id:
+            params["step_run_id"] = step_run_id
         response_body = self.get(API_TOKEN, params=params)
         if not isinstance(response_body, str):
             raise ValueError(
@@ -4518,6 +4488,40 @@ class RestZenStore(BaseZenStore):
 
         return response_model.model_validate(response_body)
 
+    def _batch_create_resources(
+        self,
+        resources: List[AnyRequest],
+        response_model: Type[AnyResponse],
+        route: str,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> List[AnyResponse]:
+        """Create a new batch of resources.
+
+        Args:
+            resources: The resources to create.
+            response_model: The response model of an individual resource.
+            route: The resource REST route to use.
+            params: Optional query parameters to pass to the endpoint.
+
+        Returns:
+            List of response models.
+        """
+        json_data = [
+            resource.model_dump(mode="json") for resource in resources
+        ]
+        response = self._request(
+            "POST",
+            self.url + API + VERSION_1 + route + BATCH,
+            json=json_data,
+            params=params,
+        )
+        assert isinstance(response, list)
+
+        return [
+            response_model.model_validate(model_data)
+            for model_data in response
+        ]
+
     def _create_workspace_scoped_resource(
         self,
         resource: AnyWorkspaceScopedRequest,

zenml/zen_stores/schemas/artifact_schemas.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of artifact table."""
 
+import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
@@ -23,6 +24,7 @@ from sqlmodel import Field, Relationship
 
 from zenml.config.source import Source
 from zenml.enums import (
+    ArtifactSaveType,
     ArtifactType,
     ExecutionStatus,
     MetadataResourceTypes,
@@ -196,6 +198,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
             overlaps="tags",
         ),
     )
+    save_type: str = Field(sa_column=Column(TEXT, nullable=False))
 
     # Foreign keys
     artifact_id: UUID = build_foreign_key_field(
@@ -300,6 +303,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
             uri=artifact_version_request.uri,
             materializer=artifact_version_request.materializer.model_dump_json(),
             data_type=artifact_version_request.data_type.model_dump_json(),
+            save_type=artifact_version_request.save_type.value,
         )
 
     def to_model(
@@ -360,6 +364,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
             updated=self.updated,
             tags=[t.tag.to_model() for t in self.tags],
             producer_pipeline_run_id=producer_pipeline_run_id,
+            save_type=ArtifactSaveType(self.save_type),
             artifact_store_id=self.artifact_store_id,
         )
 
@@ -370,7 +375,9 @@ class ArtifactVersionSchema(BaseSchema, table=True):
                 workspace=self.workspace.to_model(),
                 producer_step_run_id=producer_step_run_id,
                 visualizations=[v.to_model() for v in self.visualizations],
-                run_metadata={m.key: m.to_model() for m in self.run_metadata},
+                run_metadata={
+                    m.key: json.loads(m.value) for m in self.run_metadata
+                },
             )
 
         resources = None

zenml/zen_stores/schemas/model_schemas.py

@@ -13,12 +13,13 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of model tables."""
 
+import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 from uuid import UUID
 
 from pydantic import ConfigDict
-from sqlalchemy import BOOLEAN, INTEGER, TEXT, Column
+from sqlalchemy import BOOLEAN, INTEGER, TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.enums import MetadataResourceTypes, TaggableResourceTypes
@@ -61,6 +62,13 @@ class ModelSchema(NamedSchema, table=True):
     """SQL Model for model."""
 
     __tablename__ = "model"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_model_name_in_workspace",
+        ),
+    )
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,
@@ -219,6 +227,23 @@ class ModelVersionSchema(NamedSchema, table=True):
     """SQL Model for model version."""
 
     __tablename__ = MODEL_VERSION_TABLENAME
+    __table_args__ = (
+        # We need two 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
+        UniqueConstraint(
+            "number",
+            "model_id",
+            name="unique_version_number_for_model_id",
+        ),
+        UniqueConstraint(
+            "name",
+            "model_id",
+            name="unique_version_for_model_id",
+        ),
+    )
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,
@@ -379,8 +404,7 @@ class ModelVersionSchema(NamedSchema, table=True):
                 workspace=self.workspace.to_model(),
                 description=self.description,
                 run_metadata={
-                    rm.key: rm.to_model(include_metadata=True)
-                    for rm in self.run_metadata
+                    rm.key: json.loads(rm.value) for rm in self.run_metadata
                 },
             )
 

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -72,6 +72,11 @@ class PipelineRunSchema(NamedSchema, table=True):
             "orchestrator_run_id",
             name="unique_orchestrator_run_id_for_deployment_id",
         ),
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_run_name_in_workspace",
+        ),
     )
 
     # Fields
@@ -271,7 +276,7 @@ class PipelineRunSchema(NamedSchema, table=True):
         )
 
         run_metadata = {
-            metadata_schema.key: metadata_schema.to_model()
+            metadata_schema.key: json.loads(metadata_schema.value)
             for metadata_schema in self.run_metadata
         }
 

zenml/zen_stores/schemas/pipeline_schemas.py

@@ -17,7 +17,7 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
-from sqlalchemy import TEXT, Column
+from sqlalchemy import TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.enums import TaggableResourceTypes
@@ -50,7 +50,13 @@ class PipelineSchema(NamedSchema, table=True):
     """SQL Model for pipelines."""
 
     __tablename__ = "pipeline"
-
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_pipeline_name_in_workspace",
+        ),
+    )
     # Fields
     description: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))