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

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(
@@ -4172,7 +4172,44 @@ class RestZenStore(BaseZenStore):
                 )
 
             self._session = requests.Session()
-            retries = Retry(backoff_factor=0.1, connect=5)
+            # Retries are triggered for idempotent HTTP methods (GET, HEAD, PUT,
+            # OPTIONS and DELETE) on specific HTTP status codes:
+            #
+            #     500: Internal Server Error.
+            #     502: Bad Gateway.
+            #     503: Service Unavailable.
+            #     504: Gateway Timeout.
+            #
+            # This also handles connection level errors, if a connection attempt
+            # fails due to transient issues like:
+            #
+            #     DNS resolution errors.
+            #     Connection timeouts.
+            #     Network disruptions.
+            #
+            # Additional errors retried:
+            #
+            #     Read Timeouts: If the server does not send a response within
+            #     the timeout period.
+            #     Connection Refused: If the server refuses the connection.
+            #
+            retries = Retry(
+                connect=5,
+                read=8,
+                redirect=3,
+                status=10,
+                allowed_methods=["HEAD", "GET", "PUT", "DELETE", "OPTIONS"],
+                status_forcelist=[
+                    408,  # Request Timeout
+                    429,  # Too Many Requests
+                    500,  # Internal Server Error
+                    502,  # Bad Gateway
+                    503,  # Service Unavailable
+                    504,  # Gateway Timeout
+                ],
+                other=3,
+                backoff_factor=0.5,
+            )
             self._session.mount("https://", HTTPAdapter(max_retries=retries))
             self._session.mount("http://", HTTPAdapter(max_retries=retries))
             self._session.verify = self.config.verify_ssl

zenml/zen_stores/schemas/model_schemas.py

@@ -22,7 +22,11 @@ from pydantic import ConfigDict
 from sqlalchemy import BOOLEAN, INTEGER, TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
-from zenml.enums import MetadataResourceTypes, TaggableResourceTypes
+from zenml.enums import (
+    ArtifactType,
+    MetadataResourceTypes,
+    TaggableResourceTypes,
+)
 from zenml.models import (
     BaseResponseMetadata,
     ModelRequest,
@@ -112,14 +116,6 @@ class ModelSchema(NamedSchema, table=True):
         back_populates="model",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    artifact_links: List["ModelVersionArtifactSchema"] = Relationship(
-        back_populates="model",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
-    pipeline_run_links: List["ModelVersionPipelineRunSchema"] = Relationship(
-        back_populates="model",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
 
     @classmethod
     def from_request(cls, model_request: ModelRequest) -> "ModelSchema":
@@ -377,11 +373,14 @@ class ModelVersionSchema(NamedSchema, table=True):
             artifact_name = artifact_link.artifact_version.artifact.name
             artifact_version = str(artifact_link.artifact_version.version)
             artifact_version_id = artifact_link.artifact_version.id
-            if artifact_link.is_model_artifact:
+            if artifact_link.artifact_version.type == ArtifactType.MODEL.value:
                 model_artifact_ids.setdefault(artifact_name, {}).update(
                     {str(artifact_version): artifact_version_id}
                 )
-            elif artifact_link.is_deployment_artifact:
+            elif (
+                artifact_link.artifact_version.type
+                == ArtifactType.SERVICE.value
+            ):
                 deployment_artifact_ids.setdefault(artifact_name, {}).update(
                     {str(artifact_version): artifact_version_id}
                 )
@@ -476,39 +475,6 @@ class ModelVersionArtifactSchema(BaseSchema, table=True):
 
     __tablename__ = "model_versions_artifacts"
 
-    workspace_id: UUID = build_foreign_key_field(
-        source=__tablename__,
-        target=WorkspaceSchema.__tablename__,
-        source_column="workspace_id",
-        target_column="id",
-        ondelete="CASCADE",
-        nullable=False,
-    )
-    workspace: "WorkspaceSchema" = Relationship(
-        back_populates="model_versions_artifacts_links"
-    )
-
-    user_id: Optional[UUID] = build_foreign_key_field(
-        source=__tablename__,
-        target=UserSchema.__tablename__,
-        source_column="user_id",
-        target_column="id",
-        ondelete="SET NULL",
-        nullable=True,
-    )
-    user: Optional["UserSchema"] = Relationship(
-        back_populates="model_versions_artifacts_links"
-    )
-
-    model_id: UUID = build_foreign_key_field(
-        source=__tablename__,
-        target=ModelSchema.__tablename__,
-        source_column="model_id",
-        target_column="id",
-        ondelete="CASCADE",
-        nullable=False,
-    )
-    model: "ModelSchema" = Relationship(back_populates="artifact_links")
     model_version_id: UUID = build_foreign_key_field(
         source=__tablename__,
         target=ModelVersionSchema.__tablename__,
@@ -532,11 +498,6 @@ class ModelVersionArtifactSchema(BaseSchema, table=True):
         back_populates="model_versions_artifacts_links"
     )
 
-    is_model_artifact: bool = Field(sa_column=Column(BOOLEAN, nullable=True))
-    is_deployment_artifact: bool = Field(
-        sa_column=Column(BOOLEAN, nullable=True)
-    )
-
     # 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
@@ -559,13 +520,8 @@ class ModelVersionArtifactSchema(BaseSchema, table=True):
             The converted schema.
         """
         return cls(
-            workspace_id=model_version_artifact_request.workspace,
-            user_id=model_version_artifact_request.user,
-            model_id=model_version_artifact_request.model,
             model_version_id=model_version_artifact_request.model_version,
             artifact_version_id=model_version_artifact_request.artifact_version,
-            is_model_artifact=model_version_artifact_request.is_model_artifact,
-            is_deployment_artifact=model_version_artifact_request.is_deployment_artifact,
         )
 
     def to_model(
@@ -590,11 +546,8 @@ class ModelVersionArtifactSchema(BaseSchema, table=True):
             body=ModelVersionArtifactResponseBody(
                 created=self.created,
                 updated=self.updated,
-                model=self.model_id,
                 model_version=self.model_version_id,
                 artifact_version=self.artifact_version.to_model(),
-                is_model_artifact=self.is_model_artifact,
-                is_deployment_artifact=self.is_deployment_artifact,
             ),
             metadata=BaseResponseMetadata() if include_metadata else None,
         )
@@ -605,39 +558,6 @@ class ModelVersionPipelineRunSchema(BaseSchema, table=True):
 
     __tablename__ = "model_versions_runs"
 
-    workspace_id: UUID = build_foreign_key_field(
-        source=__tablename__,
-        target=WorkspaceSchema.__tablename__,
-        source_column="workspace_id",
-        target_column="id",
-        ondelete="CASCADE",
-        nullable=False,
-    )
-    workspace: "WorkspaceSchema" = Relationship(
-        back_populates="model_versions_pipeline_runs_links"
-    )
-
-    user_id: Optional[UUID] = build_foreign_key_field(
-        source=__tablename__,
-        target=UserSchema.__tablename__,
-        source_column="user_id",
-        target_column="id",
-        ondelete="SET NULL",
-        nullable=True,
-    )
-    user: Optional["UserSchema"] = Relationship(
-        back_populates="model_versions_pipeline_runs_links"
-    )
-
-    model_id: UUID = build_foreign_key_field(
-        source=__tablename__,
-        target=ModelSchema.__tablename__,
-        source_column="model_id",
-        target_column="id",
-        ondelete="CASCADE",
-        nullable=False,
-    )
-    model: "ModelSchema" = Relationship(back_populates="pipeline_run_links")
     model_version_id: UUID = build_foreign_key_field(
         source=__tablename__,
         target=ModelVersionSchema.__tablename__,
@@ -683,9 +603,6 @@ class ModelVersionPipelineRunSchema(BaseSchema, table=True):
             The converted schema.
         """
         return cls(
-            workspace_id=model_version_pipeline_run_request.workspace,
-            user_id=model_version_pipeline_run_request.user,
-            model_id=model_version_pipeline_run_request.model,
             model_version_id=model_version_pipeline_run_request.model_version,
             pipeline_run_id=model_version_pipeline_run_request.pipeline_run,
         )
@@ -712,7 +629,6 @@ class ModelVersionPipelineRunSchema(BaseSchema, table=True):
             body=ModelVersionPipelineRunResponseBody(
                 created=self.created,
                 updated=self.updated,
-                model=self.model_id,
                 model_version=self.model_version_id,
                 pipeline_run=self.pipeline_run.to_model(),
             ),

zenml/zen_stores/schemas/user_schemas.py

@@ -44,8 +44,6 @@ if TYPE_CHECKING:
         EventSourceSchema,
         FlavorSchema,
         ModelSchema,
-        ModelVersionArtifactSchema,
-        ModelVersionPipelineRunSchema,
         ModelVersionSchema,
         OAuthDeviceSchema,
         PipelineBuildSchema,
@@ -144,12 +142,6 @@ class UserSchema(NamedSchema, table=True):
     model_versions: List["ModelVersionSchema"] = Relationship(
         back_populates="user",
     )
-    model_versions_artifacts_links: List["ModelVersionArtifactSchema"] = (
-        Relationship(back_populates="user")
-    )
-    model_versions_pipeline_runs_links: List[
-        "ModelVersionPipelineRunSchema"
-    ] = Relationship(back_populates="user")
     auth_devices: List["OAuthDeviceSchema"] = Relationship(
         back_populates="user",
         sa_relationship_kwargs={"cascade": "delete"},

zenml/zen_stores/schemas/workspace_schemas.py

@@ -35,8 +35,6 @@ if TYPE_CHECKING:
         EventSourceSchema,
         FlavorSchema,
         ModelSchema,
-        ModelVersionArtifactSchema,
-        ModelVersionPipelineRunSchema,
         ModelVersionSchema,
         PipelineBuildSchema,
         PipelineDeploymentSchema,
@@ -142,18 +140,6 @@ class WorkspaceSchema(NamedSchema, table=True):
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    model_versions_artifacts_links: List["ModelVersionArtifactSchema"] = (
-        Relationship(
-            back_populates="workspace",
-            sa_relationship_kwargs={"cascade": "delete"},
-        )
-    )
-    model_versions_pipeline_runs_links: List[
-        "ModelVersionPipelineRunSchema"
-    ] = Relationship(
-        back_populates="workspace",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
 
     @classmethod
     def from_request(cls, workspace: WorkspaceRequest) -> "WorkspaceSchema":

{zenml_nightly-0.70.0.dev20241120.dist-info → zenml_nightly-0.70.0.dev20241125.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.70.0.dev20241120
+Version: 0.70.0.dev20241125
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0