Prune of unused artifacts links via client

src/zenml/client.py

@@ -2786,6 +2786,43 @@ def delete_artifact(
         self.zen_store.delete_artifact(artifact_id=artifact.id)
         logger.info(f"Deleted artifact '{artifact.name}'.")
 
+    def prune_artifacts(
+        self, only_artifact: bool = False, only_metadata: bool = False
+    ) -> None:
+        """Delete all unused artifacts and artifact versions.
+
+        Args:
+            only_artifact: Only delete artifacts
+            only_metadata: Only delete artifact metadata
+
+        Raises:
+            RuntimeError: If any artifact cannot be deleted.
+        """
+        unused_artifact_versions = depaginate(
+            partial(self.list_artifact_versions, only_unused=True)
+        )
+
+        if not unused_artifact_versions:
+            logger.info("No unused artifact versions found.")
+            return
+
+        for unused_artifact_version in unused_artifact_versions:
+            try:
+                self.delete_artifact_version(
+                    name_id_or_prefix=unused_artifact_version.id,
+                    delete_metadata=not only_artifact,
+                    delete_from_artifact_store=not only_metadata,
+                )
+                unused_artifact = unused_artifact_version.artifact
+                if not unused_artifact.versions and not only_artifact:
+                    self.delete_artifact(unused_artifact.id)
+
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to prune artifacts, error deleting artifact `{unused_artifact_version.id}`, error `{e}`."
+                )
+        logger.info("All unused artifacts and artifact versions deleted.")
+
     # --------------------------- Artifact Versions ---------------------------
 
     def get_artifact_version(

src/zenml/zen_stores/migrations/versions/4e1972485075_endpoint_artifact_deployment_artifact.py

@@ -1,7 +1,7 @@
 """endpoint_artifact>deployment_artifact [4e1972485075].
 
 Revision ID: 4e1972485075
-Revises: 0.52.0
+Revises: 5cc3f41cf048
 Create Date: 2023-12-12 10:51:44.177810
 
 """
@@ -10,7 +10,7 @@
 
 # revision identifiers, used by Alembic.
 revision = "4e1972485075"
-down_revision = "0.52.0"
+down_revision = "5cc3f41cf048"
 branch_labels = None
 depends_on = None
 

tests/integration/functional/test_client.py

@@ -27,6 +27,7 @@
     int_plus_one_test_step,
 )
 from tests.integration.functional.utils import sample_name
+from zenml import ExternalArtifact
 from zenml.client import Client
 from zenml.config.pipeline_spec import PipelineSpec
 from zenml.config.source import Source
@@ -1128,6 +1129,37 @@ def test_basic_crud_for_entity(
             pass
 
 
+class TestArtifact:
+    def test_prune(self, clean_client: "Client"):
+        """Test that artifact pruning works."""
+        artifact_id = ExternalArtifact(value="foo").upload_by_value()
+        artifact = clean_client.get_artifact_version(artifact_id)
+        assert artifact is not None
+        clean_client.prune_artifacts()
+        with pytest.raises(KeyError):
+            clean_client.get_artifact_version(artifact_id)
+        assert not os.path.exists(artifact.uri)
+
+    def test_prune_only_metadata(self, clean_client: "Client"):
+        """Test that artifact pruning works with only metadata flag."""
+        artifact_id = ExternalArtifact(value="foo").upload_by_value()
+        artifact = clean_client.get_artifact_version(artifact_id)
+        assert artifact is not None
+        clean_client.prune_artifacts(only_metadata=True)
+        with pytest.raises(KeyError):
+            clean_client.get_artifact_version(artifact_id)
+        assert os.path.exists(artifact.uri)
+
+    def test_prune_only_artifact(self, clean_client: "Client"):
+        """Test that artifact pruning works with only artifacts flag."""
+        artifact_id = ExternalArtifact(value="foo").upload_by_value()
+        artifact = clean_client.get_artifact_version(artifact_id)
+        assert artifact is not None
+        clean_client.prune_artifacts(only_artifact=True)
+        assert clean_client.get_artifact_version(artifact_id).id == artifact.id
+        assert not os.path.exists(artifact.uri)
+
+
 class TestModel:
     MODEL_NAME = "foo"