zenml-nightly 0.70.0.dev20241122__py3-none-any.whl → 0.70.0.dev20241201__py3-none-any.whl

zenml/zen_server/routers/workspaces_endpoints.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for workspaces."""
 
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
@@ -98,13 +98,13 @@ from zenml.zen_server.feature_gate.endpoint_utils import (
 )
 from zenml.zen_server.rbac.endpoint_utils import (
     verify_permissions_and_create_entity,
-    verify_permissions_and_delete_entity,
-    verify_permissions_and_get_entity,
     verify_permissions_and_list_entities,
-    verify_permissions_and_update_entity,
 )
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import (
+    batch_verify_permissions_for_models,
+    dehydrate_page,
+    dehydrate_response_model,
     get_allowed_resource_ids,
     verify_permission,
     verify_permission_for_model,
@@ -146,12 +146,10 @@ def list_workspaces(
     Returns:
         A list of workspaces.
     """
-    return verify_permissions_and_list_entities(
-        filter_model=workspace_filter_model,
-        resource_type=ResourceType.WORKSPACE,
-        list_method=zen_store().list_workspaces,
-        hydrate=hydrate,
+    workspaces = zen_store().list_workspaces(
+        workspace_filter_model, hydrate=hydrate
     )
+    return dehydrate_page(workspaces)
 
 
 @router.post(
@@ -160,7 +158,7 @@ def list_workspaces(
 )
 @handle_exceptions
 def create_workspace(
-    workspace: WorkspaceRequest,
+    workspace_request: WorkspaceRequest,
     _: AuthContext = Security(authorize),
 ) -> WorkspaceResponse:
     """Creates a workspace based on the requestBody.
@@ -168,16 +166,13 @@ def create_workspace(
     # noqa: DAR401
 
     Args:
-        workspace: Workspace to create.
+        workspace_request: Workspace to create.
 
     Returns:
         The created workspace.
     """
-    return verify_permissions_and_create_entity(
-        request_model=workspace,
-        resource_type=ResourceType.WORKSPACE,
-        create_method=zen_store().create_workspace,
-    )
+    workspace = zen_store().create_workspace(workspace_request)
+    return dehydrate_response_model(workspace)
 
 
 @router.get(
@@ -203,11 +198,10 @@ def get_workspace(
     Returns:
         The requested workspace.
     """
-    return verify_permissions_and_get_entity(
-        id=workspace_name_or_id,
-        get_method=zen_store().get_workspace,
-        hydrate=hydrate,
+    workspace = zen_store().get_workspace(
+        workspace_name_or_id, hydrate=hydrate
     )
+    return dehydrate_response_model(workspace)
 
 
 @router.put(
@@ -231,12 +225,11 @@ def update_workspace(
     Returns:
         The updated workspace.
     """
-    return verify_permissions_and_update_entity(
-        id=workspace_name_or_id,
-        update_model=workspace_update,
-        get_method=zen_store().get_workspace,
-        update_method=zen_store().update_workspace,
+    workspace = zen_store().get_workspace(workspace_name_or_id, hydrate=False)
+    updated_workspace = zen_store().update_workspace(
+        workspace_id=workspace.id, workspace_update=workspace_update
     )
+    return dehydrate_response_model(updated_workspace)
 
 
 @router.delete(
@@ -253,11 +246,7 @@ def delete_workspace(
     Args:
         workspace_name_or_id: Name or ID of the workspace.
     """
-    verify_permissions_and_delete_entity(
-        id=workspace_name_or_id,
-        get_method=zen_store().get_workspace,
-        delete_method=zen_store().delete_workspace,
-    )
+    zen_store().delete_workspace(workspace_name_or_id)
 
 
 @router.get(
@@ -951,20 +940,21 @@ def get_or_create_pipeline_run(
             "is not supported."
         )
 
-    verify_permission(
-        resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
-    )
+    def _pre_creation_hook() -> None:
+        verify_permission(
+            resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
+        )
+        check_entitlement(resource_type=ResourceType.PIPELINE_RUN)
 
     run, created = zen_store().get_or_create_run(
-        pipeline_run=pipeline_run,
-        pre_creation_hook=lambda: check_entitlement(
-            resource_type=ResourceType.PIPELINE_RUN
-        ),
+        pipeline_run=pipeline_run, pre_creation_hook=_pre_creation_hook
     )
     if created:
         report_usage(
             resource_type=ResourceType.PIPELINE_RUN, resource_id=run.id
         )
+    else:
+        verify_permission_for_model(run, action=Action.READ)
 
     return run, created
 
@@ -1009,24 +999,23 @@ def create_run_metadata(
             "is not supported."
         )
 
-    if run_metadata.resource_type == MetadataResourceTypes.PIPELINE_RUN:
-        run = zen_store().get_run(run_metadata.resource_id)
-        verify_permission_for_model(run, action=Action.UPDATE)
-    elif run_metadata.resource_type == MetadataResourceTypes.STEP_RUN:
-        step = zen_store().get_run_step(run_metadata.resource_id)
-        verify_permission_for_model(step, action=Action.UPDATE)
-    elif run_metadata.resource_type == MetadataResourceTypes.ARTIFACT_VERSION:
-        artifact_version = zen_store().get_artifact_version(
-            run_metadata.resource_id
-        )
-        verify_permission_for_model(artifact_version, action=Action.UPDATE)
-    elif run_metadata.resource_type == MetadataResourceTypes.MODEL_VERSION:
-        model_version = zen_store().get_model_version(run_metadata.resource_id)
-        verify_permission_for_model(model_version, action=Action.UPDATE)
-    else:
-        raise RuntimeError(
-            f"Unknown resource type: {run_metadata.resource_type}"
-        )
+    verify_models: List[Any] = []
+    for resource in run_metadata.resources:
+        if resource.type == MetadataResourceTypes.PIPELINE_RUN:
+            verify_models.append(zen_store().get_run(resource.id))
+        elif resource.type == MetadataResourceTypes.STEP_RUN:
+            verify_models.append(zen_store().get_run_step(resource.id))
+        elif resource.type == MetadataResourceTypes.ARTIFACT_VERSION:
+            verify_models.append(zen_store().get_artifact_version(resource.id))
+        elif resource.type == MetadataResourceTypes.MODEL_VERSION:
+            verify_models.append(zen_store().get_model_version(resource.id))
+        else:
+            raise RuntimeError(f"Unknown resource type: {resource.type}")
+
+    batch_verify_permissions_for_models(
+        models=verify_models,
+        action=Action.UPDATE,
+    )
 
     verify_permission(
         resource_type=ResourceType.RUN_METADATA, action=Action.CREATE

zenml/zen_server/template_execution/utils.py

@@ -363,6 +363,7 @@ def deployment_request_from_template(
                     "external_input_artifacts",
                     "model_artifacts_or_metadata",
                     "client_lazy_loaders",
+                    "substitutions",
                     "outputs",
                 }
             ),

zenml/zen_server/zen_server_api.py

@@ -180,7 +180,15 @@ class RequestBodyLimit(BaseHTTPMiddleware):
         if content_length := request.headers.get("content-length"):
             if int(content_length) > self.max_bytes:
                 return Response(status_code=413)  # Request Entity Too Large
-        return await call_next(request)
+
+        try:
+            return await call_next(request)
+        except Exception:
+            logger.exception("An error occurred while processing the request")
+            return JSONResponse(
+                status_code=500,
+                content={"detail": "An unexpected error occurred."},
+            )
 
 
 class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
@@ -220,7 +228,15 @@ class RestrictFileUploadsMiddleware(BaseHTTPMiddleware):
                         "detail": "File uploads are not allowed on this endpoint."
                     },
                 )
-        return await call_next(request)
+
+        try:
+            return await call_next(request)
+        except Exception:
+            logger.exception("An error occurred while processing the request")
+            return JSONResponse(
+                status_code=500,
+                content={"detail": "An unexpected error occurred."},
+            )
 
 
 ALLOWED_FOR_FILE_UPLOAD: Set[str] = set()
@@ -252,13 +268,21 @@ async def set_secure_headers(request: Request, call_next: Any) -> Any:
     Returns:
         The response with secure headers set.
     """
+    try:
+        response = await call_next(request)
+    except Exception:
+        logger.exception("An error occurred while processing the request")
+        response = JSONResponse(
+            status_code=500,
+            content={"detail": "An unexpected error occurred."},
+        )
+
     # If the request is for the openAPI docs, don't set secure headers
     if request.url.path.startswith("/docs") or request.url.path.startswith(
         "/redoc"
     ):
-        return await call_next(request)
+        return response
 
-    response = await call_next(request)
     secure_headers().framework.fastapi(response)
     return response
 
@@ -298,7 +322,15 @@ async def track_last_user_activity(request: Request, call_next: Any) -> Any:
         zen_store()._update_last_user_activity_timestamp(
             last_user_activity=last_user_activity
         )
-    return await call_next(request)
+
+    try:
+        return await call_next(request)
+    except Exception:
+        logger.exception("An error occurred while processing the request")
+        return JSONResponse(
+            status_code=500,
+            content={"detail": "An unexpected error occurred."},
+        )
 
 
 @app.middleware("http")
@@ -330,7 +362,14 @@ async def infer_source_context(request: Request, call_next: Any) -> Any:
         )
         source_context.set(SourceContextTypes.API)
 
-    return await call_next(request)
+    try:
+        return await call_next(request)
+    except Exception:
+        logger.exception("An error occurred while processing the request")
+        return JSONResponse(
+            status_code=500,
+            content={"detail": "An unexpected error occurred."},
+        )
 
 
 @app.on_event("startup")

zenml/zen_stores/migrations/utils.py

@@ -273,30 +273,25 @@ class MigrationUtils(BaseModel):
                         + "\n);"
                     )
 
+                # Detect self-referential foreign keys from the table schema
+                has_self_referential_foreign_keys = False
+                for fk in table.foreign_keys:
+                    # Check if the foreign key points to the same table
+                    if fk.column.table == table:
+                        has_self_referential_foreign_keys = True
+                        break
+
                 # Store the table schema
                 store_db_info(
-                    dict(table=table.name, create_stmt=create_table_stmt)
+                    dict(
+                        table=table.name,
+                        create_stmt=create_table_stmt,
+                        self_references=has_self_referential_foreign_keys,
+                    )
                 )
 
                 # 2. extract the table data in batches
-
-                # If the table has a `created` column, we use it to sort
-                # the rows in the table starting with the oldest rows.
-                # This is to ensure that the rows are inserted in the
-                # correct order, since some tables have inner foreign key
-                # constraints.
-                if "created" in table.columns:
-                    order_by = [table.columns["created"]]
-                else:
-                    order_by = []
-                if "id" in table.columns:
-                    # If the table has an `id` column, we also use it to sort
-                    # the rows in the table, even if we already use "created"
-                    # to sort the rows. We need a unique field to sort the rows,
-                    # to break the tie between rows with the same "created"
-                    # date, otherwise the same entry might end up multiple times
-                    # in subsequent pages.
-                    order_by.append(table.columns["id"])
+                order_by = [col for col in table.primary_key]
 
                 # Fetch the number of rows in the table
                 row_count = conn.scalar(
@@ -305,7 +300,7 @@ class MigrationUtils(BaseModel):
 
                 # Fetch the data from the table in batches
                 if row_count is not None:
-                    batch_size = 50
+                    batch_size = 100
                     for i in range(0, row_count, batch_size):
                         rows = conn.execute(
                             table.select()
@@ -349,6 +344,7 @@ class MigrationUtils(BaseModel):
 
         with self.engine.begin() as connection:
             # read the DB information one JSON object at a time
+            self_references: Dict[str, bool] = {}
             for table_dump in load_db_info():
                 table_name = table_dump["table"]
                 if "create_stmt" in table_dump:
@@ -356,10 +352,22 @@ class MigrationUtils(BaseModel):
                     connection.execute(text(table_dump["create_stmt"]))
                     # Reload the database metadata after creating the table
                     metadata.reflect(bind=self.engine)
+                    self_references[table_name] = table_dump.get(
+                        "self_references", False
+                    )
 
                 if "data" in table_dump:
                     # insert the data into the database
                     table = metadata.tables[table_name]
+                    if self_references.get(table_name, False):
+                        # If the table has self-referential foreign keys, we
+                        # need to disable the foreign key checks before inserting
+                        # the rows and re-enable them afterwards. This is because
+                        # the rows need to be inserted in the correct order to
+                        # satisfy the foreign key constraints and we don't sort
+                        # the rows by creation time in the backup.
+                        connection.execute(text("SET FOREIGN_KEY_CHECKS = 0"))
+
                     for row in table_dump["data"]:
                         # Convert column values to the correct type
                         for column in table.columns:
@@ -372,10 +380,18 @@ class MigrationUtils(BaseModel):
                                     row[column.name], "utf-8"
                                 )
 
-                    # Insert the rows into the table
-                    connection.execute(
-                        table.insert().values(table_dump["data"])
-                    )
+                    # Insert the rows into the table in batches
+                    batch_size = 100
+                    for i in range(0, len(table_dump["data"]), batch_size):
+                        connection.execute(
+                            table.insert().values(
+                                table_dump["data"][i : i + batch_size]
+                            )
+                        )
+
+                    if table_dump.get("self_references", False):
+                        # Re-enable the foreign key checks after inserting the rows
+                        connection.execute(text("SET FOREIGN_KEY_CHECKS = 1"))
 
     def backup_database_to_file(self, dump_file: str) -> None:
         """Backup the database to a file.

zenml/zen_stores/migrations/versions/b73bc71f1106_remove_component_spec_path.py

@@ -0,0 +1,36 @@
+"""Remove component spec path [b73bc71f1106].
+
+Revision ID: b73bc71f1106
+Revises: ec6307720f92
+Create Date: 2024-11-29 09:36:33.089945
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b73bc71f1106"
+down_revision = "ec6307720f92"
+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("stack_component", schema=None) as batch_op:
+        batch_op.drop_column("component_spec_path")
+
+    # ### 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("stack_component", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("component_spec_path", sa.VARCHAR(), nullable=True)
+        )
+
+    # ### end Alembic commands ###

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

@@ -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/__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",