Compute cascading tags server-side (#3781)

* Compute cascading tags server-side

* Linting

* More linting

Author: schustmi

src/zenml/orchestrators/step_launcher.py

@@ -292,10 +292,6 @@ def _bypass() -> None:
                             artifacts=step_run.outputs,
                             model_version=model_version,
                         )
-                    step_run_utils.cascade_tags_for_output_artifacts(
-                        artifacts=step_run.outputs,
-                        tags=pipeline_run.config.tags,
-                    )
 
         except:  # noqa: E722
             logger.error(f"Pipeline run `{pipeline_run.name}` failed.")

src/zenml/orchestrators/step_run_utils.py

@@ -14,9 +14,8 @@
 """Utilities for creating step runs."""
 
 import json
-from typing import Dict, List, Optional, Set, Tuple, Union
+from typing import Dict, List, Optional, Set, Tuple
 
-from zenml import Tag, add_tags
 from zenml.client import Client
 from zenml.config.step_configurations import Step
 from zenml.constants import CODE_HASH_PARAMETER_NAME, TEXT_FIELD_MAX_LENGTH
@@ -404,29 +403,6 @@ def link_output_artifacts_to_model_version(
             )
 
 
-def cascade_tags_for_output_artifacts(
-    artifacts: Dict[str, List[ArtifactVersionResponse]],
-    tags: Optional[List[Union[str, Tag]]] = None,
-) -> None:
-    """Tag the outputs of a step run.
-
-    Args:
-        artifacts: The step output artifacts.
-        tags: The tags to add to the artifacts.
-    """
-    if tags is None:
-        return
-
-    cascade_tags = [t for t in tags if isinstance(t, Tag) and t.cascade]
-
-    for output_artifacts in artifacts.values():
-        for output_artifact in output_artifacts:
-            add_tags(
-                tags=[t.name for t in cascade_tags],
-                artifact_version_id=output_artifact.id,
-            )
-
-
 def publish_cached_step_run(
     request: "StepRunRequest", pipeline_run: "PipelineRunResponse"
 ) -> "StepRunResponse":
@@ -447,11 +423,6 @@ def publish_cached_step_run(
             model_version=model_version,
         )
 
-    cascade_tags_for_output_artifacts(
-        artifacts=step_run.outputs,
-        tags=pipeline_run.config.tags,
-    )
-
     return step_run
 
 

src/zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -334,6 +334,29 @@ def from_request(
             trigger_execution_id=request.trigger_execution_id,
         )
 
+    def get_pipeline_configuration(self) -> PipelineConfiguration:
+        """Get the pipeline configuration for the pipeline run.
+
+        Raises:
+            RuntimeError: if the pipeline run has no deployment and no pipeline
+                configuration.
+
+        Returns:
+            The pipeline configuration.
+        """
+        if self.deployment:
+            return PipelineConfiguration.model_validate_json(
+                self.deployment.pipeline_configuration
+            )
+        elif self.pipeline_configuration:
+            return PipelineConfiguration.model_validate_json(
+                self.pipeline_configuration
+            )
+        else:
+            raise RuntimeError(
+                "Pipeline run has no deployment and no pipeline configuration."
+            )
+
     def fetch_metadata_collection(
         self, include_full_metadata: bool = False, **kwargs: Any
     ) -> Dict[str, List[RunMetadataEntry]]:

src/zenml/zen_stores/sql_zen_store.py

@@ -8835,7 +8835,6 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
                 reference_id=step_run.pipeline_run_id,
                 session=session,
             )
-
             self._get_reference_schema_by_id(
                 resource=step_run,
                 reference_schema=StepRunSchema,
@@ -8911,15 +8910,41 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
                     )
                     for link in original_metadata_links
                 ]
-                # Add all new links in a single operation
-                session.add_all(new_links)
-                # Commit the changes
-                session.commit()
-                session.refresh(step_schema)
 
-            session.commit()
-            session.refresh(step_schema)
+                if new_links:
+                    session.add_all(new_links)
+                    session.commit()
+                    session.refresh(step_schema, ["run_metadata"])
 
+            if step_run.status == ExecutionStatus.CACHED:
+                from zenml.utils.tag_utils import Tag
+
+                cascading_tags = [
+                    tag
+                    for tag in run.get_pipeline_configuration().tags or []
+                    if isinstance(tag, Tag) and tag.cascade
+                ]
+
+                if cascading_tags:
+                    output_artifact_ids = [
+                        id for ids in step_run.outputs.values() for id in ids
+                    ]
+                    output_artifacts = list(
+                        session.exec(
+                            select(ArtifactVersionSchema).where(
+                                col(ArtifactVersionSchema.id).in_(
+                                    output_artifact_ids
+                                )
+                            )
+                        ).all()
+                    )
+                    self._attach_tags_to_resources(
+                        cascading_tags,
+                        resources=output_artifacts,
+                        session=session,
+                    )
+
+            session.commit()
             step_model = step_schema.to_model(include_metadata=True)
 
             for upstream_step in step_model.spec.upstream_steps:
@@ -8977,7 +9002,9 @@ def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
                 )
 
             session.commit()
-            session.refresh(step_schema)
+            session.refresh(
+                step_schema, ["input_artifacts", "output_artifacts"]
+            )
 
             if model_version_id := self._get_or_create_model_version_for_run(
                 step_schema
@@ -9330,6 +9357,9 @@ def _update_pipeline_run_status(
         """
         from zenml.orchestrators.publish_utils import get_pipeline_run_status
 
+        # Make sure we start with a fresh transaction before locking the
+        # pipeline run
+        session.commit()
         pipeline_run = session.exec(
             select(PipelineRunSchema)
             .with_for_update()
@@ -12253,7 +12283,7 @@ def _get_tag_schema(
     def _attach_tags_to_resources(
         self,
         tags: Optional[Sequence[Union[str, tag_utils.Tag]]],
-        resources: Union[BaseSchema, List[BaseSchema]],
+        resources: Union[BaseSchema, Sequence[BaseSchema]],
         session: Session,
     ) -> None:
         """Attaches multiple tags to multiple resources.