zenml-nightly 0.70.0.dev20241125__py3-none-any.whl → 0.71.0.dev20241220__py3-none-any.whl

zenml/zen_stores/migrations/versions/cc269488e5a9_separate_run_metadata.py

@@ -0,0 +1,135 @@
+"""Separate run metadata into resource link table with new UUIDs.
+
+Revision ID: cc269488e5a9
+Revises: b73bc71f1106
+Create Date: 2024-11-12 09:46:46.587478
+"""
+
+import uuid
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "cc269488e5a9"
+down_revision = "b73bc71f1106"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Creates the 'run_metadata_resource' table and migrates data."""
+    op.create_table(
+        "run_metadata_resource",
+        sa.Column(
+            "id",
+            sqlmodel.sql.sqltypes.GUID(),
+            nullable=False,
+            primary_key=True,
+        ),
+        sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column("resource_type", sa.String(length=255), nullable=False),
+        sa.Column(
+            "run_metadata_id",
+            sqlmodel.sql.sqltypes.GUID(),
+            sa.ForeignKey("run_metadata.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+    )
+
+    connection = op.get_bind()
+
+    run_metadata_data = connection.execute(
+        sa.text("""
+            SELECT id, resource_id, resource_type
+            FROM run_metadata
+        """)
+    ).fetchall()
+
+    # Prepare data with new UUIDs for bulk insert
+    resource_data = [
+        {
+            "id": str(uuid.uuid4()),  # Generate a new UUID for each row
+            "resource_id": row.resource_id,
+            "resource_type": row.resource_type,
+            "run_metadata_id": row.id,
+        }
+        for row in run_metadata_data
+    ]
+
+    # Perform bulk insert into `run_metadata_resource`
+    if resource_data:  # Only perform insert if there's data to migrate
+        op.bulk_insert(
+            sa.table(
+                "run_metadata_resource",
+                sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+                sa.Column(
+                    "resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False
+                ),
+                sa.Column(
+                    "resource_type", sa.String(length=255), nullable=False
+                ),
+                sa.Column(
+                    "run_metadata_id",
+                    sqlmodel.sql.sqltypes.GUID(),
+                    nullable=False,
+                ),  # Changed to BIGINT
+            ),
+            resource_data,
+        )
+
+    op.drop_column("run_metadata", "resource_id")
+    op.drop_column("run_metadata", "resource_type")
+
+    op.add_column(
+        "run_metadata",
+        sa.Column(
+            "publisher_step_id", sqlmodel.sql.sqltypes.GUID(), nullable=True
+        ),
+    )
+
+
+def downgrade() -> None:
+    """Reverts the 'run_metadata_resource' table and migrates data back."""
+    # Recreate the `resource_id` and `resource_type` columns in `run_metadata`
+    op.add_column(
+        "run_metadata",
+        sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=True),
+    )
+    op.add_column(
+        "run_metadata",
+        sa.Column("resource_type", sa.String(length=255), nullable=True),
+    )
+
+    # Migrate data back from `run_metadata_resource` to `run_metadata`
+    connection = op.get_bind()
+
+    # Fetch data from `run_metadata_resource`
+    run_metadata_resource_data = connection.execute(
+        sa.text("""
+            SELECT resource_id, resource_type, run_metadata_id
+            FROM run_metadata_resource
+        """)
+    ).fetchall()
+
+    # Update `run_metadata` with the data from `run_metadata_resource`
+    for row in run_metadata_resource_data:
+        connection.execute(
+            sa.text("""
+                UPDATE run_metadata
+                SET resource_id = :resource_id, resource_type = :resource_type
+                WHERE id = :run_metadata_id
+            """),
+            {
+                "resource_id": row.resource_id,
+                "resource_type": row.resource_type,
+                "run_metadata_id": row.run_metadata_id,
+            },
+        )
+
+    # Drop the `run_metadata_resource` table
+    op.drop_table("run_metadata_resource")
+
+    # Drop the cached column
+    op.drop_column("run_metadata", "publisher_step_id")

zenml/zen_stores/migrations/versions/ec6307720f92_simplify_model_version_links.py

@@ -59,12 +59,13 @@ def _migrate_artifact_type() -> None:
                 {"id_": artifact_version_id, "type": "ServiceArtifact"}
             )
 
-    connection.execute(
-        sa.update(artifact_version_table).where(
-            artifact_version_table.c.id == sa.bindparam("id_")
-        ),
-        updates,
-    )
+    if updates:
+        connection.execute(
+            sa.update(artifact_version_table).where(
+                artifact_version_table.c.id == sa.bindparam("id_")
+            ),
+            updates,
+        )
 
 
 def upgrade() -> None:

zenml/zen_stores/rest_zen_store.py

@@ -3873,13 +3873,17 @@ class RestZenStore(BaseZenStore):
 
     def get_api_token(
         self,
+        token_type: APITokenType = APITokenType.WORKLOAD,
+        expires_in: Optional[int] = None,
         schedule_id: Optional[UUID] = None,
         pipeline_run_id: Optional[UUID] = None,
         step_run_id: Optional[UUID] = None,
     ) -> str:
-        """Get an API token for a workload.
+        """Get an API token.
 
         Args:
+            token_type: The type of the token to get.
+            expires_in: The time in seconds until the token expires.
             schedule_id: The ID of the schedule to get a token for.
             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.
@@ -3891,9 +3895,10 @@ class RestZenStore(BaseZenStore):
             ValueError: if the server response is not valid.
         """
         params: Dict[str, Any] = {
-            # Python clients may only request workload tokens.
-            "token_type": APITokenType.WORKLOAD.value,
+            "token_type": token_type.value,
         }
+        if expires_in:
+            params["expires_in"] = expires_in
         if schedule_id:
             params["schedule_id"] = schedule_id
         if pipeline_run_id:
@@ -4062,7 +4067,7 @@ class RestZenStore(BaseZenStore):
                 "you should use a service account API key to authenticate to "
                 "the server instead of temporary CLI login credentials. For "
                 "more information, see "
-                "https://docs.zenml.io/how-to/connecting-to-zenml/connect-with-a-service-account"
+                "https://docs.zenml.io/how-to/project-setup-and-management/connecting-to-zenml/connect-with-a-service-account"
             )
 
             if api_key is not None:
@@ -4344,46 +4349,74 @@ class RestZenStore(BaseZenStore):
             {source_context.name: source_context.get().value}
         )
 
-        try:
-            return self._handle_response(
-                self.session.request(
-                    method,
-                    url,
-                    params=params,
-                    verify=self.config.verify_ssl,
-                    timeout=timeout or self.config.http_timeout,
-                    **kwargs,
-                )
-            )
-        except CredentialsNotValid:
-            # NOTE: CredentialsNotValid is raised only when the server
-            # explicitly indicates that the credentials are not valid and they
-            # can be thrown away.
-
-            # We authenticate or re-authenticate here and then try the request
-            # again, this time with a valid API token in the header.
-            self.authenticate(
-                # If the last request was authenticated with an API token,
-                # we force a re-authentication to get a fresh token.
-                force=self._api_token is not None
-            )
-
-        try:
-            return self._handle_response(
-                self.session.request(
-                    method,
-                    url,
-                    params=params,
-                    verify=self.config.verify_ssl,
-                    timeout=self.config.http_timeout,
-                    **kwargs,
+        # If the server replies with a credentials validation (401 Unauthorized)
+        # error, we (re-)authenticate and retry the request here in the
+        # following cases:
+        #
+        # 1. initial authentication: the last request was not authenticated
+        # with an API token.
+        # 2. re-authentication: the last request was authenticated with an API
+        # token that was rejected by the server. This is to cover the case
+        # of expired tokens that can be refreshed by the client automatically
+        # without user intervention from other sources (e.g. API keys).
+        #
+        # NOTE: it can happen that the same request is retried here for up to
+        # two times: once after initial authentication and once after
+        # re-authentication.
+        re_authenticated = False
+        while True:
+            try:
+                return self._handle_response(
+                    self.session.request(
+                        method,
+                        url,
+                        params=params,
+                        verify=self.config.verify_ssl,
+                        timeout=timeout or self.config.http_timeout,
+                        **kwargs,
+                    )
                 )
-            )
-        except CredentialsNotValid as e:
-            raise CredentialsNotValid(
-                "The current credentials are no longer valid. Please log in "
-                "again using 'zenml login'."
-            ) from e
+            except CredentialsNotValid as e:
+                # NOTE: CredentialsNotValid is raised only when the server
+                # explicitly indicates that the credentials are not valid and
+                # they can be thrown away or when the request is not
+                # authenticated at all.
+
+                if self._api_token is None:
+                    # The last request was not authenticated with an API
+                    # token at all. We authenticate here and then try the
+                    # request again, this time with a valid API token in the
+                    # header.
+                    logger.debug(
+                        f"The last request was not authenticated: {e}\n"
+                        "Re-authenticating and retrying..."
+                    )
+                    self.authenticate()
+                elif not re_authenticated:
+                    # The last request was authenticated with an API token
+                    # that was rejected by the server. We attempt a
+                    # re-authentication here and then retry the request.
+                    logger.debug(
+                        "The last request was authenticated with an API token "
+                        f"that was rejected by the server: {e}\n"
+                        "Re-authenticating and retrying..."
+                    )
+                    re_authenticated = True
+                    self.authenticate(
+                        # Ignore the current token and force a re-authentication
+                        force=True
+                    )
+                else:
+                    # The last request was made after re-authenticating but
+                    # still failed. Bailing out.
+                    logger.debug(
+                        f"The last request failed after re-authenticating: {e}\n"
+                        "Bailing out..."
+                    )
+                    raise CredentialsNotValid(
+                        "The current credentials are no longer valid. Please "
+                        "log in again using 'zenml login'."
+                    ) from e
 
     def get(
         self,

zenml/zen_stores/schemas/__init__.py

@@ -39,7 +39,10 @@ from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
 from zenml.zen_stores.schemas.pipeline_schemas import PipelineSchema
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
-from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
+from zenml.zen_stores.schemas.run_metadata_schemas import (
+    RunMetadataResourceSchema,
+    RunMetadataSchema,
+)
 from zenml.zen_stores.schemas.schedule_schema import ScheduleSchema
 from zenml.zen_stores.schemas.secret_schemas import SecretSchema
 from zenml.zen_stores.schemas.service_schemas import ServiceSchema
@@ -90,6 +93,7 @@ __all__ = [
     "PipelineDeploymentSchema",
     "PipelineRunSchema",
     "PipelineSchema",
+    "RunMetadataResourceSchema",
     "RunMetadataSchema",
     "ScheduleSchema",
     "SecretSchema",

zenml/zen_stores/schemas/artifact_schemas.py

@@ -13,7 +13,6 @@
 #  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
@@ -50,6 +49,7 @@ from zenml.zen_stores.schemas.step_run_schemas import (
     StepRunOutputArtifactSchema,
 )
 from zenml.zen_stores.schemas.user_schemas import UserSchema
+from zenml.zen_stores.schemas.utils import RunMetadataInterface
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
@@ -59,7 +59,9 @@ if TYPE_CHECKING:
     from zenml.zen_stores.schemas.model_schemas import (
         ModelVersionArtifactSchema,
     )
-    from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
+    from zenml.zen_stores.schemas.run_metadata_schemas import (
+        RunMetadataResourceSchema,
+    )
     from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
 
 
@@ -171,7 +173,7 @@ class ArtifactSchema(NamedSchema, table=True):
         return self
 
 
-class ArtifactVersionSchema(BaseSchema, table=True):
+class ArtifactVersionSchema(BaseSchema, RunMetadataInterface, table=True):
     """SQL Model for artifact versions."""
 
     __tablename__ = "artifact_version"
@@ -242,12 +244,12 @@ class ArtifactVersionSchema(BaseSchema, table=True):
     workspace: "WorkspaceSchema" = Relationship(
         back_populates="artifact_versions"
     )
-    run_metadata: List["RunMetadataSchema"] = Relationship(
-        back_populates="artifact_version",
+    run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
+        back_populates="artifact_versions",
         sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ArtifactVersionSchema.id)",
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ArtifactVersionSchema.id)",
             cascade="delete",
-            overlaps="run_metadata",
+            overlaps="run_metadata_resources",
         ),
     )
     output_of_step_runs: List["StepRunOutputArtifactSchema"] = Relationship(
@@ -352,8 +354,9 @@ class ArtifactVersionSchema(BaseSchema, table=True):
                 producer_step_run_id = step_run.original_step_run_id
 
         # Create the body of the model
+        artifact = self.artifact.to_model()
         body = ArtifactVersionResponseBody(
-            artifact=self.artifact.to_model(),
+            artifact=artifact,
             version=self.version or str(self.version_number),
             user=self.user.to_model() if self.user else None,
             uri=self.uri,
@@ -375,9 +378,7 @@ 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: json.loads(m.value) for m in self.run_metadata
-                },
+                run_metadata=self.fetch_metadata(),
             )
 
         resources = None

zenml/zen_stores/schemas/component_schemas.py

@@ -56,7 +56,6 @@ class StackComponentSchema(NamedSchema, table=True):
     flavor: str
     configuration: bytes
     labels: Optional[bytes]
-    component_spec_path: Optional[str]
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,
@@ -135,7 +134,6 @@ class StackComponentSchema(NamedSchema, table=True):
             name=request.name,
             workspace_id=request.workspace,
             user_id=request.user,
-            component_spec_path=request.component_spec_path,
             type=request.type,
             flavor=request.flavor,
             configuration=base64.b64encode(
@@ -218,7 +216,6 @@ class StackComponentSchema(NamedSchema, table=True):
                 labels=json.loads(base64.b64decode(self.labels).decode())
                 if self.labels
                 else None,
-                component_spec_path=self.component_spec_path,
                 connector_resource_id=self.connector_resource_id,
                 connector=self.connector.to_model()
                 if self.connector

zenml/zen_stores/schemas/model_schemas.py

@@ -13,13 +13,18 @@
 #  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 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 (
@@ -51,11 +56,16 @@ from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
-from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
+from zenml.zen_stores.schemas.run_metadata_schemas import (
+    RunMetadataResourceSchema,
+)
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
 from zenml.zen_stores.schemas.user_schemas import UserSchema
-from zenml.zen_stores.schemas.utils import get_page_from_list
+from zenml.zen_stores.schemas.utils import (
+    RunMetadataInterface,
+    get_page_from_list,
+)
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
@@ -219,16 +229,18 @@ class ModelSchema(NamedSchema, table=True):
         return self
 
 
-class ModelVersionSchema(NamedSchema, table=True):
+class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
     """SQL Model for model version."""
 
     __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",
@@ -239,6 +251,11 @@ class ModelVersionSchema(NamedSchema, 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(
@@ -299,21 +316,32 @@ class ModelVersionSchema(NamedSchema, table=True):
     description: str = Field(sa_column=Column(TEXT, nullable=True))
     stage: str = Field(sa_column=Column(TEXT, nullable=True))
 
-    run_metadata: List["RunMetadataSchema"] = Relationship(
-        back_populates="model_version",
+    run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
+        back_populates="model_versions",
         sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ModelVersionSchema.id)",
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ModelVersionSchema.id)",
             cascade="delete",
-            overlaps="run_metadata",
+            overlaps="run_metadata_resources",
         ),
     )
     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
@@ -324,24 +352,36 @@ class ModelVersionSchema(NamedSchema, 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(
@@ -402,9 +442,7 @@ class ModelVersionSchema(NamedSchema, table=True):
             metadata = ModelVersionResponseMetadata(
                 workspace=self.workspace.to_model(),
                 description=self.description,
-                run_metadata={
-                    rm.key: json.loads(rm.value) for rm in self.run_metadata
-                },
+                run_metadata=self.fetch_metadata(),
             )
 
         resources = None

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,