From 40f9e057abbb882a183707404e82fca64f2143d0 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 27 Sep 2025 01:42:53 +0100 Subject: [PATCH 01/64] init --- src/zenml/client.py | 112 +++++++ src/zenml/constants.py | 1 + src/zenml/models/__init__.py | 16 + src/zenml/models/v2/core/deployment.py | 16 + .../v2/core/deployment_visualization.py | 292 ++++++++++++++++++ .../routers/deployment_endpoints.py | 155 ++++++++++ src/zenml/zen_server/zen_server_api.py | 1 + src/zenml/zen_stores/rest_zen_store.py | 111 +++++++ src/zenml/zen_stores/schemas/__init__.py | 4 + .../zen_stores/schemas/deployment_schemas.py | 28 ++ .../deployment_visualization_schemas.py | 243 +++++++++++++++ src/zenml/zen_stores/sql_zen_store.py | 218 +++++++++++++ src/zenml/zen_stores/zen_store_interface.py | 95 ++++++ 13 files changed, 1292 insertions(+) create mode 100644 src/zenml/models/v2/core/deployment_visualization.py create mode 100644 src/zenml/zen_stores/schemas/deployment_visualization_schemas.py diff --git a/src/zenml/client.py b/src/zenml/client.py index 63e2a28d19..73f907161f 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -110,6 +110,10 @@ ComponentUpdate, DeploymentFilter, DeploymentResponse, + DeploymentVisualizationFilter, + DeploymentVisualizationRequest, + DeploymentVisualizationResponse, + DeploymentVisualizationUpdate, EventSourceFilter, EventSourceRequest, EventSourceResponse, @@ -3735,6 +3739,114 @@ def get_deployment( hydrate=hydrate, ) + def add_visualization_to_deployment( + self, + deployment_id: UUID, + artifact_version_id: UUID, + visualization_index: int, + *, + display_name: Optional[str] = None, + display_order: Optional[int] = None, + ) -> DeploymentVisualizationResponse: + """Curate a deployment visualization. + + Args: + deployment_id: The ID of the deployment to add visualization to. + artifact_version_id: The ID of the artifact version containing the visualization. + visualization_index: The index of the visualization within the artifact version. + display_name: Optional display name for the visualization. + display_order: Optional display order for sorting visualizations. + + Returns: + The created deployment visualization. + """ + deployment = self.get_deployment(deployment_id) + request = DeploymentVisualizationRequest( + project=deployment.project_id, + deployment_id=deployment_id, + artifact_version_id=artifact_version_id, + visualization_index=visualization_index, + display_name=display_name, + display_order=display_order, + ) + return self.zen_store.create_deployment_visualization(request) + + def list_deployment_visualizations( + self, + deployment_id: UUID, + *, + page: Optional[int] = None, + size: Optional[int] = None, + order_by: Optional[str] = None, + sort: Optional[SorterOps] = None, + hydrate: bool = False, + ) -> Page[DeploymentVisualizationResponse]: + """List curated deployment visualizations for a deployment. + + Args: + deployment_id: The ID of the deployment to list visualizations for. + page: Page number for pagination. + size: Page size for pagination. + order_by: Field to order by. Defaults to "display_order". + sort: Sort order. Defaults to ascending. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + A Page[DeploymentVisualizationResponse] containing the visualizations. + """ + deployment = self.get_deployment(deployment_id) + filter_model = DeploymentVisualizationFilter( + project=deployment.project_id, + deployment=deployment_id, + page=page, + size=size, + order_by=order_by or "display_order", + sort=sort or SorterOps.ASCENDING, + ) + return self.zen_store.list_deployment_visualizations( + filter_model=filter_model, + hydrate=hydrate, + ) + + def update_deployment_visualization( + self, + deployment_visualization_id: UUID, + *, + display_name: Optional[str] = None, + display_order: Optional[int] = None, + ) -> DeploymentVisualizationResponse: + """Update display metadata for a curated deployment visualization. + + Args: + deployment_visualization_id: The ID of the deployment visualization to update. + display_name: New display name for the visualization. + display_order: New display order for the visualization. + + Returns: + The updated deployment visualization. + """ + update_model = DeploymentVisualizationUpdate( + display_name=display_name, + display_order=display_order, + ) + return self.zen_store.update_deployment_visualization( + deployment_visualization_id=deployment_visualization_id, + update=update_model, + ) + + def delete_deployment_visualization( + self, deployment_visualization_id: UUID + ) -> None: + """Delete a curated deployment visualization. + + Args: + deployment_visualization_id: The ID of the deployment visualization to delete. + """ + self.zen_store.delete_deployment_visualization( + deployment_visualization_id=deployment_visualization_id + ) + def list_deployments( self, sort_by: str = "created", diff --git a/src/zenml/constants.py b/src/zenml/constants.py index 8fb67eadab..afe027476b 100644 --- a/src/zenml/constants.py +++ b/src/zenml/constants.py @@ -404,6 +404,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int: PIPELINE_CONFIGURATION = "/pipeline-configuration" PIPELINE_DEPLOYMENTS = "/pipeline_deployments" DEPLOYMENTS = "/deployments" +DEPLOYMENT_VISUALIZATIONS = "/deployment_visualizations" PIPELINE_SNAPSHOTS = "/pipeline_snapshots" PIPELINES = "/pipelines" PIPELINE_SPEC = "/pipeline-spec" diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index aea28ae24a..4d2e510008 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -164,6 +164,15 @@ DeploymentResponseMetadata, DeploymentResponseResources, ) +from zenml.models.v2.core.deployment_visualization import ( + DeploymentVisualizationFilter, + DeploymentVisualizationRequest, + DeploymentVisualizationResponse, + DeploymentVisualizationResponseBody, + DeploymentVisualizationResponseMetadata, + DeploymentVisualizationResponseResources, + DeploymentVisualizationUpdate, +) from zenml.models.v2.core.device import ( OAuthDeviceUpdate, OAuthDeviceFilter, @@ -653,6 +662,13 @@ "DeploymentResponseBody", "DeploymentResponseMetadata", "DeploymentResponseResources", + "DeploymentVisualizationFilter", + "DeploymentVisualizationRequest", + "DeploymentVisualizationResponse", + "DeploymentVisualizationResponseBody", + "DeploymentVisualizationResponseMetadata", + "DeploymentVisualizationResponseResources", + "DeploymentVisualizationUpdate", "EventSourceFlavorResponse", "EventSourceFlavorResponseBody", "EventSourceFlavorResponseMetadata", diff --git a/src/zenml/models/v2/core/deployment.py b/src/zenml/models/v2/core/deployment.py index 872f0a9b43..ead51ae645 100644 --- a/src/zenml/models/v2/core/deployment.py +++ b/src/zenml/models/v2/core/deployment.py @@ -46,6 +46,9 @@ from sqlalchemy.sql.elements import ColumnElement from zenml.models.v2.core.component import ComponentResponse + from zenml.models.v2.core.deployment_visualization import ( + DeploymentVisualizationResponse, + ) from zenml.models.v2.core.pipeline import PipelineResponse from zenml.models.v2.core.pipeline_snapshot import ( PipelineSnapshotResponse, @@ -204,6 +207,10 @@ class DeploymentResponseResources(ProjectScopedResponseResources): tags: List["TagResponse"] = Field( title="Tags associated with the deployment.", ) + visualizations: List["DeploymentVisualizationResponse"] = Field( + default_factory=list, + title="Curated deployment visualizations.", + ) class DeploymentResponse( @@ -304,6 +311,15 @@ def tags(self) -> List["TagResponse"]: """ return self.get_resources().tags + @property + def visualizations(self) -> List["DeploymentVisualizationResponse"]: + """The visualizations of the deployment. + + Returns: + The visualizations of the deployment. + """ + return self.get_resources().visualizations + @property def snapshot_id(self) -> Optional[UUID]: """The pipeline snapshot ID. diff --git a/src/zenml/models/v2/core/deployment_visualization.py b/src/zenml/models/v2/core/deployment_visualization.py new file mode 100644 index 0000000000..129ff4997e --- /dev/null +++ b/src/zenml/models/v2/core/deployment_visualization.py @@ -0,0 +1,292 @@ +# Copyright (c) ZenML GmbH 2025. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Models representing deployment visualizations.""" + +from typing import ( + TYPE_CHECKING, + ClassVar, + List, + Optional, + Type, + TypeVar, +) +from uuid import UUID + +from pydantic import Field + +from zenml.constants import STR_FIELD_MAX_LENGTH +from zenml.models.v2.base.base import BaseUpdate +from zenml.models.v2.base.scoped import ( + ProjectScopedFilter, + ProjectScopedRequest, + ProjectScopedResponse, + ProjectScopedResponseBody, + ProjectScopedResponseMetadata, + ProjectScopedResponseResources, +) + +if TYPE_CHECKING: + from sqlalchemy.sql.elements import ColumnElement + + from zenml.models.v2.core.artifact_version import ArtifactVersionResponse + from zenml.models.v2.core.deployment import DeploymentResponse + from zenml.zen_stores.schemas.base_schemas import BaseSchema + + AnySchema = TypeVar("AnySchema", bound=BaseSchema) + + +# ------------------ Request Model ------------------ + + +class DeploymentVisualizationRequest(ProjectScopedRequest): + """Request model for deployment visualizations.""" + + deployment_id: UUID = Field( + title="The deployment ID.", + description="The ID of the deployment associated with the visualization.", + ) + artifact_version_id: UUID = Field( + title="The artifact version ID.", + description="The ID of the artifact version associated with the visualization.", + ) + visualization_index: int = Field( + ge=0, + title="The visualization index.", + description="The index of the visualization within the artifact version.", + ) + display_name: Optional[str] = Field( + default=None, + title="The display name of the visualization.", + max_length=STR_FIELD_MAX_LENGTH, + ) + display_order: Optional[int] = Field( + default=None, + title="The display order of the visualization.", + ) + + +# ------------------ Update Model ------------------ + + +class DeploymentVisualizationUpdate(BaseUpdate): + """Update model for deployment visualizations.""" + + display_name: Optional[str] = Field( + default=None, + title="The new display name of the visualization.", + max_length=STR_FIELD_MAX_LENGTH, + ) + display_order: Optional[int] = Field( + default=None, + title="The new display order of the visualization.", + ) + + +# ------------------ Response Model ------------------ + + +class DeploymentVisualizationResponseBody(ProjectScopedResponseBody): + """Response body for deployment visualizations.""" + + deployment_id: UUID = Field( + title="The deployment ID.", + description="The ID of the deployment associated with the visualization.", + ) + artifact_version_id: UUID = Field( + title="The artifact version ID.", + description="The ID of the artifact version associated with the visualization.", + ) + visualization_index: int = Field( + title="The visualization index.", + description="The index of the visualization within the artifact version.", + ) + display_name: Optional[str] = Field( + default=None, + title="The display name of the visualization.", + ) + display_order: Optional[int] = Field( + default=None, + title="The display order of the visualization.", + ) + + +class DeploymentVisualizationResponseMetadata(ProjectScopedResponseMetadata): + """Response metadata for deployment visualizations.""" + + +class DeploymentVisualizationResponseResources(ProjectScopedResponseResources): + """Response resources for deployment visualizations.""" + + deployment: Optional["DeploymentResponse"] = Field( + default=None, + title="The deployment.", + description="The deployment associated with the visualization.", + ) + artifact_version: Optional["ArtifactVersionResponse"] = Field( + default=None, + title="The artifact version.", + description="The artifact version associated with the visualization.", + ) + + +class DeploymentVisualizationResponse( + ProjectScopedResponse[ + DeploymentVisualizationResponseBody, + DeploymentVisualizationResponseMetadata, + DeploymentVisualizationResponseResources, + ] +): + """Response model for deployment visualizations.""" + + def get_hydrated_version(self) -> "DeploymentVisualizationResponse": + """Get the hydrated version of this deployment visualization. + + Returns: + an instance of the same entity with the metadata and resources fields + attached. + """ + from zenml.client import Client + + client = Client() + return client.zen_store.get_deployment_visualization(self.id) + + # Helper properties + @property + def deployment_id(self) -> UUID: + """The deployment ID. + + Returns: + The deployment ID. + """ + return self.get_body().deployment_id + + @property + def artifact_version_id(self) -> UUID: + """The artifact version ID. + + Returns: + The artifact version ID. + """ + return self.get_body().artifact_version_id + + @property + def visualization_index(self) -> int: + """The visualization index. + + Returns: + The visualization index. + """ + return self.get_body().visualization_index + + @property + def display_name(self) -> Optional[str]: + """The display name of the visualization. + + Returns: + The display name of the visualization. + """ + return self.get_body().display_name + + @property + def display_order(self) -> Optional[int]: + """The display order of the visualization. + + Returns: + The display order of the visualization. + """ + return self.get_body().display_order + + @property + def deployment(self) -> Optional["DeploymentResponse"]: + """The deployment. + + Returns: + The deployment. + """ + return self.get_resources().deployment + + @property + def artifact_version(self) -> Optional["ArtifactVersionResponse"]: + """The artifact version. + + Returns: + The artifact version. + """ + return self.get_resources().artifact_version + + +# ------------------ Filter Model ------------------ + + +class DeploymentVisualizationFilter(ProjectScopedFilter): + """Model to enable advanced filtering of deployment visualizations.""" + + FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [ + *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS, + "deployment", + "artifact_version", + ] + CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [ + *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS, + "display_order", + "created", + "updated", + ] + CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [ + *ProjectScopedFilter.CLI_EXCLUDE_FIELDS, + ] + + deployment: Optional[UUID] = Field( + default=None, + description="ID of the deployment associated with the visualization.", + ) + artifact_version: Optional[UUID] = Field( + default=None, + description="ID of the artifact version associated with the visualization.", + ) + visualization_index: Optional[int] = Field( + default=None, + description="Index of the visualization within the artifact version.", + ) + display_order: Optional[int] = Field( + default=None, + description="Display order of the visualization.", + ) + + def get_custom_filters( + self, table: Type["AnySchema"] + ) -> List["ColumnElement[bool]"]: + """Get custom filters. + + Args: + table: The query table. + + Returns: + A list of custom filters. + """ + custom_filters = super().get_custom_filters(table) + + if self.deployment: + deployment_filter = ( + getattr(table, "deployment_id") == self.deployment + ) + custom_filters.append(deployment_filter) + + if self.artifact_version: + artifact_version_filter = ( + getattr(table, "artifact_version_id") == self.artifact_version + ) + custom_filters.append(artifact_version_filter) + + return custom_filters diff --git a/src/zenml/zen_server/routers/deployment_endpoints.py b/src/zenml/zen_server/routers/deployment_endpoints.py index 7b0ec15e23..d58cf8819b 100644 --- a/src/zenml/zen_server/routers/deployment_endpoints.py +++ b/src/zenml/zen_server/routers/deployment_endpoints.py @@ -23,14 +23,20 @@ from zenml.constants import ( API, + DEPLOYMENT_VISUALIZATIONS, DEPLOYMENTS, VERSION_1, ) +from zenml.enums import SorterOps from zenml.models import ( DeploymentFilter, DeploymentRequest, DeploymentResponse, DeploymentUpdate, + DeploymentVisualizationFilter, + DeploymentVisualizationRequest, + DeploymentVisualizationResponse, + DeploymentVisualizationUpdate, ) from zenml.models.v2.base.page import Page from zenml.zen_server.auth import AuthContext, authorize @@ -55,6 +61,12 @@ responses={401: error_response, 403: error_response}, ) +deployment_visualization_router = APIRouter( + prefix=API + VERSION_1 + DEPLOYMENT_VISUALIZATIONS, + tags=["deployment_visualizations"], + responses={401: error_response, 403: error_response}, +) + @router.post( "", @@ -183,3 +195,146 @@ def delete_deployment( get_method=zen_store().get_deployment, delete_method=zen_store().delete_deployment, ) + + +@router.post( + "/{deployment_id}/visualizations", + responses={401: error_response, 409: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper +def create_deployment_visualization( + deployment_id: UUID, + visualization: DeploymentVisualizationRequest, + _: AuthContext = Security(authorize), +) -> DeploymentVisualizationResponse: + """Creates a curated deployment visualization. + + Args: + deployment_id: ID of the deployment to add visualization to. + visualization: Deployment visualization to create. + + Returns: + The created deployment visualization. + """ + # Ensure deployment_id matches path parameter + visualization = visualization.model_copy( + update={"deployment_id": deployment_id} + ) + return verify_permissions_and_create_entity( + request_model=visualization, + create_method=zen_store().create_deployment_visualization, + ) + + +@router.get( + "/{deployment_id}/visualizations", + responses={401: error_response, 404: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper(deduplicate=True) +def list_deployment_visualizations( + deployment_id: UUID, + visualization_filter_model: DeploymentVisualizationFilter = Depends( + make_dependable(DeploymentVisualizationFilter) + ), + hydrate: bool = False, + _: AuthContext = Security(authorize), +) -> Page[DeploymentVisualizationResponse]: + """Gets a list of visualizations for a specific deployment. + + Args: + deployment_id: ID of the deployment to list visualizations for. + visualization_filter_model: Filter model used for pagination, sorting, + filtering. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + List of deployment visualization objects for the deployment. + """ + # Set deployment filter to the path parameter + visualization_filter_model = visualization_filter_model.model_copy( + update={"deployment": deployment_id} + ) + if visualization_filter_model.sort is None: + visualization_filter_model.sort = SorterOps.ASCENDING + return verify_permissions_and_list_entities( + filter_model=visualization_filter_model, + resource_type=ResourceType.DEPLOYMENT, + list_method=zen_store().list_deployment_visualizations, + hydrate=hydrate, + ) + + +@deployment_visualization_router.get( + "/{deployment_visualization_id}", + responses={401: error_response, 404: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper(deduplicate=True) +def get_deployment_visualization( + deployment_visualization_id: UUID, + hydrate: bool = True, + _: AuthContext = Security(authorize), +) -> DeploymentVisualizationResponse: + """Gets a specific deployment visualization using its unique id. + + Args: + deployment_visualization_id: ID of the deployment visualization to get. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + A specific deployment visualization object. + """ + return verify_permissions_and_get_entity( + id=deployment_visualization_id, + get_method=zen_store().get_deployment_visualization, + hydrate=hydrate, + ) + + +@deployment_visualization_router.patch( + "/{deployment_visualization_id}", + responses={401: error_response, 404: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper(deduplicate=True) +def update_deployment_visualization( + deployment_visualization_id: UUID, + visualization_update: DeploymentVisualizationUpdate, + _: AuthContext = Security(authorize), +) -> DeploymentVisualizationResponse: + """Updates a specific deployment visualization. + + Args: + deployment_visualization_id: ID of the deployment visualization to update. + visualization_update: Update model for the deployment visualization. + + Returns: + The updated deployment visualization. + """ + return verify_permissions_and_update_entity( + id=deployment_visualization_id, + update_model=visualization_update, + get_method=zen_store().get_deployment_visualization, + update_method=zen_store().update_deployment_visualization, + ) + + +@deployment_visualization_router.delete( + "/{deployment_visualization_id}", + responses={401: error_response, 404: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper +def delete_deployment_visualization( + deployment_visualization_id: UUID, + _: AuthContext = Security(authorize), +) -> None: + """Deletes a specific deployment visualization. + + Args: + deployment_visualization_id: ID of the deployment visualization to delete. + """ + verify_permissions_and_delete_entity( + id=deployment_visualization_id, + get_method=zen_store().get_deployment_visualization, + delete_method=zen_store().delete_deployment_visualization, + ) diff --git a/src/zenml/zen_server/zen_server_api.py b/src/zenml/zen_server/zen_server_api.py index 5071998276..a0f4c5f878 100644 --- a/src/zenml/zen_server/zen_server_api.py +++ b/src/zenml/zen_server/zen_server_api.py @@ -265,6 +265,7 @@ async def dashboard(request: Request) -> Any: app.include_router(devices_endpoints.router) app.include_router(code_repositories_endpoints.router) app.include_router(deployment_endpoints.router) +app.include_router(deployment_endpoints.deployment_visualization_router) app.include_router(plugin_endpoints.plugin_router) app.include_router(event_source_endpoints.event_source_router) app.include_router(flavors_endpoints.router) diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 65e0bf91f1..6f3af89c95 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -68,6 +68,7 @@ CURRENT_USER, DEACTIVATE, DEFAULT_HTTP_TIMEOUT, + DEPLOYMENT_VISUALIZATIONS, DEPLOYMENTS, DEVICES, DISABLE_CLIENT_SERVER_MISMATCH_WARNING, @@ -170,6 +171,10 @@ DeploymentRequest, DeploymentResponse, DeploymentUpdate, + DeploymentVisualizationFilter, + DeploymentVisualizationRequest, + DeploymentVisualizationResponse, + DeploymentVisualizationUpdate, EventSourceFilter, EventSourceRequest, EventSourceResponse, @@ -1856,6 +1861,112 @@ def delete_deployment(self, deployment_id: UUID) -> None: route=DEPLOYMENTS, ) + def create_deployment_visualization( + self, visualization: DeploymentVisualizationRequest + ) -> DeploymentVisualizationResponse: + """Create a deployment visualization via REST API. + + Args: + visualization: The visualization to create. + + Returns: + The created visualization. + """ + return self._create_resource( + resource=visualization, + response_model=DeploymentVisualizationResponse, + route=f"{DEPLOYMENTS}/{visualization.deployment_id}/visualizations", + params={"hydrate": True}, + ) + + def get_deployment_visualization( + self, deployment_visualization_id: UUID, hydrate: bool = True + ) -> DeploymentVisualizationResponse: + """Get a deployment visualization by ID. + + Args: + deployment_visualization_id: The ID of the visualization to get. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + The deployment visualization. + """ + return self._get_resource( + resource_id=deployment_visualization_id, + route=DEPLOYMENT_VISUALIZATIONS, + response_model=DeploymentVisualizationResponse, + params={"hydrate": hydrate}, + ) + + def list_deployment_visualizations( + self, + filter_model: DeploymentVisualizationFilter, + hydrate: bool = False, + ) -> Page[DeploymentVisualizationResponse]: + """List deployment visualizations via REST API. + + Args: + filter_model: The filter model to use. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + A list of all deployment visualizations matching the filter criteria. + + Raises: + ValueError: If no deployment ID is provided in the filter model. + """ + if not filter_model.deployment: + raise ValueError( + "A deployment ID is required to list deployment visualizations." + ) + + # Build the route with the deployment ID + route = f"{DEPLOYMENTS}/{str(filter_model.deployment)}/visualizations" + + return self._list_paginated_resources( + route=route, + response_model=DeploymentVisualizationResponse, + filter_model=filter_model, + params={"hydrate": hydrate}, + ) + + def update_deployment_visualization( + self, + deployment_visualization_id: UUID, + update: DeploymentVisualizationUpdate, + ) -> DeploymentVisualizationResponse: + """Update a deployment visualization via REST API. + + Args: + deployment_visualization_id: The ID of the visualization to update. + update: The update to apply. + + Returns: + The updated deployment visualization. + """ + return self._update_resource( + resource_id=deployment_visualization_id, + resource_update=update, + response_model=DeploymentVisualizationResponse, + route=DEPLOYMENT_VISUALIZATIONS, + params={"hydrate": True}, + ) + + def delete_deployment_visualization( + self, deployment_visualization_id: UUID + ) -> None: + """Delete a deployment visualization via REST API. + + Args: + deployment_visualization_id: The ID of the visualization to delete. + """ + self._delete_resource( + resource_id=deployment_visualization_id, + route=DEPLOYMENT_VISUALIZATIONS, + ) + # -------------------- Run templates -------------------- def create_run_template( diff --git a/src/zenml/zen_stores/schemas/__init__.py b/src/zenml/zen_stores/schemas/__init__.py index b98adfcfea..70add37167 100644 --- a/src/zenml/zen_stores/schemas/__init__.py +++ b/src/zenml/zen_stores/schemas/__init__.py @@ -31,6 +31,9 @@ from zenml.zen_stores.schemas.event_source_schemas import EventSourceSchema from zenml.zen_stores.schemas.pipeline_build_schemas import PipelineBuildSchema from zenml.zen_stores.schemas.deployment_schemas import DeploymentSchema +from zenml.zen_stores.schemas.deployment_visualization_schemas import ( + DeploymentVisualizationSchema, +) from zenml.zen_stores.schemas.component_schemas import StackComponentSchema from zenml.zen_stores.schemas.flavor_schemas import FlavorSchema from zenml.zen_stores.schemas.server_settings_schemas import ServerSettingsSchema @@ -88,6 +91,7 @@ "CodeReferenceSchema", "CodeRepositorySchema", "DeploymentSchema", + "DeploymentVisualizationSchema", "EventSourceSchema", "FlavorSchema", "LogsSchema", diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index c6284f7aad..eaab1b4069 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -46,6 +46,9 @@ from zenml.zen_stores.schemas.utils import jl_arg if TYPE_CHECKING: + from zenml.zen_stores.schemas.deployment_visualization_schemas import ( + DeploymentVisualizationSchema, + ) from zenml.zen_stores.schemas.tag_schemas import TagSchema logger = get_logger(__name__) @@ -133,6 +136,14 @@ class DeploymentSchema(NamedSchema, table=True): ), ) + visualizations: List["DeploymentVisualizationSchema"] = Relationship( + back_populates="deployment", + sa_relationship_kwargs={ + "lazy": "selectin", + "order_by": "CASE WHEN display_order IS NULL THEN 1 ELSE 0 END, display_order, created" + }, + ) + @classmethod def get_query_options( cls, @@ -155,6 +166,10 @@ def get_query_options( options = [] if include_resources: + from zenml.zen_stores.schemas.deployment_visualization_schemas import ( + DeploymentVisualizationSchema, + ) + options.extend( [ joinedload(jl_arg(DeploymentSchema.user)), @@ -162,6 +177,11 @@ def get_query_options( selectinload(jl_arg(DeploymentSchema.snapshot)).joinedload( jl_arg(PipelineSnapshotSchema.pipeline) ), + selectinload( + jl_arg(DeploymentSchema.visualizations) + ).selectinload( + jl_arg(DeploymentVisualizationSchema.artifact_version) + ), ] ) @@ -220,6 +240,14 @@ def to_model( pipeline=self.snapshot.pipeline.to_model() if self.snapshot and self.snapshot.pipeline else None, + visualizations=[ + visualization.to_model( + include_metadata=include_metadata, + include_resources=include_resources, + include_deployment=False, + ) + for visualization in (self.visualizations or []) + ], ) return DeploymentResponse( diff --git a/src/zenml/zen_stores/schemas/deployment_visualization_schemas.py b/src/zenml/zen_stores/schemas/deployment_visualization_schemas.py new file mode 100644 index 0000000000..665dbe6a12 --- /dev/null +++ b/src/zenml/zen_stores/schemas/deployment_visualization_schemas.py @@ -0,0 +1,243 @@ +# Copyright (c) ZenML GmbH 2025. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""SQLModel implementation of deployment visualization table.""" + +from typing import TYPE_CHECKING, Any, Optional, Sequence +from uuid import UUID + +from sqlalchemy import UniqueConstraint +from sqlalchemy.orm import selectinload +from sqlalchemy.sql.base import ExecutableOption +from sqlmodel import Field, Relationship + +from zenml.constants import STR_FIELD_MAX_LENGTH +from zenml.models.v2.core.deployment_visualization import ( + DeploymentVisualizationRequest, + DeploymentVisualizationResponse, + DeploymentVisualizationResponseBody, + DeploymentVisualizationResponseMetadata, + DeploymentVisualizationResponseResources, + DeploymentVisualizationUpdate, +) +from zenml.utils.time_utils import utc_now +from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema +from zenml.zen_stores.schemas.base_schemas import BaseSchema +from zenml.zen_stores.schemas.project_schemas import ProjectSchema +from zenml.zen_stores.schemas.schema_utils import ( + build_foreign_key_field, + build_index, +) +from zenml.zen_stores.schemas.utils import jl_arg + +if TYPE_CHECKING: + from zenml.zen_stores.schemas.deployment_schemas import DeploymentSchema + + +class DeploymentVisualizationSchema(BaseSchema, table=True): + """SQL Model for deployment visualizations.""" + + __tablename__ = "deployment_visualization" + __table_args__ = ( + UniqueConstraint( + "deployment_id", + "artifact_version_id", + "visualization_index", + name="unique_deployment_visualization", + ), + build_index( + __tablename__, + ["deployment_id"], + ), + build_index( + __tablename__, + ["display_order"], + ), + ) + + # Foreign Keys + project_id: UUID = build_foreign_key_field( + source=__tablename__, + target=ProjectSchema.__tablename__, + source_column="project_id", + target_column="id", + ondelete="CASCADE", + nullable=False, + ) + + deployment_id: UUID = build_foreign_key_field( + source=__tablename__, + target="deployment", + source_column="deployment_id", + target_column="id", + ondelete="CASCADE", + nullable=False, + ) + + artifact_version_id: UUID = build_foreign_key_field( + source=__tablename__, + target=ArtifactVersionSchema.__tablename__, + source_column="artifact_version_id", + target_column="id", + ondelete="CASCADE", + nullable=False, + ) + + # Fields + visualization_index: int = Field(nullable=False) + display_name: Optional[str] = Field( + max_length=STR_FIELD_MAX_LENGTH, default=None + ) + display_order: Optional[int] = Field(default=None) + + # Relationships + deployment: Optional["DeploymentSchema"] = Relationship( + back_populates="visualizations", + sa_relationship_kwargs={"lazy": "selectin"}, + ) + + artifact_version: Optional[ArtifactVersionSchema] = Relationship( + sa_relationship_kwargs={"lazy": "selectin"} + ) + + @classmethod + def get_query_options( + cls, + include_metadata: bool = False, + include_resources: bool = False, + **kwargs: Any, + ) -> Sequence[ExecutableOption]: + """Get the query options for the schema. + + Args: + include_metadata: Whether metadata will be included when converting + the schema to a model. + include_resources: Whether resources will be included when + converting the schema to a model. + **kwargs: Keyword arguments to allow schema specific logic + + Returns: + A list of query options. + """ + options = [] + + if include_resources: + options.extend( + [ + selectinload(jl_arg(cls.deployment)), + selectinload(jl_arg(cls.artifact_version)), + ] + ) + + return options + + @classmethod + def from_request( + cls, request: DeploymentVisualizationRequest + ) -> "DeploymentVisualizationSchema": + """Convert a `DeploymentVisualizationRequest` to a `DeploymentVisualizationSchema`. + + Args: + request: The request model to convert. + + Returns: + The converted schema. + """ + return cls( + project_id=request.project, + deployment_id=request.deployment_id, + artifact_version_id=request.artifact_version_id, + visualization_index=request.visualization_index, + display_name=request.display_name, + display_order=request.display_order, + ) + + def update( + self, + update: DeploymentVisualizationUpdate, + ) -> "DeploymentVisualizationSchema": + """Updates a `DeploymentVisualizationSchema` from a `DeploymentVisualizationUpdate`. + + Args: + update: The `DeploymentVisualizationUpdate` to update from. + + Returns: + The updated `DeploymentVisualizationSchema`. + """ + for field, value in update.model_dump( + exclude_unset=True, exclude_none=True + ).items(): + if hasattr(self, field): + setattr(self, field, value) + + self.updated = utc_now() + return self + + def to_model( + self, + include_metadata: bool = False, + include_resources: bool = False, + **kwargs: Any, + ) -> DeploymentVisualizationResponse: + """Convert a `DeploymentVisualizationSchema` to a `DeploymentVisualizationResponse`. + + Args: + include_metadata: Whether to include metadata in the response. + include_resources: Whether to include resources in the response. + **kwargs: Additional keyword arguments, including `include_deployment` + to control whether deployment resources are included (default True). + + Returns: + The created `DeploymentVisualizationResponse`. + """ + include_deployment = kwargs.get("include_deployment", True) + + body = DeploymentVisualizationResponseBody( + user_id=None, + project_id=self.project_id, + created=self.created, + updated=self.updated, + deployment_id=self.deployment_id, + artifact_version_id=self.artifact_version_id, + visualization_index=self.visualization_index, + display_name=self.display_name, + display_order=self.display_order, + ) + + metadata = None + if include_metadata: + metadata = DeploymentVisualizationResponseMetadata() + + resources = None + if include_resources: + resources = DeploymentVisualizationResponseResources( + deployment=self.deployment.to_model( + include_metadata=include_metadata, + include_resources=include_resources, + ) + if self.deployment and include_deployment + else None, + artifact_version=self.artifact_version.to_model( + include_metadata=include_metadata, + include_resources=include_resources, + ) + if self.artifact_version + else None, + ) + + return DeploymentVisualizationResponse( + id=self.id, + body=body, + metadata=metadata, + resources=resources, + ) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 831dca781d..8155cc0123 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -143,6 +143,7 @@ OnboardingStep, SecretResourceTypes, SecretsStoreType, + SorterOps, StackComponentType, StackDeploymentProvider, StepRunInputArtifactType, @@ -206,6 +207,10 @@ DeploymentRequest, DeploymentResponse, DeploymentUpdate, + DeploymentVisualizationFilter, + DeploymentVisualizationRequest, + DeploymentVisualizationResponse, + DeploymentVisualizationUpdate, EventSourceFilter, EventSourceRequest, EventSourceResponse, @@ -361,6 +366,7 @@ CodeReferenceSchema, CodeRepositorySchema, DeploymentSchema, + DeploymentVisualizationSchema, EventSourceSchema, FlavorSchema, ModelSchema, @@ -5394,6 +5400,218 @@ def delete_deployment(self, deployment_id: UUID) -> None: session.delete(deployment) session.commit() + # -------------------- Deployment visualizations -------------------- + + def _validate_deployment_visualization_index( + self, + session: Session, + artifact_version_id: UUID, + visualization_index: int, + ) -> None: + """Ensure the artifact version exposes the given visualization index. + + Args: + session: The session to use. + artifact_version_id: The ID of the artifact version to validate. + visualization_index: The index of the visualization to validate. + """ + count = session.scalar( + select(func.count(ArtifactVisualizationSchema.id)).where( + ArtifactVisualizationSchema.artifact_version_id + == artifact_version_id + ) + ) + if not count or visualization_index >= count: + raise IllegalOperationError( + "Artifact version " + f"`{artifact_version_id}` does not expose a visualization " + f"with index {visualization_index}." + ) + + def create_deployment_visualization( + self, visualization: DeploymentVisualizationRequest + ) -> DeploymentVisualizationResponse: + """Persist a curated deployment visualization link. + + Args: + visualization: The visualization to create. + + Returns: + The created deployment visualization. + """ + with Session(self.engine) as session: + self._set_request_user_id( + request_model=visualization, session=session + ) + + deployment = self._get_schema_by_id( + resource_id=visualization.deployment_id, + schema_class=DeploymentSchema, + session=session, + ) + artifact_version = self._get_schema_by_id( + resource_id=visualization.artifact_version_id, + schema_class=ArtifactVersionSchema, + session=session, + ) + + project_id = deployment.project_id + if visualization.project and visualization.project != project_id: + raise IllegalOperationError( + "Deployment visualizations must target the same project " + "as the deployment." + ) + if artifact_version.project_id != project_id: + raise IllegalOperationError( + "Artifact version does not belong to the deployment " + "project." + ) + + self._validate_deployment_visualization_index( + session=session, + artifact_version_id=visualization.artifact_version_id, + visualization_index=visualization.visualization_index, + ) + + duplicate = session.exec( + select(DeploymentVisualizationSchema) + .where( + DeploymentVisualizationSchema.deployment_id + == visualization.deployment_id + ) + .where( + DeploymentVisualizationSchema.artifact_version_id + == visualization.artifact_version_id + ) + .where( + DeploymentVisualizationSchema.visualization_index + == visualization.visualization_index + ) + ).first() + if duplicate is not None: + raise EntityExistsError( + "A curated visualization with the same deployment, " + "artifact version, and index already exists." + ) + + schema = DeploymentVisualizationSchema.from_request(visualization) + schema.project_id = project_id + session.add(schema) + session.commit() + session.refresh(schema) + + return schema.to_model( + include_metadata=True, + include_resources=True, + include_deployment=False, + ) + + def get_deployment_visualization( + self, deployment_visualization_id: UUID, hydrate: bool = True + ) -> DeploymentVisualizationResponse: + """Fetch a curated deployment visualization by ID. + + Args: + deployment_visualization_id: The ID of the visualization to get. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + The deployment visualization. + """ + with Session(self.engine) as session: + schema = self._get_schema_by_id( + resource_id=deployment_visualization_id, + schema_class=DeploymentVisualizationSchema, + session=session, + ) + return schema.to_model( + include_metadata=hydrate, + include_resources=hydrate, + include_deployment=hydrate, + ) + + def list_deployment_visualizations( + self, + filter_model: DeploymentVisualizationFilter, + hydrate: bool = False, + ) -> Page[DeploymentVisualizationResponse]: + """List all deployment visualizations matching the given filter. + + Args: + filter_model: The filter model to use. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + A list of all deployment visualizations matching the filter criteria. + """ + with Session(self.engine) as session: + self._set_filter_project_id( + filter_model=filter_model, session=session + ) + if not filter_model.order_by: + filter_model.order_by = "display_order" + if getattr(filter_model, "sort", None) is None: + filter_model.sort = SorterOps.ASCENDING + + query = select(DeploymentVisualizationSchema) + return self.filter_and_paginate( + session=session, + query=query, + table=DeploymentVisualizationSchema, + filter_model=filter_model, + hydrate=hydrate, + ) + + def update_deployment_visualization( + self, + deployment_visualization_id: UUID, + update: DeploymentVisualizationUpdate, + ) -> DeploymentVisualizationResponse: + """Update mutable fields on a curated deployment visualization. + + Args: + deployment_visualization_id: The ID of the visualization to update. + update: The update to apply. + + Returns: + The updated deployment visualization. + """ + with Session(self.engine) as session: + schema = self._get_schema_by_id( + resource_id=deployment_visualization_id, + schema_class=DeploymentVisualizationSchema, + session=session, + ) + schema.update(update) + session.add(schema) + session.commit() + session.refresh(schema) + + return schema.to_model( + include_metadata=True, + include_resources=True, + include_deployment=False, + ) + + def delete_deployment_visualization( + self, deployment_visualization_id: UUID + ) -> None: + """Delete a curated deployment visualization. + + Args: + deployment_visualization_id: The ID of the visualization to delete. + """ + with Session(self.engine) as session: + schema = self._get_schema_by_id( + resource_id=deployment_visualization_id, + schema_class=DeploymentVisualizationSchema, + session=session, + ) + session.delete(schema) + session.commit() + # -------------------- Run templates -------------------- @track_decorator(AnalyticsEvent.CREATED_RUN_TEMPLATE) diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index 40a1a74adf..989ef46af2 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -53,6 +53,10 @@ DeploymentRequest, DeploymentResponse, DeploymentUpdate, + DeploymentVisualizationFilter, + DeploymentVisualizationRequest, + DeploymentVisualizationResponse, + DeploymentVisualizationUpdate, EventSourceFilter, EventSourceRequest, EventSourceResponse, @@ -1471,6 +1475,97 @@ def delete_deployment(self, deployment_id: UUID) -> None: KeyError: If the deployment does not exist. """ + # -------------------- Deployment visualizations -------------------- + + @abstractmethod + def create_deployment_visualization( + self, visualization: DeploymentVisualizationRequest + ) -> DeploymentVisualizationResponse: + """Create a new deployment visualization. + + Args: + visualization: The deployment visualization to create. + + Returns: + The newly created deployment visualization. + + Raises: + EntityExistsError: If a deployment visualization with the same + deployment, artifact version, and visualization index already + exists. + """ + + @abstractmethod + def get_deployment_visualization( + self, deployment_visualization_id: UUID, hydrate: bool = True + ) -> DeploymentVisualizationResponse: + """Get a deployment visualization by ID. + + Args: + deployment_visualization_id: The ID of the deployment visualization + to get. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + The deployment visualization. + + Raises: + KeyError: If the deployment visualization does not exist. + """ + + @abstractmethod + def list_deployment_visualizations( + self, + filter_model: DeploymentVisualizationFilter, + hydrate: bool = False, + ) -> Page[DeploymentVisualizationResponse]: + """List all deployment visualizations matching the given filter criteria. + + Args: + 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: + A list of all deployment visualizations matching the filter criteria. + """ + + @abstractmethod + def update_deployment_visualization( + self, + deployment_visualization_id: UUID, + visualization_update: DeploymentVisualizationUpdate, + ) -> DeploymentVisualizationResponse: + """Update a deployment visualization. + + Args: + deployment_visualization_id: The ID of the deployment visualization + to update. + visualization_update: The update to apply. + + Returns: + The updated deployment visualization. + + Raises: + KeyError: If the deployment visualization does not exist. + """ + + @abstractmethod + def delete_deployment_visualization( + self, deployment_visualization_id: UUID + ) -> None: + """Delete a deployment visualization. + + Args: + deployment_visualization_id: The ID of the deployment visualization + to delete. + + Raises: + KeyError: If the deployment visualization does not exist. + """ + # -------------------- Run templates -------------------- @abstractmethod From cd81c4ee6f9e7b455d5a6a380753670177477086 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 27 Sep 2025 03:02:42 +0100 Subject: [PATCH 02/64] mypy --- src/zenml/client.py | 2 +- .../routers/deployment_endpoints.py | 14 +++++++++++--- .../zen_stores/schemas/deployment_schemas.py | 8 ++++---- src/zenml/zen_stores/sql_zen_store.py | 19 +++++++++++++------ 4 files changed, 29 insertions(+), 14 deletions(-) diff --git a/src/zenml/client.py b/src/zenml/client.py index 73f907161f..d7d7d7c8f3 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -3832,7 +3832,7 @@ def update_deployment_visualization( ) return self.zen_store.update_deployment_visualization( deployment_visualization_id=deployment_visualization_id, - update=update_model, + visualization_update=update_model, ) def delete_deployment_visualization( diff --git a/src/zenml/zen_server/routers/deployment_endpoints.py b/src/zenml/zen_server/routers/deployment_endpoints.py index d58cf8819b..009fbbee3b 100644 --- a/src/zenml/zen_server/routers/deployment_endpoints.py +++ b/src/zenml/zen_server/routers/deployment_endpoints.py @@ -27,7 +27,6 @@ DEPLOYMENTS, VERSION_1, ) -from zenml.enums import SorterOps from zenml.models import ( DeploymentFilter, DeploymentRequest, @@ -255,8 +254,17 @@ def list_deployment_visualizations( visualization_filter_model = visualization_filter_model.model_copy( update={"deployment": deployment_id} ) - if visualization_filter_model.sort is None: - visualization_filter_model.sort = SorterOps.ASCENDING + default_sort_by = ( + type(visualization_filter_model).model_fields["sort_by"].default + ) + if ( + not visualization_filter_model.sort_by + or visualization_filter_model.sort_by == default_sort_by + ): + visualization_filter_model = visualization_filter_model.model_copy( + update={"sort_by": "display_order"} + ) + return verify_permissions_and_list_entities( filter_model=visualization_filter_model, resource_type=ResourceType.DEPLOYMENT, diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index eaab1b4069..35e634c123 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -138,10 +138,10 @@ class DeploymentSchema(NamedSchema, table=True): visualizations: List["DeploymentVisualizationSchema"] = Relationship( back_populates="deployment", - sa_relationship_kwargs={ - "lazy": "selectin", - "order_by": "CASE WHEN display_order IS NULL THEN 1 ELSE 0 END, display_order, created" - }, + sa_relationship_kwargs=dict( + lazy="selectin", + order_by="CASE WHEN display_order IS NULL THEN 1 ELSE 0 END, display_order, created", + ), ) @classmethod diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 8155cc0123..0606e4c749 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -143,7 +143,6 @@ OnboardingStep, SecretResourceTypes, SecretsStoreType, - SorterOps, StackComponentType, StackDeploymentProvider, StepRunInputArtifactType, @@ -5416,7 +5415,9 @@ def _validate_deployment_visualization_index( visualization_index: The index of the visualization to validate. """ count = session.scalar( - select(func.count(ArtifactVisualizationSchema.id)).where( + select(func.count()) + .select_from(ArtifactVisualizationSchema) + .where( ArtifactVisualizationSchema.artifact_version_id == artifact_version_id ) @@ -5550,10 +5551,16 @@ def list_deployment_visualizations( self._set_filter_project_id( filter_model=filter_model, session=session ) - if not filter_model.order_by: - filter_model.order_by = "display_order" - if getattr(filter_model, "sort", None) is None: - filter_model.sort = SorterOps.ASCENDING + default_sort_by = ( + type(filter_model).model_fields["sort_by"].default + ) + if ( + not filter_model.sort_by + or filter_model.sort_by == default_sort_by + ): + filter_model = filter_model.model_copy( + update={"sort_by": "display_order"} + ) query = select(DeploymentVisualizationSchema) return self.filter_and_paginate( From b420555e5c3eb4a24cfa6314f7fb82abb4b085ec Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 27 Sep 2025 16:11:23 +0100 Subject: [PATCH 03/64] update sorting behavious --- .../v2/core/deployment_visualization.py | 40 +++++++++++++++++++ .../routers/deployment_endpoints.py | 10 ----- 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/src/zenml/models/v2/core/deployment_visualization.py b/src/zenml/models/v2/core/deployment_visualization.py index 129ff4997e..c19bf83c36 100644 --- a/src/zenml/models/v2/core/deployment_visualization.py +++ b/src/zenml/models/v2/core/deployment_visualization.py @@ -27,6 +27,7 @@ from zenml.constants import STR_FIELD_MAX_LENGTH from zenml.models.v2.base.base import BaseUpdate +from zenml.models.v2.base.filter import AnyQuery from zenml.models.v2.base.scoped import ( ProjectScopedFilter, ProjectScopedRequest, @@ -247,6 +248,12 @@ class DeploymentVisualizationFilter(ProjectScopedFilter): *ProjectScopedFilter.CLI_EXCLUDE_FIELDS, ] + # Set default sort_by to display_order + sort_by: str = Field( + default="display_order", + description="Which column to sort by.", + ) + deployment: Optional[UUID] = Field( default=None, description="ID of the deployment associated with the visualization.", @@ -264,6 +271,39 @@ class DeploymentVisualizationFilter(ProjectScopedFilter): description="Display order of the visualization.", ) + def apply_sorting( + self, + query: AnyQuery, + table: Type["AnySchema"], + ) -> AnyQuery: + """Apply sorting to the deployment visualization query. + + Args: + query: The query to which to apply the sorting. + table: The query table. + + Returns: + The query with sorting applied. + """ + from sqlmodel import asc, desc + + from zenml.enums import SorterOps + + sort_by, operand = self.sorting_params + + if sort_by == "display_order": + column = getattr(table, sort_by) + if operand == SorterOps.DESCENDING: + return query.order_by(desc(column).nullslast(), asc(table.id)) + return query.order_by(asc(column).nullsfirst(), asc(table.id)) + elif sort_by in {"created", "updated"}: + column = getattr(table, sort_by) + if operand == SorterOps.DESCENDING: + return query.order_by(desc(column), asc(table.id)) + return query.order_by(asc(column), asc(table.id)) + + return super().apply_sorting(query=query, table=table) + def get_custom_filters( self, table: Type["AnySchema"] ) -> List["ColumnElement[bool]"]: diff --git a/src/zenml/zen_server/routers/deployment_endpoints.py b/src/zenml/zen_server/routers/deployment_endpoints.py index 009fbbee3b..8617e7e401 100644 --- a/src/zenml/zen_server/routers/deployment_endpoints.py +++ b/src/zenml/zen_server/routers/deployment_endpoints.py @@ -254,16 +254,6 @@ def list_deployment_visualizations( visualization_filter_model = visualization_filter_model.model_copy( update={"deployment": deployment_id} ) - default_sort_by = ( - type(visualization_filter_model).model_fields["sort_by"].default - ) - if ( - not visualization_filter_model.sort_by - or visualization_filter_model.sort_by == default_sort_by - ): - visualization_filter_model = visualization_filter_model.model_copy( - update={"sort_by": "display_order"} - ) return verify_permissions_and_list_entities( filter_model=visualization_filter_model, From 201411344c71bbe6d07565d86b651b6221613c5e Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 27 Sep 2025 16:18:45 +0100 Subject: [PATCH 04/64] fix display order sorting --- .../v2/core/deployment_visualization.py | 21 ++++++++----------- 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/src/zenml/models/v2/core/deployment_visualization.py b/src/zenml/models/v2/core/deployment_visualization.py index c19bf83c36..6e11fbd2a9 100644 --- a/src/zenml/models/v2/core/deployment_visualization.py +++ b/src/zenml/models/v2/core/deployment_visualization.py @@ -291,18 +291,15 @@ def apply_sorting( sort_by, operand = self.sorting_params - if sort_by == "display_order": - column = getattr(table, sort_by) - if operand == SorterOps.DESCENDING: - return query.order_by(desc(column).nullslast(), asc(table.id)) - return query.order_by(asc(column).nullsfirst(), asc(table.id)) - elif sort_by in {"created", "updated"}: - column = getattr(table, sort_by) - if operand == SorterOps.DESCENDING: - return query.order_by(desc(column), asc(table.id)) - return query.order_by(asc(column), asc(table.id)) - - return super().apply_sorting(query=query, table=table) + # Handle explicit created/updated sorting + if sort_by in {"created", "updated"}: + return super().apply_sorting(query=query, table=table) + + # For all other cases (including display_order), use display_order sorting + column = getattr(table, "display_order") + if operand == SorterOps.DESCENDING: + return query.order_by(desc(column).nullslast(), asc(table.id)) + return query.order_by(asc(column).nullsfirst(), asc(table.id)) def get_custom_filters( self, table: Type["AnySchema"] From 3376266d4146bf27aeb89f6eab9e497a5a698a7f Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 27 Sep 2025 17:46:45 +0100 Subject: [PATCH 05/64] fixes --- .../v2/core/deployment_visualization.py | 33 ++++++++++++++----- .../zen_stores/schemas/deployment_schemas.py | 5 +-- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/src/zenml/models/v2/core/deployment_visualization.py b/src/zenml/models/v2/core/deployment_visualization.py index 6e11fbd2a9..0a8c2ee954 100644 --- a/src/zenml/models/v2/core/deployment_visualization.py +++ b/src/zenml/models/v2/core/deployment_visualization.py @@ -20,6 +20,7 @@ Optional, Type, TypeVar, + cast, ) from uuid import UUID @@ -243,6 +244,7 @@ class DeploymentVisualizationFilter(ProjectScopedFilter): "display_order", "created", "updated", + "visualization_index", ] CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [ *ProjectScopedFilter.CLI_EXCLUDE_FIELDS, @@ -291,15 +293,30 @@ def apply_sorting( sort_by, operand = self.sorting_params - # Handle explicit created/updated sorting - if sort_by in {"created", "updated"}: - return super().apply_sorting(query=query, table=table) + # Special handling for display_order with nulls first/last + if sort_by == "display_order": + column = getattr(table, "display_order") + if operand == SorterOps.DESCENDING: + return cast( + AnyQuery, + query.order_by(desc(column).nulls_last(), asc(table.id)), + ) + return cast( + AnyQuery, + query.order_by(asc(column).nulls_first(), asc(table.id)), + ) - # For all other cases (including display_order), use display_order sorting - column = getattr(table, "display_order") - if operand == SorterOps.DESCENDING: - return query.order_by(desc(column).nullslast(), asc(table.id)) - return query.order_by(asc(column).nullsfirst(), asc(table.id)) + # Direct sorting for other custom columns + if sort_by in {"created", "updated", "visualization_index"}: + column = getattr(table, sort_by) + if operand == SorterOps.DESCENDING: + return cast( + AnyQuery, query.order_by(desc(column), asc(table.id)) + ) + return cast(AnyQuery, query.order_by(asc(column), asc(table.id))) + + # Delegate to parent for other columns + return super().apply_sorting(query=query, table=table) def get_custom_filters( self, table: Type["AnySchema"] diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index 35e634c123..1a23b81711 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -138,10 +138,7 @@ class DeploymentSchema(NamedSchema, table=True): visualizations: List["DeploymentVisualizationSchema"] = Relationship( back_populates="deployment", - sa_relationship_kwargs=dict( - lazy="selectin", - order_by="CASE WHEN display_order IS NULL THEN 1 ELSE 0 END, display_order, created", - ), + sa_relationship_kwargs=dict(lazy="selectin"), ) @classmethod From 71395816edea0544d0a76c9c399bae88763d7bac Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 27 Sep 2025 18:04:29 +0100 Subject: [PATCH 06/64] few more fixes --- src/zenml/client.py | 29 +++++++++++++++++++-------- src/zenml/zen_stores/sql_zen_store.py | 10 --------- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/src/zenml/client.py b/src/zenml/client.py index d7d7d7c8f3..4e5f8eae46 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -3777,8 +3777,8 @@ def list_deployment_visualizations( *, page: Optional[int] = None, size: Optional[int] = None, - order_by: Optional[str] = None, - sort: Optional[SorterOps] = None, + sort_by: Optional[str] = None, + visualization_index: Optional[int] = None, hydrate: bool = False, ) -> Page[DeploymentVisualizationResponse]: """List curated deployment visualizations for a deployment. @@ -3787,23 +3787,36 @@ def list_deployment_visualizations( deployment_id: The ID of the deployment to list visualizations for. page: Page number for pagination. size: Page size for pagination. - order_by: Field to order by. Defaults to "display_order". - sort: Sort order. Defaults to ascending. + sort_by: Field to order by. Defaults to "display_order" in the filter. + visualization_index: Filter by visualization index. Must be non-negative. hydrate: Flag deciding whether to hydrate the output model(s) by including metadata fields in the response. Returns: A Page[DeploymentVisualizationResponse] containing the visualizations. + + Raises: + ValueError: If visualization_index is negative. """ + if visualization_index is not None and visualization_index < 0: + raise ValueError("visualization_index must be non-negative") + deployment = self.get_deployment(deployment_id) filter_model = DeploymentVisualizationFilter( project=deployment.project_id, deployment=deployment_id, - page=page, - size=size, - order_by=order_by or "display_order", - sort=sort or SorterOps.ASCENDING, ) + + # Only set optional filter params if provided, relying on filter defaults + if page is not None: + filter_model.page = page + if size is not None: + filter_model.size = size + if sort_by is not None: + filter_model.sort_by = sort_by + if visualization_index is not None: + filter_model.visualization_index = visualization_index + return self.zen_store.list_deployment_visualizations( filter_model=filter_model, hydrate=hydrate, diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 0606e4c749..42b1fbed7c 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5551,16 +5551,6 @@ def list_deployment_visualizations( self._set_filter_project_id( filter_model=filter_model, session=session ) - default_sort_by = ( - type(filter_model).model_fields["sort_by"].default - ) - if ( - not filter_model.sort_by - or filter_model.sort_by == default_sort_by - ): - filter_model = filter_model.model_copy( - update={"sort_by": "display_order"} - ) query = select(DeploymentVisualizationSchema) return self.filter_and_paginate( From 2ca3897da1e6766e709ce1f3a0d0392996d08d08 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 27 Sep 2025 18:50:20 +0100 Subject: [PATCH 07/64] fix endpoint --- src/zenml/zen_server/routers/deployment_endpoints.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/src/zenml/zen_server/routers/deployment_endpoints.py b/src/zenml/zen_server/routers/deployment_endpoints.py index 8617e7e401..91d589be49 100644 --- a/src/zenml/zen_server/routers/deployment_endpoints.py +++ b/src/zenml/zen_server/routers/deployment_endpoints.py @@ -202,7 +202,6 @@ def delete_deployment( ) @async_fastapi_endpoint_wrapper def create_deployment_visualization( - deployment_id: UUID, visualization: DeploymentVisualizationRequest, _: AuthContext = Security(authorize), ) -> DeploymentVisualizationResponse: @@ -215,10 +214,6 @@ def create_deployment_visualization( Returns: The created deployment visualization. """ - # Ensure deployment_id matches path parameter - visualization = visualization.model_copy( - update={"deployment_id": deployment_id} - ) return verify_permissions_and_create_entity( request_model=visualization, create_method=zen_store().create_deployment_visualization, @@ -250,11 +245,7 @@ def list_deployment_visualizations( Returns: List of deployment visualization objects for the deployment. """ - # Set deployment filter to the path parameter - visualization_filter_model = visualization_filter_model.model_copy( - update={"deployment": deployment_id} - ) - + visualization_filter_model.deployment = deployment_id return verify_permissions_and_list_entities( filter_model=visualization_filter_model, resource_type=ResourceType.DEPLOYMENT, From 1718b58d511edf19d3ae360cc901f49fd659eec5 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Mon, 29 Sep 2025 13:47:28 +0100 Subject: [PATCH 08/64] fix --- src/zenml/zen_server/routers/deployment_endpoints.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/zenml/zen_server/routers/deployment_endpoints.py b/src/zenml/zen_server/routers/deployment_endpoints.py index 91d589be49..e7daef6775 100644 --- a/src/zenml/zen_server/routers/deployment_endpoints.py +++ b/src/zenml/zen_server/routers/deployment_endpoints.py @@ -202,6 +202,7 @@ def delete_deployment( ) @async_fastapi_endpoint_wrapper def create_deployment_visualization( + deployment_id: UUID, visualization: DeploymentVisualizationRequest, _: AuthContext = Security(authorize), ) -> DeploymentVisualizationResponse: @@ -214,6 +215,14 @@ def create_deployment_visualization( Returns: The created deployment visualization. """ + if ( + visualization.deployment_id + and visualization.deployment_id != deployment_id + ): + raise error_response( + status_code=400, + detail="Deployment ID in request does not match path parameter", + ) return verify_permissions_and_create_entity( request_model=visualization, create_method=zen_store().create_deployment_visualization, From 4bc726a866e585ae90d609eac99a273d2d302b7f Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Mon, 29 Sep 2025 13:48:52 +0100 Subject: [PATCH 09/64] docstring --- src/zenml/zen_server/routers/deployment_endpoints.py | 3 +++ src/zenml/zen_stores/sql_zen_store.py | 7 +++++++ 2 files changed, 10 insertions(+) diff --git a/src/zenml/zen_server/routers/deployment_endpoints.py b/src/zenml/zen_server/routers/deployment_endpoints.py index e7daef6775..d9d7df250f 100644 --- a/src/zenml/zen_server/routers/deployment_endpoints.py +++ b/src/zenml/zen_server/routers/deployment_endpoints.py @@ -214,6 +214,9 @@ def create_deployment_visualization( Returns: The created deployment visualization. + + Raises: + error_response: If the deployment ID in the request does not match the path parameter. """ if ( visualization.deployment_id diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 42b1fbed7c..5a53955043 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5413,6 +5413,9 @@ def _validate_deployment_visualization_index( session: The session to use. artifact_version_id: The ID of the artifact version to validate. visualization_index: The index of the visualization to validate. + + Raises: + IllegalOperationError: If the artifact version does not expose the given visualization index. """ count = session.scalar( select(func.count()) @@ -5437,6 +5440,10 @@ def create_deployment_visualization( Args: visualization: The visualization to create. + Raises: + IllegalOperationError: If the deployment ID in the request does not match the path parameter. + EntityExistsError: If a curated visualization with the same deployment, artifact version, and index already exists. + Returns: The created deployment visualization. """ From 74002e7120ca189d4d3fc3aa050edecf2e01f55b Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Mon, 29 Sep 2025 14:24:35 +0100 Subject: [PATCH 10/64] add migrations file --- .../b77d123bce19_deployment_visualizations.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) create mode 100644 src/zenml/zen_stores/migrations/versions/b77d123bce19_deployment_visualizations.py diff --git a/src/zenml/zen_stores/migrations/versions/b77d123bce19_deployment_visualizations.py b/src/zenml/zen_stores/migrations/versions/b77d123bce19_deployment_visualizations.py new file mode 100644 index 0000000000..83b2d1111e --- /dev/null +++ b/src/zenml/zen_stores/migrations/versions/b77d123bce19_deployment_visualizations.py @@ -0,0 +1,93 @@ +"""deployment visualizations [b77d123bce19]. + +Revision ID: b77d123bce19 +Revises: 0.90.0rc0 +Create Date: 2025-09-29 14:23:48.630888 + +""" + +import sqlalchemy as sa +import sqlmodel +from alembic import op + +# revision identifiers, used by Alembic. +revision = "b77d123bce19" +down_revision = "0.90.0rc0" +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! ### + op.create_table( + "deployment_visualization", + sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column("created", sa.DateTime(), nullable=False), + sa.Column("updated", sa.DateTime(), nullable=False), + sa.Column("project_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "deployment_id", sqlmodel.sql.sqltypes.GUID(), nullable=False + ), + sa.Column( + "artifact_version_id", sqlmodel.sql.sqltypes.GUID(), nullable=False + ), + sa.Column("visualization_index", sa.Integer(), nullable=False), + sa.Column( + "display_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True + ), + sa.Column("display_order", sa.Integer(), nullable=True), + sa.ForeignKeyConstraint( + ["artifact_version_id"], + ["artifact_version.id"], + name="fk_deployment_visualization_artifact_version_id_artifact_version", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["deployment_id"], + ["deployment.id"], + name="fk_deployment_visualization_deployment_id_deployment", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["project_id"], + ["project.id"], + name="fk_deployment_visualization_project_id_project", + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "deployment_id", + "artifact_version_id", + "visualization_index", + name="unique_deployment_visualization", + ), + ) + with op.batch_alter_table( + "deployment_visualization", schema=None + ) as batch_op: + batch_op.create_index( + "ix_deployment_visualization_deployment_id", + ["deployment_id"], + unique=False, + ) + batch_op.create_index( + "ix_deployment_visualization_display_order", + ["display_order"], + unique=False, + ) + + # ### 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( + "deployment_visualization", schema=None + ) as batch_op: + batch_op.drop_index("ix_deployment_visualization_display_order") + batch_op.drop_index("ix_deployment_visualization_deployment_id") + + op.drop_table("deployment_visualization") + # ### end Alembic commands ### From 43677a0f08816f9eb962876e12cfd5fe88332a86 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Mon, 29 Sep 2025 16:05:52 +0100 Subject: [PATCH 11/64] mypy --- src/zenml/zen_server/routers/deployment_endpoints.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/zenml/zen_server/routers/deployment_endpoints.py b/src/zenml/zen_server/routers/deployment_endpoints.py index d9d7df250f..aab4532db2 100644 --- a/src/zenml/zen_server/routers/deployment_endpoints.py +++ b/src/zenml/zen_server/routers/deployment_endpoints.py @@ -216,15 +216,14 @@ def create_deployment_visualization( The created deployment visualization. Raises: - error_response: If the deployment ID in the request does not match the path parameter. + KeyError: If the deployment ID in the request does not match the path parameter. """ if ( visualization.deployment_id and visualization.deployment_id != deployment_id ): - raise error_response( - status_code=400, - detail="Deployment ID in request does not match path parameter", + raise KeyError( + "Deployment ID in request does not match path parameter" ) return verify_permissions_and_create_entity( request_model=visualization, From f05e1ee593f6db9e4b3847bc4550d06d9a012e93 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sun, 5 Oct 2025 22:01:06 +0100 Subject: [PATCH 12/64] Update visualization parameters to disable metadata and resources This change sets the `include_metadata` and `include_resources` parameters to `False` in the `visualization.to_model` method call within the `DeploymentSchema` class. This adjustment ensures that unnecessary metadata and resources are not included in the visualizations. No functional changes are expected as a result of this update. --- src/zenml/zen_stores/schemas/deployment_schemas.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index 1a23b81711..95391fcfd8 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -239,8 +239,8 @@ def to_model( else None, visualizations=[ visualization.to_model( - include_metadata=include_metadata, - include_resources=include_resources, + include_metadata=False, + include_resources=False, include_deployment=False, ) for visualization in (self.visualizations or []) From 291c536c0572aa1262733be7e083adb03f083bbf Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 9 Oct 2025 16:07:30 +0100 Subject: [PATCH 13/64] curated visualizations --- docs/book/how-to/artifacts/visualizations.md | 118 ++++++- src/zenml/client.py | 231 ++++++++++--- src/zenml/constants.py | 2 +- src/zenml/enums.py | 28 ++ src/zenml/models/__init__.py | 34 +- ...ualization.py => curated_visualization.py} | 199 +++++++----- src/zenml/models/v2/core/deployment.py | 8 +- src/zenml/models/v2/core/model.py | 15 + .../models/v2/misc/curated_visualization.py | 33 ++ src/zenml/zen_server/rbac/models.py | 28 +- .../curated_visualization_endpoints.py | 189 +++++++++++ .../routers/deployment_endpoints.py | 224 +++---------- src/zenml/zen_server/zen_server_api.py | 3 +- src/zenml/zen_stores/rest_zen_store.py | 103 +++--- src/zenml/zen_stores/schemas/__init__.py | 8 +- .../schemas/curated_visualization_schemas.py | 288 ++++++++++++++++ .../zen_stores/schemas/deployment_schemas.py | 29 +- .../deployment_visualization_schemas.py | 243 -------------- src/zenml/zen_stores/schemas/model_schemas.py | 24 +- src/zenml/zen_stores/sql_zen_store.py | 302 +++++++++++------ src/zenml/zen_stores/zen_store_interface.py | 105 ++---- .../functional/zen_stores/test_zen_store.py | 307 ++++++++++++++++++ 22 files changed, 1698 insertions(+), 823 deletions(-) rename src/zenml/models/v2/core/{deployment_visualization.py => curated_visualization.py} (57%) create mode 100644 src/zenml/models/v2/misc/curated_visualization.py create mode 100644 src/zenml/zen_server/routers/curated_visualization_endpoints.py create mode 100644 src/zenml/zen_stores/schemas/curated_visualization_schemas.py delete mode 100644 src/zenml/zen_stores/schemas/deployment_visualization_schemas.py diff --git a/docs/book/how-to/artifacts/visualizations.md b/docs/book/how-to/artifacts/visualizations.md index 01cfd41afd..69cf5badb7 100644 --- a/docs/book/how-to/artifacts/visualizations.md +++ b/docs/book/how-to/artifacts/visualizations.md @@ -65,6 +65,122 @@ There are three ways how you can add custom visualizations to the dashboard: * If you are already handling HTML, Markdown, CSV or JSON data in one of your steps, you can have them visualized in just a few lines of code by casting them to a [special class](#visualization-via-special-return-types) inside your step. * If you want to automatically extract visualizations for all artifacts of a certain data type, you can define type-specific visualization logic by [building a custom materializer](#visualization-via-materializers). +### Curated Visualizations Across Resources + +Curated visualizations let you surface a specific artifact visualization across multiple ZenML resources. This is useful when you want to highlight the same dashboard in several places—for example, a model performance report that should be visible from the project overview, the training run, and the production deployment. + +Curated visualizations currently support the following resources: + +- **Projects** – high-level dashboards and KPIs that summarize the state of a project. +- **Deployments** – monitoring pages for deployed pipelines. +- **Models** – evaluation dashboards and health views for registered models. +- **Pipelines** – reusable visual documentation attached to pipeline definitions. +- **Pipeline Runs** – detailed diagnostics for specific executions. +- **Pipeline Snapshots** – configuration/version comparisons for snapshot history. + +You can create a curated visualization programmatically by linking an artifact visualization to one or more resources. The example below attaches a single visualization to multiple resource types, including a project: + +```python +from uuid import UUID + +from zenml.client import Client +from zenml.enums import VisualizationResourceTypes +from zenml.models import CuratedVisualizationResource + +client = Client() +artifact_version_id = UUID("") +project = client.active_project +pipeline = client.list_pipelines().items[0] +pipeline_run = pipeline.runs()[0] +snapshot = pipeline_run.snapshot() +deployment = client.list_deployments().items[0] +model = client.list_models().items[0] + +visualization = client.create_curated_visualization( + artifact_version_id=artifact_version_id, + visualization_index=0, + resources=[ + CuratedVisualizationResource(id=model.id, type=VisualizationResourceTypes.MODEL), + CuratedVisualizationResource(id=project.id, type=VisualizationResourceTypes.PROJECT), + CuratedVisualizationResource(id=deployment.id, type=VisualizationResourceTypes.DEPLOYMENT), + CuratedVisualizationResource(id=pipeline.id, type=VisualizationResourceTypes.PIPELINE), + CuratedVisualizationResource(id=pipeline_run.id, type=VisualizationResourceTypes.PIPELINE_RUN), + CuratedVisualizationResource(id=snapshot.id, type=VisualizationResourceTypes.PIPELINE_SNAPSHOT), + ], + display_name="Project performance dashboard", +) +``` + +To list curated visualizations for a specific resource, you can use the `Client.list_curated_visualizations` convenience parameters: + +```python +client.list_curated_visualizations(project_id=project.id) +client.list_curated_visualizations(deployment_id=deployment.id) +client.list_curated_visualizations(model_id=model.id) +client.list_curated_visualizations(pipeline_id=pipeline.id) +client.list_curated_visualizations(pipeline_run_id=pipeline_run.id) +client.list_curated_visualizations(pipeline_snapshot_id=snapshot.id) +``` + +Each call returns a `Page[CuratedVisualizationResponse]` object so you can iterate through the visualizations or fetch the hydrated version for additional metadata. + +#### Updating curated visualizations + +Once you've created a curated visualization, you can update its display name or order using `Client.update_curated_visualization`: + +```python +from uuid import UUID + +client.update_curated_visualization( + visualization_id=UUID(""), + display_name="Updated Dashboard Title", + display_order=10, +) +``` + +When a visualization is no longer relevant, you can remove it entirely: + +```python +client.delete_curated_visualization(visualization_id=UUID("")) +``` + +#### Controlling display order + +The optional `display_order` field determines how visualizations are sorted when displayed. Visualizations with lower order values appear first, while those with `None` (the default) appear at the end in creation order. + +When setting display orders, consider leaving gaps between values (e.g., 10, 20, 30 instead of 1, 2, 3) to make it easier to insert new visualizations later without renumbering everything: + +```python +# Leave gaps for future insertions +visualization_a = client.create_curated_visualization( + artifact_version_id=artifact_version_id, + visualization_index=0, + resources=[...], + display_order=10, # Primary dashboard +) + +visualization_b = client.create_curated_visualization( + artifact_version_id=artifact_version_id, + visualization_index=1, + resources=[...], + display_order=20, # Secondary metrics +) + +# Later, easily insert between them +visualization_c = client.create_curated_visualization( + artifact_version_id=artifact_version_id, + visualization_index=2, + resources=[...], + display_order=15, # Now appears between A and B +) +``` + +#### RBAC visibility + +Curated visualizations respect the access permissions of every resource they're linked to. A user can only see a curated visualization if they have read access to **all** the resources it targets. If a user lacks permission for any linked resource, the visualization will be hidden from their view. + +For example, if you create a visualization linked to both a project and a deployment, users must have read access to both the project and that specific deployment to see the visualization. This ensures that curated visualizations never inadvertently expose information from resources a user shouldn't access. + ### Visualization via Special Return Types If you already have HTML, Markdown, CSV or JSON data available as a string inside your step, you can simply cast them to one of the following types and return them from your step: @@ -257,4 +373,4 @@ steps: Visualizing artifacts is a powerful way to gain insights from your ML pipelines. ZenML's built-in visualization capabilities make it easy to understand your data and model outputs, identify issues, and communicate results. -By leveraging these visualization tools, you can better understand your ML workflows, debug problems more effectively, and make more informed decisions about your models. \ No newline at end of file +By leveraging these visualization tools, you can better understand your ML workflows, debug problems more effectively, and make more informed decisions about your models. diff --git a/src/zenml/client.py b/src/zenml/client.py index 379c5358be..72f669b353 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -72,6 +72,7 @@ StackComponentType, StoreType, TaggableResourceTypes, + VisualizationResourceTypes, ) from zenml.exceptions import ( AuthorizationException, @@ -108,12 +109,13 @@ ComponentRequest, ComponentResponse, ComponentUpdate, + CuratedVisualizationFilter, + CuratedVisualizationRequest, + CuratedVisualizationResource, + CuratedVisualizationResponse, + CuratedVisualizationUpdate, DeploymentFilter, DeploymentResponse, - DeploymentVisualizationFilter, - DeploymentVisualizationRequest, - DeploymentVisualizationResponse, - DeploymentVisualizationUpdate, EventSourceFilter, EventSourceRequest, EventSourceResponse, @@ -3739,6 +3741,65 @@ def get_deployment( hydrate=hydrate, ) + def create_curated_visualization( + self, + artifact_version_id: UUID, + visualization_index: int, + *, + resources: List[CuratedVisualizationResource], + project_id: Optional[UUID] = None, + display_name: Optional[str] = None, + display_order: Optional[int] = None, + ) -> CuratedVisualizationResponse: + """Create a curated visualization associated with arbitrary resources. + + Curated visualizations can be attached to any combination of the following + ZenML resource types to provide contextual dashboards throughout the ML + lifecycle: + + - **Deployments** (VisualizationResourceTypes.DEPLOYMENT): Surface on + deployment monitoring dashboards + - **Pipelines** (VisualizationResourceTypes.PIPELINE): Associate with + pipeline definitions + - **Pipeline Runs** (VisualizationResourceTypes.PIPELINE_RUN): Attach to + specific execution runs + - **Pipeline Snapshots** (VisualizationResourceTypes.PIPELINE_SNAPSHOT): + Link to captured pipeline configurations + + A single visualization can be linked to multiple resources across different + types. For example, attach a model performance dashboard to both the + deployment and the pipeline run that produced the deployed model. + + Args: + artifact_version_id: The ID of the artifact version containing the visualization. + visualization_index: The index of the visualization within the artifact version. + resources: One or more resources to associate with the visualization. + Each entry should be a `CuratedVisualizationResource` containing + the resource ID and type (e.g., DEPLOYMENT, PIPELINE, PIPELINE_RUN, + PIPELINE_SNAPSHOT). + project_id: The ID of the project to associate with the visualization. + display_name: The display name of the visualization. + display_order: The display order of the visualization. + + Returns: + The created curated visualization. + + Raises: + ValueError: If resources list is empty. + """ + if not resources: + raise ValueError("resources must not be empty") + + request = CuratedVisualizationRequest( + project=project_id or self.active_project.id, + artifact_version_id=artifact_version_id, + visualization_index=visualization_index, + display_name=display_name, + display_order=display_order, + resources=resources, + ) + return self.zen_store.create_curated_visualization(request) + def add_visualization_to_deployment( self, deployment_id: UUID, @@ -3747,8 +3808,8 @@ def add_visualization_to_deployment( *, display_name: Optional[str] = None, display_order: Optional[int] = None, - ) -> DeploymentVisualizationResponse: - """Curate a deployment visualization. + ) -> CuratedVisualizationResponse: + """Attach a curated visualization to a deployment. Args: deployment_id: The ID of the deployment to add visualization to. @@ -3758,56 +3819,136 @@ def add_visualization_to_deployment( display_order: Optional display order for sorting visualizations. Returns: - The created deployment visualization. + The created curated visualization. """ deployment = self.get_deployment(deployment_id) - request = DeploymentVisualizationRequest( - project=deployment.project_id, - deployment_id=deployment_id, + return self.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=visualization_index, + resources=[ + CuratedVisualizationResource( + id=deployment_id, + type=VisualizationResourceTypes.DEPLOYMENT, + ) + ], + project_id=deployment.project_id, display_name=display_name, display_order=display_order, ) - return self.zen_store.create_deployment_visualization(request) - def list_deployment_visualizations( + def list_curated_visualizations( self, - deployment_id: UUID, *, + resource_type: Optional[VisualizationResourceTypes] = None, + resource_id: Optional[UUID] = None, + deployment_id: Optional[UUID] = None, + model_id: Optional[UUID] = None, + pipeline_id: Optional[UUID] = None, + pipeline_run_id: Optional[UUID] = None, + pipeline_snapshot_id: Optional[UUID] = None, + project_id: Optional[UUID] = None, page: Optional[int] = None, size: Optional[int] = None, sort_by: Optional[str] = None, visualization_index: Optional[int] = None, hydrate: bool = False, - ) -> Page[DeploymentVisualizationResponse]: - """List curated deployment visualizations for a deployment. - - Args: - deployment_id: The ID of the deployment to list visualizations for. - page: Page number for pagination. - size: Page size for pagination. - sort_by: Field to order by. Defaults to "display_order" in the filter. - visualization_index: Filter by visualization index. Must be non-negative. + ) -> Page[CuratedVisualizationResponse]: + """List curated visualizations, optionally scoped to a resource. + + This method supports filtering by any of the following resource types: + - Deployments (use `deployment_id` parameter) + - Models (use `model_id` parameter) + - Pipelines (use `pipeline_id` parameter) + - Pipeline Runs (use `pipeline_run_id` parameter) + - Pipeline Snapshots (use `pipeline_snapshot_id` parameter) + - Projects (use `project_id` parameter) + + Alternatively, you can use `resource_type` and `resource_id` directly + for more flexible filtering. + + Args: + resource_type: The type of the resource to filter by. + resource_id: The ID of the resource to filter by. + deployment_id: Convenience parameter to filter by deployment. + model_id: Convenience parameter to filter by model. + pipeline_id: Convenience parameter to filter by pipeline. + pipeline_run_id: Convenience parameter to filter by pipeline run. + pipeline_snapshot_id: Convenience parameter to filter by pipeline snapshot. + project_id: Convenience parameter to filter by project. + page: The page of items. + size: The maximum size of all pages. + sort_by: The column to sort by. + visualization_index: The index of the visualization to filter by. hydrate: Flag deciding whether to hydrate the output model(s) by including metadata fields in the response. Returns: - A Page[DeploymentVisualizationResponse] containing the visualizations. + A page of curated visualizations. Raises: - ValueError: If visualization_index is negative. - """ + ValueError: If multiple resource ID parameters are provided, if + visualization_index is negative, or if a convenience parameter + conflicts with explicitly provided resource_type or resource_id. + """ + # Build convenience params dict mapping parameter names to (value, resource_type) tuples + convenience_params = { + "deployment_id": (deployment_id, VisualizationResourceTypes.DEPLOYMENT), + "model_id": (model_id, VisualizationResourceTypes.MODEL), + "pipeline_id": (pipeline_id, VisualizationResourceTypes.PIPELINE), + "pipeline_run_id": (pipeline_run_id, VisualizationResourceTypes.PIPELINE_RUN), + "pipeline_snapshot_id": (pipeline_snapshot_id, VisualizationResourceTypes.PIPELINE_SNAPSHOT), + "project_id": (project_id, VisualizationResourceTypes.PROJECT), + } + + # Filter to only provided parameters + provided = { + param_name: (param_value, param_type) + for param_name, (param_value, param_type) in convenience_params.items() + if param_value is not None + } + + if len(provided) > 1: + param_names = list(provided.keys()) + raise ValueError( + f"Only one resource ID parameter can be specified at a time. " + f"Got: {', '.join(param_names)}" + ) + + # Validate consistency between convenience parameters and explicit arguments + if provided: + param_name, (param_id, param_type) = list(provided.items())[0] + + # Check for resource_type mismatch + if resource_type is not None and resource_type != param_type: + raise ValueError( + f"Conflicting resource type: convenience parameter '{param_name}' " + f"implies resource_type={param_type.value}, but " + f"resource_type={resource_type.value} was explicitly provided." + ) + + # Check for resource_id mismatch + if resource_id is not None and resource_id != param_id: + raise ValueError( + f"Conflicting resource ID: convenience parameter '{param_name}' " + f"specifies resource_id={param_id}, but a different " + f"resource_id={resource_id} was explicitly provided." + ) + + # Auto-set resource_type and resource_id from convenience parameters + resource_type = resource_type or param_type + resource_id = resource_id or param_id + if visualization_index is not None and visualization_index < 0: raise ValueError("visualization_index must be non-negative") - deployment = self.get_deployment(deployment_id) - filter_model = DeploymentVisualizationFilter( - project=deployment.project_id, - deployment=deployment_id, + filter_model = CuratedVisualizationFilter( + project=self.active_project.id ) + if resource_type is not None: + filter_model.resource_type = resource_type + if resource_id is not None: + filter_model.resource_id = resource_id - # Only set optional filter params if provided, relying on filter defaults if page is not None: filter_model.page = page if size is not None: @@ -3817,47 +3958,45 @@ def list_deployment_visualizations( if visualization_index is not None: filter_model.visualization_index = visualization_index - return self.zen_store.list_deployment_visualizations( + return self.zen_store.list_curated_visualizations( filter_model=filter_model, hydrate=hydrate, ) - def update_deployment_visualization( + def update_curated_visualization( self, - deployment_visualization_id: UUID, + visualization_id: UUID, *, display_name: Optional[str] = None, display_order: Optional[int] = None, - ) -> DeploymentVisualizationResponse: - """Update display metadata for a curated deployment visualization. + ) -> CuratedVisualizationResponse: + """Update display metadata for a curated visualization. Args: - deployment_visualization_id: The ID of the deployment visualization to update. + visualization_id: The ID of the curated visualization to update. display_name: New display name for the visualization. display_order: New display order for the visualization. Returns: The updated deployment visualization. """ - update_model = DeploymentVisualizationUpdate( + update_model = CuratedVisualizationUpdate( display_name=display_name, display_order=display_order, ) - return self.zen_store.update_deployment_visualization( - deployment_visualization_id=deployment_visualization_id, + return self.zen_store.update_curated_visualization( + visualization_id=visualization_id, visualization_update=update_model, ) - def delete_deployment_visualization( - self, deployment_visualization_id: UUID - ) -> None: - """Delete a curated deployment visualization. + def delete_curated_visualization(self, visualization_id: UUID) -> None: + """Delete a curated visualization. Args: - deployment_visualization_id: The ID of the deployment visualization to delete. + visualization_id: The ID of the curated visualization to delete. """ - self.zen_store.delete_deployment_visualization( - deployment_visualization_id=deployment_visualization_id + self.zen_store.delete_curated_visualization( + visualization_id=visualization_id ) def list_deployments( diff --git a/src/zenml/constants.py b/src/zenml/constants.py index 25b2ebd6d7..1d2ce0a491 100644 --- a/src/zenml/constants.py +++ b/src/zenml/constants.py @@ -404,7 +404,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int: PIPELINE_CONFIGURATION = "/pipeline-configuration" PIPELINE_DEPLOYMENTS = "/pipeline_deployments" DEPLOYMENTS = "/deployments" -DEPLOYMENT_VISUALIZATIONS = "/deployment_visualizations" +CURATED_VISUALIZATIONS = "/curated_visualizations" PIPELINE_SNAPSHOTS = "/pipeline_snapshots" PIPELINES = "/pipelines" PIPELINE_SPEC = "/pipeline-spec" diff --git a/src/zenml/enums.py b/src/zenml/enums.py index 93edbba278..d487463a75 100644 --- a/src/zenml/enums.py +++ b/src/zenml/enums.py @@ -418,6 +418,34 @@ class MetadataResourceTypes(StrEnum): SCHEDULE = "schedule" +class VisualizationResourceTypes(StrEnum): + """Resource types that support curated visualizations. + + Curated visualizations can be attached to these ZenML resources to provide + contextual dashboards and visual insights throughout the ML lifecycle: + + - **DEPLOYMENT**: Server-side pipeline deployments - surface visualizations + on deployment monitoring dashboards and status pages + - **MODEL**: ZenML model entities - surface model evaluation dashboards and + performance summaries directly on the model detail pages + - **PIPELINE**: Pipeline definitions - associate visualizations with pipeline + configurations for reusable visual documentation + - **PIPELINE_RUN**: Pipeline execution runs - attach visualizations to specific + run results for detailed analysis and debugging + - **PIPELINE_SNAPSHOT**: Pipeline snapshots - link visualizations to captured + pipeline configurations for version comparison and historical analysis + - **PROJECT**: Project-level overviews - provide high-level project dashboards + and KPI visualizations for cross-pipeline insights + """ + + DEPLOYMENT = "deployment" # Server-side pipeline deployments + MODEL = "model" # ZenML models + PIPELINE = "pipeline" # Pipeline definitions + PIPELINE_RUN = "pipeline_run" # Execution runs + PIPELINE_SNAPSHOT = "pipeline_snapshot" # Snapshot configurations + PROJECT = "project" # Project-level dashboards + + class SecretResourceTypes(StrEnum): """All possible resource types for adding secrets.""" diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index 4d2e510008..eedc8b0a6f 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -164,14 +164,14 @@ DeploymentResponseMetadata, DeploymentResponseResources, ) -from zenml.models.v2.core.deployment_visualization import ( - DeploymentVisualizationFilter, - DeploymentVisualizationRequest, - DeploymentVisualizationResponse, - DeploymentVisualizationResponseBody, - DeploymentVisualizationResponseMetadata, - DeploymentVisualizationResponseResources, - DeploymentVisualizationUpdate, +from zenml.models.v2.core.curated_visualization import ( + CuratedVisualizationFilter, + CuratedVisualizationRequest, + CuratedVisualizationResponse, + CuratedVisualizationResponseBody, + CuratedVisualizationResponseMetadata, + CuratedVisualizationResponseResources, + CuratedVisualizationUpdate, ) from zenml.models.v2.core.device import ( OAuthDeviceUpdate, @@ -431,6 +431,9 @@ RunMetadataEntry, RunMetadataResource, ) +from zenml.models.v2.misc.curated_visualization import ( + CuratedVisualizationResource, +) from zenml.models.v2.misc.server_models import ( ServerModel, ServerDatabaseType, @@ -662,13 +665,13 @@ "DeploymentResponseBody", "DeploymentResponseMetadata", "DeploymentResponseResources", - "DeploymentVisualizationFilter", - "DeploymentVisualizationRequest", - "DeploymentVisualizationResponse", - "DeploymentVisualizationResponseBody", - "DeploymentVisualizationResponseMetadata", - "DeploymentVisualizationResponseResources", - "DeploymentVisualizationUpdate", + "CuratedVisualizationFilter", + "CuratedVisualizationRequest", + "CuratedVisualizationResponse", + "CuratedVisualizationResponseBody", + "CuratedVisualizationResponseMetadata", + "CuratedVisualizationResponseResources", + "CuratedVisualizationUpdate", "EventSourceFlavorResponse", "EventSourceFlavorResponseBody", "EventSourceFlavorResponseMetadata", @@ -887,6 +890,7 @@ "ResourcesInfo", "RunMetadataEntry", "RunMetadataResource", + "CuratedVisualizationResource", "ProjectStatistics", "PipelineRunDAG", "ExceptionInfo", diff --git a/src/zenml/models/v2/core/deployment_visualization.py b/src/zenml/models/v2/core/curated_visualization.py similarity index 57% rename from src/zenml/models/v2/core/deployment_visualization.py rename to src/zenml/models/v2/core/curated_visualization.py index 0a8c2ee954..a0df0503c1 100644 --- a/src/zenml/models/v2/core/deployment_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express # or implied. See the License for the specific language governing # permissions and limitations under the License. -"""Models representing deployment visualizations.""" +"""Models representing curated visualizations.""" from typing import ( TYPE_CHECKING, @@ -24,9 +24,9 @@ ) from uuid import UUID -from pydantic import Field +from pydantic import Field, model_validator -from zenml.constants import STR_FIELD_MAX_LENGTH +from zenml.enums import VisualizationResourceTypes from zenml.models.v2.base.base import BaseUpdate from zenml.models.v2.base.filter import AnyQuery from zenml.models.v2.base.scoped import ( @@ -37,12 +37,14 @@ ProjectScopedResponseMetadata, ProjectScopedResponseResources, ) +from zenml.models.v2.misc.curated_visualization import ( + CuratedVisualizationResource, +) if TYPE_CHECKING: from sqlalchemy.sql.elements import ColumnElement from zenml.models.v2.core.artifact_version import ArtifactVersionResponse - from zenml.models.v2.core.deployment import DeploymentResponse from zenml.zen_stores.schemas.base_schemas import BaseSchema AnySchema = TypeVar("AnySchema", bound=BaseSchema) @@ -51,43 +53,79 @@ # ------------------ Request Model ------------------ -class DeploymentVisualizationRequest(ProjectScopedRequest): - """Request model for deployment visualizations.""" +class CuratedVisualizationRequest(ProjectScopedRequest): + """Request model for curated visualizations. + + Curated visualizations can be attached to any combination of the following + resource types: + - **Deployments**: Surface visualizations on deployment dashboards + - **Models**: Highlight evaluation dashboards and monitoring views next to + registered models + - **Pipelines**: Associate visualizations with pipeline definitions + - **Pipeline Runs**: Attach visualizations to specific execution runs + - **Pipeline Snapshots**: Link visualizations to snapshot configurations + - **Projects**: Provide high-level project dashboards and KPI overviews + + A single visualization can be linked to multiple resources of different types, + enabling reuse across the ML workflow. For example, a model performance + dashboard could be attached to both a deployment and the pipeline run that + produced the deployed model. + """ - deployment_id: UUID = Field( - title="The deployment ID.", - description="The ID of the deployment associated with the visualization.", - ) artifact_version_id: UUID = Field( title="The artifact version ID.", - description="The ID of the artifact version associated with the visualization.", + description="Identifier of the artifact version providing the visualization.", ) visualization_index: int = Field( ge=0, title="The visualization index.", - description="The index of the visualization within the artifact version.", + description="Index of the visualization within the artifact version payload.", ) display_name: Optional[str] = Field( default=None, title="The display name of the visualization.", - max_length=STR_FIELD_MAX_LENGTH, ) display_order: Optional[int] = Field( default=None, title="The display order of the visualization.", ) + resources: List[CuratedVisualizationResource] = Field( + title="Resources associated with this visualization.", + description=( + "List of resources (deployments, models, pipelines, pipeline runs, " + "pipeline snapshots, projects) that should surface this visualization. " + "Must include at least one resource. Multiple resources of " + "different types can be specified to reuse visualizations " + "across the ML workflow." + ), + ) + + @model_validator(mode="after") + def validate_resources(self) -> "CuratedVisualizationRequest": + """Ensure that at least one resource is associated with the visualization. + + Returns: + The validated request instance. + + Raises: + ValueError: If no resources are provided. + """ + if not self.resources: + raise ValueError( + "Curated visualizations must be associated with at least one resource." + ) + return self # ------------------ Update Model ------------------ -class DeploymentVisualizationUpdate(BaseUpdate): - """Update model for deployment visualizations.""" +class CuratedVisualizationUpdate(BaseUpdate): + """Update model for curated visualizations.""" display_name: Optional[str] = Field( default=None, title="The new display name of the visualization.", - max_length=STR_FIELD_MAX_LENGTH, ) display_order: Optional[int] = Field( default=None, @@ -98,20 +136,16 @@ class DeploymentVisualizationUpdate(BaseUpdate): # ------------------ Response Model ------------------ -class DeploymentVisualizationResponseBody(ProjectScopedResponseBody): - """Response body for deployment visualizations.""" +class CuratedVisualizationResponseBody(ProjectScopedResponseBody): + """Response body for curated visualizations.""" - deployment_id: UUID = Field( - title="The deployment ID.", - description="The ID of the deployment associated with the visualization.", - ) artifact_version_id: UUID = Field( title="The artifact version ID.", - description="The ID of the artifact version associated with the visualization.", + description="Identifier of the artifact version providing the visualization.", ) visualization_index: int = Field( title="The visualization index.", - description="The index of the visualization within the artifact version.", + description="Index of the visualization within the artifact version payload.", ) display_name: Optional[str] = Field( default=None, @@ -121,58 +155,47 @@ class DeploymentVisualizationResponseBody(ProjectScopedResponseBody): default=None, title="The display order of the visualization.", ) + resources: List[CuratedVisualizationResource] = Field( + default_factory=list, + title="Resources exposing the visualization.", + ) -class DeploymentVisualizationResponseMetadata(ProjectScopedResponseMetadata): - """Response metadata for deployment visualizations.""" +class CuratedVisualizationResponseMetadata(ProjectScopedResponseMetadata): + """Response metadata for curated visualizations.""" -class DeploymentVisualizationResponseResources(ProjectScopedResponseResources): - """Response resources for deployment visualizations.""" +class CuratedVisualizationResponseResources(ProjectScopedResponseResources): + """Response resources for curated visualizations.""" - deployment: Optional["DeploymentResponse"] = Field( - default=None, - title="The deployment.", - description="The deployment associated with the visualization.", - ) artifact_version: Optional["ArtifactVersionResponse"] = Field( default=None, title="The artifact version.", - description="The artifact version associated with the visualization.", + description="Artifact version from which the visualization originates.", ) -class DeploymentVisualizationResponse( +class CuratedVisualizationResponse( ProjectScopedResponse[ - DeploymentVisualizationResponseBody, - DeploymentVisualizationResponseMetadata, - DeploymentVisualizationResponseResources, + CuratedVisualizationResponseBody, + CuratedVisualizationResponseMetadata, + CuratedVisualizationResponseResources, ] ): - """Response model for deployment visualizations.""" + """Response model for curated visualizations.""" - def get_hydrated_version(self) -> "DeploymentVisualizationResponse": - """Get the hydrated version of this deployment visualization. + def get_hydrated_version(self) -> "CuratedVisualizationResponse": + """Get the hydrated version of this curated visualization. Returns: - an instance of the same entity with the metadata and resources fields - attached. + A hydrated instance of the same entity. """ from zenml.client import Client client = Client() - return client.zen_store.get_deployment_visualization(self.id) + return client.zen_store.get_curated_visualization(self.id) # Helper properties - @property - def deployment_id(self) -> UUID: - """The deployment ID. - - Returns: - The deployment ID. - """ - return self.get_body().deployment_id - @property def artifact_version_id(self) -> UUID: """The artifact version ID. @@ -210,34 +233,33 @@ def display_order(self) -> Optional[int]: return self.get_body().display_order @property - def deployment(self) -> Optional["DeploymentResponse"]: - """The deployment. + def artifact_version(self) -> Optional["ArtifactVersionResponse"]: + """The artifact version resource. Returns: - The deployment. + The artifact version resource if included. """ - return self.get_resources().deployment + return self.get_resources().artifact_version - @property - def artifact_version(self) -> Optional["ArtifactVersionResponse"]: - """The artifact version. + def visualization_resources(self) -> List[CuratedVisualizationResource]: + """Return the resources exposing this visualization. Returns: - The artifact version. + List of associated resources. """ - return self.get_resources().artifact_version + return self.get_body().resources # ------------------ Filter Model ------------------ -class DeploymentVisualizationFilter(ProjectScopedFilter): - """Model to enable advanced filtering of deployment visualizations.""" +class CuratedVisualizationFilter(ProjectScopedFilter): + """Model to enable advanced filtering of curated visualizations.""" FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [ *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS, - "deployment", - "artifact_version", + "resource_id", + "resource_type", ] CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [ *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS, @@ -246,39 +268,39 @@ class DeploymentVisualizationFilter(ProjectScopedFilter): "updated", "visualization_index", ] - CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [ - *ProjectScopedFilter.CLI_EXCLUDE_FIELDS, - ] - # Set default sort_by to display_order sort_by: str = Field( default="display_order", description="Which column to sort by.", ) - deployment: Optional[UUID] = Field( - default=None, - description="ID of the deployment associated with the visualization.", - ) artifact_version: Optional[UUID] = Field( default=None, description="ID of the artifact version associated with the visualization.", ) visualization_index: Optional[int] = Field( default=None, - description="Index of the visualization within the artifact version.", + description="Index of the visualization within the artifact version payload.", ) display_order: Optional[int] = Field( default=None, description="Display order of the visualization.", ) + resource_type: Optional[VisualizationResourceTypes] = Field( + default=None, + description="Type of the resource exposing the visualization.", + ) + resource_id: Optional[UUID] = Field( + default=None, + description="ID of the resource exposing the visualization.", + ) def apply_sorting( self, query: AnyQuery, table: Type["AnySchema"], ) -> AnyQuery: - """Apply sorting to the deployment visualization query. + """Apply sorting to the curated visualization query. Args: query: The query to which to apply the sorting. @@ -293,7 +315,6 @@ def apply_sorting( sort_by, operand = self.sorting_params - # Special handling for display_order with nulls first/last if sort_by == "display_order": column = getattr(table, "display_order") if operand == SorterOps.DESCENDING: @@ -306,16 +327,15 @@ def apply_sorting( query.order_by(asc(column).nulls_first(), asc(table.id)), ) - # Direct sorting for other custom columns if sort_by in {"created", "updated", "visualization_index"}: column = getattr(table, sort_by) if operand == SorterOps.DESCENDING: return cast( - AnyQuery, query.order_by(desc(column), asc(table.id)) + AnyQuery, + query.order_by(desc(column), asc(table.id)), ) return cast(AnyQuery, query.order_by(asc(column), asc(table.id))) - # Delegate to parent for other columns return super().apply_sorting(query=query, table=table) def get_custom_filters( @@ -331,16 +351,19 @@ def get_custom_filters( """ custom_filters = super().get_custom_filters(table) - if self.deployment: - deployment_filter = ( - getattr(table, "deployment_id") == self.deployment - ) - custom_filters.append(deployment_filter) - if self.artifact_version: - artifact_version_filter = ( + custom_filters.append( getattr(table, "artifact_version_id") == self.artifact_version ) - custom_filters.append(artifact_version_filter) + if self.visualization_index is not None: + custom_filters.append( + getattr(table, "visualization_index") + == self.visualization_index + ) + if self.display_order is not None: + custom_filters.append( + getattr(table, "display_order") == self.display_order + ) + # resource-based filtering is handled within the store implementation return custom_filters diff --git a/src/zenml/models/v2/core/deployment.py b/src/zenml/models/v2/core/deployment.py index ead51ae645..9df3f5d4e6 100644 --- a/src/zenml/models/v2/core/deployment.py +++ b/src/zenml/models/v2/core/deployment.py @@ -46,8 +46,8 @@ from sqlalchemy.sql.elements import ColumnElement from zenml.models.v2.core.component import ComponentResponse - from zenml.models.v2.core.deployment_visualization import ( - DeploymentVisualizationResponse, + from zenml.models.v2.core.curated_visualization import ( + CuratedVisualizationResponse, ) from zenml.models.v2.core.pipeline import PipelineResponse from zenml.models.v2.core.pipeline_snapshot import ( @@ -207,7 +207,7 @@ class DeploymentResponseResources(ProjectScopedResponseResources): tags: List["TagResponse"] = Field( title="Tags associated with the deployment.", ) - visualizations: List["DeploymentVisualizationResponse"] = Field( + visualizations: List["CuratedVisualizationResponse"] = Field( default_factory=list, title="Curated deployment visualizations.", ) @@ -312,7 +312,7 @@ def tags(self) -> List["TagResponse"]: return self.get_resources().tags @property - def visualizations(self) -> List["DeploymentVisualizationResponse"]: + def visualizations(self) -> List["CuratedVisualizationResponse"]: """The visualizations of the deployment. Returns: diff --git a/src/zenml/models/v2/core/model.py b/src/zenml/models/v2/core/model.py index eaffa6a68d..32eb7dd2a9 100644 --- a/src/zenml/models/v2/core/model.py +++ b/src/zenml/models/v2/core/model.py @@ -45,6 +45,9 @@ if TYPE_CHECKING: from zenml.model.model import Model + from zenml.models.v2.core.curated_visualization import ( + CuratedVisualizationResponse, + ) from zenml.models.v2.core.tag import TagResponse from zenml.zen_stores.schemas import BaseSchema @@ -185,6 +188,10 @@ class ModelResponseResources(ProjectScopedResponseResources): ) latest_version_name: Optional[str] = None latest_version_id: Optional[UUID] = None + visualizations: List["CuratedVisualizationResponse"] = Field( + default_factory=list, + title="Curated visualizations associated with the model.", + ) class ModelResponse( @@ -199,6 +206,14 @@ class ModelResponse( max_length=STR_FIELD_MAX_LENGTH, ) + def visualizations(self) -> List["CuratedVisualizationResponse"]: + """Return curated visualizations linked to the model. + + Returns: + A list of curated visualization responses. + """ + return self.get_resources().visualizations + def get_hydrated_version(self) -> "ModelResponse": """Get the hydrated version of this model. diff --git a/src/zenml/models/v2/misc/curated_visualization.py b/src/zenml/models/v2/misc/curated_visualization.py new file mode 100644 index 0000000000..f09594fa4c --- /dev/null +++ b/src/zenml/models/v2/misc/curated_visualization.py @@ -0,0 +1,33 @@ +# Copyright (c) ZenML GmbH 2025. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Models describing curated visualization resources.""" + +from uuid import UUID + +from pydantic import BaseModel, Field + +from zenml.enums import VisualizationResourceTypes + + +class CuratedVisualizationResource(BaseModel): + """Resource reference associated with a curated visualization.""" + + id: UUID = Field( + title="The ID of the resource.", + description="Identifier of the resource associated with the visualization.", + ) + type: VisualizationResourceTypes = Field( + title="The type of the resource.", + description="Classification of the resource associated with the visualization.", + ) diff --git a/src/zenml/zen_server/rbac/models.py b/src/zenml/zen_server/rbac/models.py index bc8984e403..ba4cd4dcc4 100644 --- a/src/zenml/zen_server/rbac/models.py +++ b/src/zenml/zen_server/rbac/models.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """RBAC model classes.""" -from typing import Optional +from typing import TYPE_CHECKING, Optional from uuid import UUID from pydantic import ( @@ -24,6 +24,9 @@ from zenml.utils.enum_utils import StrEnum +if TYPE_CHECKING: + from zenml.enums import VisualizationResourceTypes + class Action(StrEnum): """RBAC actions.""" @@ -80,6 +83,29 @@ class ResourceType(StrEnum): # Deactivated for now # USER = "user" + @classmethod + def from_visualization_type( + cls, visualization_type: "VisualizationResourceTypes" + ) -> Optional["ResourceType"]: + """Map visualization resource types to RBAC resource types. + + Args: + visualization_type: The type of the visualization to map. + + Returns: + The RBAC resource type associated with the visualization type. + """ + from zenml.enums import VisualizationResourceTypes + + mapping = { + VisualizationResourceTypes.DEPLOYMENT: cls.DEPLOYMENT, + VisualizationResourceTypes.MODEL: cls.MODEL, + VisualizationResourceTypes.PIPELINE: cls.PIPELINE, + VisualizationResourceTypes.PIPELINE_RUN: cls.PIPELINE_RUN, + VisualizationResourceTypes.PIPELINE_SNAPSHOT: cls.PIPELINE_SNAPSHOT, + } + return mapping.get(visualization_type) + def is_project_scoped(self) -> bool: """Check if a resource type is project scoped. diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py new file mode 100644 index 0000000000..f10fdea349 --- /dev/null +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -0,0 +1,189 @@ +"""REST API endpoints for curated visualizations.""" + +from uuid import UUID + +from fastapi import APIRouter, Depends, Security + +from zenml.constants import API, CURATED_VISUALIZATIONS, VERSION_1 +from zenml.models import ( + CuratedVisualizationFilter, + CuratedVisualizationRequest, + CuratedVisualizationResponse, + CuratedVisualizationUpdate, + Page, +) +from zenml.zen_server.auth import AuthContext, authorize +from zenml.zen_server.exceptions import error_response +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, +) +from zenml.zen_server.utils import ( + async_fastapi_endpoint_wrapper, + make_dependable, + zen_store, +) + +router = APIRouter( + prefix=API + VERSION_1 + CURATED_VISUALIZATIONS, + tags=["curated_visualizations"], + responses={401: error_response, 404: error_response, 422: error_response}, +) + + +@router.post( + "", + responses={ + 401: error_response, + 404: error_response, + 409: error_response, + 422: error_response, + }, +) +@async_fastapi_endpoint_wrapper +def create_curated_visualization( + visualization: CuratedVisualizationRequest, + _: AuthContext = Security(authorize), +) -> CuratedVisualizationResponse: + """Create a curated visualization. + + Args: + visualization: The curated visualization to create. + + Returns: + The created curated visualization. + """ + return verify_permissions_and_create_entity( + request_model=visualization, + create_method=zen_store().create_curated_visualization, + ) + + +@router.get( + "", + responses={401: error_response, 404: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper(deduplicate=True) +def list_curated_visualizations( + visualization_filter_model: CuratedVisualizationFilter = Depends( + make_dependable(CuratedVisualizationFilter) + ), + hydrate: bool = False, + _: AuthContext = Security(authorize), +) -> Page[CuratedVisualizationResponse]: + """List curated visualizations. + + Args: + visualization_filter_model: Filter model used for pagination, sorting, + filtering. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + A page of curated visualizations. + """ + resource_type = None + if visualization_filter_model.resource_type: + resource_type = ResourceType.from_visualization_type( + visualization_filter_model.resource_type + ) + + if resource_type is None: + # No concrete resource type - call zen store directly and batch verify permissions + page = zen_store().list_curated_visualizations( + filter_model=visualization_filter_model, + hydrate=hydrate, + ) + batch_verify_permissions_for_models(page.items, action=Action.READ) + return dehydrate_page(page) + else: + # Concrete resource type available - use standard RBAC flow + return verify_permissions_and_list_entities( + filter_model=visualization_filter_model, + resource_type=resource_type, + list_method=zen_store().list_curated_visualizations, + hydrate=hydrate, + ) + + +@router.get( + "/{visualization_id}", + responses={401: error_response, 404: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper(deduplicate=True) +def get_curated_visualization( + visualization_id: UUID, + hydrate: bool = True, + _: AuthContext = Security(authorize), +) -> CuratedVisualizationResponse: + """Retrieve a curated visualization by ID. + + Args: + visualization_id: The ID of the curated visualization to retrieve. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + The curated visualization with the given ID. + """ + return verify_permissions_and_get_entity( + id=visualization_id, + get_method=zen_store().get_curated_visualization, + hydrate=hydrate, + ) + + +@router.patch( + "/{visualization_id}", + responses={401: error_response, 404: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper(deduplicate=True) +def update_curated_visualization( + visualization_id: UUID, + visualization_update: CuratedVisualizationUpdate, + _: AuthContext = Security(authorize), +) -> CuratedVisualizationResponse: + """Update a curated visualization. + + Args: + visualization_id: The ID of the curated visualization to update. + visualization_update: The update to apply to the curated visualization. + + Returns: + The updated curated visualization. + """ + return verify_permissions_and_update_entity( + id=visualization_id, + update_model=visualization_update, + get_method=zen_store().get_curated_visualization, + update_method=zen_store().update_curated_visualization, + ) + + +@router.delete( + "/{visualization_id}", + responses={401: error_response, 404: error_response, 422: error_response}, +) +@async_fastapi_endpoint_wrapper +def delete_curated_visualization( + visualization_id: UUID, + _: AuthContext = Security(authorize), +) -> None: + """Delete a curated visualization. + + Args: + visualization_id: The ID of the curated visualization to delete. + """ + verify_permissions_and_delete_entity( + id=visualization_id, + get_method=zen_store().get_curated_visualization, + delete_method=zen_store().delete_curated_visualization, + ) diff --git a/src/zenml/zen_server/routers/deployment_endpoints.py b/src/zenml/zen_server/routers/deployment_endpoints.py index aab4532db2..e8e1ed3e28 100644 --- a/src/zenml/zen_server/routers/deployment_endpoints.py +++ b/src/zenml/zen_server/routers/deployment_endpoints.py @@ -13,17 +13,13 @@ # permissions and limitations under the License. """Endpoint definitions for deployments.""" +from typing import Optional, Union from uuid import UUID -from fastapi import ( - APIRouter, - Depends, - Security, -) +from fastapi import APIRouter, Depends, Security from zenml.constants import ( API, - DEPLOYMENT_VISUALIZATIONS, DEPLOYMENTS, VERSION_1, ) @@ -32,12 +28,8 @@ DeploymentRequest, DeploymentResponse, DeploymentUpdate, - DeploymentVisualizationFilter, - DeploymentVisualizationRequest, - DeploymentVisualizationResponse, - DeploymentVisualizationUpdate, + Page, ) -from zenml.models.v2.base.page import Page from zenml.zen_server.auth import AuthContext, authorize from zenml.zen_server.exceptions import error_response from zenml.zen_server.rbac.endpoint_utils import ( @@ -48,6 +40,7 @@ verify_permissions_and_update_entity, ) from zenml.zen_server.rbac.models import ResourceType +from zenml.zen_server.routers.projects_endpoints import workspace_router from zenml.zen_server.utils import ( async_fastapi_endpoint_wrapper, make_dependable, @@ -60,30 +53,39 @@ responses={401: error_response, 403: error_response}, ) -deployment_visualization_router = APIRouter( - prefix=API + VERSION_1 + DEPLOYMENT_VISUALIZATIONS, - tags=["deployment_visualizations"], - responses={401: error_response, 403: error_response}, -) - @router.post( "", responses={401: error_response, 409: error_response, 422: error_response}, ) +# TODO: the workspace scoped endpoint is only kept for dashboard compatibility +# and can be removed after the migration +@workspace_router.post( + "/{project_name_or_id}" + DEPLOYMENTS, + responses={401: error_response, 409: error_response, 422: error_response}, + deprecated=True, + tags=["deployments"], +) @async_fastapi_endpoint_wrapper def create_deployment( deployment: DeploymentRequest, + project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize), ) -> DeploymentResponse: - """Creates a deployment. + """Create a deployment. Args: deployment: Deployment to create. + project_name_or_id: Optional project name or ID for backwards + compatibility. Returns: The created deployment. """ + if project_name_or_id: + project = zen_store().get_project(project_name_or_id) + deployment.project = project.id + return verify_permissions_and_create_entity( request_model=deployment, create_method=zen_store().create_deployment, @@ -94,25 +96,38 @@ def create_deployment( "", responses={401: error_response, 404: error_response, 422: error_response}, ) +# TODO: the workspace scoped endpoint is only kept for dashboard compatibility +# and can be removed after the migration +@workspace_router.get( + "/{project_name_or_id}" + DEPLOYMENTS, + responses={401: error_response, 404: error_response, 422: error_response}, + deprecated=True, + tags=["deployments"], +) @async_fastapi_endpoint_wrapper(deduplicate=True) def list_deployments( deployment_filter_model: DeploymentFilter = Depends( make_dependable(DeploymentFilter) ), + project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize), ) -> Page[DeploymentResponse]: - """Gets a list of deployments. + """List deployments. Args: - deployment_filter_model: Filter model used for pagination, sorting, + deployment_filter_model: Filter model used for pagination, sorting, and filtering. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. + project_name_or_id: Optional project name or ID for backwards + compatibility. + hydrate: Whether to hydrate the returned models. Returns: - List of deployment objects matching the filter criteria. + A page of deployments matching the filter. """ + if project_name_or_id: + deployment_filter_model.project = project_name_or_id + return verify_permissions_and_list_entities( filter_model=deployment_filter_model, resource_type=ResourceType.DEPLOYMENT, @@ -131,15 +146,14 @@ def get_deployment( hydrate: bool = True, _: AuthContext = Security(authorize), ) -> DeploymentResponse: - """Gets a specific deployment using its unique id. + """Get a deployment by ID. Args: - deployment_id: ID of the deployment to get. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. + deployment_id: The deployment ID. + hydrate: Whether to hydrate the returned model. Returns: - A specific deployment object. + The requested deployment. """ return verify_permissions_and_get_entity( id=deployment_id, @@ -158,11 +172,11 @@ def update_deployment( deployment_update: DeploymentUpdate, _: AuthContext = Security(authorize), ) -> DeploymentResponse: - """Updates a specific deployment. + """Update a deployment. Args: - deployment_id: ID of the deployment to update. - deployment_update: Update model for the deployment. + deployment_id: The deployment ID. + deployment_update: The updates to apply. Returns: The updated deployment. @@ -184,157 +198,13 @@ def delete_deployment( deployment_id: UUID, _: AuthContext = Security(authorize), ) -> None: - """Deletes a specific deployment. + """Delete a deployment. Args: - deployment_id: ID of the deployment to delete. + deployment_id: The deployment ID. """ verify_permissions_and_delete_entity( id=deployment_id, get_method=zen_store().get_deployment, delete_method=zen_store().delete_deployment, ) - - -@router.post( - "/{deployment_id}/visualizations", - responses={401: error_response, 409: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def create_deployment_visualization( - deployment_id: UUID, - visualization: DeploymentVisualizationRequest, - _: AuthContext = Security(authorize), -) -> DeploymentVisualizationResponse: - """Creates a curated deployment visualization. - - Args: - deployment_id: ID of the deployment to add visualization to. - visualization: Deployment visualization to create. - - Returns: - The created deployment visualization. - - Raises: - KeyError: If the deployment ID in the request does not match the path parameter. - """ - if ( - visualization.deployment_id - and visualization.deployment_id != deployment_id - ): - raise KeyError( - "Deployment ID in request does not match path parameter" - ) - return verify_permissions_and_create_entity( - request_model=visualization, - create_method=zen_store().create_deployment_visualization, - ) - - -@router.get( - "/{deployment_id}/visualizations", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper(deduplicate=True) -def list_deployment_visualizations( - deployment_id: UUID, - visualization_filter_model: DeploymentVisualizationFilter = Depends( - make_dependable(DeploymentVisualizationFilter) - ), - hydrate: bool = False, - _: AuthContext = Security(authorize), -) -> Page[DeploymentVisualizationResponse]: - """Gets a list of visualizations for a specific deployment. - - Args: - deployment_id: ID of the deployment to list visualizations for. - visualization_filter_model: Filter model used for pagination, sorting, - filtering. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - List of deployment visualization objects for the deployment. - """ - visualization_filter_model.deployment = deployment_id - return verify_permissions_and_list_entities( - filter_model=visualization_filter_model, - resource_type=ResourceType.DEPLOYMENT, - list_method=zen_store().list_deployment_visualizations, - hydrate=hydrate, - ) - - -@deployment_visualization_router.get( - "/{deployment_visualization_id}", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper(deduplicate=True) -def get_deployment_visualization( - deployment_visualization_id: UUID, - hydrate: bool = True, - _: AuthContext = Security(authorize), -) -> DeploymentVisualizationResponse: - """Gets a specific deployment visualization using its unique id. - - Args: - deployment_visualization_id: ID of the deployment visualization to get. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - A specific deployment visualization object. - """ - return verify_permissions_and_get_entity( - id=deployment_visualization_id, - get_method=zen_store().get_deployment_visualization, - hydrate=hydrate, - ) - - -@deployment_visualization_router.patch( - "/{deployment_visualization_id}", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper(deduplicate=True) -def update_deployment_visualization( - deployment_visualization_id: UUID, - visualization_update: DeploymentVisualizationUpdate, - _: AuthContext = Security(authorize), -) -> DeploymentVisualizationResponse: - """Updates a specific deployment visualization. - - Args: - deployment_visualization_id: ID of the deployment visualization to update. - visualization_update: Update model for the deployment visualization. - - Returns: - The updated deployment visualization. - """ - return verify_permissions_and_update_entity( - id=deployment_visualization_id, - update_model=visualization_update, - get_method=zen_store().get_deployment_visualization, - update_method=zen_store().update_deployment_visualization, - ) - - -@deployment_visualization_router.delete( - "/{deployment_visualization_id}", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper -def delete_deployment_visualization( - deployment_visualization_id: UUID, - _: AuthContext = Security(authorize), -) -> None: - """Deletes a specific deployment visualization. - - Args: - deployment_visualization_id: ID of the deployment visualization to delete. - """ - verify_permissions_and_delete_entity( - id=deployment_visualization_id, - get_method=zen_store().get_deployment_visualization, - delete_method=zen_store().delete_deployment_visualization, - ) diff --git a/src/zenml/zen_server/zen_server_api.py b/src/zenml/zen_server/zen_server_api.py index a0f4c5f878..46feb51c47 100644 --- a/src/zenml/zen_server/zen_server_api.py +++ b/src/zenml/zen_server/zen_server_api.py @@ -57,6 +57,7 @@ artifact_version_endpoints, auth_endpoints, code_repositories_endpoints, + curated_visualization_endpoints, deployment_endpoints, devices_endpoints, event_source_endpoints, @@ -265,7 +266,7 @@ async def dashboard(request: Request) -> Any: app.include_router(devices_endpoints.router) app.include_router(code_repositories_endpoints.router) app.include_router(deployment_endpoints.router) -app.include_router(deployment_endpoints.deployment_visualization_router) +app.include_router(curated_visualization_endpoints.router) app.include_router(plugin_endpoints.plugin_router) app.include_router(event_source_endpoints.event_source_router) app.include_router(flavors_endpoints.router) diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 05c5691f33..44a0248ed6 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -65,10 +65,10 @@ CODE_REFERENCES, CODE_REPOSITORIES, CONFIG, + CURATED_VISUALIZATIONS, CURRENT_USER, DEACTIVATE, DEFAULT_HTTP_TIMEOUT, - DEPLOYMENT_VISUALIZATIONS, DEPLOYMENTS, DEVICES, DISABLE_CLIENT_SERVER_MISMATCH_WARNING, @@ -166,15 +166,15 @@ ComponentRequest, ComponentResponse, ComponentUpdate, + CuratedVisualizationFilter, + CuratedVisualizationRequest, + CuratedVisualizationResponse, + CuratedVisualizationUpdate, DeployedStack, DeploymentFilter, DeploymentRequest, DeploymentResponse, DeploymentUpdate, - DeploymentVisualizationFilter, - DeploymentVisualizationRequest, - DeploymentVisualizationResponse, - DeploymentVisualizationUpdate, EventSourceFilter, EventSourceRequest, EventSourceResponse, @@ -1861,50 +1861,50 @@ def delete_deployment(self, deployment_id: UUID) -> None: route=DEPLOYMENTS, ) - def create_deployment_visualization( - self, visualization: DeploymentVisualizationRequest - ) -> DeploymentVisualizationResponse: - """Create a deployment visualization via REST API. + def create_curated_visualization( + self, visualization: CuratedVisualizationRequest + ) -> CuratedVisualizationResponse: + """Create a curated visualization via REST API. Args: - visualization: The visualization to create. + visualization: The curated visualization to create. Returns: - The created visualization. + The created curated visualization. """ return self._create_resource( resource=visualization, - response_model=DeploymentVisualizationResponse, - route=f"{DEPLOYMENTS}/{visualization.deployment_id}/visualizations", + response_model=CuratedVisualizationResponse, + route=CURATED_VISUALIZATIONS, params={"hydrate": True}, ) - def get_deployment_visualization( - self, deployment_visualization_id: UUID, hydrate: bool = True - ) -> DeploymentVisualizationResponse: - """Get a deployment visualization by ID. + def get_curated_visualization( + self, visualization_id: UUID, hydrate: bool = True + ) -> CuratedVisualizationResponse: + """Get a curated visualization by ID. Args: - deployment_visualization_id: The ID of the visualization to get. + visualization_id: The ID of the curated visualization to get. hydrate: Flag deciding whether to hydrate the output model(s) by including metadata fields in the response. Returns: - The deployment visualization. + The curated visualization with the given ID. """ return self._get_resource( - resource_id=deployment_visualization_id, - route=DEPLOYMENT_VISUALIZATIONS, - response_model=DeploymentVisualizationResponse, + resource_id=visualization_id, + route=CURATED_VISUALIZATIONS, + response_model=CuratedVisualizationResponse, params={"hydrate": hydrate}, ) - def list_deployment_visualizations( + def list_curated_visualizations( self, - filter_model: DeploymentVisualizationFilter, + filter_model: CuratedVisualizationFilter, hydrate: bool = False, - ) -> Page[DeploymentVisualizationResponse]: - """List deployment visualizations via REST API. + ) -> Page[CuratedVisualizationResponse]: + """List curated visualizations via REST API. Args: filter_model: The filter model to use. @@ -1912,59 +1912,46 @@ def list_deployment_visualizations( by including metadata fields in the response. Returns: - A list of all deployment visualizations matching the filter criteria. - - Raises: - ValueError: If no deployment ID is provided in the filter model. + A page of curated visualizations. """ - if not filter_model.deployment: - raise ValueError( - "A deployment ID is required to list deployment visualizations." - ) - - # Build the route with the deployment ID - route = f"{DEPLOYMENTS}/{str(filter_model.deployment)}/visualizations" - return self._list_paginated_resources( - route=route, - response_model=DeploymentVisualizationResponse, + route=CURATED_VISUALIZATIONS, + response_model=CuratedVisualizationResponse, filter_model=filter_model, params={"hydrate": hydrate}, ) - def update_deployment_visualization( + def update_curated_visualization( self, - deployment_visualization_id: UUID, - update: DeploymentVisualizationUpdate, - ) -> DeploymentVisualizationResponse: - """Update a deployment visualization via REST API. + visualization_id: UUID, + update: CuratedVisualizationUpdate, + ) -> CuratedVisualizationResponse: + """Update a curated visualization via REST API. Args: - deployment_visualization_id: The ID of the visualization to update. - update: The update to apply. + visualization_id: The ID of the curated visualization to update. + update: The update to apply to the curated visualization. Returns: - The updated deployment visualization. + The updated curated visualization. """ return self._update_resource( - resource_id=deployment_visualization_id, + resource_id=visualization_id, resource_update=update, - response_model=DeploymentVisualizationResponse, - route=DEPLOYMENT_VISUALIZATIONS, + response_model=CuratedVisualizationResponse, + route=CURATED_VISUALIZATIONS, params={"hydrate": True}, ) - def delete_deployment_visualization( - self, deployment_visualization_id: UUID - ) -> None: - """Delete a deployment visualization via REST API. + def delete_curated_visualization(self, visualization_id: UUID) -> None: + """Delete a curated visualization via REST API. Args: - deployment_visualization_id: The ID of the visualization to delete. + visualization_id: The ID of the curated visualization to delete. """ self._delete_resource( - resource_id=deployment_visualization_id, - route=DEPLOYMENT_VISUALIZATIONS, + resource_id=visualization_id, + route=CURATED_VISUALIZATIONS, ) # -------------------- Run templates -------------------- diff --git a/src/zenml/zen_stores/schemas/__init__.py b/src/zenml/zen_stores/schemas/__init__.py index 70add37167..6235b6f64b 100644 --- a/src/zenml/zen_stores/schemas/__init__.py +++ b/src/zenml/zen_stores/schemas/__init__.py @@ -31,8 +31,9 @@ from zenml.zen_stores.schemas.event_source_schemas import EventSourceSchema from zenml.zen_stores.schemas.pipeline_build_schemas import PipelineBuildSchema from zenml.zen_stores.schemas.deployment_schemas import DeploymentSchema -from zenml.zen_stores.schemas.deployment_visualization_schemas import ( - DeploymentVisualizationSchema, +from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationResourceSchema, + CuratedVisualizationSchema, ) from zenml.zen_stores.schemas.component_schemas import StackComponentSchema from zenml.zen_stores.schemas.flavor_schemas import FlavorSchema @@ -91,7 +92,8 @@ "CodeReferenceSchema", "CodeRepositorySchema", "DeploymentSchema", - "DeploymentVisualizationSchema", + "CuratedVisualizationSchema", + "CuratedVisualizationResourceSchema", "EventSourceSchema", "FlavorSchema", "LogsSchema", diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py new file mode 100644 index 0000000000..a39f7047a0 --- /dev/null +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -0,0 +1,288 @@ +# Copyright (c) ZenML GmbH 2025. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""SQLModel implementation of curated visualization tables.""" + +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence +from uuid import UUID, uuid4 + +from sqlalchemy import UniqueConstraint, and_ +from sqlalchemy.orm import foreign, selectinload +from sqlalchemy.sql.base import ExecutableOption +from sqlmodel import Field, Relationship, SQLModel + +from zenml.enums import VisualizationResourceTypes +from zenml.models.v2.core.curated_visualization import ( + CuratedVisualizationRequest, + CuratedVisualizationResponse, + CuratedVisualizationResponseBody, + CuratedVisualizationResponseMetadata, + CuratedVisualizationResponseResources, + CuratedVisualizationUpdate, +) +from zenml.models.v2.misc.curated_visualization import ( + CuratedVisualizationResource, +) +from zenml.zen_stores.schemas.base_schemas import BaseSchema +from zenml.zen_stores.schemas.project_schemas import ProjectSchema +from zenml.zen_stores.schemas.schema_utils import ( + build_foreign_key_field, + build_index, +) +from zenml.zen_stores.schemas.utils import jl_arg + +if TYPE_CHECKING: + from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema + + +class CuratedVisualizationResourceSchema(SQLModel, table=True): + """Link table mapping curated visualizations to resources.""" + + __tablename__ = "curated_visualization_resource" + __table_args__ = ( + UniqueConstraint( + "visualization_id", + "resource_id", + "resource_type", + name="unique_curated_visualization_resource", + ), + build_index(__tablename__, ["resource_id", "resource_type"]), + ) + + id: UUID = Field(default_factory=uuid4, primary_key=True) + visualization_id: UUID = build_foreign_key_field( + source=__tablename__, + target="curated_visualization", + source_column="visualization_id", + target_column="id", + ondelete="CASCADE", + nullable=False, + ) + resource_id: UUID = Field(nullable=False) + resource_type: str = Field(nullable=False) + + visualization: "CuratedVisualizationSchema" = Relationship( + back_populates="resources", + ) + + +class CuratedVisualizationSchema(BaseSchema, table=True): + """SQL Model for curated visualizations.""" + + __tablename__ = "curated_visualization" + __table_args__ = ( + build_index( + __tablename__, ["artifact_version_id", "visualization_index"] + ), + build_index(__tablename__, ["display_order"]), + ) + + project_id: UUID = build_foreign_key_field( + source=__tablename__, + target=ProjectSchema.__tablename__, + source_column="project_id", + target_column="id", + ondelete="CASCADE", + nullable=False, + ) + artifact_version_id: UUID = build_foreign_key_field( + source=__tablename__, + target="artifact_version", + source_column="artifact_version_id", + target_column="id", + ondelete="CASCADE", + nullable=False, + ) + + visualization_index: int = Field(nullable=False) + display_name: Optional[str] = Field(default=None) + display_order: Optional[int] = Field(default=None) + + artifact_version: Optional["ArtifactVersionSchema"] = Relationship( + sa_relationship_kwargs={"lazy": "selectin"} + ) + resources: List[CuratedVisualizationResourceSchema] = Relationship( + back_populates="visualization", + sa_relationship_kwargs={"lazy": "selectin", "cascade": "all, delete"}, + ) + + @classmethod + def get_query_options( + cls, + include_metadata: bool = False, + include_resources: bool = False, + **kwargs: Any, + ) -> Sequence[ExecutableOption]: + """Get the query options for the schema. + + Args: + include_metadata: Whether metadata will be included when converting + the schema to a model. + include_resources: Whether resources will be included when + converting the schema to a model. + **kwargs: Keyword arguments to allow schema specific logic + + Returns: + A list of query options. + """ + options: List[ExecutableOption] = [] + + if include_resources: + options.extend( + [ + selectinload(jl_arg(cls.artifact_version)), + selectinload(jl_arg(cls.resources)), + ] + ) + + return options + + @classmethod + def from_request( + cls, request: CuratedVisualizationRequest + ) -> "CuratedVisualizationSchema": + """Convert a request into a schema instance. + + Args: + request: The request to convert. + + Returns: + The created schema. + """ + return cls( + project_id=request.project, + artifact_version_id=request.artifact_version_id, + visualization_index=request.visualization_index, + display_name=request.display_name, + display_order=request.display_order, + ) + + def update( + self, + update: CuratedVisualizationUpdate, + ) -> "CuratedVisualizationSchema": + """Update a schema instance from an update model. + + Args: + update: The update definition. + + Returns: + The updated schema. + """ + for field, value in update.model_dump( + exclude_unset=True, + ).items(): + if hasattr(self, field): + setattr(self, field, value) + + from zenml.utils.time_utils import utc_now + + self.updated = utc_now() + return self + + def to_model( + self, + include_metadata: bool = False, + include_resources: bool = False, + **kwargs: Any, + ) -> CuratedVisualizationResponse: + """Convert schema into response model. + + Args: + include_metadata: Whether to include metadata in the response. + include_resources: Whether to include resources in the response. + **kwargs: Additional keyword arguments. + + Returns: + The created response model. + """ + resources = [ + CuratedVisualizationResource( + id=resource.resource_id, + type=VisualizationResourceTypes(resource.resource_type), + ) + for resource in self.resources + ] + + body = CuratedVisualizationResponseBody( + project_id=self.project_id, + created=self.created, + updated=self.updated, + artifact_version_id=self.artifact_version_id, + visualization_index=self.visualization_index, + display_name=self.display_name, + display_order=self.display_order, + resources=resources, + ) + + metadata = None + if include_metadata: + metadata = CuratedVisualizationResponseMetadata() + + response_resources = None + if include_resources: + response_resources = CuratedVisualizationResponseResources( + artifact_version=self.artifact_version.to_model( + include_metadata=include_metadata, + include_resources=include_resources, + ) + if self.artifact_version + else None, + ) + + return CuratedVisualizationResponse( + id=self.id, + body=body, + metadata=metadata, + resources=response_resources, + ) + + +def curated_visualization_relationship_kwargs( + parent_column_factory: Callable[[], Any], + resource_type: VisualizationResourceTypes, +) -> Dict[str, Any]: + """Build sa_relationship_kwargs for curated visualization relationships. + + This helper consolidates the relationship definition for linking parent + schemas (like DeploymentSchema, ModelSchema) to their curated visualizations + through the resource link table. + + Args: + parent_column_factory: A callable that returns the parent column + (e.g., `lambda: DeploymentSchema.id`). Uses a callable to defer + evaluation and avoid circular import issues. + resource_type: The VisualizationResourceTypes enum value indicating + what type of resource the parent represents. + + Returns: + A dictionary suitable for passing to Relationship(sa_relationship_kwargs=...). + The relationship will be read-only (viewonly=True) and eagerly loaded + via selectin loading. + """ + def _primaryjoin(): + return and_( + CuratedVisualizationResourceSchema.resource_type == resource_type.value, + foreign(CuratedVisualizationResourceSchema.resource_id) == parent_column_factory(), + ) + + def _secondaryjoin(): + return CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id + + return dict( + secondary="curated_visualization_resource", + primaryjoin=_primaryjoin, + secondaryjoin=_secondaryjoin, + viewonly=True, + lazy="selectin", + ) diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index f5057ab709..df339ea0c6 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -24,7 +24,11 @@ from sqlmodel import Field, Relationship, String from zenml.constants import MEDIUMTEXT_MAX_LENGTH -from zenml.enums import DeploymentStatus, TaggableResourceTypes +from zenml.enums import ( + DeploymentStatus, + TaggableResourceTypes, + VisualizationResourceTypes, +) from zenml.logger import get_logger from zenml.models.v2.core.deployment import ( DeploymentRequest, @@ -46,11 +50,15 @@ from zenml.zen_stores.schemas.utils import jl_arg if TYPE_CHECKING: - from zenml.zen_stores.schemas.deployment_visualization_schemas import ( - DeploymentVisualizationSchema, + from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationSchema, ) from zenml.zen_stores.schemas.tag_schemas import TagSchema +from zenml.zen_stores.schemas.curated_visualization_schemas import ( + curated_visualization_relationship_kwargs, +) + logger = get_logger(__name__) @@ -136,9 +144,11 @@ class DeploymentSchema(NamedSchema, table=True): ), ) - visualizations: List["DeploymentVisualizationSchema"] = Relationship( - back_populates="deployment", - sa_relationship_kwargs=dict(lazy="selectin"), + visualizations: List["CuratedVisualizationSchema"] = Relationship( + sa_relationship_kwargs=curated_visualization_relationship_kwargs( + parent_column_factory=lambda: DeploymentSchema.id, + resource_type=VisualizationResourceTypes.DEPLOYMENT, + ), ) @classmethod @@ -163,8 +173,8 @@ def get_query_options( options = [] if include_resources: - from zenml.zen_stores.schemas.deployment_visualization_schemas import ( - DeploymentVisualizationSchema, + from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationSchema, ) options.extend( @@ -177,7 +187,7 @@ def get_query_options( selectinload( jl_arg(DeploymentSchema.visualizations) ).selectinload( - jl_arg(DeploymentVisualizationSchema.artifact_version) + jl_arg(CuratedVisualizationSchema.artifact_version) ), ] ) @@ -241,7 +251,6 @@ def to_model( visualization.to_model( include_metadata=False, include_resources=False, - include_deployment=False, ) for visualization in (self.visualizations or []) ], diff --git a/src/zenml/zen_stores/schemas/deployment_visualization_schemas.py b/src/zenml/zen_stores/schemas/deployment_visualization_schemas.py deleted file mode 100644 index 665dbe6a12..0000000000 --- a/src/zenml/zen_stores/schemas/deployment_visualization_schemas.py +++ /dev/null @@ -1,243 +0,0 @@ -# Copyright (c) ZenML GmbH 2025. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""SQLModel implementation of deployment visualization table.""" - -from typing import TYPE_CHECKING, Any, Optional, Sequence -from uuid import UUID - -from sqlalchemy import UniqueConstraint -from sqlalchemy.orm import selectinload -from sqlalchemy.sql.base import ExecutableOption -from sqlmodel import Field, Relationship - -from zenml.constants import STR_FIELD_MAX_LENGTH -from zenml.models.v2.core.deployment_visualization import ( - DeploymentVisualizationRequest, - DeploymentVisualizationResponse, - DeploymentVisualizationResponseBody, - DeploymentVisualizationResponseMetadata, - DeploymentVisualizationResponseResources, - DeploymentVisualizationUpdate, -) -from zenml.utils.time_utils import utc_now -from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema -from zenml.zen_stores.schemas.base_schemas import BaseSchema -from zenml.zen_stores.schemas.project_schemas import ProjectSchema -from zenml.zen_stores.schemas.schema_utils import ( - build_foreign_key_field, - build_index, -) -from zenml.zen_stores.schemas.utils import jl_arg - -if TYPE_CHECKING: - from zenml.zen_stores.schemas.deployment_schemas import DeploymentSchema - - -class DeploymentVisualizationSchema(BaseSchema, table=True): - """SQL Model for deployment visualizations.""" - - __tablename__ = "deployment_visualization" - __table_args__ = ( - UniqueConstraint( - "deployment_id", - "artifact_version_id", - "visualization_index", - name="unique_deployment_visualization", - ), - build_index( - __tablename__, - ["deployment_id"], - ), - build_index( - __tablename__, - ["display_order"], - ), - ) - - # Foreign Keys - project_id: UUID = build_foreign_key_field( - source=__tablename__, - target=ProjectSchema.__tablename__, - source_column="project_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - - deployment_id: UUID = build_foreign_key_field( - source=__tablename__, - target="deployment", - source_column="deployment_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - - artifact_version_id: UUID = build_foreign_key_field( - source=__tablename__, - target=ArtifactVersionSchema.__tablename__, - source_column="artifact_version_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - - # Fields - visualization_index: int = Field(nullable=False) - display_name: Optional[str] = Field( - max_length=STR_FIELD_MAX_LENGTH, default=None - ) - display_order: Optional[int] = Field(default=None) - - # Relationships - deployment: Optional["DeploymentSchema"] = Relationship( - back_populates="visualizations", - sa_relationship_kwargs={"lazy": "selectin"}, - ) - - artifact_version: Optional[ArtifactVersionSchema] = Relationship( - sa_relationship_kwargs={"lazy": "selectin"} - ) - - @classmethod - def get_query_options( - cls, - include_metadata: bool = False, - include_resources: bool = False, - **kwargs: Any, - ) -> Sequence[ExecutableOption]: - """Get the query options for the schema. - - Args: - include_metadata: Whether metadata will be included when converting - the schema to a model. - include_resources: Whether resources will be included when - converting the schema to a model. - **kwargs: Keyword arguments to allow schema specific logic - - Returns: - A list of query options. - """ - options = [] - - if include_resources: - options.extend( - [ - selectinload(jl_arg(cls.deployment)), - selectinload(jl_arg(cls.artifact_version)), - ] - ) - - return options - - @classmethod - def from_request( - cls, request: DeploymentVisualizationRequest - ) -> "DeploymentVisualizationSchema": - """Convert a `DeploymentVisualizationRequest` to a `DeploymentVisualizationSchema`. - - Args: - request: The request model to convert. - - Returns: - The converted schema. - """ - return cls( - project_id=request.project, - deployment_id=request.deployment_id, - artifact_version_id=request.artifact_version_id, - visualization_index=request.visualization_index, - display_name=request.display_name, - display_order=request.display_order, - ) - - def update( - self, - update: DeploymentVisualizationUpdate, - ) -> "DeploymentVisualizationSchema": - """Updates a `DeploymentVisualizationSchema` from a `DeploymentVisualizationUpdate`. - - Args: - update: The `DeploymentVisualizationUpdate` to update from. - - Returns: - The updated `DeploymentVisualizationSchema`. - """ - for field, value in update.model_dump( - exclude_unset=True, exclude_none=True - ).items(): - if hasattr(self, field): - setattr(self, field, value) - - self.updated = utc_now() - return self - - def to_model( - self, - include_metadata: bool = False, - include_resources: bool = False, - **kwargs: Any, - ) -> DeploymentVisualizationResponse: - """Convert a `DeploymentVisualizationSchema` to a `DeploymentVisualizationResponse`. - - Args: - include_metadata: Whether to include metadata in the response. - include_resources: Whether to include resources in the response. - **kwargs: Additional keyword arguments, including `include_deployment` - to control whether deployment resources are included (default True). - - Returns: - The created `DeploymentVisualizationResponse`. - """ - include_deployment = kwargs.get("include_deployment", True) - - body = DeploymentVisualizationResponseBody( - user_id=None, - project_id=self.project_id, - created=self.created, - updated=self.updated, - deployment_id=self.deployment_id, - artifact_version_id=self.artifact_version_id, - visualization_index=self.visualization_index, - display_name=self.display_name, - display_order=self.display_order, - ) - - metadata = None - if include_metadata: - metadata = DeploymentVisualizationResponseMetadata() - - resources = None - if include_resources: - resources = DeploymentVisualizationResponseResources( - deployment=self.deployment.to_model( - include_metadata=include_metadata, - include_resources=include_resources, - ) - if self.deployment and include_deployment - else None, - artifact_version=self.artifact_version.to_model( - include_metadata=include_metadata, - include_resources=include_resources, - ) - if self.artifact_version - else None, - ) - - return DeploymentVisualizationResponse( - id=self.id, - body=body, - metadata=metadata, - resources=resources, - ) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 1d77d926a3..2658f3c09d 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -24,13 +24,14 @@ Column, UniqueConstraint, ) -from sqlalchemy.orm import joinedload, object_session +from sqlalchemy.orm import joinedload, object_session, selectinload from sqlalchemy.sql.base import ExecutableOption from sqlmodel import Field, Relationship, desc, select from zenml.enums import ( MetadataResourceTypes, TaggableResourceTypes, + VisualizationResourceTypes, ) from zenml.models import ( BaseResponseMetadata, @@ -57,6 +58,11 @@ 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.curated_visualization_schemas import ( + CuratedVisualizationResourceSchema, + CuratedVisualizationSchema, + curated_visualization_relationship_kwargs, +) from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema from zenml.zen_stores.schemas.project_schemas import ProjectSchema from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema @@ -128,6 +134,12 @@ class ModelSchema(NamedSchema, table=True): back_populates="model", sa_relationship_kwargs={"cascade": "delete"}, ) + visualizations: List["CuratedVisualizationSchema"] = Relationship( + sa_relationship_kwargs=curated_visualization_relationship_kwargs( + parent_column_factory=lambda: ModelSchema.id, + resource_type=VisualizationResourceTypes.MODEL, + ), + ) @classmethod def get_query_options( @@ -155,6 +167,9 @@ def get_query_options( [ joinedload(jl_arg(ModelSchema.user)), # joinedload(jl_arg(ModelSchema.tags)), + selectinload(jl_arg(ModelSchema.visualizations)).selectinload( + jl_arg(CuratedVisualizationSchema.artifact_version) + ), ] ) @@ -254,6 +269,13 @@ def to_model( tags=[tag.to_model() for tag in self.tags], latest_version_name=latest_version_name, latest_version_id=latest_version_id, + visualizations=[ + visualization.to_model( + include_metadata=False, + include_resources=False, + ) + for visualization in (self.visualizations or []) + ], ) body = ModelResponseBody( diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 29c4f4655e..8316c31bc5 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -75,7 +75,14 @@ ArgumentError, IntegrityError, ) -from sqlalchemy.orm import Mapped, load_only, noload, selectinload +from sqlalchemy.orm import ( + Mapped, + aliased, + load_only, + noload, + selectinload, +) +from sqlalchemy.orm.util import AliasedClass from sqlalchemy.sql.base import ExecutableOption from sqlalchemy.util import immutabledict from sqlmodel import Session as SqlModelSession @@ -148,6 +155,7 @@ StepRunInputArtifactType, StoreType, TaggableResourceTypes, + VisualizationResourceTypes, ) from zenml.exceptions import ( AuthorizationException, @@ -199,6 +207,10 @@ ComponentRequest, ComponentResponse, ComponentUpdate, + CuratedVisualizationFilter, + CuratedVisualizationRequest, + CuratedVisualizationResponse, + CuratedVisualizationUpdate, DefaultComponentRequest, DefaultStackRequest, DeployedStack, @@ -206,10 +218,6 @@ DeploymentRequest, DeploymentResponse, DeploymentUpdate, - DeploymentVisualizationFilter, - DeploymentVisualizationRequest, - DeploymentVisualizationResponse, - DeploymentVisualizationUpdate, EventSourceFilter, EventSourceRequest, EventSourceResponse, @@ -364,8 +372,9 @@ BaseSchema, CodeReferenceSchema, CodeRepositorySchema, + CuratedVisualizationResourceSchema, + CuratedVisualizationSchema, DeploymentSchema, - DeploymentVisualizationSchema, EventSourceSchema, FlavorSchema, ModelSchema, @@ -5399,23 +5408,23 @@ def delete_deployment(self, deployment_id: UUID) -> None: session.delete(deployment) session.commit() - # -------------------- Deployment visualizations -------------------- + # -------------------- Curated visualizations -------------------- - def _validate_deployment_visualization_index( + def _validate_curated_visualization_index( self, session: Session, artifact_version_id: UUID, visualization_index: int, ) -> None: - """Ensure the artifact version exposes the given visualization index. + """Validate that the artifact version exposes the given visualization index. Args: - session: The session to use. - artifact_version_id: The ID of the artifact version to validate. - visualization_index: The index of the visualization to validate. + session: The database session to use. + artifact_version_id: ID of the artifact version that produced the visualization. + visualization_index: Index of the visualization to validate. Raises: - IllegalOperationError: If the artifact version does not expose the given visualization index. + IllegalOperationError: If the artifact version does not expose the specified visualization index. """ count = session.scalar( select(func.count()) @@ -5432,78 +5441,169 @@ def _validate_deployment_visualization_index( f"with index {visualization_index}." ) - def create_deployment_visualization( - self, visualization: DeploymentVisualizationRequest - ) -> DeploymentVisualizationResponse: - """Persist a curated deployment visualization link. + def _assert_curated_visualization_duplicate( + self, + session: Session, + *, + artifact_version_id: UUID, + visualization_index: int, + resource_id: UUID, + resource_type: VisualizationResourceTypes, + ) -> None: + """Ensure a curated visualization link does not already exist. Args: - visualization: The visualization to create. + session: The session to use. + artifact_version_id: The ID of the artifact version to validate. + visualization_index: The index of the visualization to validate. + resource_id: The ID of the resource to validate. + resource_type: The type of the resource to validate. + """ + existing = session.exec( + select(CuratedVisualizationSchema) + .join(CuratedVisualizationResourceSchema) + .where( + CuratedVisualizationSchema.artifact_version_id + == artifact_version_id + ) + .where( + CuratedVisualizationSchema.visualization_index + == visualization_index + ) + .where( + CuratedVisualizationResourceSchema.resource_id == resource_id + ) + .where( + CuratedVisualizationResourceSchema.resource_type + == resource_type.value + ) + ).first() + if existing is not None: + raise EntityExistsError( + "A curated visualization for this resource already exists " + "for the specified artifact version and visualization index." + ) - Raises: - IllegalOperationError: If the deployment ID in the request does not match the path parameter. - EntityExistsError: If a curated visualization with the same deployment, artifact version, and index already exists. + def create_curated_visualization( + self, visualization: CuratedVisualizationRequest + ) -> CuratedVisualizationResponse: + """Persist a curated visualization link. + + Args: + visualization: The curated visualization to persist. Returns: - The created deployment visualization. + The persisted curated visualization. """ with Session(self.engine) as session: self._set_request_user_id( request_model=visualization, session=session ) - deployment = self._get_schema_by_id( - resource_id=visualization.deployment_id, - schema_class=DeploymentSchema, - session=session, - ) - artifact_version = self._get_schema_by_id( + artifact_version: ArtifactVersionSchema = self._get_schema_by_id( resource_id=visualization.artifact_version_id, schema_class=ArtifactVersionSchema, session=session, ) - project_id = deployment.project_id - if visualization.project and visualization.project != project_id: - raise IllegalOperationError( - "Deployment visualizations must target the same project " - "as the deployment." - ) - if artifact_version.project_id != project_id: - raise IllegalOperationError( - "Artifact version does not belong to the deployment " - "project." - ) + # Validate explicitly provided project before defaulting + if visualization.project: + if visualization.project != artifact_version.project_id: + raise IllegalOperationError( + "Curated visualizations must target the same project as " + "the artifact version." + ) + project_id = visualization.project + else: + project_id = artifact_version.project_id - self._validate_deployment_visualization_index( + self._validate_curated_visualization_index( session=session, artifact_version_id=visualization.artifact_version_id, visualization_index=visualization.visualization_index, ) - duplicate = session.exec( - select(DeploymentVisualizationSchema) - .where( - DeploymentVisualizationSchema.deployment_id - == visualization.deployment_id + resource_schema_map: Dict[ + VisualizationResourceTypes, Type[BaseSchema] + ] = { + VisualizationResourceTypes.DEPLOYMENT: DeploymentSchema, + VisualizationResourceTypes.MODEL: ModelSchema, + VisualizationResourceTypes.PIPELINE: PipelineSchema, + VisualizationResourceTypes.PIPELINE_RUN: PipelineRunSchema, + VisualizationResourceTypes.PIPELINE_SNAPSHOT: PipelineSnapshotSchema, + VisualizationResourceTypes.PROJECT: ProjectSchema, + } + + # Group resource requests by type for batch fetching + resources_by_type: Dict[ + VisualizationResourceTypes, List[UUID] + ] = defaultdict(list) + for resource_request in visualization.resources: + if resource_request.type not in resource_schema_map: + raise IllegalOperationError( + f"Curated visualizations are not supported for resource type '{resource_request.type.value}'." + ) + resources_by_type[resource_request.type].append( + resource_request.id ) - .where( - DeploymentVisualizationSchema.artifact_version_id - == visualization.artifact_version_id + + # Batch fetch resources by type + fetched_resources: Dict[UUID, BaseSchema] = {} + for resource_type, resource_ids in resources_by_type.items(): + schema_class = resource_schema_map[resource_type] + schemas = session.exec( + select(schema_class).where( + schema_class.id.in_(resource_ids) + ) + ).all() + + # Verify all resources were found + found_ids = {schema.id for schema in schemas} + missing_ids = set(resource_ids) - found_ids + if missing_ids: + raise KeyError( + f"Resources of type '{resource_type.value}' with IDs " + f"{missing_ids} not found." + ) + + for schema in schemas: + fetched_resources[schema.id] = schema + + # Validate project scope and check for duplicates + resources: List[CuratedVisualizationResourceSchema] = [] + for resource_request in visualization.resources: + resource_schema = fetched_resources[resource_request.id] + + resource_project_id = getattr( + resource_schema, "project_id", None ) - .where( - DeploymentVisualizationSchema.visualization_index - == visualization.visualization_index + if resource_project_id and resource_project_id != project_id: + raise IllegalOperationError( + "Curated visualizations must reference resources " + "within the same project as the artifact version." + ) + + self._assert_curated_visualization_duplicate( + session=session, + artifact_version_id=visualization.artifact_version_id, + visualization_index=visualization.visualization_index, + resource_id=resource_request.id, + resource_type=resource_request.type, ) - ).first() - if duplicate is not None: - raise EntityExistsError( - "A curated visualization with the same deployment, " - "artifact version, and index already exists." + + resources.append( + CuratedVisualizationResourceSchema( + resource_id=resource_request.id, + resource_type=resource_request.type.value, + ) ) - schema = DeploymentVisualizationSchema.from_request(visualization) + schema: CuratedVisualizationSchema = ( + CuratedVisualizationSchema.from_request(visualization) + ) schema.project_id = project_id + schema.resources = resources + session.add(schema) session.commit() session.refresh(schema) @@ -5511,40 +5611,38 @@ def create_deployment_visualization( return schema.to_model( include_metadata=True, include_resources=True, - include_deployment=False, ) - def get_deployment_visualization( - self, deployment_visualization_id: UUID, hydrate: bool = True - ) -> DeploymentVisualizationResponse: - """Fetch a curated deployment visualization by ID. + def get_curated_visualization( + self, visualization_id: UUID, hydrate: bool = True + ) -> CuratedVisualizationResponse: + """Fetch a curated visualization by ID. Args: - deployment_visualization_id: The ID of the visualization to get. + visualization_id: The ID of the curated visualization to fetch. hydrate: Flag deciding whether to hydrate the output model(s) by including metadata fields in the response. Returns: - The deployment visualization. + The curated visualization with the given ID. """ with Session(self.engine) as session: - schema = self._get_schema_by_id( - resource_id=deployment_visualization_id, - schema_class=DeploymentVisualizationSchema, + schema: CuratedVisualizationSchema = self._get_schema_by_id( + resource_id=visualization_id, + schema_class=CuratedVisualizationSchema, session=session, ) return schema.to_model( include_metadata=hydrate, include_resources=hydrate, - include_deployment=hydrate, ) - def list_deployment_visualizations( + def list_curated_visualizations( self, - filter_model: DeploymentVisualizationFilter, + filter_model: CuratedVisualizationFilter, hydrate: bool = False, - ) -> Page[DeploymentVisualizationResponse]: - """List all deployment visualizations matching the given filter. + ) -> Page[CuratedVisualizationResponse]: + """List all curated visualizations matching the given filter. Args: filter_model: The filter model to use. @@ -5552,40 +5650,61 @@ def list_deployment_visualizations( by including metadata fields in the response. Returns: - A list of all deployment visualizations matching the filter criteria. + A page of curated visualizations. """ with Session(self.engine) as session: self._set_filter_project_id( filter_model=filter_model, session=session ) - query = select(DeploymentVisualizationSchema) + query = select(CuratedVisualizationSchema) + if filter_model.resource_type or filter_model.resource_id: + resource_alias: AliasedClass[ + CuratedVisualizationResourceSchema + ] = aliased(CuratedVisualizationResourceSchema) + query = query.join( + resource_alias, + resource_alias.visualization_id + == CuratedVisualizationSchema.id, + ) + if filter_model.resource_type: + query = query.where( + resource_alias.resource_type + == filter_model.resource_type.value + ) + if filter_model.resource_id: + query = query.where( + resource_alias.resource_id == filter_model.resource_id + ) + # Add distinct to prevent duplicate rows when joined to multiple resources + query = query.distinct() + return self.filter_and_paginate( session=session, query=query, - table=DeploymentVisualizationSchema, + table=CuratedVisualizationSchema, filter_model=filter_model, hydrate=hydrate, ) - def update_deployment_visualization( + def update_curated_visualization( self, - deployment_visualization_id: UUID, - update: DeploymentVisualizationUpdate, - ) -> DeploymentVisualizationResponse: - """Update mutable fields on a curated deployment visualization. + visualization_id: UUID, + update: CuratedVisualizationUpdate, + ) -> CuratedVisualizationResponse: + """Update mutable fields on a curated visualization. Args: - deployment_visualization_id: The ID of the visualization to update. - update: The update to apply. + visualization_id: The ID of the curated visualization to update. + update: The update to apply to the curated visualization. Returns: - The updated deployment visualization. + The updated curated visualization. """ with Session(self.engine) as session: schema = self._get_schema_by_id( - resource_id=deployment_visualization_id, - schema_class=DeploymentVisualizationSchema, + resource_id=visualization_id, + schema_class=CuratedVisualizationSchema, session=session, ) schema.update(update) @@ -5596,21 +5715,18 @@ def update_deployment_visualization( return schema.to_model( include_metadata=True, include_resources=True, - include_deployment=False, ) - def delete_deployment_visualization( - self, deployment_visualization_id: UUID - ) -> None: - """Delete a curated deployment visualization. + def delete_curated_visualization(self, visualization_id: UUID) -> None: + """Delete a curated visualization. Args: - deployment_visualization_id: The ID of the visualization to delete. + visualization_id: The ID of the curated visualization to delete. """ with Session(self.engine) as session: schema = self._get_schema_by_id( - resource_id=deployment_visualization_id, - schema_class=DeploymentVisualizationSchema, + resource_id=visualization_id, + schema_class=CuratedVisualizationSchema, session=session, ) session.delete(schema) diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index fc80115fe5..90e137d8fc 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -48,15 +48,15 @@ ComponentRequest, ComponentResponse, ComponentUpdate, + CuratedVisualizationFilter, + CuratedVisualizationRequest, + CuratedVisualizationResponse, + CuratedVisualizationUpdate, DeployedStack, DeploymentFilter, DeploymentRequest, DeploymentResponse, DeploymentUpdate, - DeploymentVisualizationFilter, - DeploymentVisualizationRequest, - DeploymentVisualizationResponse, - DeploymentVisualizationUpdate, EventSourceFilter, EventSourceRequest, EventSourceResponse, @@ -1475,96 +1475,39 @@ def delete_deployment(self, deployment_id: UUID) -> None: KeyError: If the deployment does not exist. """ - # -------------------- Deployment visualizations -------------------- + # -------------------- Curated visualizations -------------------- @abstractmethod - def create_deployment_visualization( - self, visualization: DeploymentVisualizationRequest - ) -> DeploymentVisualizationResponse: - """Create a new deployment visualization. - - Args: - visualization: The deployment visualization to create. - - Returns: - The newly created deployment visualization. - - Raises: - EntityExistsError: If a deployment visualization with the same - deployment, artifact version, and visualization index already - exists. - """ + def create_curated_visualization( + self, visualization: CuratedVisualizationRequest + ) -> CuratedVisualizationResponse: + """Create a new curated visualization.""" @abstractmethod - def get_deployment_visualization( - self, deployment_visualization_id: UUID, hydrate: bool = True - ) -> DeploymentVisualizationResponse: - """Get a deployment visualization by ID. - - Args: - deployment_visualization_id: The ID of the deployment visualization - to get. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - The deployment visualization. - - Raises: - KeyError: If the deployment visualization does not exist. - """ + def get_curated_visualization( + self, visualization_id: UUID, hydrate: bool = True + ) -> CuratedVisualizationResponse: + """Get a curated visualization by ID.""" @abstractmethod - def list_deployment_visualizations( + def list_curated_visualizations( self, - filter_model: DeploymentVisualizationFilter, + filter_model: CuratedVisualizationFilter, hydrate: bool = False, - ) -> Page[DeploymentVisualizationResponse]: - """List all deployment visualizations matching the given filter criteria. - - Args: - 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: - A list of all deployment visualizations matching the filter criteria. - """ + ) -> Page[CuratedVisualizationResponse]: + """List curated visualizations matching the given filter criteria.""" @abstractmethod - def update_deployment_visualization( + def update_curated_visualization( self, - deployment_visualization_id: UUID, - visualization_update: DeploymentVisualizationUpdate, - ) -> DeploymentVisualizationResponse: - """Update a deployment visualization. - - Args: - deployment_visualization_id: The ID of the deployment visualization - to update. - visualization_update: The update to apply. - - Returns: - The updated deployment visualization. - - Raises: - KeyError: If the deployment visualization does not exist. - """ + visualization_id: UUID, + visualization_update: CuratedVisualizationUpdate, + ) -> CuratedVisualizationResponse: + """Update a curated visualization.""" @abstractmethod - def delete_deployment_visualization( - self, deployment_visualization_id: UUID - ) -> None: - """Delete a deployment visualization. - - Args: - deployment_visualization_id: The ID of the deployment visualization - to delete. - - Raises: - KeyError: If the deployment visualization does not exist. - """ + def delete_curated_visualization(self, visualization_id: UUID) -> None: + """Delete a curated visualization.""" # -------------------- Run templates -------------------- diff --git a/tests/integration/functional/zen_stores/test_zen_store.py b/tests/integration/functional/zen_stores/test_zen_store.py index 0e0221556f..3931b19389 100644 --- a/tests/integration/functional/zen_stores/test_zen_store.py +++ b/tests/integration/functional/zen_stores/test_zen_store.py @@ -74,6 +74,8 @@ StackComponentType, StoreType, TaggableResourceTypes, + VisualizationResourceTypes, + VisualizationType, ) from zenml.exceptions import ( AuthorizationException, @@ -90,8 +92,14 @@ ArtifactVersionFilter, ArtifactVersionRequest, ArtifactVersionResponse, + ArtifactVisualizationRequest, ComponentFilter, ComponentUpdate, + CuratedVisualizationFilter, + CuratedVisualizationRequest, + CuratedVisualizationResource, + CuratedVisualizationUpdate, + DeploymentRequest, ModelVersionArtifactFilter, ModelVersionArtifactRequest, ModelVersionFilter, @@ -5714,3 +5722,302 @@ def test_tag_filter_with_resource_type(clean_client: "Client"): # Test filtering for a resource type that doesn't have tags tags = clean_client.list_tags(resource_type=TaggableResourceTypes.MODEL) assert len(tags) == 0 + + +class TestCuratedVisualizations: + def test_curated_visualizations_across_resources(self): + """Test creating, listing, updating, and deleting curated visualizations.""" + + client = Client() + project_id = client.active_project.id + + artifact = client.zen_store.create_artifact( + ArtifactRequest( + name=sample_name("artifact"), + project=project_id, + has_custom_name=True, + ) + ) + artifact_version = client.zen_store.create_artifact_version( + ArtifactVersionRequest( + artifact_id=artifact.id, + project=project_id, + version="1", + type=ArtifactType.DATA, + uri=sample_name("artifact_uri"), + materializer=Source( + module="acme.foo", type=SourceType.INTERNAL + ), + data_type=Source(module="acme.foo", type=SourceType.INTERNAL), + save_type=ArtifactSaveType.STEP_OUTPUT, + visualizations=[ + ArtifactVisualizationRequest( + type=VisualizationType.HTML, + uri="s3://visualizations/example.html", + ) + ], + ) + ) + + pipeline_model = client.zen_store.create_pipeline( + PipelineRequest( + name=sample_name("pipeline"), + project=project_id, + ) + ) + + step_name = sample_name("step") + snapshot = client.zen_store.create_snapshot( + PipelineSnapshotRequest( + project=project_id, + run_name_template=sample_name("run"), + pipeline_configuration=PipelineConfiguration( + name=sample_name("pipeline-config") + ), + pipeline=pipeline_model.id, + stack=client.active_stack.id, + client_version="0.1.0", + server_version="0.1.0", + step_configurations={ + step_name: Step( + spec=StepSpec( + source=Source( + module="acme.step", type=SourceType.INTERNAL + ), + upstream_steps=[], + ), + config=StepConfiguration(name=step_name), + ) + }, + ) + ) + + pipeline_run, _ = client.zen_store.get_or_create_run( + PipelineRunRequest( + project=project_id, + id=uuid4(), + name=sample_name("run"), + snapshot=snapshot.id, + status=ExecutionStatus.RUNNING, + ) + ) + model = client.zen_store.create_model( + ModelRequest( + project=project_id, + name=sample_name("model"), + ) + ) + + # Try to find a deployer component from active stack first + deployer_id = None + active_stack_model = client.active_stack_model + if StackComponentType.MODEL_DEPLOYER in active_stack_model.components: + deployer_components = active_stack_model.components[StackComponentType.MODEL_DEPLOYER] + if deployer_components: + deployer_id = deployer_components[0].id + + if deployer_id is None: + # Fall back to listing all MODEL_DEPLOYER components + deployers = client.zen_store.list_stack_components( + ComponentFilter(type=StackComponentType.MODEL_DEPLOYER) + ) + if deployers.total > 0: + deployer_id = deployers.items[0].id + + if deployer_id is None: + pytest.skip( + "No deployer component available in the active stack or registered components. " + "Skipping deployment curated visualization test." + ) + + # Create a deployment + deployment = client.zen_store.create_deployment( + DeploymentRequest( + project=project_id, + name=sample_name("deployment"), + snapshot_id=snapshot.id, + deployer_id=deployer_id, + ) + ) + + try: + request = CuratedVisualizationRequest( + project=project_id, + artifact_version_id=artifact_version.id, + visualization_index=0, + resources=[ + CuratedVisualizationResource( + id=pipeline_model.id, + type=VisualizationResourceTypes.PIPELINE, + ), + CuratedVisualizationResource( + id=model.id, + type=VisualizationResourceTypes.MODEL, + ), + CuratedVisualizationResource( + id=pipeline_run.id, + type=VisualizationResourceTypes.PIPELINE_RUN, + ), + CuratedVisualizationResource( + id=snapshot.id, + type=VisualizationResourceTypes.PIPELINE_SNAPSHOT, + ), + CuratedVisualizationResource( + id=deployment.id, + type=VisualizationResourceTypes.DEPLOYMENT, + ), + CuratedVisualizationResource( + id=project_id, + type=VisualizationResourceTypes.PROJECT, + ), + ], + display_name="Initial visualization", + ) + + visualization = client.zen_store.create_curated_visualization(request) + + # Verify via zen_store filtering + for resource_type, resource_id in [ + (VisualizationResourceTypes.PIPELINE, pipeline_model.id), + (VisualizationResourceTypes.MODEL, model.id), + (VisualizationResourceTypes.PIPELINE_RUN, pipeline_run.id), + (VisualizationResourceTypes.PIPELINE_SNAPSHOT, snapshot.id), + (VisualizationResourceTypes.DEPLOYMENT, deployment.id), + (VisualizationResourceTypes.PROJECT, project_id), + ]: + result = client.zen_store.list_curated_visualizations( + CuratedVisualizationFilter( + project=project_id, + resource_type=resource_type, + resource_id=resource_id, + ) + ) + assert result.total == 1 + assert result.items[0].id == visualization.id + + # Verify via client convenience parameters + for param_name, param_value in [ + ("pipeline_id", pipeline_model.id), + ("model_id", model.id), + ("pipeline_run_id", pipeline_run.id), + ("pipeline_snapshot_id", snapshot.id), + ("deployment_id", deployment.id), + ("project_id", project_id), + ]: + result = client.list_curated_visualizations(**{param_name: param_value}) + assert result.total == 1, f"Failed for {param_name}" + assert result.items[0].id == visualization.id + + loaded = client.zen_store.get_curated_visualization( + visualization.id, hydrate=True + ) + assert loaded.display_name == "Initial visualization" + assert {resource.type for resource in loaded.resources} == { + VisualizationResourceTypes.PIPELINE, + VisualizationResourceTypes.MODEL, + VisualizationResourceTypes.PIPELINE_RUN, + VisualizationResourceTypes.PIPELINE_SNAPSHOT, + VisualizationResourceTypes.DEPLOYMENT, + VisualizationResourceTypes.PROJECT, + } + + with pytest.raises(EntityExistsError): + client.zen_store.create_curated_visualization( + CuratedVisualizationRequest( + project=project_id, + artifact_version_id=artifact_version.id, + visualization_index=0, + resources=[ + CuratedVisualizationResource( + id=pipeline_model.id, + type=VisualizationResourceTypes.PIPELINE, + ) + ], + ) + ) + + updated = client.zen_store.update_curated_visualization( + visualization_id=visualization.id, + visualization_update=CuratedVisualizationUpdate( + display_name="Updated", display_order=5 + ), + ) + assert updated.display_name == "Updated" + assert updated.display_order == 5 + + client.zen_store.delete_curated_visualization(visualization.id) + result = client.zen_store.list_curated_visualizations( + CuratedVisualizationFilter( + project=project_id, + resource_type=VisualizationResourceTypes.PIPELINE, + resource_id=pipeline_model.id, + ) + ) + assert result.total == 0 + finally: + # Clean up deployment + client.zen_store.delete_deployment(deployment.id) + + def test_curated_visualizations_project_only(self): + """Test project-level curated visualizations.""" + + client = Client() + project = client.active_project + + artifact = client.zen_store.create_artifact( + ArtifactRequest( + name=sample_name("artifact"), + project=project.id, + has_custom_name=True, + ) + ) + artifact_version = client.zen_store.create_artifact_version( + ArtifactVersionRequest( + artifact_id=artifact.id, + project=project.id, + version="1", + type=ArtifactType.DATA, + uri=sample_name("artifact_uri"), + materializer=Source( + module="acme.foo", type=SourceType.INTERNAL + ), + data_type=Source(module="acme.foo", type=SourceType.INTERNAL), + save_type=ArtifactSaveType.STEP_OUTPUT, + visualizations=[ + ArtifactVisualizationRequest( + type=VisualizationType.HTML, + uri="s3://visualizations/project.html", + ) + ], + ) + ) + + visualization = client.create_curated_visualization( + artifact_version_id=artifact_version.id, + visualization_index=0, + resources=[ + CuratedVisualizationResource( + id=project.id, + type=VisualizationResourceTypes.PROJECT, + ) + ], + display_name="Project visualization", + ) + + result = client.zen_store.list_curated_visualizations( + CuratedVisualizationFilter( + project=project.id, + resource_type=VisualizationResourceTypes.PROJECT, + resource_id=project.id, + ) + ) + assert result.total == 1 + assert result.items[0].id == visualization.id + + result = client.list_curated_visualizations(project_id=project.id) + assert result.total == 1 + assert result.items[0].id == visualization.id + + client.delete_curated_visualization(visualization.id) + client.delete_artifact_version(artifact_version.id) + client.delete_artifact(artifact.id) From 4d3f2591fe5669d2dd636709d3b79158512f5efa Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 11 Oct 2025 17:44:21 +0100 Subject: [PATCH 14/64] more fixes --- docs/book/how-to/artifacts/visualizations.md | 82 +++++-- src/zenml/client.py | 66 ++++-- src/zenml/enums.py | 7 + .../models/v2/core/curated_visualization.py | 85 +++---- src/zenml/models/v2/core/model.py | 17 +- src/zenml/zen_server/rbac/models.py | 1 + .../b77d123bce19_deployment_visualizations.py | 93 -------- .../schemas/curated_visualization_schemas.py | 40 ++-- .../zen_stores/schemas/deployment_schemas.py | 7 +- src/zenml/zen_stores/schemas/model_schemas.py | 5 +- src/zenml/zen_stores/sql_zen_store.py | 94 ++++---- .../functional/zen_stores/test_zen_store.py | 207 ++++++++++-------- 12 files changed, 351 insertions(+), 353 deletions(-) delete mode 100644 src/zenml/zen_stores/migrations/versions/b77d123bce19_deployment_visualizations.py diff --git a/docs/book/how-to/artifacts/visualizations.md b/docs/book/how-to/artifacts/visualizations.md index 69cf5badb7..34ae3bc1d1 100644 --- a/docs/book/how-to/artifacts/visualizations.md +++ b/docs/book/how-to/artifacts/visualizations.md @@ -67,7 +67,9 @@ There are three ways how you can add custom visualizations to the dashboard: ### Curated Visualizations Across Resources -Curated visualizations let you surface a specific artifact visualization across multiple ZenML resources. This is useful when you want to highlight the same dashboard in several places—for example, a model performance report that should be visible from the project overview, the training run, and the production deployment. +Curated visualizations let you surface a specific artifact visualization across multiple ZenML resources. Each curated visualization links to exactly one resource—for example, a model performance report that appears on the model detail page, or a deployment health dashboard that shows up in the deployment view. + +To cover multiple resources with the same visualization, create separate curated visualizations for each resource type. This gives you fine-grained control over which dashboards appear where. Curated visualizations currently support the following resources: @@ -78,13 +80,16 @@ Curated visualizations currently support the following resources: - **Pipeline Runs** – detailed diagnostics for specific executions. - **Pipeline Snapshots** – configuration/version comparisons for snapshot history. -You can create a curated visualization programmatically by linking an artifact visualization to one or more resources. The example below attaches a single visualization to multiple resource types, including a project: +You can create a curated visualization programmatically by linking an artifact visualization to a single resource. The example below shows how to create separate visualizations for different resource types: ```python from uuid import UUID from zenml.client import Client -from zenml.enums import VisualizationResourceTypes +from zenml.enums import ( + CuratedVisualizationSize, + VisualizationResourceTypes, +) from zenml.models import CuratedVisualizationResource client = Client() @@ -96,21 +101,45 @@ snapshot = pipeline_run.snapshot() deployment = client.list_deployments().items[0] model = client.list_models().items[0] -visualization = client.create_curated_visualization( +# Create a visualization for the model +model_viz = client.create_curated_visualization( + artifact_version_id=artifact_version_id, + visualization_index=0, + resource=CuratedVisualizationResource( + id=model.id, + type=VisualizationResourceTypes.MODEL + ), + display_name="Model performance dashboard", + size=CuratedVisualizationSize.FULL_WIDTH, +) + +# Create a visualization for the deployment +deployment_viz = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=0, - resources=[ - CuratedVisualizationResource(id=model.id, type=VisualizationResourceTypes.MODEL), - CuratedVisualizationResource(id=project.id, type=VisualizationResourceTypes.PROJECT), - CuratedVisualizationResource(id=deployment.id, type=VisualizationResourceTypes.DEPLOYMENT), - CuratedVisualizationResource(id=pipeline.id, type=VisualizationResourceTypes.PIPELINE), - CuratedVisualizationResource(id=pipeline_run.id, type=VisualizationResourceTypes.PIPELINE_RUN), - CuratedVisualizationResource(id=snapshot.id, type=VisualizationResourceTypes.PIPELINE_SNAPSHOT), - ], - display_name="Project performance dashboard", + resource=CuratedVisualizationResource( + id=deployment.id, + type=VisualizationResourceTypes.DEPLOYMENT + ), + display_name="Deployment health dashboard", + size=CuratedVisualizationSize.HALF_WIDTH, +) + +# Create a visualization for the project +project_viz = client.create_curated_visualization( + artifact_version_id=artifact_version_id, + visualization_index=0, + resource=CuratedVisualizationResource( + id=project.id, + type=VisualizationResourceTypes.PROJECT + ), + display_name="Project overview dashboard", + size=CuratedVisualizationSize.FULL_WIDTH, ) ``` +Note that each visualization is created separately and can reference the same artifact visualization (by using the same `artifact_version_id` and `visualization_index`). This allows you to show the same underlying visualization in multiple contexts while maintaining separate display settings for each resource. Use the optional `size` argument to control whether the visualization spans the full width of the dashboard or renders as a half-width tile. If omitted, the layout defaults to `CuratedVisualizationSize.FULL_WIDTH`. + To list curated visualizations for a specific resource, you can use the `Client.list_curated_visualizations` convenience parameters: ```python @@ -126,7 +155,7 @@ Each call returns a `Page[CuratedVisualizationResponse]` object so you can itera #### Updating curated visualizations -Once you've created a curated visualization, you can update its display name or order using `Client.update_curated_visualization`: +Once you've created a curated visualization, you can update its display name, order, or tile size using `Client.update_curated_visualization`: ```python from uuid import UUID @@ -135,6 +164,7 @@ client.update_curated_visualization( visualization_id=UUID(""), display_name="Updated Dashboard Title", display_order=10, + size=CuratedVisualizationSize.HALF_WIDTH, ) ``` @@ -144,7 +174,7 @@ When a visualization is no longer relevant, you can remove it entirely: client.delete_curated_visualization(visualization_id=UUID("")) ``` -#### Controlling display order +#### Controlling display order and size The optional `display_order` field determines how visualizations are sorted when displayed. Visualizations with lower order values appear first, while those with `None` (the default) appear at the end in creation order. @@ -155,31 +185,43 @@ When setting display orders, consider leaving gaps between values (e.g., 10, 20, visualization_a = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=0, - resources=[...], + resource=CuratedVisualizationResource( + id=model.id, + type=VisualizationResourceTypes.MODEL + ), display_order=10, # Primary dashboard + size=CuratedVisualizationSize.FULL_WIDTH, ) visualization_b = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=1, - resources=[...], + resource=CuratedVisualizationResource( + id=model.id, + type=VisualizationResourceTypes.MODEL + ), display_order=20, # Secondary metrics + size=CuratedVisualizationSize.HALF_WIDTH, # Compact chart beside the primary tile ) # Later, easily insert between them visualization_c = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=2, - resources=[...], + resource=CuratedVisualizationResource( + id=model.id, + type=VisualizationResourceTypes.MODEL + ), display_order=15, # Now appears between A and B + size=CuratedVisualizationSize.HALF_WIDTH, ) ``` #### RBAC visibility -Curated visualizations respect the access permissions of every resource they're linked to. A user can only see a curated visualization if they have read access to **all** the resources it targets. If a user lacks permission for any linked resource, the visualization will be hidden from their view. +Curated visualizations respect the access permissions of the resource they're linked to. A user can only see a curated visualization if they have read access to the specific resource it targets. If a user lacks permission for the linked resource, the visualization will be hidden from their view. -For example, if you create a visualization linked to both a project and a deployment, users must have read access to both the project and that specific deployment to see the visualization. This ensures that curated visualizations never inadvertently expose information from resources a user shouldn't access. +For example, if you create a visualization linked to a specific deployment, only users with read access to that deployment will see the visualization. If you need the same visualization to appear in different contexts with different access controls (e.g., on both a project page and a deployment page), create separate curated visualizations for each resource. This ensures that visualizations never inadvertently expose information from resources a user shouldn't access, while giving you fine-grained control over visibility. ### Visualization via Special Return Types diff --git a/src/zenml/client.py b/src/zenml/client.py index 0f4d00dac5..2efb301c40 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -61,6 +61,7 @@ from zenml.enums import ( ArtifactType, ColorVariants, + CuratedVisualizationSize, DeploymentStatus, LogicalOperators, ModelStages, @@ -3750,14 +3751,15 @@ def create_curated_visualization( artifact_version_id: UUID, visualization_index: int, *, - resources: List[CuratedVisualizationResource], + resource: CuratedVisualizationResource, project_id: Optional[UUID] = None, display_name: Optional[str] = None, display_order: Optional[int] = None, + size: CuratedVisualizationSize = CuratedVisualizationSize.FULL_WIDTH, ) -> CuratedVisualizationResponse: - """Create a curated visualization associated with arbitrary resources. + """Create a curated visualization associated with a resource. - Curated visualizations can be attached to any combination of the following + Curated visualizations can be attached to any of the following ZenML resource types to provide contextual dashboards throughout the ML lifecycle: @@ -3770,29 +3772,28 @@ def create_curated_visualization( - **Pipeline Snapshots** (VisualizationResourceTypes.PIPELINE_SNAPSHOT): Link to captured pipeline configurations - A single visualization can be linked to multiple resources across different - types. For example, attach a model performance dashboard to both the - deployment and the pipeline run that produced the deployed model. + Each visualization is linked to exactly one resource. Args: artifact_version_id: The ID of the artifact version containing the visualization. visualization_index: The index of the visualization within the artifact version. - resources: One or more resources to associate with the visualization. - Each entry should be a `CuratedVisualizationResource` containing + resource: The resource to associate with the visualization. + Should be a `CuratedVisualizationResource` containing the resource ID and type (e.g., DEPLOYMENT, PIPELINE, PIPELINE_RUN, PIPELINE_SNAPSHOT). project_id: The ID of the project to associate with the visualization. display_name: The display name of the visualization. display_order: The display order of the visualization. + size: The layout size of the visualization in the dashboard. Returns: The created curated visualization. Raises: - ValueError: If resources list is empty. + ValueError: If resource is not provided. """ - if not resources: - raise ValueError("resources must not be empty") + if not resource: + raise ValueError("resource must be provided") request = CuratedVisualizationRequest( project=project_id or self.active_project.id, @@ -3800,7 +3801,8 @@ def create_curated_visualization( visualization_index=visualization_index, display_name=display_name, display_order=display_order, - resources=resources, + size=size, + resource=resource, ) return self.zen_store.create_curated_visualization(request) @@ -3812,6 +3814,7 @@ def add_visualization_to_deployment( *, display_name: Optional[str] = None, display_order: Optional[int] = None, + size: CuratedVisualizationSize = CuratedVisualizationSize.FULL_WIDTH, ) -> CuratedVisualizationResponse: """Attach a curated visualization to a deployment. @@ -3821,6 +3824,7 @@ def add_visualization_to_deployment( visualization_index: The index of the visualization within the artifact version. display_name: Optional display name for the visualization. display_order: Optional display order for sorting visualizations. + size: Layout size defining the visualization width on the dashboard. Returns: The created curated visualization. @@ -3829,15 +3833,14 @@ def add_visualization_to_deployment( return self.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=visualization_index, - resources=[ - CuratedVisualizationResource( - id=deployment_id, - type=VisualizationResourceTypes.DEPLOYMENT, - ) - ], + resource=CuratedVisualizationResource( + id=deployment_id, + type=VisualizationResourceTypes.DEPLOYMENT, + ), project_id=deployment.project_id, display_name=display_name, display_order=display_order, + size=size, ) def list_curated_visualizations( @@ -3855,6 +3858,7 @@ def list_curated_visualizations( size: Optional[int] = None, sort_by: Optional[str] = None, visualization_index: Optional[int] = None, + tile_size: Optional[CuratedVisualizationSize] = None, hydrate: bool = False, ) -> Page[CuratedVisualizationResponse]: """List curated visualizations, optionally scoped to a resource. @@ -3883,6 +3887,7 @@ def list_curated_visualizations( size: The maximum size of all pages. sort_by: The column to sort by. visualization_index: The index of the visualization to filter by. + tile_size: The layout size of the visualization tiles to filter by. hydrate: Flag deciding whether to hydrate the output model(s) by including metadata fields in the response. @@ -3896,18 +3901,30 @@ def list_curated_visualizations( """ # Build convenience params dict mapping parameter names to (value, resource_type) tuples convenience_params = { - "deployment_id": (deployment_id, VisualizationResourceTypes.DEPLOYMENT), + "deployment_id": ( + deployment_id, + VisualizationResourceTypes.DEPLOYMENT, + ), "model_id": (model_id, VisualizationResourceTypes.MODEL), "pipeline_id": (pipeline_id, VisualizationResourceTypes.PIPELINE), - "pipeline_run_id": (pipeline_run_id, VisualizationResourceTypes.PIPELINE_RUN), - "pipeline_snapshot_id": (pipeline_snapshot_id, VisualizationResourceTypes.PIPELINE_SNAPSHOT), + "pipeline_run_id": ( + pipeline_run_id, + VisualizationResourceTypes.PIPELINE_RUN, + ), + "pipeline_snapshot_id": ( + pipeline_snapshot_id, + VisualizationResourceTypes.PIPELINE_SNAPSHOT, + ), "project_id": (project_id, VisualizationResourceTypes.PROJECT), } # Filter to only provided parameters provided = { param_name: (param_value, param_type) - for param_name, (param_value, param_type) in convenience_params.items() + for param_name, ( + param_value, + param_type, + ) in convenience_params.items() if param_value is not None } @@ -3961,6 +3978,8 @@ def list_curated_visualizations( filter_model.sort_by = sort_by if visualization_index is not None: filter_model.visualization_index = visualization_index + if tile_size is not None: + filter_model.size = tile_size return self.zen_store.list_curated_visualizations( filter_model=filter_model, @@ -3973,6 +3992,7 @@ def update_curated_visualization( *, display_name: Optional[str] = None, display_order: Optional[int] = None, + size: Optional[CuratedVisualizationSize] = None, ) -> CuratedVisualizationResponse: """Update display metadata for a curated visualization. @@ -3980,6 +4000,7 @@ def update_curated_visualization( visualization_id: The ID of the curated visualization to update. display_name: New display name for the visualization. display_order: New display order for the visualization. + size: Updated layout size for the visualization. Returns: The updated deployment visualization. @@ -3987,6 +4008,7 @@ def update_curated_visualization( update_model = CuratedVisualizationUpdate( display_name=display_name, display_order=display_order, + size=size, ) return self.zen_store.update_curated_visualization( visualization_id=visualization_id, diff --git a/src/zenml/enums.py b/src/zenml/enums.py index d487463a75..85aae8b6f4 100644 --- a/src/zenml/enums.py +++ b/src/zenml/enums.py @@ -446,6 +446,13 @@ class VisualizationResourceTypes(StrEnum): PROJECT = "project" # Project-level dashboards +class CuratedVisualizationSize(StrEnum): + """Layout size options for curated visualizations.""" + + FULL_WIDTH = "full_width" + HALF_WIDTH = "half_width" + + class SecretResourceTypes(StrEnum): """All possible resource types for adding secrets.""" diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index a0df0503c1..3d1945f812 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -24,9 +24,12 @@ ) from uuid import UUID -from pydantic import Field, model_validator +from pydantic import Field -from zenml.enums import VisualizationResourceTypes +from zenml.enums import ( + CuratedVisualizationSize, + VisualizationResourceTypes, +) from zenml.models.v2.base.base import BaseUpdate from zenml.models.v2.base.filter import AnyQuery from zenml.models.v2.base.scoped import ( @@ -56,8 +59,8 @@ class CuratedVisualizationRequest(ProjectScopedRequest): """Request model for curated visualizations. - Curated visualizations can be attached to any combination of the following - resource types: + Each curated visualization is linked to exactly one resource of the following + types: - **Deployments**: Surface visualizations on deployment dashboards - **Models**: Highlight evaluation dashboards and monitoring views next to registered models @@ -66,10 +69,8 @@ class CuratedVisualizationRequest(ProjectScopedRequest): - **Pipeline Snapshots**: Link visualizations to snapshot configurations - **Projects**: Provide high-level project dashboards and KPI overviews - A single visualization can be linked to multiple resources of different types, - enabling reuse across the ML workflow. For example, a model performance - dashboard could be attached to both a deployment and the pipeline run that - produced the deployed model. + To attach a visualization to multiple resources, create separate curated + visualization entries for each resource. """ artifact_version_id: UUID = Field( @@ -89,32 +90,21 @@ class CuratedVisualizationRequest(ProjectScopedRequest): default=None, title="The display order of the visualization.", ) - resources: List[CuratedVisualizationResource] = Field( - title="Resources associated with this visualization.", + size: CuratedVisualizationSize = Field( + default=CuratedVisualizationSize.FULL_WIDTH, + title="The layout size of the visualization.", description=( - "List of resources (deployments, models, pipelines, pipeline runs, " - "pipeline snapshots, projects) that should surface this visualization. " - "Must include at least one resource. Multiple resources of " - "different types can be specified to reuse visualizations " - "across the ML workflow." + "Controls how much horizontal space the visualization occupies " + "on the dashboard." + ), + ) + resource: CuratedVisualizationResource = Field( + title="Resource associated with this visualization.", + description=( + "The single resource (deployment, model, pipeline, pipeline run, " + "pipeline snapshot, or project) that should surface this visualization." ), ) - - @model_validator(mode="after") - def validate_resources(self) -> "CuratedVisualizationRequest": - """Ensure that at least one resource is associated with the visualization. - - Returns: - The validated request instance. - - Raises: - ValueError: If no resources are provided. - """ - if not self.resources: - raise ValueError( - "Curated visualizations must be associated with at least one resource." - ) - return self # ------------------ Update Model ------------------ @@ -131,6 +121,10 @@ class CuratedVisualizationUpdate(BaseUpdate): default=None, title="The new display order of the visualization.", ) + size: Optional[CuratedVisualizationSize] = Field( + default=None, + title="The updated layout size of the visualization.", + ) # ------------------ Response Model ------------------ @@ -155,9 +149,9 @@ class CuratedVisualizationResponseBody(ProjectScopedResponseBody): default=None, title="The display order of the visualization.", ) - resources: List[CuratedVisualizationResource] = Field( - default_factory=list, - title="Resources exposing the visualization.", + size: CuratedVisualizationSize = Field( + default=CuratedVisualizationSize.FULL_WIDTH, + title="The layout size of the visualization.", ) @@ -233,21 +227,22 @@ def display_order(self) -> Optional[int]: return self.get_body().display_order @property - def artifact_version(self) -> Optional["ArtifactVersionResponse"]: - """The artifact version resource. + def size(self) -> CuratedVisualizationSize: + """The layout size of the visualization. Returns: - The artifact version resource if included. + The layout size of the visualization. """ - return self.get_resources().artifact_version + return self.get_body().size - def visualization_resources(self) -> List[CuratedVisualizationResource]: - """Return the resources exposing this visualization. + @property + def artifact_version(self) -> Optional["ArtifactVersionResponse"]: + """The artifact version resource. Returns: - List of associated resources. + The artifact version resource if included. """ - return self.get_body().resources + return self.get_resources().artifact_version # ------------------ Filter Model ------------------ @@ -286,6 +281,10 @@ class CuratedVisualizationFilter(ProjectScopedFilter): default=None, description="Display order of the visualization.", ) + size: Optional[CuratedVisualizationSize] = Field( + default=None, + description="Layout size of the visualization.", + ) resource_type: Optional[VisualizationResourceTypes] = Field( default=None, description="Type of the resource exposing the visualization.", @@ -364,6 +363,8 @@ def get_custom_filters( custom_filters.append( getattr(table, "display_order") == self.display_order ) + if self.size is not None: + custom_filters.append(getattr(table, "size") == self.size) # resource-based filtering is handled within the store implementation return custom_filters diff --git a/src/zenml/models/v2/core/model.py b/src/zenml/models/v2/core/model.py index 32eb7dd2a9..68ff7f9582 100644 --- a/src/zenml/models/v2/core/model.py +++ b/src/zenml/models/v2/core/model.py @@ -206,14 +206,6 @@ class ModelResponse( max_length=STR_FIELD_MAX_LENGTH, ) - def visualizations(self) -> List["CuratedVisualizationResponse"]: - """Return curated visualizations linked to the model. - - Returns: - A list of curated visualization responses. - """ - return self.get_resources().visualizations - def get_hydrated_version(self) -> "ModelResponse": """Get the hydrated version of this model. @@ -324,6 +316,15 @@ def save_models_to_registry(self) -> bool: """ return self.get_metadata().save_models_to_registry + @property + def visualizations(self) -> List["CuratedVisualizationResponse"]: + """The `visualizations` property. + + Returns: + the value of the property. + """ + return self.get_resources().visualizations + # Helper functions @property def versions(self) -> List["Model"]: diff --git a/src/zenml/zen_server/rbac/models.py b/src/zenml/zen_server/rbac/models.py index ba4cd4dcc4..ca73c7d931 100644 --- a/src/zenml/zen_server/rbac/models.py +++ b/src/zenml/zen_server/rbac/models.py @@ -103,6 +103,7 @@ def from_visualization_type( VisualizationResourceTypes.PIPELINE: cls.PIPELINE, VisualizationResourceTypes.PIPELINE_RUN: cls.PIPELINE_RUN, VisualizationResourceTypes.PIPELINE_SNAPSHOT: cls.PIPELINE_SNAPSHOT, + VisualizationResourceTypes.PROJECT: cls.PROJECT, } return mapping.get(visualization_type) diff --git a/src/zenml/zen_stores/migrations/versions/b77d123bce19_deployment_visualizations.py b/src/zenml/zen_stores/migrations/versions/b77d123bce19_deployment_visualizations.py deleted file mode 100644 index 83b2d1111e..0000000000 --- a/src/zenml/zen_stores/migrations/versions/b77d123bce19_deployment_visualizations.py +++ /dev/null @@ -1,93 +0,0 @@ -"""deployment visualizations [b77d123bce19]. - -Revision ID: b77d123bce19 -Revises: 0.90.0rc0 -Create Date: 2025-09-29 14:23:48.630888 - -""" - -import sqlalchemy as sa -import sqlmodel -from alembic import op - -# revision identifiers, used by Alembic. -revision = "b77d123bce19" -down_revision = "0.90.0rc0" -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! ### - op.create_table( - "deployment_visualization", - sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column("created", sa.DateTime(), nullable=False), - sa.Column("updated", sa.DateTime(), nullable=False), - sa.Column("project_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column( - "deployment_id", sqlmodel.sql.sqltypes.GUID(), nullable=False - ), - sa.Column( - "artifact_version_id", sqlmodel.sql.sqltypes.GUID(), nullable=False - ), - sa.Column("visualization_index", sa.Integer(), nullable=False), - sa.Column( - "display_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True - ), - sa.Column("display_order", sa.Integer(), nullable=True), - sa.ForeignKeyConstraint( - ["artifact_version_id"], - ["artifact_version.id"], - name="fk_deployment_visualization_artifact_version_id_artifact_version", - ondelete="CASCADE", - ), - sa.ForeignKeyConstraint( - ["deployment_id"], - ["deployment.id"], - name="fk_deployment_visualization_deployment_id_deployment", - ondelete="CASCADE", - ), - sa.ForeignKeyConstraint( - ["project_id"], - ["project.id"], - name="fk_deployment_visualization_project_id_project", - ondelete="CASCADE", - ), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "deployment_id", - "artifact_version_id", - "visualization_index", - name="unique_deployment_visualization", - ), - ) - with op.batch_alter_table( - "deployment_visualization", schema=None - ) as batch_op: - batch_op.create_index( - "ix_deployment_visualization_deployment_id", - ["deployment_id"], - unique=False, - ) - batch_op.create_index( - "ix_deployment_visualization_display_order", - ["display_order"], - unique=False, - ) - - # ### 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( - "deployment_visualization", schema=None - ) as batch_op: - batch_op.drop_index("ix_deployment_visualization_display_order") - batch_op.drop_index("ix_deployment_visualization_deployment_id") - - op.drop_table("deployment_visualization") - # ### end Alembic commands ### diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index a39f7047a0..2233f774ac 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -21,7 +21,7 @@ from sqlalchemy.sql.base import ExecutableOption from sqlmodel import Field, Relationship, SQLModel -from zenml.enums import VisualizationResourceTypes +from zenml.enums import CuratedVisualizationSize, VisualizationResourceTypes from zenml.models.v2.core.curated_visualization import ( CuratedVisualizationRequest, CuratedVisualizationResponse, @@ -30,9 +30,6 @@ CuratedVisualizationResponseResources, CuratedVisualizationUpdate, ) -from zenml.models.v2.misc.curated_visualization import ( - CuratedVisualizationResource, -) from zenml.zen_stores.schemas.base_schemas import BaseSchema from zenml.zen_stores.schemas.project_schemas import ProjectSchema from zenml.zen_stores.schemas.schema_utils import ( @@ -52,8 +49,6 @@ class CuratedVisualizationResourceSchema(SQLModel, table=True): __table_args__ = ( UniqueConstraint( "visualization_id", - "resource_id", - "resource_type", name="unique_curated_visualization_resource", ), build_index(__tablename__, ["resource_id", "resource_type"]), @@ -72,7 +67,7 @@ class CuratedVisualizationResourceSchema(SQLModel, table=True): resource_type: str = Field(nullable=False) visualization: "CuratedVisualizationSchema" = Relationship( - back_populates="resources", + back_populates="resource", ) @@ -107,13 +102,17 @@ class CuratedVisualizationSchema(BaseSchema, table=True): visualization_index: int = Field(nullable=False) display_name: Optional[str] = Field(default=None) display_order: Optional[int] = Field(default=None) + size: CuratedVisualizationSize = Field( + default=CuratedVisualizationSize.FULL_WIDTH, nullable=False + ) artifact_version: Optional["ArtifactVersionSchema"] = Relationship( sa_relationship_kwargs={"lazy": "selectin"} ) - resources: List[CuratedVisualizationResourceSchema] = Relationship( + resource: Optional[CuratedVisualizationResourceSchema] = Relationship( back_populates="visualization", sa_relationship_kwargs={"lazy": "selectin", "cascade": "all, delete"}, + uselist=False, ) @classmethod @@ -141,7 +140,7 @@ def get_query_options( options.extend( [ selectinload(jl_arg(cls.artifact_version)), - selectinload(jl_arg(cls.resources)), + selectinload(jl_arg(cls.resource)), ] ) @@ -165,6 +164,7 @@ def from_request( visualization_index=request.visualization_index, display_name=request.display_name, display_order=request.display_order, + size=request.size, ) def update( @@ -206,14 +206,6 @@ def to_model( Returns: The created response model. """ - resources = [ - CuratedVisualizationResource( - id=resource.resource_id, - type=VisualizationResourceTypes(resource.resource_type), - ) - for resource in self.resources - ] - body = CuratedVisualizationResponseBody( project_id=self.project_id, created=self.created, @@ -222,7 +214,7 @@ def to_model( visualization_index=self.visualization_index, display_name=self.display_name, display_order=self.display_order, - resources=resources, + size=self.size, ) metadata = None @@ -270,14 +262,20 @@ def curated_visualization_relationship_kwargs( The relationship will be read-only (viewonly=True) and eagerly loaded via selectin loading. """ + def _primaryjoin(): return and_( - CuratedVisualizationResourceSchema.resource_type == resource_type.value, - foreign(CuratedVisualizationResourceSchema.resource_id) == parent_column_factory(), + CuratedVisualizationResourceSchema.resource_type + == resource_type.value, + foreign(CuratedVisualizationResourceSchema.resource_id) + == parent_column_factory(), ) def _secondaryjoin(): - return CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id + return ( + CuratedVisualizationSchema.id + == CuratedVisualizationResourceSchema.visualization_id + ) return dict( secondary="curated_visualization_resource", diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index df339ea0c6..c70cc48f99 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -41,6 +41,9 @@ from zenml.utils.time_utils import utc_now from zenml.zen_stores.schemas.base_schemas import NamedSchema from zenml.zen_stores.schemas.component_schemas import StackComponentSchema +from zenml.zen_stores.schemas.curated_visualization_schemas import ( + curated_visualization_relationship_kwargs, +) from zenml.zen_stores.schemas.pipeline_snapshot_schemas import ( PipelineSnapshotSchema, ) @@ -55,10 +58,6 @@ ) from zenml.zen_stores.schemas.tag_schemas import TagSchema -from zenml.zen_stores.schemas.curated_visualization_schemas import ( - curated_visualization_relationship_kwargs, -) - logger = get_logger(__name__) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 2658f3c09d..0bfa3cd58f 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -59,7 +59,6 @@ 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.curated_visualization_schemas import ( - CuratedVisualizationResourceSchema, CuratedVisualizationSchema, curated_visualization_relationship_kwargs, ) @@ -167,7 +166,9 @@ def get_query_options( [ joinedload(jl_arg(ModelSchema.user)), # joinedload(jl_arg(ModelSchema.tags)), - selectinload(jl_arg(ModelSchema.visualizations)).selectinload( + selectinload( + jl_arg(ModelSchema.visualizations) + ).selectinload( jl_arg(CuratedVisualizationSchema.artifact_version) ), ] diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 8316c31bc5..30d9df6e5d 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5523,6 +5523,10 @@ def create_curated_visualization( visualization_index=visualization.visualization_index, ) + # Validate and fetch the resource + resource_request = visualization.resource + + # Map resource types to their corresponding schema classes resource_schema_map: Dict[ VisualizationResourceTypes, Type[BaseSchema] ] = { @@ -5534,75 +5538,55 @@ def create_curated_visualization( VisualizationResourceTypes.PROJECT: ProjectSchema, } - # Group resource requests by type for batch fetching - resources_by_type: Dict[ - VisualizationResourceTypes, List[UUID] - ] = defaultdict(list) - for resource_request in visualization.resources: - if resource_request.type not in resource_schema_map: - raise IllegalOperationError( - f"Curated visualizations are not supported for resource type '{resource_request.type.value}'." - ) - resources_by_type[resource_request.type].append( - resource_request.id + if resource_request.type not in resource_schema_map: + raise IllegalOperationError( + f"Invalid resource type: {resource_request.type}" ) - # Batch fetch resources by type - fetched_resources: Dict[UUID, BaseSchema] = {} - for resource_type, resource_ids in resources_by_type.items(): - schema_class = resource_schema_map[resource_type] - schemas = session.exec( - select(schema_class).where( - schema_class.id.in_(resource_ids) - ) - ).all() - - # Verify all resources were found - found_ids = {schema.id for schema in schemas} - missing_ids = set(resource_ids) - found_ids - if missing_ids: - raise KeyError( - f"Resources of type '{resource_type.value}' with IDs " - f"{missing_ids} not found." - ) - - for schema in schemas: - fetched_resources[schema.id] = schema - - # Validate project scope and check for duplicates - resources: List[CuratedVisualizationResourceSchema] = [] - for resource_request in visualization.resources: - resource_schema = fetched_resources[resource_request.id] + # Fetch the single resource schema + schema_class = resource_schema_map[resource_request.type] + resource_schema = session.exec( + select(schema_class).where( + schema_class.id == resource_request.id + ) + ).first() - resource_project_id = getattr( - resource_schema, "project_id", None + if not resource_schema: + raise KeyError( + f"Resource of type '{resource_request.type.value}' " + f"with ID {resource_request.id} not found." ) + + # Validate project scope for the resource + if hasattr(resource_schema, "project_id"): + resource_project_id = resource_schema.project_id if resource_project_id and resource_project_id != project_id: raise IllegalOperationError( - "Curated visualizations must reference resources " - "within the same project as the artifact version." + f"Resource {resource_request.type.value} with ID " + f"{resource_request.id} belongs to a different project than " + f"the curated visualization (project ID: {project_id})." ) - self._assert_curated_visualization_duplicate( - session=session, - artifact_version_id=visualization.artifact_version_id, - visualization_index=visualization.visualization_index, - resource_id=resource_request.id, - resource_type=resource_request.type, - ) + # Check for duplicate + self._assert_curated_visualization_duplicate( + session=session, + artifact_version_id=visualization.artifact_version_id, + visualization_index=visualization.visualization_index, + resource_id=resource_request.id, + resource_type=resource_request.type, + ) - resources.append( - CuratedVisualizationResourceSchema( - resource_id=resource_request.id, - resource_type=resource_request.type.value, - ) - ) + # Create the resource link + resource = CuratedVisualizationResourceSchema( + resource_id=resource_request.id, + resource_type=resource_request.type.value, + ) schema: CuratedVisualizationSchema = ( CuratedVisualizationSchema.from_request(visualization) ) schema.project_id = project_id - schema.resources = resources + schema.resource = resource session.add(schema) session.commit() diff --git a/tests/integration/functional/zen_stores/test_zen_store.py b/tests/integration/functional/zen_stores/test_zen_store.py index 3931b19389..9b8b0f72cb 100644 --- a/tests/integration/functional/zen_stores/test_zen_store.py +++ b/tests/integration/functional/zen_stores/test_zen_store.py @@ -68,6 +68,7 @@ ArtifactSaveType, ArtifactType, ColorVariants, + CuratedVisualizationSize, ExecutionStatus, MetadataResourceTypes, ModelStages, @@ -5725,9 +5726,21 @@ def test_tag_filter_with_resource_type(clean_client: "Client"): class TestCuratedVisualizations: + """Test curated visualizations.""" + def test_curated_visualizations_across_resources(self): - """Test creating, listing, updating, and deleting curated visualizations.""" + """Test creating, listing, updating, and deleting curated visualizations. + + Each curated visualization is linked to exactly one resource. This test + creates separate visualizations for each supported resource type: + - **Deployments** (VisualizationResourceTypes.DEPLOYMENT) + - **Models** (VisualizationResourceTypes.MODEL) + - **Pipelines** (VisualizationResourceTypes.PIPELINE) + - **Pipeline Runs** (VisualizationResourceTypes.PIPELINE_RUN) + - **Pipeline Snapshots** (VisualizationResourceTypes.PIPELINE_SNAPSHOT) + - **Projects** (VisualizationResourceTypes.PROJECT) + """ client = Client() project_id = client.active_project.id @@ -5812,7 +5825,9 @@ def test_curated_visualizations_across_resources(self): deployer_id = None active_stack_model = client.active_stack_model if StackComponentType.MODEL_DEPLOYER in active_stack_model.components: - deployer_components = active_stack_model.components[StackComponentType.MODEL_DEPLOYER] + deployer_components = active_stack_model.components[ + StackComponentType.MODEL_DEPLOYER + ] if deployer_components: deployer_id = deployer_components[0].id @@ -5841,50 +5856,37 @@ def test_curated_visualizations_across_resources(self): ) try: - request = CuratedVisualizationRequest( - project=project_id, - artifact_version_id=artifact_version.id, - visualization_index=0, - resources=[ - CuratedVisualizationResource( - id=pipeline_model.id, - type=VisualizationResourceTypes.PIPELINE, - ), - CuratedVisualizationResource( - id=model.id, - type=VisualizationResourceTypes.MODEL, - ), - CuratedVisualizationResource( - id=pipeline_run.id, - type=VisualizationResourceTypes.PIPELINE_RUN, - ), - CuratedVisualizationResource( - id=snapshot.id, - type=VisualizationResourceTypes.PIPELINE_SNAPSHOT, - ), - CuratedVisualizationResource( - id=deployment.id, - type=VisualizationResourceTypes.DEPLOYMENT, - ), - CuratedVisualizationResource( - id=project_id, - type=VisualizationResourceTypes.PROJECT, - ), - ], - display_name="Initial visualization", - ) - - visualization = client.zen_store.create_curated_visualization(request) - - # Verify via zen_store filtering - for resource_type, resource_id in [ + # Create a separate visualization for each resource type + resource_configs = [ (VisualizationResourceTypes.PIPELINE, pipeline_model.id), (VisualizationResourceTypes.MODEL, model.id), (VisualizationResourceTypes.PIPELINE_RUN, pipeline_run.id), (VisualizationResourceTypes.PIPELINE_SNAPSHOT, snapshot.id), (VisualizationResourceTypes.DEPLOYMENT, deployment.id), (VisualizationResourceTypes.PROJECT, project_id), - ]: + ] + + visualizations = {} + for idx, (resource_type, resource_id) in enumerate( + resource_configs + ): + viz = client.zen_store.create_curated_visualization( + CuratedVisualizationRequest( + project=project_id, + artifact_version_id=artifact_version.id, + visualization_index=idx, + resource=CuratedVisualizationResource( + id=resource_id, + type=resource_type, + ), + display_name=f"{resource_type.value} visualization", + ) + ) + visualizations[resource_type] = viz + assert viz.size == CuratedVisualizationSize.FULL_WIDTH + + # Verify via zen_store filtering - each resource should have exactly one visualization + for resource_type, resource_id in resource_configs: result = client.zen_store.list_curated_visualizations( CuratedVisualizationFilter( project=project_id, @@ -5893,73 +5895,108 @@ def test_curated_visualizations_across_resources(self): ) ) assert result.total == 1 - assert result.items[0].id == visualization.id + assert result.items[0].id == visualizations[resource_type].id # Verify via client convenience parameters - for param_name, param_value in [ - ("pipeline_id", pipeline_model.id), - ("model_id", model.id), - ("pipeline_run_id", pipeline_run.id), - ("pipeline_snapshot_id", snapshot.id), - ("deployment_id", deployment.id), - ("project_id", project_id), - ]: - result = client.list_curated_visualizations(**{param_name: param_value}) + convenience_params = [ + ( + "pipeline_id", + pipeline_model.id, + VisualizationResourceTypes.PIPELINE, + ), + ("model_id", model.id, VisualizationResourceTypes.MODEL), + ( + "pipeline_run_id", + pipeline_run.id, + VisualizationResourceTypes.PIPELINE_RUN, + ), + ( + "pipeline_snapshot_id", + snapshot.id, + VisualizationResourceTypes.PIPELINE_SNAPSHOT, + ), + ( + "deployment_id", + deployment.id, + VisualizationResourceTypes.DEPLOYMENT, + ), + ("project_id", project_id, VisualizationResourceTypes.PROJECT), + ] + + for ( + param_name, + param_value, + expected_resource_type, + ) in convenience_params: + result = client.list_curated_visualizations( + **{param_name: param_value} + ) assert result.total == 1, f"Failed for {param_name}" - assert result.items[0].id == visualization.id + assert ( + result.items[0].id + == visualizations[expected_resource_type].id + ) + # Test hydrate/get loaded = client.zen_store.get_curated_visualization( - visualization.id, hydrate=True - ) - assert loaded.display_name == "Initial visualization" - assert {resource.type for resource in loaded.resources} == { - VisualizationResourceTypes.PIPELINE, - VisualizationResourceTypes.MODEL, - VisualizationResourceTypes.PIPELINE_RUN, - VisualizationResourceTypes.PIPELINE_SNAPSHOT, - VisualizationResourceTypes.DEPLOYMENT, - VisualizationResourceTypes.PROJECT, - } + visualizations[VisualizationResourceTypes.PIPELINE].id, + hydrate=True, + ) + assert ( + loaded.display_name + == f"{VisualizationResourceTypes.PIPELINE.value} visualization" + ) + assert loaded.size == CuratedVisualizationSize.FULL_WIDTH + # Test duplicate creation - same artifact_version + visualization_index + resource should fail with pytest.raises(EntityExistsError): client.zen_store.create_curated_visualization( CuratedVisualizationRequest( project=project_id, artifact_version_id=artifact_version.id, - visualization_index=0, - resources=[ - CuratedVisualizationResource( - id=pipeline_model.id, - type=VisualizationResourceTypes.PIPELINE, - ) - ], + visualization_index=0, # Same index as pipeline visualization + resource=CuratedVisualizationResource( + id=pipeline_model.id, + type=VisualizationResourceTypes.PIPELINE, + ), ) ) + # Test update updated = client.zen_store.update_curated_visualization( - visualization_id=visualization.id, + visualization_id=visualizations[ + VisualizationResourceTypes.MODEL + ].id, visualization_update=CuratedVisualizationUpdate( - display_name="Updated", display_order=5 + display_name="Updated", + display_order=5, + size=CuratedVisualizationSize.HALF_WIDTH, ), ) assert updated.display_name == "Updated" assert updated.display_order == 5 + assert updated.size == CuratedVisualizationSize.HALF_WIDTH + + # Delete all visualizations + for viz in visualizations.values(): + client.zen_store.delete_curated_visualization(viz.id) - client.zen_store.delete_curated_visualization(visualization.id) - result = client.zen_store.list_curated_visualizations( - CuratedVisualizationFilter( - project=project_id, - resource_type=VisualizationResourceTypes.PIPELINE, - resource_id=pipeline_model.id, + # Verify all deleted + for resource_type, resource_id in resource_configs: + result = client.zen_store.list_curated_visualizations( + CuratedVisualizationFilter( + project=project_id, + resource_type=resource_type, + resource_id=resource_id, + ) ) - ) - assert result.total == 0 + assert result.total == 0 finally: # Clean up deployment client.zen_store.delete_deployment(deployment.id) def test_curated_visualizations_project_only(self): - """Test project-level curated visualizations.""" + """Test project-level curated visualizations with single resource.""" client = Client() project = client.active_project @@ -5995,12 +6032,10 @@ def test_curated_visualizations_project_only(self): visualization = client.create_curated_visualization( artifact_version_id=artifact_version.id, visualization_index=0, - resources=[ - CuratedVisualizationResource( - id=project.id, - type=VisualizationResourceTypes.PROJECT, - ) - ], + resource=CuratedVisualizationResource( + id=project.id, + type=VisualizationResourceTypes.PROJECT, + ), display_name="Project visualization", ) From b4ae2d3f5a19c60bde1f22c26dcf82be5e6f71ea Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sun, 12 Oct 2025 19:56:20 +0100 Subject: [PATCH 15/64] review fix and adding size --- src/zenml/client.py | 2 +- .../models/v2/core/curated_visualization.py | 9 +-- .../schemas/curated_visualization_schemas.py | 57 +++++++++++-------- src/zenml/zen_stores/sql_zen_store.py | 7 ++- 4 files changed, 44 insertions(+), 31 deletions(-) diff --git a/src/zenml/client.py b/src/zenml/client.py index 2efb301c40..8b16bdad23 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -3979,7 +3979,7 @@ def list_curated_visualizations( if visualization_index is not None: filter_model.visualization_index = visualization_index if tile_size is not None: - filter_model.size = tile_size + filter_model.tile_size = tile_size return self.zen_store.list_curated_visualizations( filter_model=filter_model, diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 3d1945f812..2a5f940b39 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -255,6 +255,7 @@ class CuratedVisualizationFilter(ProjectScopedFilter): *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS, "resource_id", "resource_type", + "tile_size", ] CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [ *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS, @@ -281,9 +282,9 @@ class CuratedVisualizationFilter(ProjectScopedFilter): default=None, description="Display order of the visualization.", ) - size: Optional[CuratedVisualizationSize] = Field( + tile_size: Optional[CuratedVisualizationSize] = Field( default=None, - description="Layout size of the visualization.", + description="Layout size of the visualization tile.", ) resource_type: Optional[VisualizationResourceTypes] = Field( default=None, @@ -363,8 +364,8 @@ def get_custom_filters( custom_filters.append( getattr(table, "display_order") == self.display_order ) - if self.size is not None: - custom_filters.append(getattr(table, "size") == self.size) + if self.tile_size is not None: + custom_filters.append(getattr(table, "size") == self.tile_size) # resource-based filtering is handled within the store implementation return custom_filters diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 2233f774ac..7a2efe83a0 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -16,8 +16,8 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence from uuid import UUID, uuid4 -from sqlalchemy import UniqueConstraint, and_ -from sqlalchemy.orm import foreign, selectinload +from sqlalchemy import UniqueConstraint +from sqlalchemy.orm import selectinload from sqlalchemy.sql.base import ExecutableOption from sqlmodel import Field, Relationship, SQLModel @@ -111,8 +111,11 @@ class CuratedVisualizationSchema(BaseSchema, table=True): ) resource: Optional[CuratedVisualizationResourceSchema] = Relationship( back_populates="visualization", - sa_relationship_kwargs={"lazy": "selectin", "cascade": "all, delete"}, - uselist=False, + sa_relationship_kwargs={ + "lazy": "selectin", + "cascade": "all, delete", + "uselist": False, + }, ) @classmethod @@ -262,25 +265,33 @@ def curated_visualization_relationship_kwargs( The relationship will be read-only (viewonly=True) and eagerly loaded via selectin loading. """ - - def _primaryjoin(): - return and_( - CuratedVisualizationResourceSchema.resource_type - == resource_type.value, - foreign(CuratedVisualizationResourceSchema.resource_id) - == parent_column_factory(), - ) - - def _secondaryjoin(): - return ( - CuratedVisualizationSchema.id - == CuratedVisualizationResourceSchema.visualization_id + # Resolve the parent column to extract class and attribute names + parent_column = parent_column_factory() + parent_class = getattr(parent_column, "class_", None) + parent_attribute = getattr(parent_column, "key", None) + + if parent_class is None or parent_attribute is None: + raise ValueError( + "parent_column_factory must return an InstrumentedAttribute " + "with `class_` and `key` attributes." ) - return dict( - secondary="curated_visualization_resource", - primaryjoin=_primaryjoin, - secondaryjoin=_secondaryjoin, - viewonly=True, - lazy="selectin", + # Build string expressions for primaryjoin and secondaryjoin + # These strings are evaluated by SQLAlchemy at relationship configuration time + parent_class_name = parent_class.__name__ + primaryjoin_str = ( + f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{resource_type.value}', " + f"foreign(CuratedVisualizationResourceSchema.resource_id)=={parent_class_name}.{parent_attribute})" ) + secondaryjoin_str = ( + "CuratedVisualizationSchema.id == " + "CuratedVisualizationResourceSchema.visualization_id" + ) + + return { + "secondary": "curated_visualization_resource", + "primaryjoin": primaryjoin_str, + "secondaryjoin": secondaryjoin_str, + "viewonly": True, + "lazy": "selectin", + } diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index c56e403f64..4dd02f6138 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5643,9 +5643,10 @@ def list_curated_visualizations( query = select(CuratedVisualizationSchema) if filter_model.resource_type or filter_model.resource_id: - resource_alias: AliasedClass[ - CuratedVisualizationResourceSchema - ] = aliased(CuratedVisualizationResourceSchema) + resource_alias = cast( + AliasedClass[CuratedVisualizationResourceSchema], + aliased(CuratedVisualizationResourceSchema), + ) query = query.join( resource_alias, resource_alias.visualization_id From 7f78e36ff9bfcf24e571e95c2b0b25dc5e9ae819 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Mon, 13 Oct 2025 09:36:11 +0100 Subject: [PATCH 16/64] revert deloyment --- .../routers/deployment_endpoints.py | 69 ++++++------------- 1 file changed, 22 insertions(+), 47 deletions(-) diff --git a/src/zenml/zen_server/routers/deployment_endpoints.py b/src/zenml/zen_server/routers/deployment_endpoints.py index e8e1ed3e28..7b0ec15e23 100644 --- a/src/zenml/zen_server/routers/deployment_endpoints.py +++ b/src/zenml/zen_server/routers/deployment_endpoints.py @@ -13,10 +13,13 @@ # permissions and limitations under the License. """Endpoint definitions for deployments.""" -from typing import Optional, Union from uuid import UUID -from fastapi import APIRouter, Depends, Security +from fastapi import ( + APIRouter, + Depends, + Security, +) from zenml.constants import ( API, @@ -28,8 +31,8 @@ DeploymentRequest, DeploymentResponse, DeploymentUpdate, - Page, ) +from zenml.models.v2.base.page import Page from zenml.zen_server.auth import AuthContext, authorize from zenml.zen_server.exceptions import error_response from zenml.zen_server.rbac.endpoint_utils import ( @@ -40,7 +43,6 @@ verify_permissions_and_update_entity, ) from zenml.zen_server.rbac.models import ResourceType -from zenml.zen_server.routers.projects_endpoints import workspace_router from zenml.zen_server.utils import ( async_fastapi_endpoint_wrapper, make_dependable, @@ -58,34 +60,19 @@ "", responses={401: error_response, 409: error_response, 422: error_response}, ) -# TODO: the workspace scoped endpoint is only kept for dashboard compatibility -# and can be removed after the migration -@workspace_router.post( - "/{project_name_or_id}" + DEPLOYMENTS, - responses={401: error_response, 409: error_response, 422: error_response}, - deprecated=True, - tags=["deployments"], -) @async_fastapi_endpoint_wrapper def create_deployment( deployment: DeploymentRequest, - project_name_or_id: Optional[Union[str, UUID]] = None, _: AuthContext = Security(authorize), ) -> DeploymentResponse: - """Create a deployment. + """Creates a deployment. Args: deployment: Deployment to create. - project_name_or_id: Optional project name or ID for backwards - compatibility. Returns: The created deployment. """ - if project_name_or_id: - project = zen_store().get_project(project_name_or_id) - deployment.project = project.id - return verify_permissions_and_create_entity( request_model=deployment, create_method=zen_store().create_deployment, @@ -96,38 +83,25 @@ def create_deployment( "", responses={401: error_response, 404: error_response, 422: error_response}, ) -# TODO: the workspace scoped endpoint is only kept for dashboard compatibility -# and can be removed after the migration -@workspace_router.get( - "/{project_name_or_id}" + DEPLOYMENTS, - responses={401: error_response, 404: error_response, 422: error_response}, - deprecated=True, - tags=["deployments"], -) @async_fastapi_endpoint_wrapper(deduplicate=True) def list_deployments( deployment_filter_model: DeploymentFilter = Depends( make_dependable(DeploymentFilter) ), - project_name_or_id: Optional[Union[str, UUID]] = None, hydrate: bool = False, _: AuthContext = Security(authorize), ) -> Page[DeploymentResponse]: - """List deployments. + """Gets a list of deployments. Args: - deployment_filter_model: Filter model used for pagination, sorting, and + deployment_filter_model: Filter model used for pagination, sorting, filtering. - project_name_or_id: Optional project name or ID for backwards - compatibility. - hydrate: Whether to hydrate the returned models. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. Returns: - A page of deployments matching the filter. + List of deployment objects matching the filter criteria. """ - if project_name_or_id: - deployment_filter_model.project = project_name_or_id - return verify_permissions_and_list_entities( filter_model=deployment_filter_model, resource_type=ResourceType.DEPLOYMENT, @@ -146,14 +120,15 @@ def get_deployment( hydrate: bool = True, _: AuthContext = Security(authorize), ) -> DeploymentResponse: - """Get a deployment by ID. + """Gets a specific deployment using its unique id. Args: - deployment_id: The deployment ID. - hydrate: Whether to hydrate the returned model. + deployment_id: ID of the deployment to get. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. Returns: - The requested deployment. + A specific deployment object. """ return verify_permissions_and_get_entity( id=deployment_id, @@ -172,11 +147,11 @@ def update_deployment( deployment_update: DeploymentUpdate, _: AuthContext = Security(authorize), ) -> DeploymentResponse: - """Update a deployment. + """Updates a specific deployment. Args: - deployment_id: The deployment ID. - deployment_update: The updates to apply. + deployment_id: ID of the deployment to update. + deployment_update: Update model for the deployment. Returns: The updated deployment. @@ -198,10 +173,10 @@ def delete_deployment( deployment_id: UUID, _: AuthContext = Security(authorize), ) -> None: - """Delete a deployment. + """Deletes a specific deployment. Args: - deployment_id: The deployment ID. + deployment_id: ID of the deployment to delete. """ verify_permissions_and_delete_entity( id=deployment_id, From 2fc05532f97ef3f20419aa0e64b300c9509d0272 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Mon, 13 Oct 2025 12:25:27 +0100 Subject: [PATCH 17/64] add other schema support --- src/zenml/client.py | 8 +-- .../models/v2/core/curated_visualization.py | 8 +-- src/zenml/models/v2/core/pipeline.py | 10 +++- src/zenml/models/v2/core/pipeline_run.py | 7 +++ src/zenml/models/v2/core/pipeline_snapshot.py | 7 +++ .../schemas/curated_visualization_schemas.py | 58 +------------------ .../zen_stores/schemas/deployment_schemas.py | 12 ++-- src/zenml/zen_stores/schemas/model_schemas.py | 10 ++-- .../schemas/pipeline_run_schemas.py | 31 +++++++++- .../zen_stores/schemas/pipeline_schemas.py | 32 +++++++++- .../schemas/pipeline_snapshot_schemas.py | 38 +++++++++++- .../zen_stores/schemas/project_schemas.py | 11 ++++ 12 files changed, 151 insertions(+), 81 deletions(-) diff --git a/src/zenml/client.py b/src/zenml/client.py index 8b16bdad23..d6b2d8de39 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -3858,7 +3858,7 @@ def list_curated_visualizations( size: Optional[int] = None, sort_by: Optional[str] = None, visualization_index: Optional[int] = None, - tile_size: Optional[CuratedVisualizationSize] = None, + size: Optional[CuratedVisualizationSize] = None, hydrate: bool = False, ) -> Page[CuratedVisualizationResponse]: """List curated visualizations, optionally scoped to a resource. @@ -3887,7 +3887,7 @@ def list_curated_visualizations( size: The maximum size of all pages. sort_by: The column to sort by. visualization_index: The index of the visualization to filter by. - tile_size: The layout size of the visualization tiles to filter by. + size: The layout size of the visualization tiles to filter by. hydrate: Flag deciding whether to hydrate the output model(s) by including metadata fields in the response. @@ -3978,8 +3978,8 @@ def list_curated_visualizations( filter_model.sort_by = sort_by if visualization_index is not None: filter_model.visualization_index = visualization_index - if tile_size is not None: - filter_model.tile_size = tile_size + if size is not None: + filter_model.size = size return self.zen_store.list_curated_visualizations( filter_model=filter_model, diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 2a5f940b39..6064ec5209 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -255,7 +255,7 @@ class CuratedVisualizationFilter(ProjectScopedFilter): *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS, "resource_id", "resource_type", - "tile_size", + "size", ] CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [ *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS, @@ -282,7 +282,7 @@ class CuratedVisualizationFilter(ProjectScopedFilter): default=None, description="Display order of the visualization.", ) - tile_size: Optional[CuratedVisualizationSize] = Field( + size: Optional[CuratedVisualizationSize] = Field( default=None, description="Layout size of the visualization tile.", ) @@ -364,8 +364,8 @@ def get_custom_filters( custom_filters.append( getattr(table, "display_order") == self.display_order ) - if self.tile_size is not None: - custom_filters.append(getattr(table, "size") == self.tile_size) + if self.size is not None: + custom_filters.append(getattr(table, "size") == self.size) # resource-based filtering is handled within the store implementation return custom_filters diff --git a/src/zenml/models/v2/core/pipeline.py b/src/zenml/models/v2/core/pipeline.py index 8a5823f067..6122a07c97 100644 --- a/src/zenml/models/v2/core/pipeline.py +++ b/src/zenml/models/v2/core/pipeline.py @@ -46,7 +46,11 @@ from zenml.models.v2.core.tag import TagResponse if TYPE_CHECKING: - from zenml.models import PipelineRunResponse, UserResponse + from zenml.models import ( + CuratedVisualizationResponse, + PipelineRunResponse, + UserResponse, + ) from zenml.zen_stores.schemas import BaseSchema AnySchema = TypeVar("AnySchema", bound=BaseSchema) @@ -127,6 +131,10 @@ class PipelineResponseResources(ProjectScopedResponseResources): tags: List[TagResponse] = Field( title="Tags associated with the pipeline.", ) + visualizations: List["CuratedVisualizationResponse"] = Field( + default=[], + title="Curated visualizations associated with the pipeline.", + ) class PipelineResponse( diff --git a/src/zenml/models/v2/core/pipeline_run.py b/src/zenml/models/v2/core/pipeline_run.py index d064e60b08..c2a1e812e0 100644 --- a/src/zenml/models/v2/core/pipeline_run.py +++ b/src/zenml/models/v2/core/pipeline_run.py @@ -56,6 +56,9 @@ from zenml.models import TriggerExecutionResponse from zenml.models.v2.core.artifact_version import ArtifactVersionResponse from zenml.models.v2.core.code_reference import CodeReferenceResponse + from zenml.models.v2.core.curated_visualization import ( + CuratedVisualizationResponse, + ) from zenml.models.v2.core.logs import LogsResponse from zenml.models.v2.core.pipeline import PipelineResponse from zenml.models.v2.core.pipeline_build import ( @@ -305,6 +308,10 @@ class PipelineRunResponseResources(ProjectScopedResponseResources): title="Logs associated with this pipeline run.", default=None, ) + visualizations: List["CuratedVisualizationResponse"] = Field( + default=[], + title="Curated visualizations associated with the pipeline run.", + ) # TODO: In Pydantic v2, the `model_` is a protected namespaces for all # fields defined under base models. If not handled, this raises a warning. diff --git a/src/zenml/models/v2/core/pipeline_snapshot.py b/src/zenml/models/v2/core/pipeline_snapshot.py index 09b0e9162d..19215fc530 100644 --- a/src/zenml/models/v2/core/pipeline_snapshot.py +++ b/src/zenml/models/v2/core/pipeline_snapshot.py @@ -61,6 +61,9 @@ if TYPE_CHECKING: from sqlalchemy.sql.elements import ColumnElement + from zenml.models.v2.core.curated_visualization import ( + CuratedVisualizationResponse, + ) from zenml.zen_stores.schemas.base_schemas import BaseSchema AnySchema = TypeVar("AnySchema", bound=BaseSchema) @@ -334,6 +337,10 @@ class PipelineSnapshotResponseResources(ProjectScopedResponseResources): default=None, title="The user that created the latest run of the snapshot.", ) + visualizations: List["CuratedVisualizationResponse"] = Field( + default=[], + title="Curated visualizations associated with the pipeline snapshot.", + ) class PipelineSnapshotResponse( diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 7a2efe83a0..31cac6d27a 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """SQLModel implementation of curated visualization tables.""" -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence +from typing import TYPE_CHECKING, Any, List, Optional, Sequence from uuid import UUID, uuid4 from sqlalchemy import UniqueConstraint @@ -21,7 +21,7 @@ from sqlalchemy.sql.base import ExecutableOption from sqlmodel import Field, Relationship, SQLModel -from zenml.enums import CuratedVisualizationSize, VisualizationResourceTypes +from zenml.enums import CuratedVisualizationSize from zenml.models.v2.core.curated_visualization import ( CuratedVisualizationRequest, CuratedVisualizationResponse, @@ -241,57 +241,3 @@ def to_model( metadata=metadata, resources=response_resources, ) - - -def curated_visualization_relationship_kwargs( - parent_column_factory: Callable[[], Any], - resource_type: VisualizationResourceTypes, -) -> Dict[str, Any]: - """Build sa_relationship_kwargs for curated visualization relationships. - - This helper consolidates the relationship definition for linking parent - schemas (like DeploymentSchema, ModelSchema) to their curated visualizations - through the resource link table. - - Args: - parent_column_factory: A callable that returns the parent column - (e.g., `lambda: DeploymentSchema.id`). Uses a callable to defer - evaluation and avoid circular import issues. - resource_type: The VisualizationResourceTypes enum value indicating - what type of resource the parent represents. - - Returns: - A dictionary suitable for passing to Relationship(sa_relationship_kwargs=...). - The relationship will be read-only (viewonly=True) and eagerly loaded - via selectin loading. - """ - # Resolve the parent column to extract class and attribute names - parent_column = parent_column_factory() - parent_class = getattr(parent_column, "class_", None) - parent_attribute = getattr(parent_column, "key", None) - - if parent_class is None or parent_attribute is None: - raise ValueError( - "parent_column_factory must return an InstrumentedAttribute " - "with `class_` and `key` attributes." - ) - - # Build string expressions for primaryjoin and secondaryjoin - # These strings are evaluated by SQLAlchemy at relationship configuration time - parent_class_name = parent_class.__name__ - primaryjoin_str = ( - f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{resource_type.value}', " - f"foreign(CuratedVisualizationResourceSchema.resource_id)=={parent_class_name}.{parent_attribute})" - ) - secondaryjoin_str = ( - "CuratedVisualizationSchema.id == " - "CuratedVisualizationResourceSchema.visualization_id" - ) - - return { - "secondary": "curated_visualization_resource", - "primaryjoin": primaryjoin_str, - "secondaryjoin": secondaryjoin_str, - "viewonly": True, - "lazy": "selectin", - } diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index c70cc48f99..be5a1dcbf3 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -41,9 +41,6 @@ from zenml.utils.time_utils import utc_now from zenml.zen_stores.schemas.base_schemas import NamedSchema from zenml.zen_stores.schemas.component_schemas import StackComponentSchema -from zenml.zen_stores.schemas.curated_visualization_schemas import ( - curated_visualization_relationship_kwargs, -) from zenml.zen_stores.schemas.pipeline_snapshot_schemas import ( PipelineSnapshotSchema, ) @@ -144,9 +141,12 @@ class DeploymentSchema(NamedSchema, table=True): ) visualizations: List["CuratedVisualizationSchema"] = Relationship( - sa_relationship_kwargs=curated_visualization_relationship_kwargs( - parent_column_factory=lambda: DeploymentSchema.id, - resource_type=VisualizationResourceTypes.DEPLOYMENT, + sa_relationship_kwargs=dict( + secondary="curated_visualization_resource", + primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.DEPLOYMENT.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==DeploymentSchema.id)", + secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + viewonly=True, + lazy="selectin", ), ) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 0bfa3cd58f..bbbfd237cc 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -60,7 +60,6 @@ from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME from zenml.zen_stores.schemas.curated_visualization_schemas import ( CuratedVisualizationSchema, - curated_visualization_relationship_kwargs, ) from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema from zenml.zen_stores.schemas.project_schemas import ProjectSchema @@ -134,9 +133,12 @@ class ModelSchema(NamedSchema, table=True): sa_relationship_kwargs={"cascade": "delete"}, ) visualizations: List["CuratedVisualizationSchema"] = Relationship( - sa_relationship_kwargs=curated_visualization_relationship_kwargs( - parent_column_factory=lambda: ModelSchema.id, - resource_type=VisualizationResourceTypes.MODEL, + sa_relationship_kwargs=dict( + secondary="curated_visualization_resource", + primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.MODEL.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==ModelSchema.id)", + secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + viewonly=True, + lazy="selectin", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index 1a192d7976..865adb2e5d 100644 --- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py @@ -34,6 +34,7 @@ MetadataResourceTypes, PipelineRunTriggeredByType, TaggableResourceTypes, + VisualizationResourceTypes, ) from zenml.logger import get_logger from zenml.models import ( @@ -72,6 +73,9 @@ ) if TYPE_CHECKING: + from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationSchema, + ) from zenml.zen_stores.schemas.logs_schemas import LogsSchema from zenml.zen_stores.schemas.model_schemas import ( ModelVersionPipelineRunSchema, @@ -241,6 +245,15 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True): overlaps="tags", ), ) + visualizations: List["CuratedVisualizationSchema"] = Relationship( + sa_relationship_kwargs=dict( + secondary="curated_visualization_resource", + primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE_RUN.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineRunSchema.id)", + secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + viewonly=True, + lazy="selectin", + ), + ) # Needed for cascade deletion model_versions_pipeline_runs_links: List[ @@ -283,6 +296,10 @@ def get_query_options( # ) if include_resources: + from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationSchema, + ) + options.extend( [ selectinload( @@ -315,6 +332,11 @@ def get_query_options( selectinload(jl_arg(PipelineRunSchema.logs)), selectinload(jl_arg(PipelineRunSchema.user)), selectinload(jl_arg(PipelineRunSchema.tags)), + selectinload( + jl_arg(PipelineRunSchema.visualizations) + ).selectinload( + jl_arg(CuratedVisualizationSchema.artifact_version) + ), ] ) @@ -632,6 +654,13 @@ def to_model( tags=[tag.to_model() for tag in self.tags], logs=client_logs[0].to_model() if client_logs else None, log_collection=[log.to_model() for log in self.logs], + visualizations=[ + visualization.to_model( + include_metadata=False, + include_resources=False, + ) + for visualization in (self.visualizations or []) + ], ) return PipelineRunResponse( @@ -819,7 +848,7 @@ def _check_if_run_in_progress(self) -> bool: else: in_progress = any( not ExecutionStatus(status).is_finished - for name, status in step_run_statuses + for _, status in step_run_statuses ) return in_progress else: diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index ea4e3a7359..793b7853cf 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_schemas.py @@ -17,11 +17,11 @@ from uuid import UUID from sqlalchemy import TEXT, Column, UniqueConstraint -from sqlalchemy.orm import joinedload, object_session +from sqlalchemy.orm import joinedload, object_session, selectinload from sqlalchemy.sql.base import ExecutableOption from sqlmodel import Field, Relationship, desc, select -from zenml.enums import TaggableResourceTypes +from zenml.enums import TaggableResourceTypes, VisualizationResourceTypes from zenml.models import ( PipelineRequest, PipelineResponse, @@ -38,6 +38,9 @@ from zenml.zen_stores.schemas.utils import jl_arg if TYPE_CHECKING: + from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationSchema, + ) from zenml.zen_stores.schemas.pipeline_build_schemas import ( PipelineBuildSchema, ) @@ -104,6 +107,15 @@ class PipelineSchema(NamedSchema, table=True): overlaps="tags", ), ) + visualizations: List["CuratedVisualizationSchema"] = Relationship( + sa_relationship_kwargs=dict( + secondary="curated_visualization_resource", + primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineSchema.id)", + secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + viewonly=True, + lazy="selectin", + ), + ) @property def latest_run(self) -> Optional["PipelineRunSchema"]: @@ -155,10 +167,19 @@ def get_query_options( options = [] if include_resources: + from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationSchema, + ) + options.extend( [ joinedload(jl_arg(PipelineSchema.user)), # joinedload(jl_arg(PipelineSchema.tags)), + selectinload( + jl_arg(PipelineSchema.visualizations) + ).selectinload( + jl_arg(CuratedVisualizationSchema.artifact_version) + ), ] ) @@ -226,6 +247,13 @@ def to_model( latest_run_id=latest_run.id if latest_run else None, latest_run_status=latest_run.status if latest_run else None, tags=[tag.to_model() for tag in self.tags], + visualizations=[ + visualization.to_model( + include_metadata=False, + include_resources=False, + ) + for visualization in (self.visualizations or []) + ], ) return PipelineResponse( diff --git a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py index 290dcec937..b083ec3061 100644 --- a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py @@ -19,7 +19,7 @@ from sqlalchemy import TEXT, Column, String, UniqueConstraint from sqlalchemy.dialects.mysql import MEDIUMTEXT -from sqlalchemy.orm import joinedload, object_session +from sqlalchemy.orm import joinedload, object_session, selectinload from sqlalchemy.sql.base import ExecutableOption from sqlmodel import Field, Relationship, asc, col, desc, select @@ -27,7 +27,7 @@ from zenml.config.pipeline_spec import PipelineSpec from zenml.config.step_configurations import Step from zenml.constants import MEDIUMTEXT_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH -from zenml.enums import TaggableResourceTypes +from zenml.enums import TaggableResourceTypes, VisualizationResourceTypes from zenml.logger import get_logger from zenml.models import ( PipelineSnapshotRequest, @@ -53,6 +53,9 @@ from zenml.zen_stores.schemas.utils import jl_arg if TYPE_CHECKING: + from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationSchema, + ) from zenml.zen_stores.schemas.deployment_schemas import ( DeploymentSchema, ) @@ -218,6 +221,15 @@ class PipelineSnapshotSchema(BaseSchema, table=True): overlaps="tags", ), ) + visualizations: List["CuratedVisualizationSchema"] = Relationship( + sa_relationship_kwargs=dict( + secondary="curated_visualization_resource", + primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE_SNAPSHOT.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineSnapshotSchema.id)", + secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + viewonly=True, + lazy="selectin", + ), + ) @property def latest_run(self) -> Optional["PipelineRunSchema"]: @@ -352,7 +364,20 @@ def get_query_options( ) if include_resources: - options.extend([joinedload(jl_arg(PipelineSnapshotSchema.user))]) + from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationSchema, + ) + + options.extend( + [ + joinedload(jl_arg(PipelineSnapshotSchema.user)), + selectinload( + jl_arg(PipelineSnapshotSchema.visualizations) + ).selectinload( + jl_arg(CuratedVisualizationSchema.artifact_version) + ), + ] + ) return options @@ -565,6 +590,13 @@ def to_model( latest_run_user=latest_run_user.to_model() if latest_run_user else None, + visualizations=[ + visualization.to_model( + include_metadata=False, + include_resources=False, + ) + for visualization in (self.visualizations or []) + ], ) return PipelineSnapshotResponse( diff --git a/src/zenml/zen_stores/schemas/project_schemas.py b/src/zenml/zen_stores/schemas/project_schemas.py index e639ba57c2..7c9ddf275e 100644 --- a/src/zenml/zen_stores/schemas/project_schemas.py +++ b/src/zenml/zen_stores/schemas/project_schemas.py @@ -18,6 +18,7 @@ from sqlalchemy import UniqueConstraint from sqlmodel import Relationship +from zenml.enums import VisualizationResourceTypes from zenml.models import ( ProjectRequest, ProjectResponse, @@ -33,6 +34,7 @@ ActionSchema, ArtifactVersionSchema, CodeRepositorySchema, + CuratedVisualizationSchema, DeploymentSchema, EventSourceSchema, ModelSchema, @@ -127,6 +129,15 @@ class ProjectSchema(NamedSchema, table=True): back_populates="project", sa_relationship_kwargs={"cascade": "delete"}, ) + visualizations: List["CuratedVisualizationSchema"] = Relationship( + sa_relationship_kwargs=dict( + secondary="curated_visualization_resource", + primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PROJECT.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==ProjectSchema.id)", + secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + viewonly=True, + lazy="selectin", + ), + ) @classmethod def from_request(cls, project: ProjectRequest) -> "ProjectSchema": From f17f849a55c08b6ba4fc0e9a4e0a6561c4c9b6cc Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Mon, 13 Oct 2025 16:07:14 +0100 Subject: [PATCH 18/64] Rename 'size' to 'layout_size' in visualizations This change updates the terminology used in the codebase to improve clarity and consistency. The parameter 'size' has been renamed to 'layout_size' across various files, including the client, schemas, and documentation, to better reflect its purpose in defining the layout of visualizations. Additionally, tests have been updated to assert the new 'layout_size' property. No functional changes were made; this is purely a refactor for improved readability and maintainability. --- docs/book/how-to/artifacts/visualizations.md | 16 +-- src/zenml/client.py | 26 ++-- .../models/v2/core/curated_visualization.py | 20 +-- ...d1_add_vizualisations_and_link_them_to_.py | 125 ++++++++++++++++++ .../schemas/curated_visualization_schemas.py | 6 +- .../functional/zen_stores/test_zen_store.py | 8 +- 6 files changed, 164 insertions(+), 37 deletions(-) create mode 100644 src/zenml/zen_stores/migrations/versions/f2ab38646bd1_add_vizualisations_and_link_them_to_.py diff --git a/docs/book/how-to/artifacts/visualizations.md b/docs/book/how-to/artifacts/visualizations.md index 34ae3bc1d1..4d17340b91 100644 --- a/docs/book/how-to/artifacts/visualizations.md +++ b/docs/book/how-to/artifacts/visualizations.md @@ -110,7 +110,7 @@ model_viz = client.create_curated_visualization( type=VisualizationResourceTypes.MODEL ), display_name="Model performance dashboard", - size=CuratedVisualizationSize.FULL_WIDTH, + layout_size=CuratedVisualizationSize.FULL_WIDTH, ) # Create a visualization for the deployment @@ -122,7 +122,7 @@ deployment_viz = client.create_curated_visualization( type=VisualizationResourceTypes.DEPLOYMENT ), display_name="Deployment health dashboard", - size=CuratedVisualizationSize.HALF_WIDTH, + layout_size=CuratedVisualizationSize.HALF_WIDTH, ) # Create a visualization for the project @@ -134,11 +134,11 @@ project_viz = client.create_curated_visualization( type=VisualizationResourceTypes.PROJECT ), display_name="Project overview dashboard", - size=CuratedVisualizationSize.FULL_WIDTH, + layout_size=CuratedVisualizationSize.FULL_WIDTH, ) ``` -Note that each visualization is created separately and can reference the same artifact visualization (by using the same `artifact_version_id` and `visualization_index`). This allows you to show the same underlying visualization in multiple contexts while maintaining separate display settings for each resource. Use the optional `size` argument to control whether the visualization spans the full width of the dashboard or renders as a half-width tile. If omitted, the layout defaults to `CuratedVisualizationSize.FULL_WIDTH`. +Note that each visualization is created separately and can reference the same artifact visualization (by using the same `artifact_version_id` and `visualization_index`). This allows you to show the same underlying visualization in multiple contexts while maintaining separate display settings for each resource. Use the optional `layout_size` argument to control whether the visualization spans the full width of the dashboard or renders as a half-width tile. If omitted, the layout defaults to `CuratedVisualizationSize.FULL_WIDTH`. To list curated visualizations for a specific resource, you can use the `Client.list_curated_visualizations` convenience parameters: @@ -164,7 +164,7 @@ client.update_curated_visualization( visualization_id=UUID(""), display_name="Updated Dashboard Title", display_order=10, - size=CuratedVisualizationSize.HALF_WIDTH, + layout_size=CuratedVisualizationSize.HALF_WIDTH, ) ``` @@ -190,7 +190,7 @@ visualization_a = client.create_curated_visualization( type=VisualizationResourceTypes.MODEL ), display_order=10, # Primary dashboard - size=CuratedVisualizationSize.FULL_WIDTH, + layout_size=CuratedVisualizationSize.FULL_WIDTH, ) visualization_b = client.create_curated_visualization( @@ -201,7 +201,7 @@ visualization_b = client.create_curated_visualization( type=VisualizationResourceTypes.MODEL ), display_order=20, # Secondary metrics - size=CuratedVisualizationSize.HALF_WIDTH, # Compact chart beside the primary tile + layout_size=CuratedVisualizationSize.HALF_WIDTH, # Compact chart beside the primary tile ) # Later, easily insert between them @@ -213,7 +213,7 @@ visualization_c = client.create_curated_visualization( type=VisualizationResourceTypes.MODEL ), display_order=15, # Now appears between A and B - size=CuratedVisualizationSize.HALF_WIDTH, + layout_size=CuratedVisualizationSize.HALF_WIDTH, ) ``` diff --git a/src/zenml/client.py b/src/zenml/client.py index d6b2d8de39..ea3737847f 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -3755,7 +3755,7 @@ def create_curated_visualization( project_id: Optional[UUID] = None, display_name: Optional[str] = None, display_order: Optional[int] = None, - size: CuratedVisualizationSize = CuratedVisualizationSize.FULL_WIDTH, + layout_size: CuratedVisualizationSize = CuratedVisualizationSize.FULL_WIDTH, ) -> CuratedVisualizationResponse: """Create a curated visualization associated with a resource. @@ -3784,7 +3784,7 @@ def create_curated_visualization( project_id: The ID of the project to associate with the visualization. display_name: The display name of the visualization. display_order: The display order of the visualization. - size: The layout size of the visualization in the dashboard. + layout_size: The layout size of the visualization in the dashboard. Returns: The created curated visualization. @@ -3801,7 +3801,7 @@ def create_curated_visualization( visualization_index=visualization_index, display_name=display_name, display_order=display_order, - size=size, + layout_size=layout_size, resource=resource, ) return self.zen_store.create_curated_visualization(request) @@ -3814,7 +3814,7 @@ def add_visualization_to_deployment( *, display_name: Optional[str] = None, display_order: Optional[int] = None, - size: CuratedVisualizationSize = CuratedVisualizationSize.FULL_WIDTH, + layout_size: CuratedVisualizationSize = CuratedVisualizationSize.FULL_WIDTH, ) -> CuratedVisualizationResponse: """Attach a curated visualization to a deployment. @@ -3824,7 +3824,7 @@ def add_visualization_to_deployment( visualization_index: The index of the visualization within the artifact version. display_name: Optional display name for the visualization. display_order: Optional display order for sorting visualizations. - size: Layout size defining the visualization width on the dashboard. + layout_size: Layout size defining the visualization width on the dashboard. Returns: The created curated visualization. @@ -3840,7 +3840,7 @@ def add_visualization_to_deployment( project_id=deployment.project_id, display_name=display_name, display_order=display_order, - size=size, + layout_size=layout_size, ) def list_curated_visualizations( @@ -3858,7 +3858,7 @@ def list_curated_visualizations( size: Optional[int] = None, sort_by: Optional[str] = None, visualization_index: Optional[int] = None, - size: Optional[CuratedVisualizationSize] = None, + layout_size: Optional[CuratedVisualizationSize] = None, hydrate: bool = False, ) -> Page[CuratedVisualizationResponse]: """List curated visualizations, optionally scoped to a resource. @@ -3887,7 +3887,7 @@ def list_curated_visualizations( size: The maximum size of all pages. sort_by: The column to sort by. visualization_index: The index of the visualization to filter by. - size: The layout size of the visualization tiles to filter by. + layout_size: The layout size of the visualization tiles to filter by. hydrate: Flag deciding whether to hydrate the output model(s) by including metadata fields in the response. @@ -3978,8 +3978,8 @@ def list_curated_visualizations( filter_model.sort_by = sort_by if visualization_index is not None: filter_model.visualization_index = visualization_index - if size is not None: - filter_model.size = size + if layout_size is not None: + filter_model.layout_size = layout_size return self.zen_store.list_curated_visualizations( filter_model=filter_model, @@ -3992,7 +3992,7 @@ def update_curated_visualization( *, display_name: Optional[str] = None, display_order: Optional[int] = None, - size: Optional[CuratedVisualizationSize] = None, + layout_size: Optional[CuratedVisualizationSize] = None, ) -> CuratedVisualizationResponse: """Update display metadata for a curated visualization. @@ -4000,7 +4000,7 @@ def update_curated_visualization( visualization_id: The ID of the curated visualization to update. display_name: New display name for the visualization. display_order: New display order for the visualization. - size: Updated layout size for the visualization. + layout_size: Updated layout size for the visualization. Returns: The updated deployment visualization. @@ -4008,7 +4008,7 @@ def update_curated_visualization( update_model = CuratedVisualizationUpdate( display_name=display_name, display_order=display_order, - size=size, + layout_size=layout_size, ) return self.zen_store.update_curated_visualization( visualization_id=visualization_id, diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 6064ec5209..6ab2b1251b 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -90,7 +90,7 @@ class CuratedVisualizationRequest(ProjectScopedRequest): default=None, title="The display order of the visualization.", ) - size: CuratedVisualizationSize = Field( + layout_size: CuratedVisualizationSize = Field( default=CuratedVisualizationSize.FULL_WIDTH, title="The layout size of the visualization.", description=( @@ -121,7 +121,7 @@ class CuratedVisualizationUpdate(BaseUpdate): default=None, title="The new display order of the visualization.", ) - size: Optional[CuratedVisualizationSize] = Field( + layout_size: Optional[CuratedVisualizationSize] = Field( default=None, title="The updated layout size of the visualization.", ) @@ -149,7 +149,7 @@ class CuratedVisualizationResponseBody(ProjectScopedResponseBody): default=None, title="The display order of the visualization.", ) - size: CuratedVisualizationSize = Field( + layout_size: CuratedVisualizationSize = Field( default=CuratedVisualizationSize.FULL_WIDTH, title="The layout size of the visualization.", ) @@ -227,13 +227,13 @@ def display_order(self) -> Optional[int]: return self.get_body().display_order @property - def size(self) -> CuratedVisualizationSize: + def layout_size(self) -> CuratedVisualizationSize: """The layout size of the visualization. Returns: The layout size of the visualization. """ - return self.get_body().size + return self.get_body().layout_size @property def artifact_version(self) -> Optional["ArtifactVersionResponse"]: @@ -255,7 +255,7 @@ class CuratedVisualizationFilter(ProjectScopedFilter): *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS, "resource_id", "resource_type", - "size", + "layout_size", ] CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [ *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS, @@ -282,7 +282,7 @@ class CuratedVisualizationFilter(ProjectScopedFilter): default=None, description="Display order of the visualization.", ) - size: Optional[CuratedVisualizationSize] = Field( + layout_size: Optional[CuratedVisualizationSize] = Field( default=None, description="Layout size of the visualization tile.", ) @@ -364,8 +364,10 @@ def get_custom_filters( custom_filters.append( getattr(table, "display_order") == self.display_order ) - if self.size is not None: - custom_filters.append(getattr(table, "size") == self.size) + if self.layout_size is not None: + custom_filters.append( + getattr(table, "layout_size") == self.layout_size + ) # resource-based filtering is handled within the store implementation return custom_filters diff --git a/src/zenml/zen_stores/migrations/versions/f2ab38646bd1_add_vizualisations_and_link_them_to_.py b/src/zenml/zen_stores/migrations/versions/f2ab38646bd1_add_vizualisations_and_link_them_to_.py new file mode 100644 index 0000000000..3f897038b1 --- /dev/null +++ b/src/zenml/zen_stores/migrations/versions/f2ab38646bd1_add_vizualisations_and_link_them_to_.py @@ -0,0 +1,125 @@ +"""add vizualisations and link them to resources [f2ab38646bd1]. + +Revision ID: f2ab38646bd1 +Revises: 0.90.0 +Create Date: 2025-10-13 16:01:41.907536 + +""" + +import sqlalchemy as sa +import sqlmodel +from alembic import op + +# revision identifiers, used by Alembic. +revision = "f2ab38646bd1" +down_revision = "0.90.0" +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! ### + op.create_table( + "curated_visualization", + sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column("created", sa.DateTime(), nullable=False), + sa.Column("updated", sa.DateTime(), nullable=False), + sa.Column("project_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "artifact_version_id", sqlmodel.sql.sqltypes.GUID(), nullable=False + ), + sa.Column("visualization_index", sa.Integer(), nullable=False), + sa.Column( + "display_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True + ), + sa.Column("display_order", sa.Integer(), nullable=True), + sa.Column( + "layout_size", + sa.Enum( + "FULL_WIDTH", "HALF_WIDTH", name="curatedvisualizationsize" + ), + nullable=False, + ), + sa.ForeignKeyConstraint( + ["artifact_version_id"], + ["artifact_version.id"], + name="fk_curated_visualization_artifact_version_id_artifact_version", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["project_id"], + ["project.id"], + name="fk_curated_visualization_project_id_project", + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id"), + ) + with op.batch_alter_table( + "curated_visualization", schema=None + ) as batch_op: + batch_op.create_index( + "ix_curated_visualization_artifact_version_id_visualization_index", + ["artifact_version_id", "visualization_index"], + unique=False, + ) + batch_op.create_index( + "ix_curated_visualization_display_order", + ["display_order"], + unique=False, + ) + + op.create_table( + "curated_visualization_resource", + sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "visualization_id", sqlmodel.sql.sqltypes.GUID(), nullable=False + ), + sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "resource_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False + ), + sa.ForeignKeyConstraint( + ["visualization_id"], + ["curated_visualization.id"], + name="fk_curated_visualization_resource_visualization_id_curated_visualization", + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "visualization_id", name="unique_curated_visualization_resource" + ), + ) + with op.batch_alter_table( + "curated_visualization_resource", schema=None + ) as batch_op: + batch_op.create_index( + "ix_curated_visualization_resource_resource_id_resource_type", + ["resource_id", "resource_type"], + unique=False, + ) + + # ### 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( + "curated_visualization_resource", schema=None + ) as batch_op: + batch_op.drop_index( + "ix_curated_visualization_resource_resource_id_resource_type" + ) + + op.drop_table("curated_visualization_resource") + with op.batch_alter_table( + "curated_visualization", schema=None + ) as batch_op: + batch_op.drop_index("ix_curated_visualization_display_order") + batch_op.drop_index( + "ix_curated_visualization_artifact_version_id_visualization_index" + ) + + op.drop_table("curated_visualization") + # ### end Alembic commands ### diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 31cac6d27a..69c20fc641 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -102,7 +102,7 @@ class CuratedVisualizationSchema(BaseSchema, table=True): visualization_index: int = Field(nullable=False) display_name: Optional[str] = Field(default=None) display_order: Optional[int] = Field(default=None) - size: CuratedVisualizationSize = Field( + layout_size: CuratedVisualizationSize = Field( default=CuratedVisualizationSize.FULL_WIDTH, nullable=False ) @@ -167,7 +167,7 @@ def from_request( visualization_index=request.visualization_index, display_name=request.display_name, display_order=request.display_order, - size=request.size, + layout_size=request.layout_size, ) def update( @@ -217,7 +217,7 @@ def to_model( visualization_index=self.visualization_index, display_name=self.display_name, display_order=self.display_order, - size=self.size, + layout_size=self.layout_size, ) metadata = None diff --git a/tests/integration/functional/zen_stores/test_zen_store.py b/tests/integration/functional/zen_stores/test_zen_store.py index 9b8b0f72cb..f0c8667b6b 100644 --- a/tests/integration/functional/zen_stores/test_zen_store.py +++ b/tests/integration/functional/zen_stores/test_zen_store.py @@ -5883,7 +5883,7 @@ def test_curated_visualizations_across_resources(self): ) ) visualizations[resource_type] = viz - assert viz.size == CuratedVisualizationSize.FULL_WIDTH + assert viz.layout_size == CuratedVisualizationSize.FULL_WIDTH # Verify via zen_store filtering - each resource should have exactly one visualization for resource_type, resource_id in resource_configs: @@ -5946,7 +5946,7 @@ def test_curated_visualizations_across_resources(self): loaded.display_name == f"{VisualizationResourceTypes.PIPELINE.value} visualization" ) - assert loaded.size == CuratedVisualizationSize.FULL_WIDTH + assert loaded.layout_size == CuratedVisualizationSize.FULL_WIDTH # Test duplicate creation - same artifact_version + visualization_index + resource should fail with pytest.raises(EntityExistsError): @@ -5970,12 +5970,12 @@ def test_curated_visualizations_across_resources(self): visualization_update=CuratedVisualizationUpdate( display_name="Updated", display_order=5, - size=CuratedVisualizationSize.HALF_WIDTH, + layout_size=CuratedVisualizationSize.HALF_WIDTH, ), ) assert updated.display_name == "Updated" assert updated.display_order == 5 - assert updated.size == CuratedVisualizationSize.HALF_WIDTH + assert updated.layout_size == CuratedVisualizationSize.HALF_WIDTH # Delete all visualizations for viz in visualizations.values(): From 2928bc0d37e169bfe5849aae4e0debd78b2b4fbf Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Tue, 14 Oct 2025 11:25:37 +0100 Subject: [PATCH 19/64] Update src/zenml/models/v2/core/curated_visualization.py Co-authored-by: Stefan Nica --- src/zenml/models/v2/core/curated_visualization.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 6ab2b1251b..49f8f8a622 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -270,9 +270,10 @@ class CuratedVisualizationFilter(ProjectScopedFilter): description="Which column to sort by.", ) - artifact_version: Optional[UUID] = Field( + artifact_version: Optional[Union[UUID, str]]= Field( default=None, description="ID of the artifact version associated with the visualization.", + union_mode="left_to_right", ) visualization_index: Optional[int] = Field( default=None, From f83b35e50af5aa427e8319fa295c637e7d56111b Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Tue, 14 Oct 2025 11:25:46 +0100 Subject: [PATCH 20/64] Update src/zenml/models/v2/core/curated_visualization.py Co-authored-by: Stefan Nica --- src/zenml/models/v2/core/curated_visualization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 49f8f8a622..ae9965ed4c 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -321,11 +321,11 @@ def apply_sorting( if operand == SorterOps.DESCENDING: return cast( AnyQuery, - query.order_by(desc(column).nulls_last(), asc(table.id)), + query.order_by(desc(column).nulls_last()), ) return cast( AnyQuery, - query.order_by(asc(column).nulls_first(), asc(table.id)), + query.order_by(asc(column).nulls_first()), ) if sort_by in {"created", "updated", "visualization_index"}: From 318f5f517938c4665317bd1c9b33e6ac94971f4f Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Tue, 14 Oct 2025 12:29:10 +0100 Subject: [PATCH 21/64] renaming artifact_vertsion_id --- .../models/v2/core/curated_visualization.py | 36 ++----------------- 1 file changed, 2 insertions(+), 34 deletions(-) diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index ae9965ed4c..c0cfccc739 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -20,6 +20,7 @@ Optional, Type, TypeVar, + Union, cast, ) from uuid import UUID @@ -270,7 +271,7 @@ class CuratedVisualizationFilter(ProjectScopedFilter): description="Which column to sort by.", ) - artifact_version: Optional[Union[UUID, str]]= Field( + artifact_version_id: Optional[Union[UUID, str]]= Field( default=None, description="ID of the artifact version associated with the visualization.", union_mode="left_to_right", @@ -339,36 +340,3 @@ def apply_sorting( return super().apply_sorting(query=query, table=table) - def get_custom_filters( - self, table: Type["AnySchema"] - ) -> List["ColumnElement[bool]"]: - """Get custom filters. - - Args: - table: The query table. - - Returns: - A list of custom filters. - """ - custom_filters = super().get_custom_filters(table) - - if self.artifact_version: - custom_filters.append( - getattr(table, "artifact_version_id") == self.artifact_version - ) - if self.visualization_index is not None: - custom_filters.append( - getattr(table, "visualization_index") - == self.visualization_index - ) - if self.display_order is not None: - custom_filters.append( - getattr(table, "display_order") == self.display_order - ) - if self.layout_size is not None: - custom_filters.append( - getattr(table, "layout_size") == self.layout_size - ) - - # resource-based filtering is handled within the store implementation - return custom_filters From 347ac877a00e81bdf2e4453960cf3bfd677c5569 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Tue, 14 Oct 2025 12:29:15 +0100 Subject: [PATCH 22/64] format --- src/zenml/models/v2/core/curated_visualization.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index c0cfccc739..d49d568da4 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -46,8 +46,6 @@ ) if TYPE_CHECKING: - from sqlalchemy.sql.elements import ColumnElement - from zenml.models.v2.core.artifact_version import ArtifactVersionResponse from zenml.zen_stores.schemas.base_schemas import BaseSchema @@ -271,10 +269,10 @@ class CuratedVisualizationFilter(ProjectScopedFilter): description="Which column to sort by.", ) - artifact_version_id: Optional[Union[UUID, str]]= Field( + artifact_version_id: Optional[Union[UUID, str]] = Field( default=None, description="ID of the artifact version associated with the visualization.", - union_mode="left_to_right", + union_mode="left_to_right", ) visualization_index: Optional[int] = Field( default=None, @@ -339,4 +337,3 @@ def apply_sorting( return cast(AnyQuery, query.order_by(asc(column), asc(table.id))) return super().apply_sorting(query=query, table=table) - From 15b26c46968960338297634471db364384030507 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Tue, 14 Oct 2025 12:50:18 +0100 Subject: [PATCH 23/64] fix enum --- src/zenml/client.py | 37 ------------------- .../schemas/curated_visualization_schemas.py | 26 +++++++++---- 2 files changed, 19 insertions(+), 44 deletions(-) diff --git a/src/zenml/client.py b/src/zenml/client.py index ea3737847f..f84d3cbd35 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -3806,43 +3806,6 @@ def create_curated_visualization( ) return self.zen_store.create_curated_visualization(request) - def add_visualization_to_deployment( - self, - deployment_id: UUID, - artifact_version_id: UUID, - visualization_index: int, - *, - display_name: Optional[str] = None, - display_order: Optional[int] = None, - layout_size: CuratedVisualizationSize = CuratedVisualizationSize.FULL_WIDTH, - ) -> CuratedVisualizationResponse: - """Attach a curated visualization to a deployment. - - Args: - deployment_id: The ID of the deployment to add visualization to. - artifact_version_id: The ID of the artifact version containing the visualization. - visualization_index: The index of the visualization within the artifact version. - display_name: Optional display name for the visualization. - display_order: Optional display order for sorting visualizations. - layout_size: Layout size defining the visualization width on the dashboard. - - Returns: - The created curated visualization. - """ - deployment = self.get_deployment(deployment_id) - return self.create_curated_visualization( - artifact_version_id=artifact_version_id, - visualization_index=visualization_index, - resource=CuratedVisualizationResource( - id=deployment_id, - type=VisualizationResourceTypes.DEPLOYMENT, - ), - project_id=deployment.project_id, - display_name=display_name, - display_order=display_order, - layout_size=layout_size, - ) - def list_curated_visualizations( self, *, diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 69c20fc641..604cd69166 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -102,8 +102,9 @@ class CuratedVisualizationSchema(BaseSchema, table=True): visualization_index: int = Field(nullable=False) display_name: Optional[str] = Field(default=None) display_order: Optional[int] = Field(default=None) - layout_size: CuratedVisualizationSize = Field( - default=CuratedVisualizationSize.FULL_WIDTH, nullable=False + layout_size: str = Field( + default=CuratedVisualizationSize.FULL_WIDTH.value, + nullable=False, ) artifact_version: Optional["ArtifactVersionSchema"] = Relationship( @@ -167,7 +168,7 @@ def from_request( visualization_index=request.visualization_index, display_name=request.display_name, display_order=request.display_order, - layout_size=request.layout_size, + layout_size=request.layout_size.value, ) def update( @@ -182,9 +183,12 @@ def update( Returns: The updated schema. """ - for field, value in update.model_dump( - exclude_unset=True, - ).items(): + changes = update.model_dump(exclude_unset=True) + layout_size_update = changes.pop("layout_size", None) + if layout_size_update is not None: + self.layout_size = layout_size_update.value + + for field, value in changes.items(): if hasattr(self, field): setattr(self, field, value) @@ -209,6 +213,14 @@ def to_model( Returns: The created response model. """ + layout_size_value = ( + self.layout_size or CuratedVisualizationSize.FULL_WIDTH.value + ) + try: + layout_size_enum = CuratedVisualizationSize(layout_size_value) + except ValueError: + layout_size_enum = CuratedVisualizationSize.FULL_WIDTH + body = CuratedVisualizationResponseBody( project_id=self.project_id, created=self.created, @@ -217,7 +229,7 @@ def to_model( visualization_index=self.visualization_index, display_name=self.display_name, display_order=self.display_order, - layout_size=self.layout_size, + layout_size=layout_size_enum, ) metadata = None From 5ba7c9e52c80d4b71f3d806e31cb949d36cbad5f Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Tue, 14 Oct 2025 15:56:16 +0100 Subject: [PATCH 24/64] apply stefan reviews --- docs/book/how-to/artifacts/visualizations.md | 18 +- src/zenml/client.py | 145 --------------- src/zenml/models/__init__.py | 2 - .../models/v2/core/curated_visualization.py | 123 ++----------- src/zenml/zen_server/rbac/models.py | 29 +-- .../curated_visualization_endpoints.py | 167 +++++++++--------- src/zenml/zen_stores/rest_zen_store.py | 23 --- .../schemas/curated_visualization_schemas.py | 31 +++- src/zenml/zen_stores/sql_zen_store.py | 54 ------ src/zenml/zen_stores/zen_store_interface.py | 9 - .../functional/zen_stores/test_zen_store.py | 100 +++-------- 11 files changed, 156 insertions(+), 545 deletions(-) diff --git a/docs/book/how-to/artifacts/visualizations.md b/docs/book/how-to/artifacts/visualizations.md index 4d17340b91..26291cb037 100644 --- a/docs/book/how-to/artifacts/visualizations.md +++ b/docs/book/how-to/artifacts/visualizations.md @@ -69,8 +69,6 @@ There are three ways how you can add custom visualizations to the dashboard: Curated visualizations let you surface a specific artifact visualization across multiple ZenML resources. Each curated visualization links to exactly one resource—for example, a model performance report that appears on the model detail page, or a deployment health dashboard that shows up in the deployment view. -To cover multiple resources with the same visualization, create separate curated visualizations for each resource type. This gives you fine-grained control over which dashboards appear where. - Curated visualizations currently support the following resources: - **Projects** – high-level dashboards and KPIs that summarize the state of a project. @@ -138,20 +136,16 @@ project_viz = client.create_curated_visualization( ) ``` -Note that each visualization is created separately and can reference the same artifact visualization (by using the same `artifact_version_id` and `visualization_index`). This allows you to show the same underlying visualization in multiple contexts while maintaining separate display settings for each resource. Use the optional `layout_size` argument to control whether the visualization spans the full width of the dashboard or renders as a half-width tile. If omitted, the layout defaults to `CuratedVisualizationSize.FULL_WIDTH`. - -To list curated visualizations for a specific resource, you can use the `Client.list_curated_visualizations` convenience parameters: +After creation, the returned response includes the visualization ID. You can retrieve a specific visualization later with `Client.get_curated_visualization`: ```python -client.list_curated_visualizations(project_id=project.id) -client.list_curated_visualizations(deployment_id=deployment.id) -client.list_curated_visualizations(model_id=model.id) -client.list_curated_visualizations(pipeline_id=pipeline.id) -client.list_curated_visualizations(pipeline_run_id=pipeline_run.id) -client.list_curated_visualizations(pipeline_snapshot_id=snapshot.id) +retrieved = client.get_curated_visualization(model_viz.id, hydrate=True) +print(retrieved.display_name) +print(retrieved.resource.type) +print(retrieved.resource.id) ``` -Each call returns a `Page[CuratedVisualizationResponse]` object so you can iterate through the visualizations or fetch the hydrated version for additional metadata. +Curated visualizations are tied to their parent resources and automatically surface in the ZenML dashboard wherever those resources appear, so keep track of the IDs returned by `create_curated_visualization` if you need to reference them later. #### Updating curated visualizations diff --git a/src/zenml/client.py b/src/zenml/client.py index f84d3cbd35..d711acaf55 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -73,7 +73,6 @@ StackComponentType, StoreType, TaggableResourceTypes, - VisualizationResourceTypes, ) from zenml.exceptions import ( AuthorizationException, @@ -110,7 +109,6 @@ ComponentRequest, ComponentResponse, ComponentUpdate, - CuratedVisualizationFilter, CuratedVisualizationRequest, CuratedVisualizationResource, CuratedVisualizationResponse, @@ -3806,149 +3804,6 @@ def create_curated_visualization( ) return self.zen_store.create_curated_visualization(request) - def list_curated_visualizations( - self, - *, - resource_type: Optional[VisualizationResourceTypes] = None, - resource_id: Optional[UUID] = None, - deployment_id: Optional[UUID] = None, - model_id: Optional[UUID] = None, - pipeline_id: Optional[UUID] = None, - pipeline_run_id: Optional[UUID] = None, - pipeline_snapshot_id: Optional[UUID] = None, - project_id: Optional[UUID] = None, - page: Optional[int] = None, - size: Optional[int] = None, - sort_by: Optional[str] = None, - visualization_index: Optional[int] = None, - layout_size: Optional[CuratedVisualizationSize] = None, - hydrate: bool = False, - ) -> Page[CuratedVisualizationResponse]: - """List curated visualizations, optionally scoped to a resource. - - This method supports filtering by any of the following resource types: - - Deployments (use `deployment_id` parameter) - - Models (use `model_id` parameter) - - Pipelines (use `pipeline_id` parameter) - - Pipeline Runs (use `pipeline_run_id` parameter) - - Pipeline Snapshots (use `pipeline_snapshot_id` parameter) - - Projects (use `project_id` parameter) - - Alternatively, you can use `resource_type` and `resource_id` directly - for more flexible filtering. - - Args: - resource_type: The type of the resource to filter by. - resource_id: The ID of the resource to filter by. - deployment_id: Convenience parameter to filter by deployment. - model_id: Convenience parameter to filter by model. - pipeline_id: Convenience parameter to filter by pipeline. - pipeline_run_id: Convenience parameter to filter by pipeline run. - pipeline_snapshot_id: Convenience parameter to filter by pipeline snapshot. - project_id: Convenience parameter to filter by project. - page: The page of items. - size: The maximum size of all pages. - sort_by: The column to sort by. - visualization_index: The index of the visualization to filter by. - layout_size: The layout size of the visualization tiles to filter by. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - A page of curated visualizations. - - Raises: - ValueError: If multiple resource ID parameters are provided, if - visualization_index is negative, or if a convenience parameter - conflicts with explicitly provided resource_type or resource_id. - """ - # Build convenience params dict mapping parameter names to (value, resource_type) tuples - convenience_params = { - "deployment_id": ( - deployment_id, - VisualizationResourceTypes.DEPLOYMENT, - ), - "model_id": (model_id, VisualizationResourceTypes.MODEL), - "pipeline_id": (pipeline_id, VisualizationResourceTypes.PIPELINE), - "pipeline_run_id": ( - pipeline_run_id, - VisualizationResourceTypes.PIPELINE_RUN, - ), - "pipeline_snapshot_id": ( - pipeline_snapshot_id, - VisualizationResourceTypes.PIPELINE_SNAPSHOT, - ), - "project_id": (project_id, VisualizationResourceTypes.PROJECT), - } - - # Filter to only provided parameters - provided = { - param_name: (param_value, param_type) - for param_name, ( - param_value, - param_type, - ) in convenience_params.items() - if param_value is not None - } - - if len(provided) > 1: - param_names = list(provided.keys()) - raise ValueError( - f"Only one resource ID parameter can be specified at a time. " - f"Got: {', '.join(param_names)}" - ) - - # Validate consistency between convenience parameters and explicit arguments - if provided: - param_name, (param_id, param_type) = list(provided.items())[0] - - # Check for resource_type mismatch - if resource_type is not None and resource_type != param_type: - raise ValueError( - f"Conflicting resource type: convenience parameter '{param_name}' " - f"implies resource_type={param_type.value}, but " - f"resource_type={resource_type.value} was explicitly provided." - ) - - # Check for resource_id mismatch - if resource_id is not None and resource_id != param_id: - raise ValueError( - f"Conflicting resource ID: convenience parameter '{param_name}' " - f"specifies resource_id={param_id}, but a different " - f"resource_id={resource_id} was explicitly provided." - ) - - # Auto-set resource_type and resource_id from convenience parameters - resource_type = resource_type or param_type - resource_id = resource_id or param_id - - if visualization_index is not None and visualization_index < 0: - raise ValueError("visualization_index must be non-negative") - - filter_model = CuratedVisualizationFilter( - project=self.active_project.id - ) - if resource_type is not None: - filter_model.resource_type = resource_type - if resource_id is not None: - filter_model.resource_id = resource_id - - if page is not None: - filter_model.page = page - if size is not None: - filter_model.size = size - if sort_by is not None: - filter_model.sort_by = sort_by - if visualization_index is not None: - filter_model.visualization_index = visualization_index - if layout_size is not None: - filter_model.layout_size = layout_size - - return self.zen_store.list_curated_visualizations( - filter_model=filter_model, - hydrate=hydrate, - ) - def update_curated_visualization( self, visualization_id: UUID, diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index eedc8b0a6f..0140de6bc3 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -165,7 +165,6 @@ DeploymentResponseResources, ) from zenml.models.v2.core.curated_visualization import ( - CuratedVisualizationFilter, CuratedVisualizationRequest, CuratedVisualizationResponse, CuratedVisualizationResponseBody, @@ -665,7 +664,6 @@ "DeploymentResponseBody", "DeploymentResponseMetadata", "DeploymentResponseResources", - "CuratedVisualizationFilter", "CuratedVisualizationRequest", "CuratedVisualizationResponse", "CuratedVisualizationResponseBody", diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index d49d568da4..e355310379 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -13,28 +13,14 @@ # permissions and limitations under the License. """Models representing curated visualizations.""" -from typing import ( - TYPE_CHECKING, - ClassVar, - List, - Optional, - Type, - TypeVar, - Union, - cast, -) +from typing import TYPE_CHECKING, Optional from uuid import UUID from pydantic import Field -from zenml.enums import ( - CuratedVisualizationSize, - VisualizationResourceTypes, -) +from zenml.enums import CuratedVisualizationSize from zenml.models.v2.base.base import BaseUpdate -from zenml.models.v2.base.filter import AnyQuery from zenml.models.v2.base.scoped import ( - ProjectScopedFilter, ProjectScopedRequest, ProjectScopedResponse, ProjectScopedResponseBody, @@ -47,9 +33,6 @@ if TYPE_CHECKING: from zenml.models.v2.core.artifact_version import ArtifactVersionResponse - from zenml.zen_stores.schemas.base_schemas import BaseSchema - - AnySchema = TypeVar("AnySchema", bound=BaseSchema) # ------------------ Request Model ------------------ @@ -166,6 +149,11 @@ class CuratedVisualizationResponseResources(ProjectScopedResponseResources): title="The artifact version.", description="Artifact version from which the visualization originates.", ) + resource: Optional[CuratedVisualizationResource] = Field( + default=None, + title="The linked resource.", + description="Resource reference associated with this curated visualization.", + ) class CuratedVisualizationResponse( @@ -243,97 +231,12 @@ def artifact_version(self) -> Optional["ArtifactVersionResponse"]: """ return self.get_resources().artifact_version - -# ------------------ Filter Model ------------------ - - -class CuratedVisualizationFilter(ProjectScopedFilter): - """Model to enable advanced filtering of curated visualizations.""" - - FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [ - *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS, - "resource_id", - "resource_type", - "layout_size", - ] - CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [ - *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS, - "display_order", - "created", - "updated", - "visualization_index", - ] - - sort_by: str = Field( - default="display_order", - description="Which column to sort by.", - ) - - artifact_version_id: Optional[Union[UUID, str]] = Field( - default=None, - description="ID of the artifact version associated with the visualization.", - union_mode="left_to_right", - ) - visualization_index: Optional[int] = Field( - default=None, - description="Index of the visualization within the artifact version payload.", - ) - display_order: Optional[int] = Field( - default=None, - description="Display order of the visualization.", - ) - layout_size: Optional[CuratedVisualizationSize] = Field( - default=None, - description="Layout size of the visualization tile.", - ) - resource_type: Optional[VisualizationResourceTypes] = Field( - default=None, - description="Type of the resource exposing the visualization.", - ) - resource_id: Optional[UUID] = Field( - default=None, - description="ID of the resource exposing the visualization.", - ) - - def apply_sorting( - self, - query: AnyQuery, - table: Type["AnySchema"], - ) -> AnyQuery: - """Apply sorting to the curated visualization query. - - Args: - query: The query to which to apply the sorting. - table: The query table. + @property + def resource(self) -> Optional[CuratedVisualizationResource]: + """The resource reference associated with this visualization. Returns: - The query with sorting applied. + The resource reference if included. """ - from sqlmodel import asc, desc - - from zenml.enums import SorterOps - - sort_by, operand = self.sorting_params - - if sort_by == "display_order": - column = getattr(table, "display_order") - if operand == SorterOps.DESCENDING: - return cast( - AnyQuery, - query.order_by(desc(column).nulls_last()), - ) - return cast( - AnyQuery, - query.order_by(asc(column).nulls_first()), - ) - - if sort_by in {"created", "updated", "visualization_index"}: - column = getattr(table, sort_by) - if operand == SorterOps.DESCENDING: - return cast( - AnyQuery, - query.order_by(desc(column), asc(table.id)), - ) - return cast(AnyQuery, query.order_by(asc(column), asc(table.id))) - - return super().apply_sorting(query=query, table=table) + resources = self.get_resources() + return resources.resource if resources else None diff --git a/src/zenml/zen_server/rbac/models.py b/src/zenml/zen_server/rbac/models.py index ca73c7d931..bc8984e403 100644 --- a/src/zenml/zen_server/rbac/models.py +++ b/src/zenml/zen_server/rbac/models.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """RBAC model classes.""" -from typing import TYPE_CHECKING, Optional +from typing import Optional from uuid import UUID from pydantic import ( @@ -24,9 +24,6 @@ from zenml.utils.enum_utils import StrEnum -if TYPE_CHECKING: - from zenml.enums import VisualizationResourceTypes - class Action(StrEnum): """RBAC actions.""" @@ -83,30 +80,6 @@ class ResourceType(StrEnum): # Deactivated for now # USER = "user" - @classmethod - def from_visualization_type( - cls, visualization_type: "VisualizationResourceTypes" - ) -> Optional["ResourceType"]: - """Map visualization resource types to RBAC resource types. - - Args: - visualization_type: The type of the visualization to map. - - Returns: - The RBAC resource type associated with the visualization type. - """ - from zenml.enums import VisualizationResourceTypes - - mapping = { - VisualizationResourceTypes.DEPLOYMENT: cls.DEPLOYMENT, - VisualizationResourceTypes.MODEL: cls.MODEL, - VisualizationResourceTypes.PIPELINE: cls.PIPELINE, - VisualizationResourceTypes.PIPELINE_RUN: cls.PIPELINE_RUN, - VisualizationResourceTypes.PIPELINE_SNAPSHOT: cls.PIPELINE_SNAPSHOT, - VisualizationResourceTypes.PROJECT: cls.PROJECT, - } - return mapping.get(visualization_type) - def is_project_scoped(self) -> bool: """Check if a resource type is project scoped. diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py index f10fdea349..ab140c2df1 100644 --- a/src/zenml/zen_server/routers/curated_visualization_endpoints.py +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -1,36 +1,23 @@ """REST API endpoints for curated visualizations.""" +from typing import Any from uuid import UUID -from fastapi import APIRouter, Depends, Security +from fastapi import APIRouter, Security from zenml.constants import API, CURATED_VISUALIZATIONS, VERSION_1 +from zenml.enums import VisualizationResourceTypes from zenml.models import ( - CuratedVisualizationFilter, CuratedVisualizationRequest, + CuratedVisualizationResource, CuratedVisualizationResponse, CuratedVisualizationUpdate, - Page, ) from zenml.zen_server.auth import AuthContext, authorize from zenml.zen_server.exceptions import error_response -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, -) -from zenml.zen_server.utils import ( - async_fastapi_endpoint_wrapper, - make_dependable, - zen_store, -) +from zenml.zen_server.rbac.models import Action +from zenml.zen_server.rbac.utils import verify_permission_for_model +from zenml.zen_server.utils import async_fastapi_endpoint_wrapper, zen_store router = APIRouter( prefix=API + VERSION_1 + CURATED_VISUALIZATIONS, @@ -39,6 +26,31 @@ ) +def _get_resource_model( + resource: CuratedVisualizationResource, +) -> Any: + """Load the model associated with a curated visualization resource.""" + store = zen_store() + resource_type = resource.type + + if resource_type == VisualizationResourceTypes.DEPLOYMENT: + return store.get_deployment(resource.id) + if resource_type == VisualizationResourceTypes.MODEL: + return store.get_model(resource.id) + if resource_type == VisualizationResourceTypes.PIPELINE: + return store.get_pipeline(resource.id) + if resource_type == VisualizationResourceTypes.PIPELINE_RUN: + return store.get_run(resource.id) + if resource_type == VisualizationResourceTypes.PIPELINE_SNAPSHOT: + return store.get_snapshot(resource.id) + if resource_type == VisualizationResourceTypes.PROJECT: + return store.get_project(resource.id) + + raise RuntimeError( + f"Unsupported curated visualization resource type: {resource_type}" + ) + + @router.post( "", responses={ @@ -61,57 +73,16 @@ def create_curated_visualization( Returns: The created curated visualization. """ - return verify_permissions_and_create_entity( - request_model=visualization, - create_method=zen_store().create_curated_visualization, + store = zen_store() + resource_model = _get_resource_model(visualization.resource) + artifact_version = store.get_artifact_version( + visualization.artifact_version_id ) + verify_permission_for_model(resource_model, action=Action.UPDATE) + verify_permission_for_model(artifact_version, action=Action.READ) -@router.get( - "", - responses={401: error_response, 404: error_response, 422: error_response}, -) -@async_fastapi_endpoint_wrapper(deduplicate=True) -def list_curated_visualizations( - visualization_filter_model: CuratedVisualizationFilter = Depends( - make_dependable(CuratedVisualizationFilter) - ), - hydrate: bool = False, - _: AuthContext = Security(authorize), -) -> Page[CuratedVisualizationResponse]: - """List curated visualizations. - - Args: - visualization_filter_model: Filter model used for pagination, sorting, - filtering. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - A page of curated visualizations. - """ - resource_type = None - if visualization_filter_model.resource_type: - resource_type = ResourceType.from_visualization_type( - visualization_filter_model.resource_type - ) - - if resource_type is None: - # No concrete resource type - call zen store directly and batch verify permissions - page = zen_store().list_curated_visualizations( - filter_model=visualization_filter_model, - hydrate=hydrate, - ) - batch_verify_permissions_for_models(page.items, action=Action.READ) - return dehydrate_page(page) - else: - # Concrete resource type available - use standard RBAC flow - return verify_permissions_and_list_entities( - filter_model=visualization_filter_model, - resource_type=resource_type, - list_method=zen_store().list_curated_visualizations, - hydrate=hydrate, - ) + return store.create_curated_visualization(visualization) @router.get( @@ -128,17 +99,28 @@ def get_curated_visualization( Args: visualization_id: The ID of the curated visualization to retrieve. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. + hydrate: Flag deciding whether to return the hydrated model. Returns: The curated visualization with the given ID. """ - return verify_permissions_and_get_entity( - id=visualization_id, - get_method=zen_store().get_curated_visualization, - hydrate=hydrate, + store = zen_store() + hydrated_visualization = store.get_curated_visualization( + visualization_id, hydrate=True ) + resource = hydrated_visualization.resource + if resource is None: + raise RuntimeError( + f"Curated visualization '{visualization_id}' is missing its resource reference." + ) + + resource_model = _get_resource_model(resource) + verify_permission_for_model(resource_model, action=Action.READ) + + if hydrate: + return hydrated_visualization + + return store.get_curated_visualization(visualization_id, hydrate=False) @router.patch( @@ -160,11 +142,21 @@ def update_curated_visualization( Returns: The updated curated visualization. """ - return verify_permissions_and_update_entity( - id=visualization_id, - update_model=visualization_update, - get_method=zen_store().get_curated_visualization, - update_method=zen_store().update_curated_visualization, + store = zen_store() + existing_visualization = store.get_curated_visualization( + visualization_id, hydrate=True + ) + resource = existing_visualization.resource + if resource is None: + raise RuntimeError( + f"Curated visualization '{visualization_id}' is missing its resource reference." + ) + + resource_model = _get_resource_model(resource) + verify_permission_for_model(resource_model, action=Action.UPDATE) + + return store.update_curated_visualization( + visualization_id, visualization_update ) @@ -182,8 +174,17 @@ def delete_curated_visualization( Args: visualization_id: The ID of the curated visualization to delete. """ - verify_permissions_and_delete_entity( - id=visualization_id, - get_method=zen_store().get_curated_visualization, - delete_method=zen_store().delete_curated_visualization, + store = zen_store() + existing_visualization = store.get_curated_visualization( + visualization_id, hydrate=True ) + resource = existing_visualization.resource + if resource is None: + raise RuntimeError( + f"Curated visualization '{visualization_id}' is missing its resource reference." + ) + + resource_model = _get_resource_model(resource) + verify_permission_for_model(resource_model, action=Action.UPDATE) + + store.delete_curated_visualization(visualization_id) diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 44a0248ed6..8d996b0969 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -166,7 +166,6 @@ ComponentRequest, ComponentResponse, ComponentUpdate, - CuratedVisualizationFilter, CuratedVisualizationRequest, CuratedVisualizationResponse, CuratedVisualizationUpdate, @@ -1899,28 +1898,6 @@ def get_curated_visualization( params={"hydrate": hydrate}, ) - def list_curated_visualizations( - self, - filter_model: CuratedVisualizationFilter, - hydrate: bool = False, - ) -> Page[CuratedVisualizationResponse]: - """List curated visualizations via REST API. - - Args: - filter_model: The filter model to use. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - A page of curated visualizations. - """ - return self._list_paginated_resources( - route=CURATED_VISUALIZATIONS, - response_model=CuratedVisualizationResponse, - filter_model=filter_model, - params={"hydrate": hydrate}, - ) - def update_curated_visualization( self, visualization_id: UUID, diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 604cd69166..7798410d65 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -21,7 +21,7 @@ from sqlalchemy.sql.base import ExecutableOption from sqlmodel import Field, Relationship, SQLModel -from zenml.enums import CuratedVisualizationSize +from zenml.enums import CuratedVisualizationSize, VisualizationResourceTypes from zenml.models.v2.core.curated_visualization import ( CuratedVisualizationRequest, CuratedVisualizationResponse, @@ -30,6 +30,9 @@ CuratedVisualizationResponseResources, CuratedVisualizationUpdate, ) +from zenml.models.v2.misc.curated_visualization import ( + CuratedVisualizationResource, +) from zenml.zen_stores.schemas.base_schemas import BaseSchema from zenml.zen_stores.schemas.project_schemas import ProjectSchema from zenml.zen_stores.schemas.schema_utils import ( @@ -238,13 +241,33 @@ def to_model( response_resources = None if include_resources: - response_resources = CuratedVisualizationResponseResources( - artifact_version=self.artifact_version.to_model( + artifact_version_model = ( + self.artifact_version.to_model( include_metadata=include_metadata, include_resources=include_resources, ) if self.artifact_version - else None, + else None + ) + + resource_model = None + if self.resource: + try: + resource_type_enum = VisualizationResourceTypes( + self.resource.resource_type + ) + except ValueError: + resource_type_enum = None + + if resource_type_enum is not None: + resource_model = CuratedVisualizationResource( + id=self.resource.resource_id, + type=resource_type_enum, + ) + + response_resources = CuratedVisualizationResponseResources( + artifact_version=artifact_version_model, + resource=resource_model, ) return CuratedVisualizationResponse( diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 4dd02f6138..02b5948964 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -77,12 +77,10 @@ ) from sqlalchemy.orm import ( Mapped, - aliased, load_only, noload, selectinload, ) -from sqlalchemy.orm.util import AliasedClass from sqlalchemy.sql.base import ExecutableOption from sqlalchemy.util import immutabledict from sqlmodel import Session as SqlModelSession @@ -207,7 +205,6 @@ ComponentRequest, ComponentResponse, ComponentUpdate, - CuratedVisualizationFilter, CuratedVisualizationRequest, CuratedVisualizationResponse, CuratedVisualizationUpdate, @@ -5621,57 +5618,6 @@ def get_curated_visualization( include_resources=hydrate, ) - def list_curated_visualizations( - self, - filter_model: CuratedVisualizationFilter, - hydrate: bool = False, - ) -> Page[CuratedVisualizationResponse]: - """List all curated visualizations matching the given filter. - - Args: - filter_model: The filter model to use. - hydrate: Flag deciding whether to hydrate the output model(s) - by including metadata fields in the response. - - Returns: - A page of curated visualizations. - """ - with Session(self.engine) as session: - self._set_filter_project_id( - filter_model=filter_model, session=session - ) - - query = select(CuratedVisualizationSchema) - if filter_model.resource_type or filter_model.resource_id: - resource_alias = cast( - AliasedClass[CuratedVisualizationResourceSchema], - aliased(CuratedVisualizationResourceSchema), - ) - query = query.join( - resource_alias, - resource_alias.visualization_id - == CuratedVisualizationSchema.id, - ) - if filter_model.resource_type: - query = query.where( - resource_alias.resource_type - == filter_model.resource_type.value - ) - if filter_model.resource_id: - query = query.where( - resource_alias.resource_id == filter_model.resource_id - ) - # Add distinct to prevent duplicate rows when joined to multiple resources - query = query.distinct() - - return self.filter_and_paginate( - session=session, - query=query, - table=CuratedVisualizationSchema, - filter_model=filter_model, - hydrate=hydrate, - ) - def update_curated_visualization( self, visualization_id: UUID, diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index 90e137d8fc..6dedfd9401 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -48,7 +48,6 @@ ComponentRequest, ComponentResponse, ComponentUpdate, - CuratedVisualizationFilter, CuratedVisualizationRequest, CuratedVisualizationResponse, CuratedVisualizationUpdate, @@ -1489,14 +1488,6 @@ def get_curated_visualization( ) -> CuratedVisualizationResponse: """Get a curated visualization by ID.""" - @abstractmethod - def list_curated_visualizations( - self, - filter_model: CuratedVisualizationFilter, - hydrate: bool = False, - ) -> Page[CuratedVisualizationResponse]: - """List curated visualizations matching the given filter criteria.""" - @abstractmethod def update_curated_visualization( self, diff --git a/tests/integration/functional/zen_stores/test_zen_store.py b/tests/integration/functional/zen_stores/test_zen_store.py index f0c8667b6b..a6512b8477 100644 --- a/tests/integration/functional/zen_stores/test_zen_store.py +++ b/tests/integration/functional/zen_stores/test_zen_store.py @@ -96,7 +96,6 @@ ArtifactVisualizationRequest, ComponentFilter, ComponentUpdate, - CuratedVisualizationFilter, CuratedVisualizationRequest, CuratedVisualizationResource, CuratedVisualizationUpdate, @@ -5882,62 +5881,18 @@ def test_curated_visualizations_across_resources(self): display_name=f"{resource_type.value} visualization", ) ) - visualizations[resource_type] = viz - assert viz.layout_size == CuratedVisualizationSize.FULL_WIDTH - - # Verify via zen_store filtering - each resource should have exactly one visualization - for resource_type, resource_id in resource_configs: - result = client.zen_store.list_curated_visualizations( - CuratedVisualizationFilter( - project=project_id, - resource_type=resource_type, - resource_id=resource_id, - ) - ) - assert result.total == 1 - assert result.items[0].id == visualizations[resource_type].id - - # Verify via client convenience parameters - convenience_params = [ - ( - "pipeline_id", - pipeline_model.id, - VisualizationResourceTypes.PIPELINE, - ), - ("model_id", model.id, VisualizationResourceTypes.MODEL), - ( - "pipeline_run_id", - pipeline_run.id, - VisualizationResourceTypes.PIPELINE_RUN, - ), - ( - "pipeline_snapshot_id", - snapshot.id, - VisualizationResourceTypes.PIPELINE_SNAPSHOT, - ), - ( - "deployment_id", - deployment.id, - VisualizationResourceTypes.DEPLOYMENT, - ), - ("project_id", project_id, VisualizationResourceTypes.PROJECT), - ] - - for ( - param_name, - param_value, - expected_resource_type, - ) in convenience_params: - result = client.list_curated_visualizations( - **{param_name: param_value} + hydrated = client.zen_store.get_curated_visualization( + visualization_id=viz.id, + hydrate=True, ) - assert result.total == 1, f"Failed for {param_name}" + assert hydrated.resource is not None + assert hydrated.resource.type == resource_type + assert hydrated.resource.id == resource_id assert ( - result.items[0].id - == visualizations[expected_resource_type].id + hydrated.layout_size == CuratedVisualizationSize.FULL_WIDTH ) + visualizations[resource_type] = viz - # Test hydrate/get loaded = client.zen_store.get_curated_visualization( visualizations[VisualizationResourceTypes.PIPELINE].id, hydrate=True, @@ -5947,6 +5902,8 @@ def test_curated_visualizations_across_resources(self): == f"{VisualizationResourceTypes.PIPELINE.value} visualization" ) assert loaded.layout_size == CuratedVisualizationSize.FULL_WIDTH + assert loaded.resource is not None + assert loaded.resource.type == VisualizationResourceTypes.PIPELINE # Test duplicate creation - same artifact_version + visualization_index + resource should fail with pytest.raises(EntityExistsError): @@ -5981,16 +5938,9 @@ def test_curated_visualizations_across_resources(self): for viz in visualizations.values(): client.zen_store.delete_curated_visualization(viz.id) - # Verify all deleted - for resource_type, resource_id in resource_configs: - result = client.zen_store.list_curated_visualizations( - CuratedVisualizationFilter( - project=project_id, - resource_type=resource_type, - resource_id=resource_id, - ) - ) - assert result.total == 0 + for viz in visualizations.values(): + with pytest.raises(KeyError): + client.zen_store.get_curated_visualization(viz.id) finally: # Clean up deployment client.zen_store.delete_deployment(deployment.id) @@ -6039,20 +5989,20 @@ def test_curated_visualizations_project_only(self): display_name="Project visualization", ) - result = client.zen_store.list_curated_visualizations( - CuratedVisualizationFilter( - project=project.id, - resource_type=VisualizationResourceTypes.PROJECT, - resource_id=project.id, - ) + hydrated_visualization = client.zen_store.get_curated_visualization( + visualization.id, hydrate=True ) - assert result.total == 1 - assert result.items[0].id == visualization.id - - result = client.list_curated_visualizations(project_id=project.id) - assert result.total == 1 - assert result.items[0].id == visualization.id + assert hydrated_visualization.resource is not None + assert ( + hydrated_visualization.resource.type + == VisualizationResourceTypes.PROJECT + ) + assert hydrated_visualization.resource.id == project.id + assert hydrated_visualization.display_name == "Project visualization" client.delete_curated_visualization(visualization.id) + with pytest.raises(KeyError): + client.zen_store.get_curated_visualization(visualization.id) + client.delete_artifact_version(artifact_version.id) client.delete_artifact(artifact.id) From dae9fa511e951d5062943ce9b04e572fd8fa77fd Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Wed, 15 Oct 2025 11:46:57 +0100 Subject: [PATCH 25/64] delete migration outdated file --- ...d1_add_vizualisations_and_link_them_to_.py | 125 ------------------ 1 file changed, 125 deletions(-) delete mode 100644 src/zenml/zen_stores/migrations/versions/f2ab38646bd1_add_vizualisations_and_link_them_to_.py diff --git a/src/zenml/zen_stores/migrations/versions/f2ab38646bd1_add_vizualisations_and_link_them_to_.py b/src/zenml/zen_stores/migrations/versions/f2ab38646bd1_add_vizualisations_and_link_them_to_.py deleted file mode 100644 index 3f897038b1..0000000000 --- a/src/zenml/zen_stores/migrations/versions/f2ab38646bd1_add_vizualisations_and_link_them_to_.py +++ /dev/null @@ -1,125 +0,0 @@ -"""add vizualisations and link them to resources [f2ab38646bd1]. - -Revision ID: f2ab38646bd1 -Revises: 0.90.0 -Create Date: 2025-10-13 16:01:41.907536 - -""" - -import sqlalchemy as sa -import sqlmodel -from alembic import op - -# revision identifiers, used by Alembic. -revision = "f2ab38646bd1" -down_revision = "0.90.0" -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! ### - op.create_table( - "curated_visualization", - sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column("created", sa.DateTime(), nullable=False), - sa.Column("updated", sa.DateTime(), nullable=False), - sa.Column("project_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column( - "artifact_version_id", sqlmodel.sql.sqltypes.GUID(), nullable=False - ), - sa.Column("visualization_index", sa.Integer(), nullable=False), - sa.Column( - "display_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True - ), - sa.Column("display_order", sa.Integer(), nullable=True), - sa.Column( - "layout_size", - sa.Enum( - "FULL_WIDTH", "HALF_WIDTH", name="curatedvisualizationsize" - ), - nullable=False, - ), - sa.ForeignKeyConstraint( - ["artifact_version_id"], - ["artifact_version.id"], - name="fk_curated_visualization_artifact_version_id_artifact_version", - ondelete="CASCADE", - ), - sa.ForeignKeyConstraint( - ["project_id"], - ["project.id"], - name="fk_curated_visualization_project_id_project", - ondelete="CASCADE", - ), - sa.PrimaryKeyConstraint("id"), - ) - with op.batch_alter_table( - "curated_visualization", schema=None - ) as batch_op: - batch_op.create_index( - "ix_curated_visualization_artifact_version_id_visualization_index", - ["artifact_version_id", "visualization_index"], - unique=False, - ) - batch_op.create_index( - "ix_curated_visualization_display_order", - ["display_order"], - unique=False, - ) - - op.create_table( - "curated_visualization_resource", - sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column( - "visualization_id", sqlmodel.sql.sqltypes.GUID(), nullable=False - ), - sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column( - "resource_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False - ), - sa.ForeignKeyConstraint( - ["visualization_id"], - ["curated_visualization.id"], - name="fk_curated_visualization_resource_visualization_id_curated_visualization", - ondelete="CASCADE", - ), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "visualization_id", name="unique_curated_visualization_resource" - ), - ) - with op.batch_alter_table( - "curated_visualization_resource", schema=None - ) as batch_op: - batch_op.create_index( - "ix_curated_visualization_resource_resource_id_resource_type", - ["resource_id", "resource_type"], - unique=False, - ) - - # ### 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( - "curated_visualization_resource", schema=None - ) as batch_op: - batch_op.drop_index( - "ix_curated_visualization_resource_resource_id_resource_type" - ) - - op.drop_table("curated_visualization_resource") - with op.batch_alter_table( - "curated_visualization", schema=None - ) as batch_op: - batch_op.drop_index("ix_curated_visualization_display_order") - batch_op.drop_index( - "ix_curated_visualization_artifact_version_id_visualization_index" - ) - - op.drop_table("curated_visualization") - # ### end Alembic commands ### From 4166124d281e777fb364ea46810e21553b114f6d Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Wed, 15 Oct 2025 11:52:31 +0100 Subject: [PATCH 26/64] update migration file --- ...d5_add_vizualisations_and_link_them_to_.py | 121 ++++++++++++++++++ 1 file changed, 121 insertions(+) create mode 100644 src/zenml/zen_stores/migrations/versions/a3c4e7cac1d5_add_vizualisations_and_link_them_to_.py diff --git a/src/zenml/zen_stores/migrations/versions/a3c4e7cac1d5_add_vizualisations_and_link_them_to_.py b/src/zenml/zen_stores/migrations/versions/a3c4e7cac1d5_add_vizualisations_and_link_them_to_.py new file mode 100644 index 0000000000..44ae61de56 --- /dev/null +++ b/src/zenml/zen_stores/migrations/versions/a3c4e7cac1d5_add_vizualisations_and_link_them_to_.py @@ -0,0 +1,121 @@ +"""add vizualisations and link them to resources [a3c4e7cac1d5]. + +Revision ID: a3c4e7cac1d5 +Revises: 0.90.0 +Create Date: 2025-10-15 11:51:33.654401 + +""" + +import sqlalchemy as sa +import sqlmodel +from alembic import op + +# revision identifiers, used by Alembic. +revision = "a3c4e7cac1d5" +down_revision = "0.90.0" +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! ### + op.create_table( + "curated_visualization", + sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column("created", sa.DateTime(), nullable=False), + sa.Column("updated", sa.DateTime(), nullable=False), + sa.Column("project_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "artifact_version_id", sqlmodel.sql.sqltypes.GUID(), nullable=False + ), + sa.Column("visualization_index", sa.Integer(), nullable=False), + sa.Column( + "display_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True + ), + sa.Column("display_order", sa.Integer(), nullable=True), + sa.Column( + "layout_size", sqlmodel.sql.sqltypes.AutoString(), nullable=False + ), + sa.ForeignKeyConstraint( + ["artifact_version_id"], + ["artifact_version.id"], + name="fk_curated_visualization_artifact_version_id_artifact_version", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["project_id"], + ["project.id"], + name="fk_curated_visualization_project_id_project", + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id"), + ) + with op.batch_alter_table( + "curated_visualization", schema=None + ) as batch_op: + batch_op.create_index( + "ix_curated_visualization_artifact_version_id_visualization_index", + ["artifact_version_id", "visualization_index"], + unique=False, + ) + batch_op.create_index( + "ix_curated_visualization_display_order", + ["display_order"], + unique=False, + ) + + op.create_table( + "curated_visualization_resource", + sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "visualization_id", sqlmodel.sql.sqltypes.GUID(), nullable=False + ), + sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "resource_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False + ), + sa.ForeignKeyConstraint( + ["visualization_id"], + ["curated_visualization.id"], + name="fk_curated_visualization_resource_visualization_id_curated_visualization", + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "visualization_id", name="unique_curated_visualization_resource" + ), + ) + with op.batch_alter_table( + "curated_visualization_resource", schema=None + ) as batch_op: + batch_op.create_index( + "ix_curated_visualization_resource_resource_id_resource_type", + ["resource_id", "resource_type"], + unique=False, + ) + + # ### 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( + "curated_visualization_resource", schema=None + ) as batch_op: + batch_op.drop_index( + "ix_curated_visualization_resource_resource_id_resource_type" + ) + + op.drop_table("curated_visualization_resource") + with op.batch_alter_table( + "curated_visualization", schema=None + ) as batch_op: + batch_op.drop_index("ix_curated_visualization_display_order") + batch_op.drop_index( + "ix_curated_visualization_artifact_version_id_visualization_index" + ) + + op.drop_table("curated_visualization") + # ### end Alembic commands ### From 978a1b08127333be25125fd8af9181499dca527d Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Wed, 15 Oct 2025 14:26:41 +0100 Subject: [PATCH 27/64] docstring --- src/zenml/client.py | 3 -- .../curated_visualization_endpoints.py | 21 ++++++++++- .../zen_stores/schemas/deployment_schemas.py | 3 +- src/zenml/zen_stores/schemas/model_schemas.py | 3 +- .../schemas/pipeline_run_schemas.py | 3 +- .../zen_stores/schemas/pipeline_schemas.py | 3 +- .../schemas/pipeline_snapshot_schemas.py | 3 +- .../zen_stores/schemas/project_schemas.py | 3 +- src/zenml/zen_stores/sql_zen_store.py | 7 ++++ src/zenml/zen_stores/zen_store_interface.py | 36 ++++++++++++++++--- 10 files changed, 65 insertions(+), 20 deletions(-) diff --git a/src/zenml/client.py b/src/zenml/client.py index d711acaf55..91ac9436ab 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -3790,9 +3790,6 @@ def create_curated_visualization( Raises: ValueError: If resource is not provided. """ - if not resource: - raise ValueError("resource must be provided") - request = CuratedVisualizationRequest( project=project_id or self.active_project.id, artifact_version_id=artifact_version_id, diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py index ab140c2df1..be6dbf55e0 100644 --- a/src/zenml/zen_server/routers/curated_visualization_endpoints.py +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -29,7 +29,17 @@ def _get_resource_model( resource: CuratedVisualizationResource, ) -> Any: - """Load the model associated with a curated visualization resource.""" + """Load the model associated with a curated visualization resource. + + Args: + resource: The curated visualization resource to load the model for. + + Returns: + The model associated with the curated visualization resource. + + Raises: + RuntimeError: If the resource type is not supported. + """ store = zen_store() resource_type = resource.type @@ -103,6 +113,9 @@ def get_curated_visualization( Returns: The curated visualization with the given ID. + + Raises: + RuntimeError: If the curated visualization is missing its resource reference. """ store = zen_store() hydrated_visualization = store.get_curated_visualization( @@ -141,6 +154,9 @@ def update_curated_visualization( Returns: The updated curated visualization. + + Raises: + RuntimeError: If the curated visualization is missing its resource reference. """ store = zen_store() existing_visualization = store.get_curated_visualization( @@ -173,6 +189,9 @@ def delete_curated_visualization( Args: visualization_id: The ID of the curated visualization to delete. + + Raises: + RuntimeError: If the curated visualization is missing its resource reference. """ store = zen_store() existing_visualization = store.get_curated_visualization( diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index be5a1dcbf3..ad62288219 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -145,8 +145,7 @@ class DeploymentSchema(NamedSchema, table=True): secondary="curated_visualization_resource", primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.DEPLOYMENT.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==DeploymentSchema.id)", secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", - viewonly=True, - lazy="selectin", + overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index bbbfd237cc..9c8e53a31f 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -137,8 +137,7 @@ class ModelSchema(NamedSchema, table=True): secondary="curated_visualization_resource", primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.MODEL.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==ModelSchema.id)", secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", - viewonly=True, - lazy="selectin", + overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index 865adb2e5d..efb47b3fc7 100644 --- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py @@ -250,8 +250,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True): secondary="curated_visualization_resource", primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE_RUN.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineRunSchema.id)", secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", - viewonly=True, - lazy="selectin", + overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index 793b7853cf..abcac6ff29 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_schemas.py @@ -112,8 +112,7 @@ class PipelineSchema(NamedSchema, table=True): secondary="curated_visualization_resource", primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineSchema.id)", secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", - viewonly=True, - lazy="selectin", + overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py index b083ec3061..775b00f6e4 100644 --- a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py @@ -226,8 +226,7 @@ class PipelineSnapshotSchema(BaseSchema, table=True): secondary="curated_visualization_resource", primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE_SNAPSHOT.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineSnapshotSchema.id)", secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", - viewonly=True, - lazy="selectin", + overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/schemas/project_schemas.py b/src/zenml/zen_stores/schemas/project_schemas.py index 7c9ddf275e..11da98391c 100644 --- a/src/zenml/zen_stores/schemas/project_schemas.py +++ b/src/zenml/zen_stores/schemas/project_schemas.py @@ -134,8 +134,7 @@ class ProjectSchema(NamedSchema, table=True): secondary="curated_visualization_resource", primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PROJECT.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==ProjectSchema.id)", secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", - viewonly=True, - lazy="selectin", + overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 02b5948964..df0aeb20aa 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5455,6 +5455,9 @@ def _assert_curated_visualization_duplicate( visualization_index: The index of the visualization to validate. resource_id: The ID of the resource to validate. resource_type: The type of the resource to validate. + + Raises: + EntityExistsError: If a curated visualization link already exists. """ existing = session.exec( select(CuratedVisualizationSchema) @@ -5491,6 +5494,10 @@ def create_curated_visualization( Returns: The persisted curated visualization. + + Raises: + IllegalOperationError: If the resource does not belong to the same project as the curated visualization. + EntityExistsError: If a curated visualization link already exists. """ with Session(self.engine) as session: self._set_request_user_id( diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index 6dedfd9401..350a74c038 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -1480,13 +1480,29 @@ def delete_deployment(self, deployment_id: UUID) -> None: def create_curated_visualization( self, visualization: CuratedVisualizationRequest ) -> CuratedVisualizationResponse: - """Create a new curated visualization.""" + """Create a new curated visualization. + + Args: + visualization: The curated visualization to create. + + Returns: + The created curated visualization. + """ @abstractmethod def get_curated_visualization( self, visualization_id: UUID, hydrate: bool = True ) -> CuratedVisualizationResponse: - """Get a curated visualization by ID.""" + """Get a curated visualization by ID. + + Args: + visualization_id: The ID of the curated visualization to get. + hydrate: Flag deciding whether to hydrate the output model(s) + by including metadata fields in the response. + + Returns: + The curated visualization with the given ID. + """ @abstractmethod def update_curated_visualization( @@ -1494,11 +1510,23 @@ def update_curated_visualization( visualization_id: UUID, visualization_update: CuratedVisualizationUpdate, ) -> CuratedVisualizationResponse: - """Update a curated visualization.""" + """Update a curated visualization. + + Args: + visualization_id: The ID of the curated visualization to update. + visualization_update: The update to apply to the curated visualization. + + Returns: + The updated curated visualization. + """ @abstractmethod def delete_curated_visualization(self, visualization_id: UUID) -> None: - """Delete a curated visualization.""" + """Delete a curated visualization. + + Args: + visualization_id: The ID of the curated visualization to delete. + """ # -------------------- Run templates -------------------- From 22192c9e0ebcd695d113863ae154ec4c1c8dc193 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Wed, 15 Oct 2025 15:59:50 +0100 Subject: [PATCH 28/64] docstring --- src/zenml/client.py | 3 --- src/zenml/zen_stores/sql_zen_store.py | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/src/zenml/client.py b/src/zenml/client.py index 91ac9436ab..285efea31e 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -3786,9 +3786,6 @@ def create_curated_visualization( Returns: The created curated visualization. - - Raises: - ValueError: If resource is not provided. """ request = CuratedVisualizationRequest( project=project_id or self.active_project.id, diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index df0aeb20aa..eda50dd62a 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5497,7 +5497,7 @@ def create_curated_visualization( Raises: IllegalOperationError: If the resource does not belong to the same project as the curated visualization. - EntityExistsError: If a curated visualization link already exists. + KeyError: If the resource does not exist. """ with Session(self.engine) as session: self._set_request_user_id( From 7d53da945dad47d492734b7a48921486dcdc26ef Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Wed, 15 Oct 2025 21:01:42 +0100 Subject: [PATCH 29/64] Update src/zenml/zen_stores/schemas/curated_visualization_schemas.py Co-authored-by: Stefan Nica --- src/zenml/zen_stores/schemas/curated_visualization_schemas.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 7798410d65..4a39192d34 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -110,7 +110,7 @@ class CuratedVisualizationSchema(BaseSchema, table=True): nullable=False, ) - artifact_version: Optional["ArtifactVersionSchema"] = Relationship( + artifact_version: "ArtifactVersionSchema" = Relationship( sa_relationship_kwargs={"lazy": "selectin"} ) resource: Optional[CuratedVisualizationResourceSchema] = Relationship( From e718e433cf0bb7df1e237db63c52220d2386b51c Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Wed, 15 Oct 2025 21:01:55 +0100 Subject: [PATCH 30/64] Update src/zenml/zen_stores/schemas/curated_visualization_schemas.py Co-authored-by: Stefan Nica --- src/zenml/zen_stores/schemas/curated_visualization_schemas.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 4a39192d34..da15e00065 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -246,8 +246,6 @@ def to_model( include_metadata=include_metadata, include_resources=include_resources, ) - if self.artifact_version - else None ) resource_model = None From 176a83835e852b08b44f3da67745c9390bcfa60e Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Wed, 15 Oct 2025 21:02:05 +0100 Subject: [PATCH 31/64] Update src/zenml/models/v2/core/curated_visualization.py Co-authored-by: Stefan Nica --- src/zenml/models/v2/core/curated_visualization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index e355310379..df4834825b 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -223,7 +223,7 @@ def layout_size(self) -> CuratedVisualizationSize: return self.get_body().layout_size @property - def artifact_version(self) -> Optional["ArtifactVersionResponse"]: + def artifact_version(self) -> "ArtifactVersionResponse": """The artifact version resource. Returns: From 6c7906b5707403398180e9d1ff2e3c790209b84e Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Wed, 15 Oct 2025 21:02:19 +0100 Subject: [PATCH 32/64] Update src/zenml/models/v2/core/curated_visualization.py Co-authored-by: Stefan Nica --- src/zenml/models/v2/core/curated_visualization.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index df4834825b..5afb9d8175 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -144,8 +144,7 @@ class CuratedVisualizationResponseMetadata(ProjectScopedResponseMetadata): class CuratedVisualizationResponseResources(ProjectScopedResponseResources): """Response resources for curated visualizations.""" - artifact_version: Optional["ArtifactVersionResponse"] = Field( - default=None, + artifact_version: "ArtifactVersionResponse" = Field( title="The artifact version.", description="Artifact version from which the visualization originates.", ) From b8ddcf92a2656808de7be315aa008f6f8d5a7b60 Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Wed, 15 Oct 2025 21:02:31 +0100 Subject: [PATCH 33/64] Update src/zenml/zen_stores/schemas/pipeline_run_schemas.py Co-authored-by: Stefan Nica --- src/zenml/zen_stores/schemas/pipeline_run_schemas.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index efb47b3fc7..a42ac78b28 100644 --- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py @@ -251,6 +251,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True): primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE_RUN.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineRunSchema.id)", secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", overlaps="visualizations", + cascade="all, delete-orphan", ), ) From efc8eff2aeccd746bb7ab17d62702b1f8a453244 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Wed, 15 Oct 2025 23:15:34 +0100 Subject: [PATCH 34/64] final round of review --- docs/book/how-to/artifacts/visualizations.md | 39 +++----- src/zenml/client.py | 14 +-- src/zenml/models/__init__.py | 4 - .../models/v2/core/curated_visualization.py | 51 ++++++---- .../models/v2/misc/curated_visualization.py | 33 ------- .../curated_visualization_endpoints.py | 63 ++++++------ src/zenml/zen_stores/schemas/__init__.py | 2 - .../schemas/curated_visualization_schemas.py | 99 ++++++------------- .../zen_stores/schemas/deployment_schemas.py | 8 +- src/zenml/zen_stores/schemas/model_schemas.py | 8 +- .../schemas/pipeline_run_schemas.py | 10 +- .../zen_stores/schemas/pipeline_schemas.py | 8 +- .../schemas/pipeline_snapshot_schemas.py | 8 +- .../zen_stores/schemas/project_schemas.py | 8 +- src/zenml/zen_stores/sql_zen_store.py | 43 +++----- .../functional/zen_stores/test_zen_store.py | 34 +++---- 16 files changed, 172 insertions(+), 260 deletions(-) delete mode 100644 src/zenml/models/v2/misc/curated_visualization.py diff --git a/docs/book/how-to/artifacts/visualizations.md b/docs/book/how-to/artifacts/visualizations.md index 26291cb037..39c68ae0c6 100644 --- a/docs/book/how-to/artifacts/visualizations.md +++ b/docs/book/how-to/artifacts/visualizations.md @@ -78,7 +78,7 @@ Curated visualizations currently support the following resources: - **Pipeline Runs** – detailed diagnostics for specific executions. - **Pipeline Snapshots** – configuration/version comparisons for snapshot history. -You can create a curated visualization programmatically by linking an artifact visualization to a single resource. The example below shows how to create separate visualizations for different resource types: +You can create a curated visualization programmatically by linking an artifact visualization to a single resource. Provide the resource identifier and resource type directly when creating the visualization. The example below shows how to create separate visualizations for different resource types: ```python from uuid import UUID @@ -88,7 +88,6 @@ from zenml.enums import ( CuratedVisualizationSize, VisualizationResourceTypes, ) -from zenml.models import CuratedVisualizationResource client = Client() artifact_version_id = UUID("") @@ -103,10 +102,8 @@ model = client.list_models().items[0] model_viz = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=0, - resource=CuratedVisualizationResource( - id=model.id, - type=VisualizationResourceTypes.MODEL - ), + resource_id=model.id, + resource_type=VisualizationResourceTypes.MODEL, display_name="Model performance dashboard", layout_size=CuratedVisualizationSize.FULL_WIDTH, ) @@ -115,10 +112,8 @@ model_viz = client.create_curated_visualization( deployment_viz = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=0, - resource=CuratedVisualizationResource( - id=deployment.id, - type=VisualizationResourceTypes.DEPLOYMENT - ), + resource_id=deployment.id, + resource_type=VisualizationResourceTypes.DEPLOYMENT, display_name="Deployment health dashboard", layout_size=CuratedVisualizationSize.HALF_WIDTH, ) @@ -127,10 +122,8 @@ deployment_viz = client.create_curated_visualization( project_viz = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=0, - resource=CuratedVisualizationResource( - id=project.id, - type=VisualizationResourceTypes.PROJECT - ), + resource_id=project.id, + resource_type=VisualizationResourceTypes.PROJECT, display_name="Project overview dashboard", layout_size=CuratedVisualizationSize.FULL_WIDTH, ) @@ -179,10 +172,8 @@ When setting display orders, consider leaving gaps between values (e.g., 10, 20, visualization_a = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=0, - resource=CuratedVisualizationResource( - id=model.id, - type=VisualizationResourceTypes.MODEL - ), + resource_id=model.id, + resource_type=VisualizationResourceTypes.MODEL, display_order=10, # Primary dashboard layout_size=CuratedVisualizationSize.FULL_WIDTH, ) @@ -190,10 +181,8 @@ visualization_a = client.create_curated_visualization( visualization_b = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=1, - resource=CuratedVisualizationResource( - id=model.id, - type=VisualizationResourceTypes.MODEL - ), + resource_id=model.id, + resource_type=VisualizationResourceTypes.MODEL, display_order=20, # Secondary metrics layout_size=CuratedVisualizationSize.HALF_WIDTH, # Compact chart beside the primary tile ) @@ -202,10 +191,8 @@ visualization_b = client.create_curated_visualization( visualization_c = client.create_curated_visualization( artifact_version_id=artifact_version_id, visualization_index=2, - resource=CuratedVisualizationResource( - id=model.id, - type=VisualizationResourceTypes.MODEL - ), + resource_id=model.id, + resource_type=VisualizationResourceTypes.MODEL, display_order=15, # Now appears between A and B layout_size=CuratedVisualizationSize.HALF_WIDTH, ) diff --git a/src/zenml/client.py b/src/zenml/client.py index 285efea31e..1f1a45d9d0 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -73,6 +73,7 @@ StackComponentType, StoreType, TaggableResourceTypes, + VisualizationResourceTypes, ) from zenml.exceptions import ( AuthorizationException, @@ -110,7 +111,6 @@ ComponentResponse, ComponentUpdate, CuratedVisualizationRequest, - CuratedVisualizationResource, CuratedVisualizationResponse, CuratedVisualizationUpdate, DeploymentFilter, @@ -3749,7 +3749,8 @@ def create_curated_visualization( artifact_version_id: UUID, visualization_index: int, *, - resource: CuratedVisualizationResource, + resource_id: UUID, + resource_type: VisualizationResourceTypes, project_id: Optional[UUID] = None, display_name: Optional[str] = None, display_order: Optional[int] = None, @@ -3775,10 +3776,8 @@ def create_curated_visualization( Args: artifact_version_id: The ID of the artifact version containing the visualization. visualization_index: The index of the visualization within the artifact version. - resource: The resource to associate with the visualization. - Should be a `CuratedVisualizationResource` containing - the resource ID and type (e.g., DEPLOYMENT, PIPELINE, PIPELINE_RUN, - PIPELINE_SNAPSHOT). + resource_id: The identifier of the resource tied to the visualization. + resource_type: The type of resource referenced by the visualization. project_id: The ID of the project to associate with the visualization. display_name: The display name of the visualization. display_order: The display order of the visualization. @@ -3794,7 +3793,8 @@ def create_curated_visualization( display_name=display_name, display_order=display_order, layout_size=layout_size, - resource=resource, + resource_id=resource_id, + resource_type=resource_type, ) return self.zen_store.create_curated_visualization(request) diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index 0140de6bc3..bd7c229a1e 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -430,9 +430,6 @@ RunMetadataEntry, RunMetadataResource, ) -from zenml.models.v2.misc.curated_visualization import ( - CuratedVisualizationResource, -) from zenml.models.v2.misc.server_models import ( ServerModel, ServerDatabaseType, @@ -888,7 +885,6 @@ "ResourcesInfo", "RunMetadataEntry", "RunMetadataResource", - "CuratedVisualizationResource", "ProjectStatistics", "PipelineRunDAG", "ExceptionInfo", diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 5afb9d8175..987a419e68 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -18,7 +18,7 @@ from pydantic import Field -from zenml.enums import CuratedVisualizationSize +from zenml.enums import CuratedVisualizationSize, VisualizationResourceTypes from zenml.models.v2.base.base import BaseUpdate from zenml.models.v2.base.scoped import ( ProjectScopedRequest, @@ -27,9 +27,6 @@ ProjectScopedResponseMetadata, ProjectScopedResponseResources, ) -from zenml.models.v2.misc.curated_visualization import ( - CuratedVisualizationResource, -) if TYPE_CHECKING: from zenml.models.v2.core.artifact_version import ArtifactVersionResponse @@ -80,13 +77,18 @@ class CuratedVisualizationRequest(ProjectScopedRequest): "on the dashboard." ), ) - resource: CuratedVisualizationResource = Field( - title="Resource associated with this visualization.", + resource_id: UUID = Field( + title="The linked resource ID.", description=( - "The single resource (deployment, model, pipeline, pipeline run, " - "pipeline snapshot, or project) that should surface this visualization." + "Identifier of the resource (deployment, model, pipeline, pipeline " + "run, pipeline snapshot, or project) that should surface this " + "visualization." ), ) + resource_type: VisualizationResourceTypes = Field( + title="The linked resource type.", + description="Type of the resource associated with this visualization.", + ) # ------------------ Update Model ------------------ @@ -135,6 +137,14 @@ class CuratedVisualizationResponseBody(ProjectScopedResponseBody): default=CuratedVisualizationSize.FULL_WIDTH, title="The layout size of the visualization.", ) + resource_id: UUID = Field( + title="The linked resource ID.", + description="Identifier of the resource associated with this visualization.", + ) + resource_type: VisualizationResourceTypes = Field( + title="The linked resource type.", + description="Type of the resource associated with this visualization.", + ) class CuratedVisualizationResponseMetadata(ProjectScopedResponseMetadata): @@ -142,17 +152,12 @@ class CuratedVisualizationResponseMetadata(ProjectScopedResponseMetadata): class CuratedVisualizationResponseResources(ProjectScopedResponseResources): - """Response resources for curated visualizations.""" + """Response resources included for curated visualizations.""" artifact_version: "ArtifactVersionResponse" = Field( title="The artifact version.", description="Artifact version from which the visualization originates.", ) - resource: Optional[CuratedVisualizationResource] = Field( - default=None, - title="The linked resource.", - description="Resource reference associated with this curated visualization.", - ) class CuratedVisualizationResponse( @@ -231,11 +236,19 @@ def artifact_version(self) -> "ArtifactVersionResponse": return self.get_resources().artifact_version @property - def resource(self) -> Optional[CuratedVisualizationResource]: - """The resource reference associated with this visualization. + def resource_id(self) -> UUID: + """The identifier of the linked resource. + + Returns: + The resource identifier associated with this visualization. + """ + return self.get_body().resource_id + + @property + def resource_type(self) -> VisualizationResourceTypes: + """The type of the linked resource. Returns: - The resource reference if included. + The resource type associated with this visualization. """ - resources = self.get_resources() - return resources.resource if resources else None + return self.get_body().resource_type diff --git a/src/zenml/models/v2/misc/curated_visualization.py b/src/zenml/models/v2/misc/curated_visualization.py deleted file mode 100644 index f09594fa4c..0000000000 --- a/src/zenml/models/v2/misc/curated_visualization.py +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) ZenML GmbH 2025. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at: -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express -# or implied. See the License for the specific language governing -# permissions and limitations under the License. -"""Models describing curated visualization resources.""" - -from uuid import UUID - -from pydantic import BaseModel, Field - -from zenml.enums import VisualizationResourceTypes - - -class CuratedVisualizationResource(BaseModel): - """Resource reference associated with a curated visualization.""" - - id: UUID = Field( - title="The ID of the resource.", - description="Identifier of the resource associated with the visualization.", - ) - type: VisualizationResourceTypes = Field( - title="The type of the resource.", - description="Classification of the resource associated with the visualization.", - ) diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py index be6dbf55e0..689607d607 100644 --- a/src/zenml/zen_server/routers/curated_visualization_endpoints.py +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -9,7 +9,6 @@ from zenml.enums import VisualizationResourceTypes from zenml.models import ( CuratedVisualizationRequest, - CuratedVisualizationResource, CuratedVisualizationResponse, CuratedVisualizationUpdate, ) @@ -27,34 +26,35 @@ def _get_resource_model( - resource: CuratedVisualizationResource, + resource_type: VisualizationResourceTypes, + resource_id: UUID, ) -> Any: - """Load the model associated with a curated visualization resource. + """Fetch the concrete resource model for a curated visualization. Args: - resource: The curated visualization resource to load the model for. + resource_type: The type of resource linked to the curated visualization. + resource_id: The unique identifier of the linked resource. Returns: - The model associated with the curated visualization resource. + The hydrated resource model retrieved from the Zen store. Raises: - RuntimeError: If the resource type is not supported. + RuntimeError: If the provided resource type is not supported. """ store = zen_store() - resource_type = resource.type if resource_type == VisualizationResourceTypes.DEPLOYMENT: - return store.get_deployment(resource.id) + return store.get_deployment(resource_id) if resource_type == VisualizationResourceTypes.MODEL: - return store.get_model(resource.id) + return store.get_model(resource_id) if resource_type == VisualizationResourceTypes.PIPELINE: - return store.get_pipeline(resource.id) + return store.get_pipeline(resource_id) if resource_type == VisualizationResourceTypes.PIPELINE_RUN: - return store.get_run(resource.id) + return store.get_run(resource_id) if resource_type == VisualizationResourceTypes.PIPELINE_SNAPSHOT: - return store.get_snapshot(resource.id) + return store.get_snapshot(resource_id) if resource_type == VisualizationResourceTypes.PROJECT: - return store.get_project(resource.id) + return store.get_project(resource_id) raise RuntimeError( f"Unsupported curated visualization resource type: {resource_type}" @@ -84,7 +84,9 @@ def create_curated_visualization( The created curated visualization. """ store = zen_store() - resource_model = _get_resource_model(visualization.resource) + resource_model = _get_resource_model( + visualization.resource_type, visualization.resource_id + ) artifact_version = store.get_artifact_version( visualization.artifact_version_id ) @@ -115,19 +117,20 @@ def get_curated_visualization( The curated visualization with the given ID. Raises: - RuntimeError: If the curated visualization is missing its resource reference. + RuntimeError: If the curated visualization is missing its resource identifier or type. """ store = zen_store() hydrated_visualization = store.get_curated_visualization( visualization_id, hydrate=True ) - resource = hydrated_visualization.resource - if resource is None: + resource_type = hydrated_visualization.resource_type + resource_id = hydrated_visualization.resource_id + if resource_type is None or resource_id is None: raise RuntimeError( - f"Curated visualization '{visualization_id}' is missing its resource reference." + f"Curated visualization '{visualization_id}' is missing its resource identifier or type." ) - resource_model = _get_resource_model(resource) + resource_model = _get_resource_model(resource_type, resource_id) verify_permission_for_model(resource_model, action=Action.READ) if hydrate: @@ -156,19 +159,20 @@ def update_curated_visualization( The updated curated visualization. Raises: - RuntimeError: If the curated visualization is missing its resource reference. + RuntimeError: If the curated visualization is missing its resource identifier or type. """ store = zen_store() existing_visualization = store.get_curated_visualization( visualization_id, hydrate=True ) - resource = existing_visualization.resource - if resource is None: + resource_type = existing_visualization.resource_type + resource_id = existing_visualization.resource_id + if resource_type is None or resource_id is None: raise RuntimeError( - f"Curated visualization '{visualization_id}' is missing its resource reference." + f"Curated visualization '{visualization_id}' is missing its resource identifier or type." ) - resource_model = _get_resource_model(resource) + resource_model = _get_resource_model(resource_type, resource_id) verify_permission_for_model(resource_model, action=Action.UPDATE) return store.update_curated_visualization( @@ -191,19 +195,20 @@ def delete_curated_visualization( visualization_id: The ID of the curated visualization to delete. Raises: - RuntimeError: If the curated visualization is missing its resource reference. + RuntimeError: If the curated visualization is missing its resource identifier or type. """ store = zen_store() existing_visualization = store.get_curated_visualization( visualization_id, hydrate=True ) - resource = existing_visualization.resource - if resource is None: + resource_type = existing_visualization.resource_type + resource_id = existing_visualization.resource_id + if resource_type is None or resource_id is None: raise RuntimeError( - f"Curated visualization '{visualization_id}' is missing its resource reference." + f"Curated visualization '{visualization_id}' is missing its resource identifier or type." ) - resource_model = _get_resource_model(resource) + resource_model = _get_resource_model(resource_type, resource_id) verify_permission_for_model(resource_model, action=Action.UPDATE) store.delete_curated_visualization(visualization_id) diff --git a/src/zenml/zen_stores/schemas/__init__.py b/src/zenml/zen_stores/schemas/__init__.py index 6235b6f64b..24a3c88c8f 100644 --- a/src/zenml/zen_stores/schemas/__init__.py +++ b/src/zenml/zen_stores/schemas/__init__.py @@ -32,7 +32,6 @@ from zenml.zen_stores.schemas.pipeline_build_schemas import PipelineBuildSchema from zenml.zen_stores.schemas.deployment_schemas import DeploymentSchema from zenml.zen_stores.schemas.curated_visualization_schemas import ( - CuratedVisualizationResourceSchema, CuratedVisualizationSchema, ) from zenml.zen_stores.schemas.component_schemas import StackComponentSchema @@ -93,7 +92,6 @@ "CodeRepositorySchema", "DeploymentSchema", "CuratedVisualizationSchema", - "CuratedVisualizationResourceSchema", "EventSourceSchema", "FlavorSchema", "LogsSchema", diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index da15e00065..0c537f0924 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -14,12 +14,12 @@ """SQLModel implementation of curated visualization tables.""" from typing import TYPE_CHECKING, Any, List, Optional, Sequence -from uuid import UUID, uuid4 +from uuid import UUID from sqlalchemy import UniqueConstraint from sqlalchemy.orm import selectinload from sqlalchemy.sql.base import ExecutableOption -from sqlmodel import Field, Relationship, SQLModel +from sqlmodel import Field, Relationship from zenml.enums import CuratedVisualizationSize, VisualizationResourceTypes from zenml.models.v2.core.curated_visualization import ( @@ -30,9 +30,6 @@ CuratedVisualizationResponseResources, CuratedVisualizationUpdate, ) -from zenml.models.v2.misc.curated_visualization import ( - CuratedVisualizationResource, -) from zenml.zen_stores.schemas.base_schemas import BaseSchema from zenml.zen_stores.schemas.project_schemas import ProjectSchema from zenml.zen_stores.schemas.schema_utils import ( @@ -45,35 +42,6 @@ from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema -class CuratedVisualizationResourceSchema(SQLModel, table=True): - """Link table mapping curated visualizations to resources.""" - - __tablename__ = "curated_visualization_resource" - __table_args__ = ( - UniqueConstraint( - "visualization_id", - name="unique_curated_visualization_resource", - ), - build_index(__tablename__, ["resource_id", "resource_type"]), - ) - - id: UUID = Field(default_factory=uuid4, primary_key=True) - visualization_id: UUID = build_foreign_key_field( - source=__tablename__, - target="curated_visualization", - source_column="visualization_id", - target_column="id", - ondelete="CASCADE", - nullable=False, - ) - resource_id: UUID = Field(nullable=False) - resource_type: str = Field(nullable=False) - - visualization: "CuratedVisualizationSchema" = Relationship( - back_populates="resource", - ) - - class CuratedVisualizationSchema(BaseSchema, table=True): """SQL Model for curated visualizations.""" @@ -83,6 +51,13 @@ class CuratedVisualizationSchema(BaseSchema, table=True): __tablename__, ["artifact_version_id", "visualization_index"] ), build_index(__tablename__, ["display_order"]), + UniqueConstraint( + "artifact_version_id", + "visualization_index", + "resource_id", + "resource_type", + name="unique_curated_visualization_resource_link", + ), ) project_id: UUID = build_foreign_key_field( @@ -109,18 +84,12 @@ class CuratedVisualizationSchema(BaseSchema, table=True): default=CuratedVisualizationSize.FULL_WIDTH.value, nullable=False, ) + resource_id: UUID = Field(nullable=False) + resource_type: str = Field(nullable=False) artifact_version: "ArtifactVersionSchema" = Relationship( sa_relationship_kwargs={"lazy": "selectin"} ) - resource: Optional[CuratedVisualizationResourceSchema] = Relationship( - back_populates="visualization", - sa_relationship_kwargs={ - "lazy": "selectin", - "cascade": "all, delete", - "uselist": False, - }, - ) @classmethod def get_query_options( @@ -144,12 +113,7 @@ def get_query_options( options: List[ExecutableOption] = [] if include_resources: - options.extend( - [ - selectinload(jl_arg(cls.artifact_version)), - selectinload(jl_arg(cls.resource)), - ] - ) + options.append(selectinload(jl_arg(cls.artifact_version))) return options @@ -172,6 +136,8 @@ def from_request( display_name=request.display_name, display_order=request.display_order, layout_size=request.layout_size.value, + resource_id=request.resource_id, + resource_type=request.resource_type.value, ) def update( @@ -224,6 +190,17 @@ def to_model( except ValueError: layout_size_enum = CuratedVisualizationSize.FULL_WIDTH + resource_type_str = self.resource_type + if resource_type_str: + try: + resource_type_enum = VisualizationResourceTypes( + resource_type_str + ) + except ValueError: + resource_type_enum = VisualizationResourceTypes.PROJECT + else: + resource_type_enum = VisualizationResourceTypes.PROJECT + body = CuratedVisualizationResponseBody( project_id=self.project_id, created=self.created, @@ -233,6 +210,8 @@ def to_model( display_name=self.display_name, display_order=self.display_order, layout_size=layout_size_enum, + resource_id=self.resource_id, + resource_type=resource_type_enum, ) metadata = None @@ -241,31 +220,13 @@ def to_model( response_resources = None if include_resources: - artifact_version_model = ( - self.artifact_version.to_model( - include_metadata=include_metadata, - include_resources=include_resources, - ) + artifact_version_model = self.artifact_version.to_model( + include_metadata=include_metadata, + include_resources=include_resources, ) - resource_model = None - if self.resource: - try: - resource_type_enum = VisualizationResourceTypes( - self.resource.resource_type - ) - except ValueError: - resource_type_enum = None - - if resource_type_enum is not None: - resource_model = CuratedVisualizationResource( - id=self.resource.resource_id, - type=resource_type_enum, - ) - response_resources = CuratedVisualizationResponseResources( artifact_version=artifact_version_model, - resource=resource_model, ) return CuratedVisualizationResponse( diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index ad62288219..4dfb2461ff 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -142,9 +142,11 @@ class DeploymentSchema(NamedSchema, table=True): visualizations: List["CuratedVisualizationSchema"] = Relationship( sa_relationship_kwargs=dict( - secondary="curated_visualization_resource", - primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.DEPLOYMENT.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==DeploymentSchema.id)", - secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + primaryjoin=( + "and_(CuratedVisualizationSchema.resource_type" + f"=='{VisualizationResourceTypes.DEPLOYMENT.value}', " + "foreign(CuratedVisualizationSchema.resource_id)==DeploymentSchema.id)" + ), overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 9c8e53a31f..66577b7b06 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -134,9 +134,11 @@ class ModelSchema(NamedSchema, table=True): ) visualizations: List["CuratedVisualizationSchema"] = Relationship( sa_relationship_kwargs=dict( - secondary="curated_visualization_resource", - primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.MODEL.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==ModelSchema.id)", - secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + primaryjoin=( + "and_(CuratedVisualizationSchema.resource_type" + f"=='{VisualizationResourceTypes.MODEL.value}', " + "foreign(CuratedVisualizationSchema.resource_id)==ModelSchema.id)" + ), overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index a42ac78b28..4740c1bb1c 100644 --- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py @@ -247,11 +247,13 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True): ) visualizations: List["CuratedVisualizationSchema"] = Relationship( sa_relationship_kwargs=dict( - secondary="curated_visualization_resource", - primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE_RUN.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineRunSchema.id)", - secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + primaryjoin=( + "and_(CuratedVisualizationSchema.resource_type" + f"=='{VisualizationResourceTypes.PIPELINE_RUN.value}', " + "foreign(CuratedVisualizationSchema.resource_id)==PipelineRunSchema.id)" + ), overlaps="visualizations", - cascade="all, delete-orphan", + cascade="all, delete-orphan", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index abcac6ff29..27f5d7b6b4 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_schemas.py @@ -109,9 +109,11 @@ class PipelineSchema(NamedSchema, table=True): ) visualizations: List["CuratedVisualizationSchema"] = Relationship( sa_relationship_kwargs=dict( - secondary="curated_visualization_resource", - primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineSchema.id)", - secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + primaryjoin=( + "and_(CuratedVisualizationSchema.resource_type" + f"=='{VisualizationResourceTypes.PIPELINE.value}', " + "foreign(CuratedVisualizationSchema.resource_id)==PipelineSchema.id)" + ), overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py index 775b00f6e4..dea0eb08b7 100644 --- a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py @@ -223,9 +223,11 @@ class PipelineSnapshotSchema(BaseSchema, table=True): ) visualizations: List["CuratedVisualizationSchema"] = Relationship( sa_relationship_kwargs=dict( - secondary="curated_visualization_resource", - primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PIPELINE_SNAPSHOT.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==PipelineSnapshotSchema.id)", - secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + primaryjoin=( + "and_(CuratedVisualizationSchema.resource_type" + f"=='{VisualizationResourceTypes.PIPELINE_SNAPSHOT.value}', " + "foreign(CuratedVisualizationSchema.resource_id)==PipelineSnapshotSchema.id)" + ), overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/schemas/project_schemas.py b/src/zenml/zen_stores/schemas/project_schemas.py index 11da98391c..84284cc2fb 100644 --- a/src/zenml/zen_stores/schemas/project_schemas.py +++ b/src/zenml/zen_stores/schemas/project_schemas.py @@ -131,9 +131,11 @@ class ProjectSchema(NamedSchema, table=True): ) visualizations: List["CuratedVisualizationSchema"] = Relationship( sa_relationship_kwargs=dict( - secondary="curated_visualization_resource", - primaryjoin=f"and_(foreign(CuratedVisualizationResourceSchema.resource_type)=='{VisualizationResourceTypes.PROJECT.value}', foreign(CuratedVisualizationResourceSchema.resource_id)==ProjectSchema.id)", - secondaryjoin="CuratedVisualizationSchema.id == CuratedVisualizationResourceSchema.visualization_id", + primaryjoin=( + "and_(CuratedVisualizationSchema.resource_type" + f"=='{VisualizationResourceTypes.PROJECT.value}', " + "foreign(CuratedVisualizationSchema.resource_id)==ProjectSchema.id)" + ), overlaps="visualizations", ), ) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index eda50dd62a..a007f79c7a 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -369,7 +369,6 @@ BaseSchema, CodeReferenceSchema, CodeRepositorySchema, - CuratedVisualizationResourceSchema, CuratedVisualizationSchema, DeploymentSchema, EventSourceSchema, @@ -5461,7 +5460,6 @@ def _assert_curated_visualization_duplicate( """ existing = session.exec( select(CuratedVisualizationSchema) - .join(CuratedVisualizationResourceSchema) .where( CuratedVisualizationSchema.artifact_version_id == artifact_version_id @@ -5470,12 +5468,9 @@ def _assert_curated_visualization_duplicate( CuratedVisualizationSchema.visualization_index == visualization_index ) + .where(CuratedVisualizationSchema.resource_id == resource_id) .where( - CuratedVisualizationResourceSchema.resource_id == resource_id - ) - .where( - CuratedVisualizationResourceSchema.resource_type - == resource_type.value + CuratedVisualizationSchema.resource_type == resource_type.value ) ).first() if existing is not None: @@ -5527,10 +5522,6 @@ def create_curated_visualization( visualization_index=visualization.visualization_index, ) - # Validate and fetch the resource - resource_request = visualization.resource - - # Map resource types to their corresponding schema classes resource_schema_map: Dict[ VisualizationResourceTypes, Type[BaseSchema] ] = { @@ -5542,55 +5533,45 @@ def create_curated_visualization( VisualizationResourceTypes.PROJECT: ProjectSchema, } - if resource_request.type not in resource_schema_map: + if visualization.resource_type not in resource_schema_map: raise IllegalOperationError( - f"Invalid resource type: {resource_request.type}" + f"Invalid resource type: {visualization.resource_type}" ) - # Fetch the single resource schema - schema_class = resource_schema_map[resource_request.type] + schema_class = resource_schema_map[visualization.resource_type] resource_schema = session.exec( select(schema_class).where( - schema_class.id == resource_request.id + schema_class.id == visualization.resource_id ) ).first() if not resource_schema: raise KeyError( - f"Resource of type '{resource_request.type.value}' " - f"with ID {resource_request.id} not found." + f"Resource of type '{visualization.resource_type.value}' " + f"with ID {visualization.resource_id} not found." ) - # Validate project scope for the resource if hasattr(resource_schema, "project_id"): resource_project_id = resource_schema.project_id if resource_project_id and resource_project_id != project_id: raise IllegalOperationError( - f"Resource {resource_request.type.value} with ID " - f"{resource_request.id} belongs to a different project than " + f"Resource {visualization.resource_type.value} with ID " + f"{visualization.resource_id} belongs to a different project than " f"the curated visualization (project ID: {project_id})." ) - # Check for duplicate self._assert_curated_visualization_duplicate( session=session, artifact_version_id=visualization.artifact_version_id, visualization_index=visualization.visualization_index, - resource_id=resource_request.id, - resource_type=resource_request.type, - ) - - # Create the resource link - resource = CuratedVisualizationResourceSchema( - resource_id=resource_request.id, - resource_type=resource_request.type.value, + resource_id=visualization.resource_id, + resource_type=visualization.resource_type, ) schema: CuratedVisualizationSchema = ( CuratedVisualizationSchema.from_request(visualization) ) schema.project_id = project_id - schema.resource = resource session.add(schema) session.commit() diff --git a/tests/integration/functional/zen_stores/test_zen_store.py b/tests/integration/functional/zen_stores/test_zen_store.py index a6512b8477..c47ef6e51c 100644 --- a/tests/integration/functional/zen_stores/test_zen_store.py +++ b/tests/integration/functional/zen_stores/test_zen_store.py @@ -97,7 +97,6 @@ ComponentFilter, ComponentUpdate, CuratedVisualizationRequest, - CuratedVisualizationResource, CuratedVisualizationUpdate, DeploymentRequest, ModelVersionArtifactFilter, @@ -5874,10 +5873,8 @@ def test_curated_visualizations_across_resources(self): project=project_id, artifact_version_id=artifact_version.id, visualization_index=idx, - resource=CuratedVisualizationResource( - id=resource_id, - type=resource_type, - ), + resource_id=resource_id, + resource_type=resource_type, display_name=f"{resource_type.value} visualization", ) ) @@ -5885,9 +5882,8 @@ def test_curated_visualizations_across_resources(self): visualization_id=viz.id, hydrate=True, ) - assert hydrated.resource is not None - assert hydrated.resource.type == resource_type - assert hydrated.resource.id == resource_id + assert hydrated.resource_id == resource_id + assert hydrated.resource_type == resource_type assert ( hydrated.layout_size == CuratedVisualizationSize.FULL_WIDTH ) @@ -5902,8 +5898,8 @@ def test_curated_visualizations_across_resources(self): == f"{VisualizationResourceTypes.PIPELINE.value} visualization" ) assert loaded.layout_size == CuratedVisualizationSize.FULL_WIDTH - assert loaded.resource is not None - assert loaded.resource.type == VisualizationResourceTypes.PIPELINE + assert loaded.resource_id == pipeline_model.id + assert loaded.resource_type == VisualizationResourceTypes.PIPELINE # Test duplicate creation - same artifact_version + visualization_index + resource should fail with pytest.raises(EntityExistsError): @@ -5912,10 +5908,8 @@ def test_curated_visualizations_across_resources(self): project=project_id, artifact_version_id=artifact_version.id, visualization_index=0, # Same index as pipeline visualization - resource=CuratedVisualizationResource( - id=pipeline_model.id, - type=VisualizationResourceTypes.PIPELINE, - ), + resource_id=pipeline_model.id, + resource_type=VisualizationResourceTypes.PIPELINE, ) ) @@ -5982,22 +5976,20 @@ def test_curated_visualizations_project_only(self): visualization = client.create_curated_visualization( artifact_version_id=artifact_version.id, visualization_index=0, - resource=CuratedVisualizationResource( - id=project.id, - type=VisualizationResourceTypes.PROJECT, - ), + resource_id=project.id, + resource_type=VisualizationResourceTypes.PROJECT, + project_id=project.id, display_name="Project visualization", ) hydrated_visualization = client.zen_store.get_curated_visualization( visualization.id, hydrate=True ) - assert hydrated_visualization.resource is not None + assert hydrated_visualization.resource_id == project.id assert ( - hydrated_visualization.resource.type + hydrated_visualization.resource_type == VisualizationResourceTypes.PROJECT ) - assert hydrated_visualization.resource.id == project.id assert hydrated_visualization.display_name == "Project visualization" client.delete_curated_visualization(visualization.id) From b4e7b4400db816024b28fc83030b512a67a6e238 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Wed, 15 Oct 2025 23:46:59 +0100 Subject: [PATCH 35/64] docstring --- src/zenml/zen_stores/schemas/curated_visualization_schemas.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 0c537f0924..77c2903a88 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -87,9 +87,7 @@ class CuratedVisualizationSchema(BaseSchema, table=True): resource_id: UUID = Field(nullable=False) resource_type: str = Field(nullable=False) - artifact_version: "ArtifactVersionSchema" = Relationship( - sa_relationship_kwargs={"lazy": "selectin"} - ) + artifact_version: "ArtifactVersionSchema" = Relationship() @classmethod def get_query_options( From 24f920e00525d32480e135e8ae1806a9db17cc9d Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Wed, 15 Oct 2025 23:49:36 +0100 Subject: [PATCH 36/64] update migration --- ...a_add_vizualisations_and_link_them_to_.py} | 57 +++++-------------- 1 file changed, 15 insertions(+), 42 deletions(-) rename src/zenml/zen_stores/migrations/versions/{a3c4e7cac1d5_add_vizualisations_and_link_them_to_.py => 0ede689b266a_add_vizualisations_and_link_them_to_.py} (68%) diff --git a/src/zenml/zen_stores/migrations/versions/a3c4e7cac1d5_add_vizualisations_and_link_them_to_.py b/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations_and_link_them_to_.py similarity index 68% rename from src/zenml/zen_stores/migrations/versions/a3c4e7cac1d5_add_vizualisations_and_link_them_to_.py rename to src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations_and_link_them_to_.py index 44ae61de56..18abce20c7 100644 --- a/src/zenml/zen_stores/migrations/versions/a3c4e7cac1d5_add_vizualisations_and_link_them_to_.py +++ b/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations_and_link_them_to_.py @@ -1,8 +1,8 @@ -"""add vizualisations and link them to resources [a3c4e7cac1d5]. +"""add vizualisations and link them to resources [0ede689b266a]. -Revision ID: a3c4e7cac1d5 +Revision ID: 0ede689b266a Revises: 0.90.0 -Create Date: 2025-10-15 11:51:33.654401 +Create Date: 2025-10-15 23:48:54.468181 """ @@ -11,7 +11,7 @@ from alembic import op # revision identifiers, used by Alembic. -revision = "a3c4e7cac1d5" +revision = "0ede689b266a" down_revision = "0.90.0" branch_labels = None depends_on = None @@ -37,6 +37,10 @@ def upgrade() -> None: sa.Column( "layout_size", sqlmodel.sql.sqltypes.AutoString(), nullable=False ), + sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "resource_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False + ), sa.ForeignKeyConstraint( ["artifact_version_id"], ["artifact_version.id"], @@ -50,6 +54,13 @@ def upgrade() -> None: ondelete="CASCADE", ), sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "artifact_version_id", + "visualization_index", + "resource_id", + "resource_type", + name="unique_curated_visualization_resource_link", + ), ) with op.batch_alter_table( "curated_visualization", schema=None @@ -65,50 +76,12 @@ def upgrade() -> None: unique=False, ) - op.create_table( - "curated_visualization_resource", - sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column( - "visualization_id", sqlmodel.sql.sqltypes.GUID(), nullable=False - ), - sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column( - "resource_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False - ), - sa.ForeignKeyConstraint( - ["visualization_id"], - ["curated_visualization.id"], - name="fk_curated_visualization_resource_visualization_id_curated_visualization", - ondelete="CASCADE", - ), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "visualization_id", name="unique_curated_visualization_resource" - ), - ) - with op.batch_alter_table( - "curated_visualization_resource", schema=None - ) as batch_op: - batch_op.create_index( - "ix_curated_visualization_resource_resource_id_resource_type", - ["resource_id", "resource_type"], - unique=False, - ) - # ### 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( - "curated_visualization_resource", schema=None - ) as batch_op: - batch_op.drop_index( - "ix_curated_visualization_resource_resource_id_resource_type" - ) - - op.drop_table("curated_visualization_resource") with op.batch_alter_table( "curated_visualization", schema=None ) as batch_op: From 18fee070e52d9a55377f46c2a87294ce16166e98 Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Thu, 16 Oct 2025 13:34:16 +0100 Subject: [PATCH 37/64] Update src/zenml/zen_server/routers/curated_visualization_endpoints.py Co-authored-by: Stefan Nica --- .../routers/curated_visualization_endpoints.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py index 689607d607..c8d8ce50ee 100644 --- a/src/zenml/zen_server/routers/curated_visualization_endpoints.py +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -121,22 +121,15 @@ def get_curated_visualization( """ store = zen_store() hydrated_visualization = store.get_curated_visualization( - visualization_id, hydrate=True + visualization_id, hydrate=hydrate ) resource_type = hydrated_visualization.resource_type resource_id = hydrated_visualization.resource_id - if resource_type is None or resource_id is None: - raise RuntimeError( - f"Curated visualization '{visualization_id}' is missing its resource identifier or type." - ) resource_model = _get_resource_model(resource_type, resource_id) verify_permission_for_model(resource_model, action=Action.READ) - if hydrate: - return hydrated_visualization - - return store.get_curated_visualization(visualization_id, hydrate=False) + return hydrated_visualization @router.patch( From ba3d5fbbb729197b1db16569e64c9e6e2c33ea1e Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Thu, 16 Oct 2025 13:34:28 +0100 Subject: [PATCH 38/64] Update src/zenml/zen_server/routers/curated_visualization_endpoints.py Co-authored-by: Stefan Nica --- .../zen_server/routers/curated_visualization_endpoints.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py index c8d8ce50ee..159b902e8c 100644 --- a/src/zenml/zen_server/routers/curated_visualization_endpoints.py +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -160,10 +160,6 @@ def update_curated_visualization( ) resource_type = existing_visualization.resource_type resource_id = existing_visualization.resource_id - if resource_type is None or resource_id is None: - raise RuntimeError( - f"Curated visualization '{visualization_id}' is missing its resource identifier or type." - ) resource_model = _get_resource_model(resource_type, resource_id) verify_permission_for_model(resource_model, action=Action.UPDATE) From 960f354c9889ca0938da6b65e90aad7de54a21ac Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Thu, 16 Oct 2025 13:34:37 +0100 Subject: [PATCH 39/64] Update src/zenml/zen_server/routers/curated_visualization_endpoints.py Co-authored-by: Stefan Nica --- .../zen_server/routers/curated_visualization_endpoints.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py index 159b902e8c..8e7abcedbc 100644 --- a/src/zenml/zen_server/routers/curated_visualization_endpoints.py +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -192,10 +192,6 @@ def delete_curated_visualization( ) resource_type = existing_visualization.resource_type resource_id = existing_visualization.resource_id - if resource_type is None or resource_id is None: - raise RuntimeError( - f"Curated visualization '{visualization_id}' is missing its resource identifier or type." - ) resource_model = _get_resource_model(resource_type, resource_id) verify_permission_for_model(resource_model, action=Action.UPDATE) From c58b0b2e7e0a622de4030f6d6566b7fa6b878432 Mon Sep 17 00:00:00 2001 From: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Date: Fri, 17 Oct 2025 10:57:38 +0100 Subject: [PATCH 40/64] Update migration file description for visualizations --- ..._and_link_them_to_.py => 0ede689b266a_add_vizualisations.py} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename src/zenml/zen_stores/migrations/versions/{0ede689b266a_add_vizualisations_and_link_them_to_.py => 0ede689b266a_add_vizualisations.py} (97%) diff --git a/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations_and_link_them_to_.py b/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations.py similarity index 97% rename from src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations_and_link_them_to_.py rename to src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations.py index 18abce20c7..fef954811b 100644 --- a/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations_and_link_them_to_.py +++ b/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations.py @@ -1,4 +1,4 @@ -"""add vizualisations and link them to resources [0ede689b266a]. +"""add vizualisations [0ede689b266a]. Revision ID: 0ede689b266a Revises: 0.90.0 From cfd35b0684454bcee3e8ca1d7ba7d8a56ca95e3b Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Fri, 17 Oct 2025 11:04:28 +0100 Subject: [PATCH 41/64] renaming of visualisation --- ...add_vizualisations.py => 0ede689b266a_add_visualisations.py} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename src/zenml/zen_stores/migrations/versions/{0ede689b266a_add_vizualisations.py => 0ede689b266a_add_visualisations.py} (98%) diff --git a/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations.py b/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_visualisations.py similarity index 98% rename from src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations.py rename to src/zenml/zen_stores/migrations/versions/0ede689b266a_add_visualisations.py index fef954811b..c152d8f990 100644 --- a/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_vizualisations.py +++ b/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_visualisations.py @@ -1,4 +1,4 @@ -"""add vizualisations [0ede689b266a]. +"""add visualisations [0ede689b266a]. Revision ID: 0ede689b266a Revises: 0.90.0 From a41ffa1e7614e948614cfad0fadcf2087b033dba Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Fri, 17 Oct 2025 11:46:21 +0100 Subject: [PATCH 42/64] migration --- ...sations.py => 553964d69699_add_visualisations.py} | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) rename src/zenml/zen_stores/migrations/versions/{0ede689b266a_add_visualisations.py => 553964d69699_add_visualisations.py} (94%) diff --git a/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_visualisations.py b/src/zenml/zen_stores/migrations/versions/553964d69699_add_visualisations.py similarity index 94% rename from src/zenml/zen_stores/migrations/versions/0ede689b266a_add_visualisations.py rename to src/zenml/zen_stores/migrations/versions/553964d69699_add_visualisations.py index c152d8f990..157796af76 100644 --- a/src/zenml/zen_stores/migrations/versions/0ede689b266a_add_visualisations.py +++ b/src/zenml/zen_stores/migrations/versions/553964d69699_add_visualisations.py @@ -1,8 +1,8 @@ -"""add visualisations [0ede689b266a]. +"""add visualisations [553964d69699]. -Revision ID: 0ede689b266a -Revises: 0.90.0 -Create Date: 2025-10-15 23:48:54.468181 +Revision ID: 553964d69699 +Revises: 7497d2ff5731 +Create Date: 2025-10-17 11:44:15.358521 """ @@ -11,8 +11,8 @@ from alembic import op # revision identifiers, used by Alembic. -revision = "0ede689b266a" -down_revision = "0.90.0" +revision = "553964d69699" +down_revision = "7497d2ff5731" branch_labels = None depends_on = None From 4a958d65d5981d965f027553e70707b898d0b2d9 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Fri, 17 Oct 2025 14:46:52 +0100 Subject: [PATCH 43/64] add model rebuild --- src/zenml/models/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index bd7c229a1e..75f8b455c3 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -486,6 +486,10 @@ DeploymentResponseBody.model_rebuild() DeploymentResponseMetadata.model_rebuild() DeploymentResponseResources.model_rebuild() +CuratedVisualizationResponseBody.model_rebuild() +CuratedVisualizationResponseMetadata.model_rebuild() +CuratedVisualizationResponseResources.model_rebuild() +CuratedVisualizationResponse.model_rebuild() EventSourceResponseBody.model_rebuild() EventSourceResponseMetadata.model_rebuild() EventSourceResponseResources.model_rebuild() From edde4b67ae1e24f0df7d315dd3d6b559984b8c40 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sun, 19 Oct 2025 01:53:25 +0100 Subject: [PATCH 44/64] update tests --- .../functional/zen_stores/test_zen_store.py | 80 +++++++++++++++---- 1 file changed, 66 insertions(+), 14 deletions(-) diff --git a/tests/integration/functional/zen_stores/test_zen_store.py b/tests/integration/functional/zen_stores/test_zen_store.py index c47ef6e51c..3664e5a9e7 100644 --- a/tests/integration/functional/zen_stores/test_zen_store.py +++ b/tests/integration/functional/zen_stores/test_zen_store.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express # or implied. See the License for the specific language governing # permissions and limitations under the License. +import atexit import json import os import random @@ -26,7 +27,7 @@ import pytest from pydantic import ValidationError -from sqlalchemy.exc import IntegrityError +from sqlalchemy.exc import IntegrityError, OperationalError, ProgrammingError from tests.integration.functional.utils import sample_name from tests.integration.functional.zen_stores.utils import ( @@ -142,6 +143,33 @@ from zenml.zen_stores.rest_zen_store import RestZenStore from zenml.zen_stores.sql_zen_store import SqlZenStore +_ORIGINAL_INITIALIZE_DATABASE = SqlZenStore._initialize_database + + +def _patched_initialize_database(self): + try: + _ORIGINAL_INITIALIZE_DATABASE(self) + except (OperationalError, ProgrammingError) as error: + message = str(error).lower() + if ( + "no such column: stack.environment" in message + or "unknown column 'stack.environment'" in message + ): + self.migrate_database() + _ORIGINAL_INITIALIZE_DATABASE(self) + else: + raise + + +SqlZenStore._initialize_database = _patched_initialize_database + + +def _restore_sql_zen_store_initialize_database() -> None: + SqlZenStore._initialize_database = _ORIGINAL_INITIALIZE_DATABASE + + +atexit.register(_restore_sql_zen_store_initialize_database) + DEFAULT_NAME = "default" # .--------------. @@ -5749,6 +5777,32 @@ def test_curated_visualizations_across_resources(self): has_custom_name=True, ) ) + resource_configs = [ + { + "resource_type": VisualizationResourceTypes.PIPELINE, + "resource_id": None, + }, + { + "resource_type": VisualizationResourceTypes.MODEL, + "resource_id": None, + }, + { + "resource_type": VisualizationResourceTypes.PIPELINE_RUN, + "resource_id": None, + }, + { + "resource_type": VisualizationResourceTypes.PIPELINE_SNAPSHOT, + "resource_id": None, + }, + { + "resource_type": VisualizationResourceTypes.DEPLOYMENT, + "resource_id": None, + }, + { + "resource_type": VisualizationResourceTypes.PROJECT, + "resource_id": None, + }, + ] artifact_version = client.zen_store.create_artifact_version( ArtifactVersionRequest( artifact_id=artifact.id, @@ -5764,8 +5818,9 @@ def test_curated_visualizations_across_resources(self): visualizations=[ ArtifactVisualizationRequest( type=VisualizationType.HTML, - uri="s3://visualizations/example.html", + uri=f"s3://visualizations/{config['resource_type'].value}_{index}.html", ) + for index, config in enumerate(resource_configs) ], ) ) @@ -5854,20 +5909,17 @@ def test_curated_visualizations_across_resources(self): ) try: - # Create a separate visualization for each resource type - resource_configs = [ - (VisualizationResourceTypes.PIPELINE, pipeline_model.id), - (VisualizationResourceTypes.MODEL, model.id), - (VisualizationResourceTypes.PIPELINE_RUN, pipeline_run.id), - (VisualizationResourceTypes.PIPELINE_SNAPSHOT, snapshot.id), - (VisualizationResourceTypes.DEPLOYMENT, deployment.id), - (VisualizationResourceTypes.PROJECT, project_id), - ] + resource_configs[0]["resource_id"] = pipeline_model.id + resource_configs[1]["resource_id"] = model.id + resource_configs[2]["resource_id"] = pipeline_run.id + resource_configs[3]["resource_id"] = snapshot.id + resource_configs[4]["resource_id"] = deployment.id + resource_configs[5]["resource_id"] = project_id visualizations = {} - for idx, (resource_type, resource_id) in enumerate( - resource_configs - ): + for idx, config in enumerate(resource_configs): + resource_type = config["resource_type"] + resource_id = config["resource_id"] viz = client.zen_store.create_curated_visualization( CuratedVisualizationRequest( project=project_id, From a7b45b5d12d97748af59c74fc71b940be9d6085d Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Mon, 20 Oct 2025 12:32:48 +0200 Subject: [PATCH 45/64] Fix schema relationships, docstrings and unit tests. --- .../curated_visualization_endpoints.py | 9 - .../zen_stores/schemas/deployment_schemas.py | 1 + src/zenml/zen_stores/schemas/model_schemas.py | 1 + .../zen_stores/schemas/pipeline_schemas.py | 1 + .../schemas/pipeline_snapshot_schemas.py | 1 + .../zen_stores/schemas/project_schemas.py | 1 + src/zenml/zen_stores/sql_zen_store.py | 16 +- .../functional/zen_stores/test_zen_store.py | 222 ++++++++++++------ 8 files changed, 173 insertions(+), 79 deletions(-) diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py index 8e7abcedbc..d0cb161304 100644 --- a/src/zenml/zen_server/routers/curated_visualization_endpoints.py +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -115,9 +115,6 @@ def get_curated_visualization( Returns: The curated visualization with the given ID. - - Raises: - RuntimeError: If the curated visualization is missing its resource identifier or type. """ store = zen_store() hydrated_visualization = store.get_curated_visualization( @@ -150,9 +147,6 @@ def update_curated_visualization( Returns: The updated curated visualization. - - Raises: - RuntimeError: If the curated visualization is missing its resource identifier or type. """ store = zen_store() existing_visualization = store.get_curated_visualization( @@ -182,9 +176,6 @@ def delete_curated_visualization( Args: visualization_id: The ID of the curated visualization to delete. - - Raises: - RuntimeError: If the curated visualization is missing its resource identifier or type. """ store = zen_store() existing_visualization = store.get_curated_visualization( diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index 4dfb2461ff..16b4295e65 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -148,6 +148,7 @@ class DeploymentSchema(NamedSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==DeploymentSchema.id)" ), overlaps="visualizations", + cascade="all, delete-orphan", ), ) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 66577b7b06..330f212b2d 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -140,6 +140,7 @@ class ModelSchema(NamedSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==ModelSchema.id)" ), overlaps="visualizations", + cascade="all, delete-orphan", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index 27f5d7b6b4..ad4e0b8d79 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_schemas.py @@ -115,6 +115,7 @@ class PipelineSchema(NamedSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==PipelineSchema.id)" ), overlaps="visualizations", + cascade="all, delete-orphan", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py index dea0eb08b7..0a028edec3 100644 --- a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py @@ -229,6 +229,7 @@ class PipelineSnapshotSchema(BaseSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==PipelineSnapshotSchema.id)" ), overlaps="visualizations", + cascade="all, delete-orphan", ), ) diff --git a/src/zenml/zen_stores/schemas/project_schemas.py b/src/zenml/zen_stores/schemas/project_schemas.py index 84284cc2fb..e4794277f0 100644 --- a/src/zenml/zen_stores/schemas/project_schemas.py +++ b/src/zenml/zen_stores/schemas/project_schemas.py @@ -137,6 +137,7 @@ class ProjectSchema(NamedSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==ProjectSchema.id)" ), overlaps="visualizations", + cascade="all, delete-orphan", ), ) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index a007f79c7a..b01320ad19 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -69,7 +69,7 @@ field_validator, model_validator, ) -from sqlalchemy import QueuePool, func, update +from sqlalchemy import QueuePool, event, func, update from sqlalchemy.engine import URL, Engine, make_url from sqlalchemy.exc import ( ArgumentError, @@ -1262,6 +1262,14 @@ def _initialize(self) -> None: engine_args=engine_args, ) + if self.config.driver == SQLDatabaseDriver.SQLITE: + # Enable foreign key checks at the SQLite database level + @event.listens_for(self._engine, "connect") + def _(dbapi_connection: Any, connection_record: Any) -> None: + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + # SQLite: As long as the parent directory exists, SQLAlchemy will # automatically create the database. if ( @@ -5609,13 +5617,13 @@ def get_curated_visualization( def update_curated_visualization( self, visualization_id: UUID, - update: CuratedVisualizationUpdate, + visualization_update: CuratedVisualizationUpdate, ) -> CuratedVisualizationResponse: """Update mutable fields on a curated visualization. Args: visualization_id: The ID of the curated visualization to update. - update: The update to apply to the curated visualization. + visualization_update: The update to apply to the curated visualization. Returns: The updated curated visualization. @@ -5626,7 +5634,7 @@ def update_curated_visualization( schema_class=CuratedVisualizationSchema, session=session, ) - schema.update(update) + schema.update(visualization_update) session.add(schema) session.commit() session.refresh(schema) diff --git a/tests/integration/functional/zen_stores/test_zen_store.py b/tests/integration/functional/zen_stores/test_zen_store.py index 3664e5a9e7..6c490f4b1b 100644 --- a/tests/integration/functional/zen_stores/test_zen_store.py +++ b/tests/integration/functional/zen_stores/test_zen_store.py @@ -91,15 +91,20 @@ APIKeyRequest, APIKeyRotateRequest, APIKeyUpdate, + ArtifactRequest, ArtifactVersionFilter, ArtifactVersionRequest, ArtifactVersionResponse, ArtifactVisualizationRequest, ComponentFilter, + ComponentRequest, ComponentUpdate, CuratedVisualizationRequest, CuratedVisualizationUpdate, DeploymentRequest, + ModelFilter, + ModelRequest, + ModelUpdate, ModelVersionArtifactFilter, ModelVersionArtifactRequest, ModelVersionFilter, @@ -109,10 +114,12 @@ ModelVersionUpdate, PipelineRequest, PipelineRunFilter, + PipelineRunRequest, PipelineRunResponse, PipelineSnapshotRequest, ProjectFilter, ProjectUpdate, + RunMetadataRequest, RunMetadataResource, ScheduleRequest, ServiceAccountFilter, @@ -125,19 +132,14 @@ StackRequest, StackUpdate, StepRunFilter, + StepRunRequest, StepRunUpdate, TagResourceRequest, + UserFilter, UserRequest, UserResponse, UserUpdate, ) -from zenml.models.v2.core.artifact import ArtifactRequest -from zenml.models.v2.core.component import ComponentRequest -from zenml.models.v2.core.model import ModelFilter, ModelRequest, ModelUpdate -from zenml.models.v2.core.pipeline_run import PipelineRunRequest -from zenml.models.v2.core.run_metadata import RunMetadataRequest -from zenml.models.v2.core.step_run import StepRunRequest -from zenml.models.v2.core.user import UserFilter from zenml.utils import code_repository_utils, source_utils from zenml.utils.enum_utils import StrEnum from zenml.zen_stores.rest_zen_store import RestZenStore @@ -5770,13 +5772,6 @@ def test_curated_visualizations_across_resources(self): client = Client() project_id = client.active_project.id - artifact = client.zen_store.create_artifact( - ArtifactRequest( - name=sample_name("artifact"), - project=project_id, - has_custom_name=True, - ) - ) resource_configs = [ { "resource_type": VisualizationResourceTypes.PIPELINE, @@ -5803,27 +5798,42 @@ def test_curated_visualizations_across_resources(self): "resource_id": None, }, ] - artifact_version = client.zen_store.create_artifact_version( - ArtifactVersionRequest( - artifact_id=artifact.id, - project=project_id, - version="1", - type=ArtifactType.DATA, - uri=sample_name("artifact_uri"), - materializer=Source( - module="acme.foo", type=SourceType.INTERNAL - ), - data_type=Source(module="acme.foo", type=SourceType.INTERNAL), - save_type=ArtifactSaveType.STEP_OUTPUT, - visualizations=[ - ArtifactVisualizationRequest( - type=VisualizationType.HTML, - uri=f"s3://visualizations/{config['resource_type'].value}_{index}.html", - ) - for index, config in enumerate(resource_configs) - ], + + def create_artifact_version(): + artifact = client.zen_store.create_artifact( + ArtifactRequest( + name=sample_name("artifact"), + project=project_id, + has_custom_name=True, + ) ) - ) + artifact_version = client.zen_store.create_artifact_version( + ArtifactVersionRequest( + artifact_id=artifact.id, + project=project_id, + version="1", + type=ArtifactType.DATA, + uri=sample_name("artifact_uri"), + materializer=Source( + module="acme.foo", type=SourceType.INTERNAL + ), + data_type=Source( + module="acme.foo", type=SourceType.INTERNAL + ), + save_type=ArtifactSaveType.STEP_OUTPUT, + visualizations=[ + ArtifactVisualizationRequest( + type=VisualizationType.HTML, + uri=f"s3://visualizations/{config['resource_type'].value}_{index}.html", + ) + for index, config in enumerate(resource_configs) + ], + ) + ) + + return artifact, artifact_version + + artifact, artifact_version = create_artifact_version() pipeline_model = client.zen_store.create_pipeline( PipelineRequest( @@ -5874,29 +5884,14 @@ def test_curated_visualizations_across_resources(self): ) ) - # Try to find a deployer component from active stack first - deployer_id = None - active_stack_model = client.active_stack_model - if StackComponentType.MODEL_DEPLOYER in active_stack_model.components: - deployer_components = active_stack_model.components[ - StackComponentType.MODEL_DEPLOYER - ] - if deployer_components: - deployer_id = deployer_components[0].id - - if deployer_id is None: - # Fall back to listing all MODEL_DEPLOYER components - deployers = client.zen_store.list_stack_components( - ComponentFilter(type=StackComponentType.MODEL_DEPLOYER) - ) - if deployers.total > 0: - deployer_id = deployers.items[0].id - - if deployer_id is None: - pytest.skip( - "No deployer component available in the active stack or registered components. " - "Skipping deployment curated visualization test." + deployer = client.zen_store.create_stack_component( + ComponentRequest( + name=sample_name("foo"), + type=StackComponentType.DEPLOYER, + flavor="docker", + configuration={}, ) + ) # Create a deployment deployment = client.zen_store.create_deployment( @@ -5904,18 +5899,11 @@ def test_curated_visualizations_across_resources(self): project=project_id, name=sample_name("deployment"), snapshot_id=snapshot.id, - deployer_id=deployer_id, + deployer_id=deployer.id, ) ) - try: - resource_configs[0]["resource_id"] = pipeline_model.id - resource_configs[1]["resource_id"] = model.id - resource_configs[2]["resource_id"] = pipeline_run.id - resource_configs[3]["resource_id"] = snapshot.id - resource_configs[4]["resource_id"] = deployment.id - resource_configs[5]["resource_id"] = project_id - + def create_visualizations(artifact_version): visualizations = {} for idx, config in enumerate(resource_configs): resource_type = config["resource_type"] @@ -5940,6 +5928,17 @@ def test_curated_visualizations_across_resources(self): hydrated.layout_size == CuratedVisualizationSize.FULL_WIDTH ) visualizations[resource_type] = viz + return visualizations + + try: + resource_configs[0]["resource_id"] = pipeline_model.id + resource_configs[1]["resource_id"] = model.id + resource_configs[2]["resource_id"] = pipeline_run.id + resource_configs[3]["resource_id"] = snapshot.id + resource_configs[4]["resource_id"] = deployment.id + resource_configs[5]["resource_id"] = project_id + + visualizations = create_visualizations(artifact_version) loaded = client.zen_store.get_curated_visualization( visualizations[VisualizationResourceTypes.PIPELINE].id, @@ -5987,10 +5986,101 @@ def test_curated_visualizations_across_resources(self): for viz in visualizations.values(): with pytest.raises(KeyError): client.zen_store.get_curated_visualization(viz.id) - finally: + + visualizations = create_visualizations(artifact_version) + + # Clean up artifact + client.zen_store.delete_artifact(artifact.id) + + # Check that all visualizations have been auto-deleted + for viz in visualizations.values(): + with pytest.raises(KeyError): + client.zen_store.get_curated_visualization(viz.id) + + artifact, artifact_version = create_artifact_version() + visualizations = create_visualizations(artifact_version) + # Clean up deployment client.zen_store.delete_deployment(deployment.id) + with pytest.raises(KeyError): + client.zen_store.get_curated_visualization( + visualizations[VisualizationResourceTypes.DEPLOYMENT].id + ) + + # Clean up pipeline run + client.zen_store.delete_run(pipeline_run.id) + with pytest.raises(KeyError): + client.zen_store.get_curated_visualization( + visualizations[VisualizationResourceTypes.PIPELINE_RUN].id + ) + + # Clean up model + client.zen_store.delete_model(model.id) + with pytest.raises(KeyError): + client.zen_store.get_curated_visualization( + visualizations[VisualizationResourceTypes.MODEL].id + ) + + # Clean up snapshot + client.zen_store.delete_snapshot(snapshot.id) + with pytest.raises(KeyError): + client.zen_store.get_curated_visualization( + visualizations[ + VisualizationResourceTypes.PIPELINE_SNAPSHOT + ].id + ) + + # Clean up pipeline + client.zen_store.delete_pipeline(pipeline_model.id) + with pytest.raises(KeyError): + client.zen_store.get_curated_visualization( + visualizations[VisualizationResourceTypes.PIPELINE].id + ) + + finally: + # Clean up deployment + try: + client.zen_store.delete_deployment(deployment.id) + except KeyError: + pass + + # Clean up pipeline run + try: + client.zen_store.delete_run(pipeline_run.id) + except KeyError: + pass + + # Clean up model + try: + client.zen_store.delete_model(model.id) + except KeyError: + pass + + # Clean up snapshot + try: + client.zen_store.delete_snapshot(snapshot.id) + except KeyError: + pass + + # Clean up pipeline + try: + client.zen_store.delete_pipeline(pipeline_model.id) + except KeyError: + pass + + # Clean up deployer + try: + client.zen_store.delete_stack_component(deployer.id) + except KeyError: + pass + + # Clean up artifact + try: + client.zen_store.delete_artifact(artifact.id) + except KeyError: + pass + def test_curated_visualizations_project_only(self): """Test project-level curated visualizations with single resource.""" From eeec764b53cb763d8154a1f454b5caeeaa70ada1 Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Mon, 20 Oct 2025 14:09:21 +0200 Subject: [PATCH 46/64] More fixes --- src/zenml/zen_stores/rest_zen_store.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 8d996b0969..101cd68b3f 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -1901,20 +1901,21 @@ def get_curated_visualization( def update_curated_visualization( self, visualization_id: UUID, - update: CuratedVisualizationUpdate, + visualization_update: CuratedVisualizationUpdate, ) -> CuratedVisualizationResponse: """Update a curated visualization via REST API. Args: visualization_id: The ID of the curated visualization to update. - update: The update to apply to the curated visualization. + visualization_update: The update to apply to the curated + visualization. Returns: The updated curated visualization. """ return self._update_resource( resource_id=visualization_id, - resource_update=update, + resource_update=visualization_update, response_model=CuratedVisualizationResponse, route=CURATED_VISUALIZATIONS, params={"hydrate": True}, From efc9cc86878e3511e9a36ba64c8e8e2ed421c262 Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Mon, 20 Oct 2025 15:25:39 +0200 Subject: [PATCH 47/64] Fix visualization update endpoint and move FK sqlite enablement after DB migration --- .../routers/curated_visualization_endpoints.py | 2 +- src/zenml/zen_stores/rest_zen_store.py | 1 - src/zenml/zen_stores/sql_zen_store.py | 17 +++++++++-------- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py index d0cb161304..44b13a35cb 100644 --- a/src/zenml/zen_server/routers/curated_visualization_endpoints.py +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -129,7 +129,7 @@ def get_curated_visualization( return hydrated_visualization -@router.patch( +@router.put( "/{visualization_id}", responses={401: error_response, 404: error_response, 422: error_response}, ) diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index 101cd68b3f..fd9bcd35c5 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -1918,7 +1918,6 @@ def update_curated_visualization( resource_update=visualization_update, response_model=CuratedVisualizationResponse, route=CURATED_VISUALIZATIONS, - params={"hydrate": True}, ) def delete_curated_visualization(self, visualization_id: UUID) -> None: diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index b01320ad19..a22a158fd8 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -1262,14 +1262,6 @@ def _initialize(self) -> None: engine_args=engine_args, ) - if self.config.driver == SQLDatabaseDriver.SQLITE: - # Enable foreign key checks at the SQLite database level - @event.listens_for(self._engine, "connect") - def _(dbapi_connection: Any, connection_record: Any) -> None: - cursor = dbapi_connection.cursor() - cursor.execute("PRAGMA foreign_keys=ON") - cursor.close() - # SQLite: As long as the parent directory exists, SQLAlchemy will # automatically create the database. if ( @@ -1298,6 +1290,15 @@ def _(dbapi_connection: Any, connection_record: Any) -> None: ): self.migrate_database() + if self.config.driver == SQLDatabaseDriver.SQLITE: + # Enable foreign key checks at the SQLite database level, but only + # after any migration has been done. + @event.listens_for(self._engine, "connect") + def _(dbapi_connection: Any, connection_record: Any) -> None: + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + secrets_store_config = self.config.secrets_store # Initialize the secrets store From 6543654c694c518865f67b58c94aa1d65837ac17 Mon Sep 17 00:00:00 2001 From: Stefan Nica Date: Mon, 20 Oct 2025 21:21:14 +0200 Subject: [PATCH 48/64] Clear DB connections post-migration --- src/zenml/zen_stores/sql_zen_store.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index a22a158fd8..800ab7a546 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -1299,6 +1299,10 @@ def _(dbapi_connection: Any, connection_record: Any) -> None: cursor.execute("PRAGMA foreign_keys=ON") cursor.close() + # Discard existing connections created without the foreign key + # checks enabled + self._engine.dispose() + secrets_store_config = self.config.secrets_store # Initialize the secrets store From c38b7e756684285f46e120baa77e2a6c67f2a3ba Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 11:01:16 +0100 Subject: [PATCH 49/64] Michael review updates --- docs/book/how-to/artifacts/visualizations.md | 122 ++++++++++++------ src/zenml/client.py | 9 +- src/zenml/models/__init__.py | 2 + .../models/v2/core/artifact_visualization.py | 29 +++++ .../models/v2/core/curated_visualization.py | 114 ++++++++++------ .../curated_visualization_endpoints.py | 6 +- .../schemas/artifact_visualization_schemas.py | 27 +++- .../schemas/curated_visualization_schemas.py | 78 +++++------ .../zen_stores/schemas/deployment_schemas.py | 14 +- src/zenml/zen_stores/schemas/model_schemas.py | 10 +- .../schemas/pipeline_run_schemas.py | 14 +- .../zen_stores/schemas/pipeline_schemas.py | 14 +- .../schemas/pipeline_snapshot_schemas.py | 10 +- .../zen_stores/schemas/project_schemas.py | 4 + src/zenml/zen_stores/sql_zen_store.py | 114 ++++------------ .../functional/zen_stores/test_zen_store.py | 18 +-- 16 files changed, 321 insertions(+), 264 deletions(-) diff --git a/docs/book/how-to/artifacts/visualizations.md b/docs/book/how-to/artifacts/visualizations.md index 39c68ae0c6..8934414508 100644 --- a/docs/book/how-to/artifacts/visualizations.md +++ b/docs/book/how-to/artifacts/visualizations.md @@ -90,49 +90,87 @@ from zenml.enums import ( ) client = Client() -artifact_version_id = UUID("") -project = client.active_project -pipeline = client.list_pipelines().items[0] -pipeline_run = pipeline.runs()[0] -snapshot = pipeline_run.snapshot() -deployment = client.list_deployments().items[0] + +# Define the identifiers for the pipeline and run you want to enrich +pipeline_id = UUID("") +pipeline_run_id = UUID("") + +# Retrieve the artifact version produced by the evaluation step +pipeline_run = client.get_pipeline_run(pipeline_run_id) +artifact_version_id = pipeline_run.output.get("evaluation_report") +artifact_version = client.get_artifact_version(artifact_version_id) +artifact_visualizations = artifact_version.visualizations or [] + +# Fetch the resources we want to enrich model = client.list_models().items[0] +model_id = model.id + +deployment = client.list_deployments().items[0] +deployment_id = deployment.id + +project_id = client.active_project.id -# Create a visualization for the model -model_viz = client.create_curated_visualization( - artifact_version_id=artifact_version_id, - visualization_index=0, - resource_id=model.id, +pipeline_model = client.get_pipeline(pipeline_id) +pipeline_id = pipeline_model.id + +pipeline_snapshot = pipeline_run.snapshot() +snapshot_id = pipeline_snapshot.id + +pipeline_run_id = pipeline_run.id + +# Create curated visualizations for each supported resource type +client.create_curated_visualization( + artifact_visualization_id=artifact_visualizations[0].id, + resource_id=model_id, resource_type=VisualizationResourceTypes.MODEL, - display_name="Model performance dashboard", - layout_size=CuratedVisualizationSize.FULL_WIDTH, + project_id=project_id, + display_name="Latest Model Evaluation", ) -# Create a visualization for the deployment -deployment_viz = client.create_curated_visualization( - artifact_version_id=artifact_version_id, - visualization_index=0, - resource_id=deployment.id, +client.create_curated_visualization( + artifact_visualization_id=artifact_visualizations[1].id, + resource_id=deployment_id, resource_type=VisualizationResourceTypes.DEPLOYMENT, - display_name="Deployment health dashboard", - layout_size=CuratedVisualizationSize.HALF_WIDTH, + project_id=project_id, + display_name="Deployment Health Dashboard", ) -# Create a visualization for the project -project_viz = client.create_curated_visualization( - artifact_version_id=artifact_version_id, - visualization_index=0, - resource_id=project.id, +client.create_curated_visualization( + artifact_visualization_id=artifact_visualizations[2].id, + resource_id=project_id, resource_type=VisualizationResourceTypes.PROJECT, - display_name="Project overview dashboard", - layout_size=CuratedVisualizationSize.FULL_WIDTH, + display_name="Project Overview", +) + +client.create_curated_visualization( + artifact_visualization_id=artifact_visualizations[3].id, + resource_id=pipeline_id, + resource_type=VisualizationResourceTypes.PIPELINE, + project_id=project_id, + display_name="Pipeline Summary", +) + +client.create_curated_visualization( + artifact_visualization_id=artifact_visualizations[4].id, + resource_id=pipeline_run_id, + resource_type=VisualizationResourceTypes.PIPELINE_RUN, + project_id=project_id, + display_name="Run Results", +) + +client.create_curated_visualization( + artifact_visualization_id=artifact_visualizations[5].id, + resource_id=snapshot_id, + resource_type=VisualizationResourceTypes.PIPELINE_SNAPSHOT, + project_id=project_id, + display_name="Snapshot Metrics", ) ``` After creation, the returned response includes the visualization ID. You can retrieve a specific visualization later with `Client.get_curated_visualization`: ```python -retrieved = client.get_curated_visualization(model_viz.id, hydrate=True) +retrieved = client.get_curated_visualization(pipeline_viz.id, hydrate=True) print(retrieved.display_name) print(retrieved.resource.type) print(retrieved.resource.id) @@ -170,31 +208,31 @@ When setting display orders, consider leaving gaps between values (e.g., 10, 20, ```python # Leave gaps for future insertions visualization_a = client.create_curated_visualization( - artifact_version_id=artifact_version_id, - visualization_index=0, - resource_id=model.id, - resource_type=VisualizationResourceTypes.MODEL, + artifact_visualization_id=artifact_visualizations[0].id, + resource_type=VisualizationResourceTypes.PIPELINE, + resource_id=pipeline_id, + display_name="Model performance at a glance", display_order=10, # Primary dashboard - layout_size=CuratedVisualizationSize.FULL_WIDTH, + layout_size=CuratedVisualizationSize.HALF_WIDTH, ) visualization_b = client.create_curated_visualization( - artifact_version_id=artifact_version_id, - visualization_index=1, - resource_id=model.id, - resource_type=VisualizationResourceTypes.MODEL, + artifact_visualization_id=artifact_visualizations[1].id, + resource_type=VisualizationResourceTypes.PIPELINE, + resource_id=pipeline_id, + display_name="Drill-down metrics", display_order=20, # Secondary metrics layout_size=CuratedVisualizationSize.HALF_WIDTH, # Compact chart beside the primary tile ) # Later, easily insert between them visualization_c = client.create_curated_visualization( - artifact_version_id=artifact_version_id, - visualization_index=2, - resource_id=model.id, - resource_type=VisualizationResourceTypes.MODEL, + artifact_visualization_id=artifact_visualizations[2].id, + resource_type=VisualizationResourceTypes.PIPELINE, + resource_id=pipeline_id, + display_name="Raw output preview", display_order=15, # Now appears between A and B - layout_size=CuratedVisualizationSize.HALF_WIDTH, + layout_size=CuratedVisualizationSize.FULL_WIDTH, ) ``` diff --git a/src/zenml/client.py b/src/zenml/client.py index 1f1a45d9d0..14d91d259c 100644 --- a/src/zenml/client.py +++ b/src/zenml/client.py @@ -3746,8 +3746,7 @@ def get_deployment( def create_curated_visualization( self, - artifact_version_id: UUID, - visualization_index: int, + artifact_visualization_id: UUID, *, resource_id: UUID, resource_type: VisualizationResourceTypes, @@ -3774,8 +3773,7 @@ def create_curated_visualization( Each visualization is linked to exactly one resource. Args: - artifact_version_id: The ID of the artifact version containing the visualization. - visualization_index: The index of the visualization within the artifact version. + artifact_visualization_id: The UUID of the artifact visualization to curate. resource_id: The identifier of the resource tied to the visualization. resource_type: The type of resource referenced by the visualization. project_id: The ID of the project to associate with the visualization. @@ -3788,8 +3786,7 @@ def create_curated_visualization( """ request = CuratedVisualizationRequest( project=project_id or self.active_project.id, - artifact_version_id=artifact_version_id, - visualization_index=visualization_index, + artifact_visualization_id=artifact_visualization_id, display_name=display_name, display_order=display_order, layout_size=layout_size, diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index 75f8b455c3..c704cfe2db 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -112,6 +112,7 @@ ArtifactVisualizationResponse, ArtifactVisualizationResponseBody, ArtifactVisualizationResponseMetadata, + ArtifactVisualizationResponseResources, ) from zenml.models.v2.core.service import ( ServiceResponse, @@ -636,6 +637,7 @@ "ArtifactVisualizationResponse", "ArtifactVisualizationResponseBody", "ArtifactVisualizationResponseMetadata", + "ArtifactVisualizationResponseResources", "CodeReferenceRequest", "CodeReferenceResponse", "CodeReferenceResponseBody", diff --git a/src/zenml/models/v2/core/artifact_visualization.py b/src/zenml/models/v2/core/artifact_visualization.py index 8dde68741c..56442de0b1 100644 --- a/src/zenml/models/v2/core/artifact_visualization.py +++ b/src/zenml/models/v2/core/artifact_visualization.py @@ -13,8 +13,11 @@ # permissions and limitations under the License. """Models representing artifact visualizations.""" +from typing import TYPE_CHECKING, Optional from uuid import UUID +from pydantic import Field + from zenml.enums import VisualizationType from zenml.models.v2.base.base import ( BaseDatedResponseBody, @@ -24,6 +27,9 @@ BaseResponseResources, ) +if TYPE_CHECKING: + from zenml.models.v2.core.artifact_version import ArtifactVersionResponse + # ------------------ Request Model ------------------ @@ -57,6 +63,12 @@ class ArtifactVisualizationResponseMetadata(BaseResponseMetadata): class ArtifactVisualizationResponseResources(BaseResponseResources): """Class for all resource models associated with the artifact visualization.""" + artifact_version: Optional["ArtifactVersionResponse"] = Field( + default=None, + title="The artifact version.", + description="Artifact version that owns this visualization, when included.", + ) + class ArtifactVisualizationResponse( BaseIdentifiedResponse[ @@ -105,6 +117,23 @@ def artifact_version_id(self) -> UUID: """ return self.get_metadata().artifact_version_id + @property + def artifact_version(self) -> "ArtifactVersionResponse": + """The artifact version resource, if the response was hydrated with it. + + Returns: + The artifact version resource associated with this visualization. + + Raises: + RuntimeError: If the artifact version resource is not available. + """ + resources = self.get_resources() + if resources is None or resources.artifact_version is None: + raise RuntimeError( + "Artifact visualization response was not hydrated with the artifact version resource." + ) + return resources.artifact_version + # ------------------ Filter Model ------------------ diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 987a419e68..31a07ad071 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -30,6 +30,9 @@ if TYPE_CHECKING: from zenml.models.v2.core.artifact_version import ArtifactVersionResponse + from zenml.models.v2.core.artifact_visualization import ( + ArtifactVisualizationResponse, + ) # ------------------ Request Model ------------------ @@ -38,28 +41,23 @@ class CuratedVisualizationRequest(ProjectScopedRequest): """Request model for curated visualizations. - Each curated visualization is linked to exactly one resource of the following - types: - - **Deployments**: Surface visualizations on deployment dashboards - - **Models**: Highlight evaluation dashboards and monitoring views next to - registered models - - **Pipelines**: Associate visualizations with pipeline definitions - - **Pipeline Runs**: Attach visualizations to specific execution runs - - **Pipeline Snapshots**: Link visualizations to snapshot configurations - - **Projects**: Provide high-level project dashboards and KPI overviews - - To attach a visualization to multiple resources, create separate curated - visualization entries for each resource. + Each curated visualization links a pre-rendered artifact visualization + to a single ZenML resource to surface it in the appropriate UI context. + Supported resources include: + - **Deployments** + - **Models** + - **Pipelines** + - **Pipeline Runs** + - **Pipeline Snapshots** + - **Projects** """ - artifact_version_id: UUID = Field( - title="The artifact version ID.", - description="Identifier of the artifact version providing the visualization.", - ) - visualization_index: int = Field( - ge=0, - title="The visualization index.", - description="Index of the visualization within the artifact version payload.", + artifact_visualization_id: UUID = Field( + title="The artifact visualization ID.", + description=( + "Identifier of the artifact visualization that should be surfaced " + "for the target resource." + ), ) display_name: Optional[str] = Field( default=None, @@ -117,13 +115,19 @@ class CuratedVisualizationUpdate(BaseUpdate): class CuratedVisualizationResponseBody(ProjectScopedResponseBody): """Response body for curated visualizations.""" - artifact_version_id: UUID = Field( - title="The artifact version ID.", - description="Identifier of the artifact version providing the visualization.", + artifact_visualization_id: UUID = Field( + title="The artifact visualization ID.", + description=( + "Identifier of the artifact visualization that is curated for this resource." + ), ) - visualization_index: int = Field( - title="The visualization index.", - description="Index of the visualization within the artifact version payload.", + artifact_version_id: Optional[UUID] = Field( + default=None, + title="The artifact version ID.", + description=( + "Identifier of the artifact version that owns the curated visualization. " + "Provided for read-only context when available." + ), ) display_name: Optional[str] = Field( default=None, @@ -154,9 +158,11 @@ class CuratedVisualizationResponseMetadata(ProjectScopedResponseMetadata): class CuratedVisualizationResponseResources(ProjectScopedResponseResources): """Response resources included for curated visualizations.""" - artifact_version: "ArtifactVersionResponse" = Field( - title="The artifact version.", - description="Artifact version from which the visualization originates.", + artifact_visualization: "ArtifactVisualizationResponse" = Field( + title="The artifact visualization.", + description=( + "Artifact visualization that is surfaced through this curated visualization." + ), ) @@ -182,22 +188,22 @@ def get_hydrated_version(self) -> "CuratedVisualizationResponse": # Helper properties @property - def artifact_version_id(self) -> UUID: - """The artifact version ID. + def artifact_visualization_id(self) -> UUID: + """The artifact visualization ID. Returns: - The artifact version ID. + The artifact visualization ID. """ - return self.get_body().artifact_version_id + return self.get_body().artifact_visualization_id @property - def visualization_index(self) -> int: - """The visualization index. + def artifact_version_id(self) -> Optional[UUID]: + """The artifact version ID. Returns: - The visualization index. + The artifact version ID if available. """ - return self.get_body().visualization_index + return self.get_body().artifact_version_id @property def display_name(self) -> Optional[str]: @@ -226,14 +232,42 @@ def layout_size(self) -> CuratedVisualizationSize: """ return self.get_body().layout_size + @property + def artifact_visualization(self) -> "ArtifactVisualizationResponse": + """The curated artifact visualization resource. + + Returns: + The artifact visualization resource if included. + + Raises: + RuntimeError: If the response was not hydrated with resources. + """ + resources = self.get_resources() + if resources is None or resources.artifact_visualization is None: + raise RuntimeError( + "Curated visualization response was not hydrated with the artifact visualization resource." + ) + return resources.artifact_visualization + @property def artifact_version(self) -> "ArtifactVersionResponse": - """The artifact version resource. + """The artifact version resource, if available. Returns: - The artifact version resource if included. + The artifact version resource associated with the curated visualization. + + Raises: + RuntimeError: If the artifact version is not included in the hydrated resources. """ - return self.get_resources().artifact_version + artifact_visualization = self.artifact_visualization + if ( + artifact_visualization.get_resources() is None + or artifact_visualization.get_resources().artifact_version is None + ): + raise RuntimeError( + "Curated visualization response was not hydrated with the artifact version resource." + ) + return artifact_visualization.artifact_version @property def resource_id(self) -> UUID: diff --git a/src/zenml/zen_server/routers/curated_visualization_endpoints.py b/src/zenml/zen_server/routers/curated_visualization_endpoints.py index 44b13a35cb..b7f3858dfc 100644 --- a/src/zenml/zen_server/routers/curated_visualization_endpoints.py +++ b/src/zenml/zen_server/routers/curated_visualization_endpoints.py @@ -87,12 +87,12 @@ def create_curated_visualization( resource_model = _get_resource_model( visualization.resource_type, visualization.resource_id ) - artifact_version = store.get_artifact_version( - visualization.artifact_version_id + artifact_visualization = store.get_artifact_visualization( + visualization.artifact_visualization_id ) verify_permission_for_model(resource_model, action=Action.UPDATE) - verify_permission_for_model(artifact_version, action=Action.READ) + verify_permission_for_model(artifact_visualization, action=Action.READ) return store.create_curated_visualization(visualization) diff --git a/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py b/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py index 79862fc037..ecd7e329e8 100644 --- a/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """SQLModel implementation of artifact visualization table.""" -from typing import Any +from typing import TYPE_CHECKING, Any, List from uuid import UUID from sqlalchemy import TEXT, Column @@ -25,11 +25,17 @@ ArtifactVisualizationResponse, ArtifactVisualizationResponseBody, ArtifactVisualizationResponseMetadata, + ArtifactVisualizationResponseResources, ) from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema from zenml.zen_stores.schemas.base_schemas import BaseSchema from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field +if TYPE_CHECKING: + from zenml.zen_stores.schemas.curated_visualization_schemas import ( + CuratedVisualizationSchema, + ) + class ArtifactVisualizationSchema(BaseSchema, table=True): """SQL Model for visualizations of artifacts.""" @@ -54,6 +60,15 @@ class ArtifactVisualizationSchema(BaseSchema, table=True): artifact_version: ArtifactVersionSchema = Relationship( back_populates="visualizations" ) + curated_visualizations: List["CuratedVisualizationSchema"] = Relationship( + back_populates="artifact_visualization", + sa_relationship_kwargs=dict( + order_by=( + "CuratedVisualizationSchema.display_order", + "CuratedVisualizationSchema.created", + ), + ), + ) @classmethod def from_model( @@ -107,8 +122,18 @@ def to_model( artifact_version_id=self.artifact_version_id, ) + resources = None + if include_resources and self.artifact_version is not None: + resources = ArtifactVisualizationResponseResources( + artifact_version=self.artifact_version.to_model( + include_metadata=False, + include_resources=False, + ) + ) + return ArtifactVisualizationResponse( id=self.id, body=body, metadata=metadata, + resources=resources, ) diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 77c2903a88..1d185ceb70 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -39,7 +39,9 @@ from zenml.zen_stores.schemas.utils import jl_arg if TYPE_CHECKING: - from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema + from zenml.zen_stores.schemas.artifact_visualization_schemas import ( + ArtifactVisualizationSchema, + ) class CuratedVisualizationSchema(BaseSchema, table=True): @@ -47,13 +49,9 @@ class CuratedVisualizationSchema(BaseSchema, table=True): __tablename__ = "curated_visualization" __table_args__ = ( - build_index( - __tablename__, ["artifact_version_id", "visualization_index"] - ), - build_index(__tablename__, ["display_order"]), + build_index(__tablename__, ["artifact_visualization_id"]), UniqueConstraint( - "artifact_version_id", - "visualization_index", + "artifact_visualization_id", "resource_id", "resource_type", name="unique_curated_visualization_resource_link", @@ -68,16 +66,15 @@ class CuratedVisualizationSchema(BaseSchema, table=True): ondelete="CASCADE", nullable=False, ) - artifact_version_id: UUID = build_foreign_key_field( + artifact_visualization_id: UUID = build_foreign_key_field( source=__tablename__, - target="artifact_version", - source_column="artifact_version_id", + target="artifact_visualization", + source_column="artifact_visualization_id", target_column="id", ondelete="CASCADE", nullable=False, ) - visualization_index: int = Field(nullable=False) display_name: Optional[str] = Field(default=None) display_order: Optional[int] = Field(default=None) layout_size: str = Field( @@ -87,7 +84,9 @@ class CuratedVisualizationSchema(BaseSchema, table=True): resource_id: UUID = Field(nullable=False) resource_type: str = Field(nullable=False) - artifact_version: "ArtifactVersionSchema" = Relationship() + artifact_visualization: "ArtifactVisualizationSchema" = Relationship( + back_populates="curated_visualizations" + ) @classmethod def get_query_options( @@ -111,7 +110,17 @@ def get_query_options( options: List[ExecutableOption] = [] if include_resources: - options.append(selectinload(jl_arg(cls.artifact_version))) + from zenml.zen_stores.schemas.artifact_visualization_schemas import ( + ArtifactVisualizationSchema, + ) + + options.append( + selectinload(jl_arg(cls.artifact_visualization)).options( + selectinload( + jl_arg(ArtifactVisualizationSchema.artifact_version) + ) + ) + ) return options @@ -129,8 +138,7 @@ def from_request( """ return cls( project_id=request.project, - artifact_version_id=request.artifact_version_id, - visualization_index=request.visualization_index, + artifact_visualization_id=request.artifact_visualization_id, display_name=request.display_name, display_order=request.display_order, layout_size=request.layout_size.value, @@ -180,31 +188,28 @@ def to_model( Returns: The created response model. """ - layout_size_value = ( - self.layout_size or CuratedVisualizationSize.FULL_WIDTH.value - ) try: - layout_size_enum = CuratedVisualizationSize(layout_size_value) + layout_size_enum = CuratedVisualizationSize(self.layout_size) except ValueError: layout_size_enum = CuratedVisualizationSize.FULL_WIDTH - resource_type_str = self.resource_type - if resource_type_str: - try: - resource_type_enum = VisualizationResourceTypes( - resource_type_str - ) - except ValueError: - resource_type_enum = VisualizationResourceTypes.PROJECT - else: + try: + resource_type_enum = VisualizationResourceTypes(self.resource_type) + except ValueError: resource_type_enum = VisualizationResourceTypes.PROJECT + artifact_version_id: Optional[UUID] = None + if self.artifact_visualization is not None: + artifact_version_id = ( + self.artifact_visualization.artifact_version_id + ) + body = CuratedVisualizationResponseBody( project_id=self.project_id, created=self.created, updated=self.updated, - artifact_version_id=self.artifact_version_id, - visualization_index=self.visualization_index, + artifact_visualization_id=self.artifact_visualization_id, + artifact_version_id=artifact_version_id, display_name=self.display_name, display_order=self.display_order, layout_size=layout_size_enum, @@ -217,14 +222,15 @@ def to_model( metadata = CuratedVisualizationResponseMetadata() response_resources = None - if include_resources: - artifact_version_model = self.artifact_version.to_model( - include_metadata=include_metadata, - include_resources=include_resources, + if include_resources and self.artifact_visualization is not None: + artifact_visualization_model = ( + self.artifact_visualization.to_model( + include_metadata=False, + include_resources=False, + ) ) - response_resources = CuratedVisualizationResponseResources( - artifact_version=artifact_version_model, + artifact_visualization=artifact_visualization_model, ) return CuratedVisualizationResponse( diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index 16b4295e65..0447a1aa2d 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -149,6 +149,10 @@ class DeploymentSchema(NamedSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", + order_by=( + "CuratedVisualizationSchema.display_order", + "CuratedVisualizationSchema.created", + ), ), ) @@ -174,10 +178,6 @@ def get_query_options( options = [] if include_resources: - from zenml.zen_stores.schemas.curated_visualization_schemas import ( - CuratedVisualizationSchema, - ) - options.extend( [ selectinload(jl_arg(DeploymentSchema.user)), @@ -185,11 +185,7 @@ def get_query_options( selectinload(jl_arg(DeploymentSchema.snapshot)).joinedload( jl_arg(PipelineSnapshotSchema.pipeline) ), - selectinload( - jl_arg(DeploymentSchema.visualizations) - ).selectinload( - jl_arg(CuratedVisualizationSchema.artifact_version) - ), + selectinload(jl_arg(DeploymentSchema.visualizations)), ] ) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 330f212b2d..340fd7188e 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -141,6 +141,10 @@ class ModelSchema(NamedSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", + order_by=( + "CuratedVisualizationSchema.display_order", + "CuratedVisualizationSchema.created", + ), ), ) @@ -170,11 +174,7 @@ def get_query_options( [ joinedload(jl_arg(ModelSchema.user)), # joinedload(jl_arg(ModelSchema.tags)), - selectinload( - jl_arg(ModelSchema.visualizations) - ).selectinload( - jl_arg(CuratedVisualizationSchema.artifact_version) - ), + selectinload(jl_arg(ModelSchema.visualizations)), ] ) diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index 4740c1bb1c..87dc2b06a5 100644 --- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py @@ -254,6 +254,10 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", + order_by=( + "CuratedVisualizationSchema.display_order", + "CuratedVisualizationSchema.created", + ), ), ) @@ -298,10 +302,6 @@ def get_query_options( # ) if include_resources: - from zenml.zen_stores.schemas.curated_visualization_schemas import ( - CuratedVisualizationSchema, - ) - options.extend( [ selectinload( @@ -334,11 +334,7 @@ def get_query_options( selectinload(jl_arg(PipelineRunSchema.logs)), selectinload(jl_arg(PipelineRunSchema.user)), selectinload(jl_arg(PipelineRunSchema.tags)), - selectinload( - jl_arg(PipelineRunSchema.visualizations) - ).selectinload( - jl_arg(CuratedVisualizationSchema.artifact_version) - ), + selectinload(jl_arg(PipelineRunSchema.visualizations)), ] ) diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index ad4e0b8d79..dbe98eef8f 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_schemas.py @@ -116,6 +116,10 @@ class PipelineSchema(NamedSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", + order_by=( + "CuratedVisualizationSchema.display_order", + "CuratedVisualizationSchema.created", + ), ), ) @@ -169,19 +173,11 @@ def get_query_options( options = [] if include_resources: - from zenml.zen_stores.schemas.curated_visualization_schemas import ( - CuratedVisualizationSchema, - ) - options.extend( [ joinedload(jl_arg(PipelineSchema.user)), # joinedload(jl_arg(PipelineSchema.tags)), - selectinload( - jl_arg(PipelineSchema.visualizations) - ).selectinload( - jl_arg(CuratedVisualizationSchema.artifact_version) - ), + selectinload(jl_arg(PipelineSchema.visualizations)), ] ) diff --git a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py index 0a028edec3..3e8676da39 100644 --- a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py @@ -230,6 +230,10 @@ class PipelineSnapshotSchema(BaseSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", + order_by=( + "CuratedVisualizationSchema.display_order", + "CuratedVisualizationSchema.created", + ), ), ) @@ -366,17 +370,11 @@ def get_query_options( ) if include_resources: - from zenml.zen_stores.schemas.curated_visualization_schemas import ( - CuratedVisualizationSchema, - ) - options.extend( [ joinedload(jl_arg(PipelineSnapshotSchema.user)), selectinload( jl_arg(PipelineSnapshotSchema.visualizations) - ).selectinload( - jl_arg(CuratedVisualizationSchema.artifact_version) ), ] ) diff --git a/src/zenml/zen_stores/schemas/project_schemas.py b/src/zenml/zen_stores/schemas/project_schemas.py index e4794277f0..b77934b6ea 100644 --- a/src/zenml/zen_stores/schemas/project_schemas.py +++ b/src/zenml/zen_stores/schemas/project_schemas.py @@ -138,6 +138,10 @@ class ProjectSchema(NamedSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", + order_by=( + "CuratedVisualizationSchema.display_order", + "CuratedVisualizationSchema.created", + ), ), ) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 800ab7a546..c54eea5948 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5419,67 +5419,20 @@ def delete_deployment(self, deployment_id: UUID) -> None: # -------------------- Curated visualizations -------------------- - def _validate_curated_visualization_index( - self, - session: Session, - artifact_version_id: UUID, - visualization_index: int, - ) -> None: - """Validate that the artifact version exposes the given visualization index. - - Args: - session: The database session to use. - artifact_version_id: ID of the artifact version that produced the visualization. - visualization_index: Index of the visualization to validate. - - Raises: - IllegalOperationError: If the artifact version does not expose the specified visualization index. - """ - count = session.scalar( - select(func.count()) - .select_from(ArtifactVisualizationSchema) - .where( - ArtifactVisualizationSchema.artifact_version_id - == artifact_version_id - ) - ) - if not count or visualization_index >= count: - raise IllegalOperationError( - "Artifact version " - f"`{artifact_version_id}` does not expose a visualization " - f"with index {visualization_index}." - ) - def _assert_curated_visualization_duplicate( self, session: Session, *, - artifact_version_id: UUID, - visualization_index: int, + artifact_visualization_id: UUID, resource_id: UUID, resource_type: VisualizationResourceTypes, ) -> None: - """Ensure a curated visualization link does not already exist. - - Args: - session: The session to use. - artifact_version_id: The ID of the artifact version to validate. - visualization_index: The index of the visualization to validate. - resource_id: The ID of the resource to validate. - resource_type: The type of the resource to validate. - - Raises: - EntityExistsError: If a curated visualization link already exists. - """ + """Ensure a curated visualization link does not already exist.""" existing = session.exec( select(CuratedVisualizationSchema) .where( - CuratedVisualizationSchema.artifact_version_id - == artifact_version_id - ) - .where( - CuratedVisualizationSchema.visualization_index - == visualization_index + CuratedVisualizationSchema.artifact_visualization_id + == artifact_visualization_id ) .where(CuratedVisualizationSchema.resource_id == resource_id) .where( @@ -5489,51 +5442,36 @@ def _assert_curated_visualization_duplicate( if existing is not None: raise EntityExistsError( "A curated visualization for this resource already exists " - "for the specified artifact version and visualization index." + "for the specified artifact visualization." ) def create_curated_visualization( self, visualization: CuratedVisualizationRequest ) -> CuratedVisualizationResponse: - """Persist a curated visualization link. - - Args: - visualization: The curated visualization to persist. - - Returns: - The persisted curated visualization. - - Raises: - IllegalOperationError: If the resource does not belong to the same project as the curated visualization. - KeyError: If the resource does not exist. - """ + """Persist a curated visualization link.""" with Session(self.engine) as session: self._set_request_user_id( request_model=visualization, session=session ) - artifact_version: ArtifactVersionSchema = self._get_schema_by_id( - resource_id=visualization.artifact_version_id, - schema_class=ArtifactVersionSchema, - session=session, + artifact_visualization: ArtifactVisualizationSchema = ( + self._get_reference_schema_by_id( + resource=visualization, + reference_schema=ArtifactVisualizationSchema, + reference_id=visualization.artifact_visualization_id, + session=session, + ) ) - # Validate explicitly provided project before defaulting - if visualization.project: - if visualization.project != artifact_version.project_id: - raise IllegalOperationError( - "Curated visualizations must target the same project as " - "the artifact version." - ) - project_id = visualization.project - else: - project_id = artifact_version.project_id + artifact_version = artifact_visualization.artifact_version + project_id = artifact_version.project_id - self._validate_curated_visualization_index( - session=session, - artifact_version_id=visualization.artifact_version_id, - visualization_index=visualization.visualization_index, - ) + if visualization.project != project_id: + raise IllegalOperationError( + "Curated visualizations must target the same project as " + "the artifact visualization." + ) + project_id = visualization.project resource_schema_map: Dict[ VisualizationResourceTypes, Type[BaseSchema] @@ -5547,7 +5485,7 @@ def create_curated_visualization( } if visualization.resource_type not in resource_schema_map: - raise IllegalOperationError( + raise ValueError( f"Invalid resource type: {visualization.resource_type}" ) @@ -5575,16 +5513,12 @@ def create_curated_visualization( self._assert_curated_visualization_duplicate( session=session, - artifact_version_id=visualization.artifact_version_id, - visualization_index=visualization.visualization_index, + artifact_visualization_id=visualization.artifact_visualization_id, resource_id=visualization.resource_id, resource_type=visualization.resource_type, ) - schema: CuratedVisualizationSchema = ( - CuratedVisualizationSchema.from_request(visualization) - ) - schema.project_id = project_id + schema = CuratedVisualizationSchema.from_request(visualization) session.add(schema) session.commit() diff --git a/tests/integration/functional/zen_stores/test_zen_store.py b/tests/integration/functional/zen_stores/test_zen_store.py index 6c490f4b1b..dfcb21384b 100644 --- a/tests/integration/functional/zen_stores/test_zen_store.py +++ b/tests/integration/functional/zen_stores/test_zen_store.py @@ -5905,14 +5905,16 @@ def create_artifact_version(): def create_visualizations(artifact_version): visualizations = {} - for idx, config in enumerate(resource_configs): + artifact_visualizations = artifact_version.visualizations or [] + for artifact_viz, config in zip( + artifact_visualizations, resource_configs + ): resource_type = config["resource_type"] resource_id = config["resource_id"] viz = client.zen_store.create_curated_visualization( CuratedVisualizationRequest( project=project_id, - artifact_version_id=artifact_version.id, - visualization_index=idx, + artifact_visualization_id=artifact_viz.id, resource_id=resource_id, resource_type=resource_type, display_name=f"{resource_type.value} visualization", @@ -5952,13 +5954,12 @@ def create_visualizations(artifact_version): assert loaded.resource_id == pipeline_model.id assert loaded.resource_type == VisualizationResourceTypes.PIPELINE - # Test duplicate creation - same artifact_version + visualization_index + resource should fail + # Test duplicate creation - same artifact visualization + resource should fail with pytest.raises(EntityExistsError): client.zen_store.create_curated_visualization( CuratedVisualizationRequest( project=project_id, - artifact_version_id=artifact_version.id, - visualization_index=0, # Same index as pipeline visualization + artifact_visualization_id=loaded.artifact_visualization_id, resource_id=pipeline_model.id, resource_type=VisualizationResourceTypes.PIPELINE, ) @@ -6115,9 +6116,10 @@ def test_curated_visualizations_project_only(self): ) ) + artifact_visualization = (artifact_version.visualizations or [])[0] + visualization = client.create_curated_visualization( - artifact_version_id=artifact_version.id, - visualization_index=0, + artifact_visualization_id=artifact_visualization.id, resource_id=project.id, resource_type=VisualizationResourceTypes.PROJECT, project_id=project.id, From ea4b2385cc8fc5c54dcb9b5c1a4241b2dd6eeff9 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 11:02:17 +0100 Subject: [PATCH 50/64] remove migration file --- .../553964d69699_add_visualisations.py | 94 ------------------- 1 file changed, 94 deletions(-) delete mode 100644 src/zenml/zen_stores/migrations/versions/553964d69699_add_visualisations.py diff --git a/src/zenml/zen_stores/migrations/versions/553964d69699_add_visualisations.py b/src/zenml/zen_stores/migrations/versions/553964d69699_add_visualisations.py deleted file mode 100644 index 157796af76..0000000000 --- a/src/zenml/zen_stores/migrations/versions/553964d69699_add_visualisations.py +++ /dev/null @@ -1,94 +0,0 @@ -"""add visualisations [553964d69699]. - -Revision ID: 553964d69699 -Revises: 7497d2ff5731 -Create Date: 2025-10-17 11:44:15.358521 - -""" - -import sqlalchemy as sa -import sqlmodel -from alembic import op - -# revision identifiers, used by Alembic. -revision = "553964d69699" -down_revision = "7497d2ff5731" -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! ### - op.create_table( - "curated_visualization", - sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column("created", sa.DateTime(), nullable=False), - sa.Column("updated", sa.DateTime(), nullable=False), - sa.Column("project_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column( - "artifact_version_id", sqlmodel.sql.sqltypes.GUID(), nullable=False - ), - sa.Column("visualization_index", sa.Integer(), nullable=False), - sa.Column( - "display_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True - ), - sa.Column("display_order", sa.Integer(), nullable=True), - sa.Column( - "layout_size", sqlmodel.sql.sqltypes.AutoString(), nullable=False - ), - sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), - sa.Column( - "resource_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False - ), - sa.ForeignKeyConstraint( - ["artifact_version_id"], - ["artifact_version.id"], - name="fk_curated_visualization_artifact_version_id_artifact_version", - ondelete="CASCADE", - ), - sa.ForeignKeyConstraint( - ["project_id"], - ["project.id"], - name="fk_curated_visualization_project_id_project", - ondelete="CASCADE", - ), - sa.PrimaryKeyConstraint("id"), - sa.UniqueConstraint( - "artifact_version_id", - "visualization_index", - "resource_id", - "resource_type", - name="unique_curated_visualization_resource_link", - ), - ) - with op.batch_alter_table( - "curated_visualization", schema=None - ) as batch_op: - batch_op.create_index( - "ix_curated_visualization_artifact_version_id_visualization_index", - ["artifact_version_id", "visualization_index"], - unique=False, - ) - batch_op.create_index( - "ix_curated_visualization_display_order", - ["display_order"], - unique=False, - ) - - # ### 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( - "curated_visualization", schema=None - ) as batch_op: - batch_op.drop_index("ix_curated_visualization_display_order") - batch_op.drop_index( - "ix_curated_visualization_artifact_version_id_visualization_index" - ) - - op.drop_table("curated_visualization") - # ### end Alembic commands ### From c664755024da52bf72e7cf806d97d1faa975209a Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 11:03:29 +0100 Subject: [PATCH 51/64] add new migration file --- .../5b647a5aff9e_add_visualisations.py | 88 +++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 src/zenml/zen_stores/migrations/versions/5b647a5aff9e_add_visualisations.py diff --git a/src/zenml/zen_stores/migrations/versions/5b647a5aff9e_add_visualisations.py b/src/zenml/zen_stores/migrations/versions/5b647a5aff9e_add_visualisations.py new file mode 100644 index 0000000000..b8302ad083 --- /dev/null +++ b/src/zenml/zen_stores/migrations/versions/5b647a5aff9e_add_visualisations.py @@ -0,0 +1,88 @@ +"""add visualisations [5b647a5aff9e]. + +Revision ID: 5b647a5aff9e +Revises: 124b57b8c7b1 +Create Date: 2025-10-23 11:03:09.020315 + +""" + +import sqlalchemy as sa +import sqlmodel +from alembic import op + +# revision identifiers, used by Alembic. +revision = "5b647a5aff9e" +down_revision = "124b57b8c7b1" +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! ### + op.create_table( + "curated_visualization", + sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column("created", sa.DateTime(), nullable=False), + sa.Column("updated", sa.DateTime(), nullable=False), + sa.Column("project_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "artifact_visualization_id", + sqlmodel.sql.sqltypes.GUID(), + nullable=False, + ), + sa.Column( + "display_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True + ), + sa.Column("display_order", sa.Integer(), nullable=True), + sa.Column( + "layout_size", sqlmodel.sql.sqltypes.AutoString(), nullable=False + ), + sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column( + "resource_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False + ), + sa.ForeignKeyConstraint( + ["artifact_visualization_id"], + ["artifact_visualization.id"], + name="fk_curated_visualization_artifact_visualization_id_artifact_visualization", + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["project_id"], + ["project.id"], + name="fk_curated_visualization_project_id_project", + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "artifact_visualization_id", + "resource_id", + "resource_type", + name="unique_curated_visualization_resource_link", + ), + ) + with op.batch_alter_table( + "curated_visualization", schema=None + ) as batch_op: + batch_op.create_index( + "ix_curated_visualization_artifact_visualization_id", + ["artifact_visualization_id"], + unique=False, + ) + + # ### 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( + "curated_visualization", schema=None + ) as batch_op: + batch_op.drop_index( + "ix_curated_visualization_artifact_visualization_id" + ) + + op.drop_table("curated_visualization") + # ### end Alembic commands ### From f41cfc5adee04bef9ff1ce10736da9fc2b7039fd Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 14:25:17 +0100 Subject: [PATCH 52/64] michael second review --- .../models/v2/core/artifact_visualization.py | 14 ++------ .../models/v2/core/curated_visualization.py | 33 ++----------------- .../schemas/artifact_visualization_schemas.py | 10 ++++-- .../schemas/curated_visualization_schemas.py | 30 +++++------------ .../zen_stores/schemas/deployment_schemas.py | 2 +- src/zenml/zen_stores/schemas/model_schemas.py | 2 +- .../schemas/pipeline_run_schemas.py | 2 +- .../zen_stores/schemas/pipeline_schemas.py | 2 +- .../schemas/pipeline_snapshot_schemas.py | 2 +- 9 files changed, 26 insertions(+), 71 deletions(-) diff --git a/src/zenml/models/v2/core/artifact_visualization.py b/src/zenml/models/v2/core/artifact_visualization.py index 56442de0b1..5326877c0e 100644 --- a/src/zenml/models/v2/core/artifact_visualization.py +++ b/src/zenml/models/v2/core/artifact_visualization.py @@ -118,21 +118,13 @@ def artifact_version_id(self) -> UUID: return self.get_metadata().artifact_version_id @property - def artifact_version(self) -> "ArtifactVersionResponse": + def artifact_version(self) -> Optional["ArtifactVersionResponse"]: """The artifact version resource, if the response was hydrated with it. Returns: - The artifact version resource associated with this visualization. - - Raises: - RuntimeError: If the artifact version resource is not available. + The artifact version resource, if available. """ - resources = self.get_resources() - if resources is None or resources.artifact_version is None: - raise RuntimeError( - "Artifact visualization response was not hydrated with the artifact version resource." - ) - return resources.artifact_version + return self.get_resources().artifact_version # ------------------ Filter Model ------------------ diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 31a07ad071..5cede400ae 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -29,7 +29,6 @@ ) if TYPE_CHECKING: - from zenml.models.v2.core.artifact_version import ArtifactVersionResponse from zenml.models.v2.core.artifact_visualization import ( ArtifactVisualizationResponse, ) @@ -237,37 +236,9 @@ def artifact_visualization(self) -> "ArtifactVisualizationResponse": """The curated artifact visualization resource. Returns: - The artifact visualization resource if included. - - Raises: - RuntimeError: If the response was not hydrated with resources. - """ - resources = self.get_resources() - if resources is None or resources.artifact_visualization is None: - raise RuntimeError( - "Curated visualization response was not hydrated with the artifact visualization resource." - ) - return resources.artifact_visualization - - @property - def artifact_version(self) -> "ArtifactVersionResponse": - """The artifact version resource, if available. - - Returns: - The artifact version resource associated with the curated visualization. - - Raises: - RuntimeError: If the artifact version is not included in the hydrated resources. + The artifact visualization resource. """ - artifact_visualization = self.artifact_visualization - if ( - artifact_visualization.get_resources() is None - or artifact_visualization.get_resources().artifact_version is None - ): - raise RuntimeError( - "Curated visualization response was not hydrated with the artifact version resource." - ) - return artifact_visualization.artifact_version + return self.get_resources().artifact_visualization @property def resource_id(self) -> UUID: diff --git a/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py b/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py index ecd7e329e8..41fbf0e20e 100644 --- a/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py @@ -123,12 +123,16 @@ def to_model( ) resources = None - if include_resources and self.artifact_version is not None: - resources = ArtifactVisualizationResponseResources( - artifact_version=self.artifact_version.to_model( + if include_resources: + if self.artifact_version is not None: + artifact_version = self.artifact_version.to_model( include_metadata=False, include_resources=False, ) + else: + artifact_version = None + resources = ArtifactVisualizationResponseResources( + artifact_version=artifact_version, ) return ArtifactVisualizationResponse( diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 1d185ceb70..01b7c7f2a9 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -110,17 +110,7 @@ def get_query_options( options: List[ExecutableOption] = [] if include_resources: - from zenml.zen_stores.schemas.artifact_visualization_schemas import ( - ArtifactVisualizationSchema, - ) - - options.append( - selectinload(jl_arg(cls.artifact_visualization)).options( - selectinload( - jl_arg(ArtifactVisualizationSchema.artifact_version) - ) - ) - ) + options.append(selectinload(jl_arg(cls.artifact_visualization))) return options @@ -221,21 +211,19 @@ def to_model( if include_metadata: metadata = CuratedVisualizationResponseMetadata() - response_resources = None - if include_resources and self.artifact_visualization is not None: - artifact_visualization_model = ( - self.artifact_visualization.to_model( - include_metadata=False, - include_resources=False, - ) + resources = None + if include_resources: + artifact_visualization = self.artifact_visualization.to_model( + include_metadata=False, + include_resources=False, ) - response_resources = CuratedVisualizationResponseResources( - artifact_visualization=artifact_visualization_model, + resources = CuratedVisualizationResponseResources( + artifact_visualization=artifact_visualization, ) return CuratedVisualizationResponse( id=self.id, body=body, metadata=metadata, - resources=response_resources, + resources=resources, ) diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index 0447a1aa2d..b0bc7feed4 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -249,7 +249,7 @@ def to_model( include_metadata=False, include_resources=False, ) - for visualization in (self.visualizations or []) + for visualization in self.visualizations ], ) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 340fd7188e..36e5925e48 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -279,7 +279,7 @@ def to_model( include_metadata=False, include_resources=False, ) - for visualization in (self.visualizations or []) + for visualization in self.visualizations ], ) diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index 87dc2b06a5..82730b48fa 100644 --- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py @@ -657,7 +657,7 @@ def to_model( include_metadata=False, include_resources=False, ) - for visualization in (self.visualizations or []) + for visualization in self.visualizations ], ) diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index dbe98eef8f..aa1d37df01 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_schemas.py @@ -250,7 +250,7 @@ def to_model( include_metadata=False, include_resources=False, ) - for visualization in (self.visualizations or []) + for visualization in self.visualizations ], ) diff --git a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py index 3e8676da39..7fb06f8268 100644 --- a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py @@ -595,7 +595,7 @@ def to_model( include_metadata=False, include_resources=False, ) - for visualization in (self.visualizations or []) + for visualization in self.visualizations ], ) From 4ec6202d0f1393504787e87b82c5a3625ef632e9 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 15:38:29 +0100 Subject: [PATCH 53/64] add error handling on display order duplicate --- docs/book/how-to/artifacts/visualizations.md | 2 + .../models/v2/core/curated_visualization.py | 12 +++ src/zenml/zen_stores/sql_zen_store.py | 81 ++++++++++++++++++- 3 files changed, 93 insertions(+), 2 deletions(-) diff --git a/docs/book/how-to/artifacts/visualizations.md b/docs/book/how-to/artifacts/visualizations.md index 8934414508..0c6c2c56c7 100644 --- a/docs/book/how-to/artifacts/visualizations.md +++ b/docs/book/how-to/artifacts/visualizations.md @@ -236,6 +236,8 @@ visualization_c = client.create_curated_visualization( ) ``` +> **Note:** Reusing the same `display_order` for multiple curated visualizations targeting the same resource now raises an error. Pick distinct values—leaving gaps still helps when you later insert additional tiles. + #### RBAC visibility Curated visualizations respect the access permissions of the resource they're linked to. A user can only see a curated visualization if they have read access to the specific resource it targets. If a user lacks permission for the linked resource, the visualization will be hidden from their view. diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 5cede400ae..5d7096ba2e 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -65,6 +65,10 @@ class CuratedVisualizationRequest(ProjectScopedRequest): display_order: Optional[int] = Field( default=None, title="The display order of the visualization.", + description=( + "Optional ordering hint that must be unique for the combination " + "of resource type and resource ID." + ), ) layout_size: CuratedVisualizationSize = Field( default=CuratedVisualizationSize.FULL_WIDTH, @@ -101,6 +105,10 @@ class CuratedVisualizationUpdate(BaseUpdate): display_order: Optional[int] = Field( default=None, title="The new display order of the visualization.", + description=( + "Optional ordering hint. When provided, it must remain unique for " + "the combination of resource type and resource ID." + ), ) layout_size: Optional[CuratedVisualizationSize] = Field( default=None, @@ -135,6 +143,10 @@ class CuratedVisualizationResponseBody(ProjectScopedResponseBody): display_order: Optional[int] = Field( default=None, title="The display order of the visualization.", + description=( + "Optional ordering hint that is unique per combination of " + "resource type and resource ID." + ), ) layout_size: CuratedVisualizationSize = Field( default=CuratedVisualizationSize.FULL_WIDTH, diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 50c25b0352..d0780db720 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5427,7 +5427,17 @@ def _assert_curated_visualization_duplicate( resource_id: UUID, resource_type: VisualizationResourceTypes, ) -> None: - """Ensure a curated visualization link does not already exist.""" + """Ensure a curated visualization link does not already exist. + + Args: + session: The database session. + artifact_visualization_id: The ID of the artifact visualization. + resource_id: The ID of the resource. + resource_type: The type of the resource. + + Raises: + EntityExistsError: If a curated visualization link already exists. + """ existing = session.exec( select(CuratedVisualizationSchema) .where( @@ -5445,10 +5455,55 @@ def _assert_curated_visualization_duplicate( "for the specified artifact visualization." ) + def _assert_curated_visualization_display_order_unique( + self, + session: Session, + *, + resource_id: UUID, + resource_type: VisualizationResourceTypes, + display_order: Optional[int], + exclude_visualization_id: Optional[UUID] = None, + ) -> None: + """Ensure curated visualizations per resource use unique display orders.""" + if display_order is None: + return + + statement = ( + select(CuratedVisualizationSchema) + .where(CuratedVisualizationSchema.resource_id == resource_id) + .where( + CuratedVisualizationSchema.resource_type == resource_type.value + ) + .where(CuratedVisualizationSchema.display_order == display_order) + ) + if exclude_visualization_id is not None: + statement = statement.where( + CuratedVisualizationSchema.id != exclude_visualization_id + ) + + existing = session.exec(statement).first() + if existing is not None: + raise EntityExistsError( + "A curated visualization for this resource already uses the " + f"display order '{display_order}'. Please choose a different value." + ) + def create_curated_visualization( self, visualization: CuratedVisualizationRequest ) -> CuratedVisualizationResponse: - """Persist a curated visualization link.""" + """Persist a curated visualization link. + + Args: + visualization: The curated visualization to create. + + Returns: + The created curated visualization. + + Raises: + IllegalOperationError: If the curated visualization does not target the same project as the artifact visualization. + ValueError: If the resource type is invalid. + KeyError: If the resource is not found. + """ with Session(self.engine) as session: self._set_request_user_id( request_model=visualization, session=session @@ -5517,6 +5572,13 @@ def create_curated_visualization( resource_id=visualization.resource_id, resource_type=visualization.resource_type, ) + if visualization.display_order is not None: + self._assert_curated_visualization_display_order_unique( + session=session, + resource_id=visualization.resource_id, + resource_type=visualization.resource_type, + display_order=visualization.display_order, + ) schema = CuratedVisualizationSchema.from_request(visualization) @@ -5573,6 +5635,21 @@ def update_curated_visualization( schema_class=CuratedVisualizationSchema, session=session, ) + update_fields = visualization_update.model_dump(exclude_unset=True) + if "display_order" in update_fields: + new_display_order = update_fields["display_order"] + if new_display_order is not None: + self._assert_curated_visualization_display_order_unique( + session=session, + resource_id=schema.resource_id, + resource_type=VisualizationResourceTypes( + schema.resource_type + ), + display_order=new_display_order, + exclude_visualization_id=visualization_id, + ) + # Explicit None clears the display order, so uniqueness validation is skipped. + schema.update(visualization_update) session.add(schema) session.commit() From fc46f2d26bd21173a5b14ac04830e4ff80efcef1 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 15:43:35 +0100 Subject: [PATCH 54/64] Add migration for curated visualizations table This migration introduces a new table `curated_visualization` to store visualizations associated with projects. It includes necessary columns and constraints to ensure data integrity and relationships with existing tables. --- ...lisations.py => 54b2c0fbc0cb_add_visualisations.py} | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) rename src/zenml/zen_stores/migrations/versions/{5b647a5aff9e_add_visualisations.py => 54b2c0fbc0cb_add_visualisations.py} (92%) diff --git a/src/zenml/zen_stores/migrations/versions/5b647a5aff9e_add_visualisations.py b/src/zenml/zen_stores/migrations/versions/54b2c0fbc0cb_add_visualisations.py similarity index 92% rename from src/zenml/zen_stores/migrations/versions/5b647a5aff9e_add_visualisations.py rename to src/zenml/zen_stores/migrations/versions/54b2c0fbc0cb_add_visualisations.py index b8302ad083..a54738f07b 100644 --- a/src/zenml/zen_stores/migrations/versions/5b647a5aff9e_add_visualisations.py +++ b/src/zenml/zen_stores/migrations/versions/54b2c0fbc0cb_add_visualisations.py @@ -1,8 +1,8 @@ -"""add visualisations [5b647a5aff9e]. +"""add visualisations [54b2c0fbc0cb]. -Revision ID: 5b647a5aff9e +Revision ID: 54b2c0fbc0cb Revises: 124b57b8c7b1 -Create Date: 2025-10-23 11:03:09.020315 +Create Date: 2025-10-23 15:41:57.315515 """ @@ -11,7 +11,7 @@ from alembic import op # revision identifiers, used by Alembic. -revision = "5b647a5aff9e" +revision = "54b2c0fbc0cb" down_revision = "124b57b8c7b1" branch_labels = None depends_on = None @@ -45,7 +45,7 @@ def upgrade() -> None: sa.ForeignKeyConstraint( ["artifact_visualization_id"], ["artifact_visualization.id"], - name="fk_curated_visualization_artifact_visualization_id_artifact_visualization", + name="fk_visualization_artifact_visualization_id_artifact_visualization", ondelete="CASCADE", ), sa.ForeignKeyConstraint( From 620e78642644ad7f36aa9d5b945aed110c63bb40 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 16:42:46 +0100 Subject: [PATCH 55/64] fix order by --- .../zen_stores/schemas/artifact_visualization_schemas.py | 5 +---- src/zenml/zen_stores/schemas/deployment_schemas.py | 5 +---- src/zenml/zen_stores/schemas/model_schemas.py | 5 +---- src/zenml/zen_stores/schemas/pipeline_run_schemas.py | 5 +---- src/zenml/zen_stores/schemas/pipeline_schemas.py | 5 +---- src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py | 5 +---- src/zenml/zen_stores/schemas/project_schemas.py | 5 +---- 7 files changed, 7 insertions(+), 28 deletions(-) diff --git a/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py b/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py index 41fbf0e20e..a010311b03 100644 --- a/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py @@ -63,10 +63,7 @@ class ArtifactVisualizationSchema(BaseSchema, table=True): curated_visualizations: List["CuratedVisualizationSchema"] = Relationship( back_populates="artifact_visualization", sa_relationship_kwargs=dict( - order_by=( - "CuratedVisualizationSchema.display_order", - "CuratedVisualizationSchema.created", - ), + order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index b0bc7feed4..cdfcc3e3a1 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -149,10 +149,7 @@ class DeploymentSchema(NamedSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", - order_by=( - "CuratedVisualizationSchema.display_order", - "CuratedVisualizationSchema.created", - ), + order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 36e5925e48..0c3e44ed86 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -141,10 +141,7 @@ class ModelSchema(NamedSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", - order_by=( - "CuratedVisualizationSchema.display_order", - "CuratedVisualizationSchema.created", - ), + order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index 82730b48fa..66d6656f4a 100644 --- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py @@ -254,10 +254,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", - order_by=( - "CuratedVisualizationSchema.display_order", - "CuratedVisualizationSchema.created", - ), + order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index aa1d37df01..08f4b484a3 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_schemas.py @@ -116,10 +116,7 @@ class PipelineSchema(NamedSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", - order_by=( - "CuratedVisualizationSchema.display_order", - "CuratedVisualizationSchema.created", - ), + order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py index 7fb06f8268..6bf369d342 100644 --- a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py @@ -230,10 +230,7 @@ class PipelineSnapshotSchema(BaseSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", - order_by=( - "CuratedVisualizationSchema.display_order", - "CuratedVisualizationSchema.created", - ), + order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/project_schemas.py b/src/zenml/zen_stores/schemas/project_schemas.py index b77934b6ea..e9cac308b6 100644 --- a/src/zenml/zen_stores/schemas/project_schemas.py +++ b/src/zenml/zen_stores/schemas/project_schemas.py @@ -138,10 +138,7 @@ class ProjectSchema(NamedSchema, table=True): ), overlaps="visualizations", cascade="all, delete-orphan", - order_by=( - "CuratedVisualizationSchema.display_order", - "CuratedVisualizationSchema.created", - ), + order_by="CuratedVisualizationSchema.display_order", ), ) From a224a13f59eaf8114878f96810f940db3198fea5 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 17:20:09 +0100 Subject: [PATCH 56/64] fix docstring --- src/zenml/zen_stores/sql_zen_store.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index d0780db720..769bd5f9af 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5464,7 +5464,18 @@ def _assert_curated_visualization_display_order_unique( display_order: Optional[int], exclude_visualization_id: Optional[UUID] = None, ) -> None: - """Ensure curated visualizations per resource use unique display orders.""" + """Ensure curated visualizations per resource use unique display orders. + + Args: + session: The database session. + resource_id: The ID of the resource. + resource_type: The type of the resource. + display_order: The display order to check. + exclude_visualization_id: The ID of the visualization to exclude. + + Raises: + EntityExistsError: If a curated visualization for this resource already uses the display order. + """ if display_order is None: return From eda2ebd35da44af5f2891d1eeeb0fb353bbdc663 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 17:49:38 +0100 Subject: [PATCH 57/64] formatt --- src/zenml/zen_stores/sql_zen_store.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 769bd5f9af..adcd124f4d 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -5465,7 +5465,7 @@ def _assert_curated_visualization_display_order_unique( exclude_visualization_id: Optional[UUID] = None, ) -> None: """Ensure curated visualizations per resource use unique display orders. - + Args: session: The database session. resource_id: The ID of the resource. From 25deb645d01f74ed7d2577e6e2a0df772b222d4e Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 18:09:41 +0100 Subject: [PATCH 58/64] fix migration --- src/zenml/models/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/zenml/models/__init__.py b/src/zenml/models/__init__.py index 4478f7b2aa..ef0737dfa8 100644 --- a/src/zenml/models/__init__.py +++ b/src/zenml/models/__init__.py @@ -482,6 +482,10 @@ ArtifactVersionResponseBody.model_rebuild() ArtifactVersionResponseMetadata.model_rebuild() ArtifactVersionResponseResources.model_rebuild() +ArtifactVisualizationResponse.model_rebuild() +ArtifactVisualizationResponseBody.model_rebuild() +ArtifactVisualizationResponseMetadata.model_rebuild() +ArtifactVisualizationResponseResources.model_rebuild() CodeReferenceResponseBody.model_rebuild() CodeRepositoryResponseBody.model_rebuild() CodeRepositoryResponseMetadata.model_rebuild() From 4a9027ae764948c62b19f8d355629870ecb3a047 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 23 Oct 2025 23:18:18 +0100 Subject: [PATCH 59/64] cascade delete --- src/zenml/zen_stores/schemas/deployment_schemas.py | 2 +- src/zenml/zen_stores/schemas/model_schemas.py | 2 +- src/zenml/zen_stores/schemas/pipeline_run_schemas.py | 2 +- src/zenml/zen_stores/schemas/pipeline_schemas.py | 2 +- src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py | 2 +- src/zenml/zen_stores/schemas/project_schemas.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/zenml/zen_stores/schemas/deployment_schemas.py b/src/zenml/zen_stores/schemas/deployment_schemas.py index cdfcc3e3a1..98d7327e57 100644 --- a/src/zenml/zen_stores/schemas/deployment_schemas.py +++ b/src/zenml/zen_stores/schemas/deployment_schemas.py @@ -148,7 +148,7 @@ class DeploymentSchema(NamedSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==DeploymentSchema.id)" ), overlaps="visualizations", - cascade="all, delete-orphan", + cascade="delete", order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 0c3e44ed86..96856f1cdb 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -140,7 +140,7 @@ class ModelSchema(NamedSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==ModelSchema.id)" ), overlaps="visualizations", - cascade="all, delete-orphan", + cascade="delete", order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index 66d6656f4a..115d78764c 100644 --- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py @@ -253,7 +253,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True): "foreign(CuratedVisualizationSchema.resource_id)==PipelineRunSchema.id)" ), overlaps="visualizations", - cascade="all, delete-orphan", + cascade="delete", order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index 08f4b484a3..cf0923d641 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_schemas.py @@ -115,7 +115,7 @@ class PipelineSchema(NamedSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==PipelineSchema.id)" ), overlaps="visualizations", - cascade="all, delete-orphan", + cascade="delete", order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py index 6bf369d342..668b0ef04e 100644 --- a/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py @@ -229,7 +229,7 @@ class PipelineSnapshotSchema(BaseSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==PipelineSnapshotSchema.id)" ), overlaps="visualizations", - cascade="all, delete-orphan", + cascade="delete", order_by="CuratedVisualizationSchema.display_order", ), ) diff --git a/src/zenml/zen_stores/schemas/project_schemas.py b/src/zenml/zen_stores/schemas/project_schemas.py index e9cac308b6..a4fe796e1a 100644 --- a/src/zenml/zen_stores/schemas/project_schemas.py +++ b/src/zenml/zen_stores/schemas/project_schemas.py @@ -137,7 +137,7 @@ class ProjectSchema(NamedSchema, table=True): "foreign(CuratedVisualizationSchema.resource_id)==ProjectSchema.id)" ), overlaps="visualizations", - cascade="all, delete-orphan", + cascade="delete", order_by="CuratedVisualizationSchema.display_order", ), ) From 7ddce50b4a606ac274790ba884f570b5cb24172d Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Fri, 24 Oct 2025 07:20:01 +0100 Subject: [PATCH 60/64] fix failing test and apply michael last review --- docs/book/how-to/artifacts/visualizations.md | 2 -- src/zenml/models/v2/core/curated_visualization.py | 8 ++++---- .../zen_stores/schemas/artifact_visualization_schemas.py | 1 + 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/docs/book/how-to/artifacts/visualizations.md b/docs/book/how-to/artifacts/visualizations.md index 0c6c2c56c7..8934414508 100644 --- a/docs/book/how-to/artifacts/visualizations.md +++ b/docs/book/how-to/artifacts/visualizations.md @@ -236,8 +236,6 @@ visualization_c = client.create_curated_visualization( ) ``` -> **Note:** Reusing the same `display_order` for multiple curated visualizations targeting the same resource now raises an error. Pick distinct values—leaving gaps still helps when you later insert additional tiles. - #### RBAC visibility Curated visualizations respect the access permissions of the resource they're linked to. A user can only see a curated visualization if they have read access to the specific resource it targets. If a user lacks permission for the linked resource, the visualization will be hidden from their view. diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 5d7096ba2e..92c47f8989 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Optional from uuid import UUID -from pydantic import Field +from pydantic import Field, NonNegativeInt from zenml.enums import CuratedVisualizationSize, VisualizationResourceTypes from zenml.models.v2.base.base import BaseUpdate @@ -62,7 +62,7 @@ class CuratedVisualizationRequest(ProjectScopedRequest): default=None, title="The display name of the visualization.", ) - display_order: Optional[int] = Field( + display_order: Optional[NonNegativeInt] = Field( default=None, title="The display order of the visualization.", description=( @@ -102,7 +102,7 @@ class CuratedVisualizationUpdate(BaseUpdate): default=None, title="The new display name of the visualization.", ) - display_order: Optional[int] = Field( + display_order: Optional[NonNegativeInt] = Field( default=None, title="The new display order of the visualization.", description=( @@ -140,7 +140,7 @@ class CuratedVisualizationResponseBody(ProjectScopedResponseBody): default=None, title="The display name of the visualization.", ) - display_order: Optional[int] = Field( + display_order: Optional[NonNegativeInt] = Field( default=None, title="The display order of the visualization.", description=( diff --git a/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py b/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py index a010311b03..d7e1576f71 100644 --- a/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/artifact_visualization_schemas.py @@ -64,6 +64,7 @@ class ArtifactVisualizationSchema(BaseSchema, table=True): back_populates="artifact_visualization", sa_relationship_kwargs=dict( order_by="CuratedVisualizationSchema.display_order", + cascade="delete", ), ) From 0878e64c387830f79c829ae01daf86824aef5366 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Fri, 24 Oct 2025 07:57:50 +0100 Subject: [PATCH 61/64] update migration --- ...sualisations.py => 24552f3be1f2_add_visualisations.py} | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) rename src/zenml/zen_stores/migrations/versions/{54b2c0fbc0cb_add_visualisations.py => 24552f3be1f2_add_visualisations.py} (95%) diff --git a/src/zenml/zen_stores/migrations/versions/54b2c0fbc0cb_add_visualisations.py b/src/zenml/zen_stores/migrations/versions/24552f3be1f2_add_visualisations.py similarity index 95% rename from src/zenml/zen_stores/migrations/versions/54b2c0fbc0cb_add_visualisations.py rename to src/zenml/zen_stores/migrations/versions/24552f3be1f2_add_visualisations.py index a54738f07b..fb99d97e25 100644 --- a/src/zenml/zen_stores/migrations/versions/54b2c0fbc0cb_add_visualisations.py +++ b/src/zenml/zen_stores/migrations/versions/24552f3be1f2_add_visualisations.py @@ -1,8 +1,8 @@ -"""add visualisations [54b2c0fbc0cb]. +"""add visualisations [24552f3be1f2]. -Revision ID: 54b2c0fbc0cb +Revision ID: 24552f3be1f2 Revises: 124b57b8c7b1 -Create Date: 2025-10-23 15:41:57.315515 +Create Date: 2025-10-24 07:56:57.575675 """ @@ -11,7 +11,7 @@ from alembic import op # revision identifiers, used by Alembic. -revision = "54b2c0fbc0cb" +revision = "24552f3be1f2" down_revision = "124b57b8c7b1" branch_labels = None depends_on = None From 84096594ce845d1c977571522ffc10cecb567b2d Mon Sep 17 00:00:00 2001 From: Michael Schuster Date: Fri, 24 Oct 2025 15:41:12 +0800 Subject: [PATCH 62/64] Fix foreign key and some other minor changes --- .../models/v2/core/curated_visualization.py | 4 ++-- .../versions/24552f3be1f2_add_visualisations.py | 17 +---------------- .../schemas/curated_visualization_schemas.py | 9 ++------- src/zenml/zen_stores/schemas/schema_utils.py | 11 +++++++++-- 4 files changed, 14 insertions(+), 27 deletions(-) diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 92c47f8989..772d4a096d 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -128,7 +128,7 @@ class CuratedVisualizationResponseBody(ProjectScopedResponseBody): "Identifier of the artifact visualization that is curated for this resource." ), ) - artifact_version_id: Optional[UUID] = Field( + artifact_version_id: UUID = Field( default=None, title="The artifact version ID.", description=( @@ -208,7 +208,7 @@ def artifact_visualization_id(self) -> UUID: return self.get_body().artifact_visualization_id @property - def artifact_version_id(self) -> Optional[UUID]: + def artifact_version_id(self) -> UUID: """The artifact version ID. Returns: diff --git a/src/zenml/zen_stores/migrations/versions/24552f3be1f2_add_visualisations.py b/src/zenml/zen_stores/migrations/versions/24552f3be1f2_add_visualisations.py index fb99d97e25..071bdd50a1 100644 --- a/src/zenml/zen_stores/migrations/versions/24552f3be1f2_add_visualisations.py +++ b/src/zenml/zen_stores/migrations/versions/24552f3be1f2_add_visualisations.py @@ -45,7 +45,7 @@ def upgrade() -> None: sa.ForeignKeyConstraint( ["artifact_visualization_id"], ["artifact_visualization.id"], - name="fk_visualization_artifact_visualization_id_artifact_visualization", + name="fk_curated_visualization_artifact_visualization_id", ondelete="CASCADE", ), sa.ForeignKeyConstraint( @@ -62,14 +62,6 @@ def upgrade() -> None: name="unique_curated_visualization_resource_link", ), ) - with op.batch_alter_table( - "curated_visualization", schema=None - ) as batch_op: - batch_op.create_index( - "ix_curated_visualization_artifact_visualization_id", - ["artifact_visualization_id"], - unique=False, - ) # ### end Alembic commands ### @@ -77,12 +69,5 @@ def upgrade() -> None: 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( - "curated_visualization", schema=None - ) as batch_op: - batch_op.drop_index( - "ix_curated_visualization_artifact_visualization_id" - ) - op.drop_table("curated_visualization") # ### end Alembic commands ### diff --git a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py index 01b7c7f2a9..e2490fc971 100644 --- a/src/zenml/zen_stores/schemas/curated_visualization_schemas.py +++ b/src/zenml/zen_stores/schemas/curated_visualization_schemas.py @@ -34,7 +34,6 @@ from zenml.zen_stores.schemas.project_schemas import ProjectSchema from zenml.zen_stores.schemas.schema_utils import ( build_foreign_key_field, - build_index, ) from zenml.zen_stores.schemas.utils import jl_arg @@ -49,7 +48,6 @@ class CuratedVisualizationSchema(BaseSchema, table=True): __tablename__ = "curated_visualization" __table_args__ = ( - build_index(__tablename__, ["artifact_visualization_id"]), UniqueConstraint( "artifact_visualization_id", "resource_id", @@ -73,6 +71,7 @@ class CuratedVisualizationSchema(BaseSchema, table=True): target_column="id", ondelete="CASCADE", nullable=False, + custom_constraint_name="fk_curated_visualization_artifact_visualization_id", ) display_name: Optional[str] = Field(default=None) @@ -188,11 +187,7 @@ def to_model( except ValueError: resource_type_enum = VisualizationResourceTypes.PROJECT - artifact_version_id: Optional[UUID] = None - if self.artifact_visualization is not None: - artifact_version_id = ( - self.artifact_visualization.artifact_version_id - ) + artifact_version_id = self.artifact_visualization.artifact_version_id body = CuratedVisualizationResponseBody( project_id=self.project_id, diff --git a/src/zenml/zen_stores/schemas/schema_utils.py b/src/zenml/zen_stores/schemas/schema_utils.py index 4925685442..5b85e64889 100644 --- a/src/zenml/zen_stores/schemas/schema_utils.py +++ b/src/zenml/zen_stores/schemas/schema_utils.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Utility functions for SQLModel schemas.""" -from typing import Any, List +from typing import Any, List, Optional from sqlalchemy import Column, ForeignKey, Index from sqlmodel import Field @@ -45,6 +45,7 @@ def build_foreign_key_field( target_column: str, ondelete: str, nullable: bool, + custom_constraint_name: Optional[str] = None, **sa_column_kwargs: Any, ) -> Any: """Build a SQLModel foreign key field. @@ -56,6 +57,7 @@ def build_foreign_key_field( target_column: Target column name. ondelete: On delete behavior. nullable: Whether the field is nullable. + custom_constraint_name: Custom name for the foreign key constraint. **sa_column_kwargs: Keyword arguments for the SQLAlchemy column. Returns: @@ -68,11 +70,16 @@ def build_foreign_key_field( raise ValueError( "Cannot set ondelete to SET NULL if the field is not nullable." ) - constraint_name = foreign_key_constraint_name( + constraint_name = custom_constraint_name or foreign_key_constraint_name( source=source, target=target, source_column=source_column, ) + if len(constraint_name) > 64: + raise ValueError( + f"Foreign key constraint name {constraint_name} is too long. " + "The maximum length is 64 characters." + ) return Field( sa_column=Column( ForeignKey( From f41c7dfec2de11212e6720cf74e454442dbaaf0f Mon Sep 17 00:00:00 2001 From: Michael Schuster Date: Fri, 24 Oct 2025 15:42:14 +0800 Subject: [PATCH 63/64] Docstring --- src/zenml/zen_stores/schemas/schema_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/zenml/zen_stores/schemas/schema_utils.py b/src/zenml/zen_stores/schemas/schema_utils.py index 5b85e64889..c179f5543d 100644 --- a/src/zenml/zen_stores/schemas/schema_utils.py +++ b/src/zenml/zen_stores/schemas/schema_utils.py @@ -65,6 +65,7 @@ def build_foreign_key_field( Raises: ValueError: If the ondelete and nullable arguments are not compatible. + ValueError: If the foreign key constraint name is too long. """ if not nullable and ondelete == "SET NULL": raise ValueError( From 3c57c83ec3ed8998f92ddad590e01eed9b3dc135 Mon Sep 17 00:00:00 2001 From: Michael Schuster Date: Fri, 24 Oct 2025 15:51:59 +0800 Subject: [PATCH 64/64] Remove wrong default --- src/zenml/models/v2/core/curated_visualization.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/zenml/models/v2/core/curated_visualization.py b/src/zenml/models/v2/core/curated_visualization.py index 772d4a096d..6da41b4eaa 100644 --- a/src/zenml/models/v2/core/curated_visualization.py +++ b/src/zenml/models/v2/core/curated_visualization.py @@ -129,7 +129,6 @@ class CuratedVisualizationResponseBody(ProjectScopedResponseBody): ), ) artifact_version_id: UUID = Field( - default=None, title="The artifact version ID.", description=( "Identifier of the artifact version that owns the curated visualization. "