Deployment custom visualizations

src/zenml/client.py

@@ -110,6 +110,10 @@
     ComponentUpdate,
     DeploymentFilter,
     DeploymentResponse,
+    DeploymentVisualizationFilter,
+    DeploymentVisualizationRequest,
+    DeploymentVisualizationResponse,
+    DeploymentVisualizationUpdate,
     EventSourceFilter,
     EventSourceRequest,
     EventSourceResponse,
@@ -3735,6 +3739,127 @@ 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,
+        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.
+            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,
+        )
+
+        # 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,
+        )
+
+    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,
+            visualization_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",

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"

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",

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.