Utilizing cascading tags for cached step runs (#3655)

* cascading tags for cached step runs

* moved tags out

* fixed the tests

* another pipeline

Author: bcdurak

src/zenml/orchestrators/step_launcher.py

@@ -292,6 +292,10 @@ 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

@@ -13,8 +13,9 @@
 #  permissions and limitations under the License.
 """Utilities for creating step runs."""
 
-from typing import Dict, List, Optional, Set, Tuple
+from typing import Dict, List, Optional, Set, Tuple, Union
 
+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
@@ -333,6 +334,11 @@ def create_cached_step_runs(
                     model_version=model_version,
                 )
 
+            cascade_tags_for_output_artifacts(
+                artifacts=step_run.outputs,
+                tags=pipeline_run.config.tags,
+            )
+
             logger.info("Using cached version of step `%s`.", invocation_id)
             cached_invocations.add(invocation_id)
 
@@ -382,3 +388,26 @@ def link_output_artifacts_to_model_version(
                 artifact_version=output_artifact,
                 model_version=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,
+            )

src/zenml/utils/tag_utils.py

@@ -348,10 +348,10 @@ def add_tags(
             if isinstance(tag, Tag):
                 tag_model = client.get_tag(tag.name)
 
-                if tag.exclusive != tag_model.exclusive:
+                if bool(tag.exclusive) != tag_model.exclusive:
                     raise ValueError(
-                        f"The tag `{tag.name}` is an "
-                        f"{'exclusive' if tag_model.exclusive else 'non-exclusive'} "
+                        f"The tag `{tag.name}` is "
+                        f"{'an exclusive' if tag_model.exclusive else 'a non-exclusive'} "
                         "tag. Please update it before attaching it to a resource."
                     )
                 if tag.cascade is not None:

src/zenml/zen_stores/sql_zen_store.py

@@ -11354,13 +11354,10 @@ def _attach_tags_to_resources(
             except EntityExistsError:
                 if isinstance(tag, tag_utils.Tag):
                     tag_schema = self._get_tag_schema(tag.name, session)
-                    if (
-                        tag.exclusive is not None
-                        and tag.exclusive != tag_schema.exclusive
-                    ):
+                    if bool(tag.exclusive) != tag_schema.exclusive:
                         raise ValueError(
-                            f"Tag `{tag_schema.name}` has been defined as a "
-                            f"{'exclusive' if tag_schema.exclusive else 'non-exclusive'} "
+                            f"Tag `{tag_schema.name}` has been defined as "
+                            f"{'an exclusive' if tag_schema.exclusive else 'a non-exclusive'} "
                             "tag. Please update it before attaching it to resources."
                         )
                 else:

tests/integration/functional/utils/test_tag_utils.py

@@ -13,11 +13,14 @@
 #  permissions and limitations under the License.
 
 
+import os
 from typing import Annotated, Tuple
 
 import pytest
 
 from zenml import ArtifactConfig, Tag, add_tags, pipeline, remove_tags, step
+from zenml.constants import ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING
+from zenml.enums import ExecutionStatus
 
 
 @step
@@ -139,3 +142,101 @@ def test_tag_utils(clean_client):
         clean_client.update_tag(
             tag_name_or_id=non_exclusive_tag.id, exclusive=True
         )
+
+
+@pipeline(
+    tags=[Tag(name="cascade_tag", cascade=True)],
+    enable_cache=False,
+)
+def pipeline_with_cascade_tag():
+    """Pipeline definition to test the tag utils."""
+    _ = step_single_output()
+
+
+def test_cascade_tags_for_output_artifacts_of_cached_pipeline_run(
+    clean_client,
+):
+    """Test that the cascade tags are added to the output artifacts of a cached step."""
+    # Run the pipeline once without caching
+    pipeline_with_cascade_tag()
+
+    pipeline_runs = clean_client.list_pipeline_runs(sort_by="created")
+    assert len(pipeline_runs.items) == 1
+    assert (
+        pipeline_runs.items[0].steps["step_single_output"].status
+        == ExecutionStatus.COMPLETED
+    )
+    assert "cascade_tag" in [
+        t.name
+        for t in pipeline_runs.items[0]
+        .steps["step_single_output"]
+        .outputs["single"][0]
+        .tags
+    ]
+
+    # Run it once again with caching
+    pipeline_with_cascade_tag.configure(enable_cache=True)
+    pipeline_with_cascade_tag()
+    pipeline_runs = clean_client.list_pipeline_runs(sort_by="created")
+    assert len(pipeline_runs.items) == 2
+    assert (
+        pipeline_runs.items[1].steps["step_single_output"].status
+        == ExecutionStatus.CACHED
+    )
+
+    # Run it once again with caching and a new cascade tag
+    pipeline_with_cascade_tag.configure(
+        tags=[Tag(name="second_cascade_tag", cascade=True)]
+    )
+    pipeline_with_cascade_tag()
+    pipeline_runs = clean_client.list_pipeline_runs(sort_by="created")
+    assert len(pipeline_runs.items) == 3
+    assert (
+        pipeline_runs.items[2].steps["step_single_output"].status
+        == ExecutionStatus.CACHED
+    )
+
+    assert "second_cascade_tag" in [
+        t.name
+        for t in pipeline_runs.items[0]
+        .steps["step_single_output"]
+        .outputs["single"][0]
+        .tags
+    ]
+    assert "second_cascade_tag" in [
+        t.name
+        for t in pipeline_runs.items[2]
+        .steps["step_single_output"]
+        .outputs["single"][0]
+        .tags
+    ]
+
+    # Run it once again with caching (preventing client side caching) and a new cascade tag
+    os.environ[ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING] = "true"
+    pipeline_with_cascade_tag.configure(
+        tags=[Tag(name="third_cascade_tag", cascade=True)]
+    )
+    pipeline_with_cascade_tag()
+
+    pipeline_runs = clean_client.list_pipeline_runs(sort_by="created")
+    assert len(pipeline_runs.items) == 4
+    assert (
+        pipeline_runs.items[3].steps["step_single_output"].status
+        == ExecutionStatus.CACHED
+    )
+
+    assert "third_cascade_tag" in [
+        t.name
+        for t in pipeline_runs.items[0]
+        .steps["step_single_output"]
+        .outputs["single"][0]
+        .tags
+    ]
+    assert "third_cascade_tag" in [
+        t.name
+        for t in pipeline_runs.items[3]
+        .steps["step_single_output"]
+        .outputs["single"][0]
+        .tags
+    ]
+    os.environ.pop(ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING, None)