Add Delete runs API

mlflow/protos/service.proto

@@ -836,6 +836,15 @@ message DeleteRun {
   message Response {}
 }
 
+message DeleteRuns {
+  option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";
+
+  // IDs of the runs to delete.
+  repeated string run_id = 1 [(validate_required) = true];
+
+  message Response {}
+}
+
 message RestoreRun {
   option (scalapb.message).extends = "com.databricks.rpc.RPC[$this.Response]";
 

mlflow/server/handlers.py

@@ -67,6 +67,7 @@
     CreateRun,
     DeleteExperiment,
     DeleteRun,
+    DeleteRuns,
     DeleteTag,
     GetExperiment,
     GetExperimentByName,
@@ -312,6 +313,12 @@ def _assert_string(x):
     assert isinstance(x, str)
 
 
+def _assert_list_of_strings(x):
+    assert isinstance(x, list)
+    for item in x:
+        _assert_string(item)
+
+
 def _assert_intlike(x):
     try:
         x = int(x)
@@ -765,6 +772,18 @@ def _delete_run():
     return response
 
 
+@catch_mlflow_exception
+@_disable_if_artifacts_only
+def _delete_runs():
+    request_message = _get_request_message(
+        DeleteRuns(), schema={"run_id": [_assert_required, _assert_list_of_strings]}
+    )
+    _get_tracking_store().delete_runs(request_message.run_id)
+    response_message = DeleteRuns.Response()
+    response = Response(mimetype="application/json")
+    response.set_data(message_to_json(response_message))
+
+
 @catch_mlflow_exception
 @_disable_if_artifacts_only
 def _restore_run():
@@ -2292,6 +2311,7 @@ def get_endpoints(get_handler=get_handler):
     CreateRun: _create_run,
     UpdateRun: _update_run,
     DeleteRun: _delete_run,
+    DeleteRuns: _delete_runs,
     RestoreRun: _restore_run,
     LogParam: _log_param,
     LogMetric: _log_metric,

mlflow/store/tracking/abstract_store.py

@@ -1,4 +1,4 @@
 from abc import ABCMeta, abstractmethod
 from typing import List, Optional
 
 from mlflow.entities import DatasetInput, ViewType
@@ -223,6 +223,17 @@
         """
         pass
 
+    @abstractmethod
+    def delete_runs(self, run_ids):
+        """
+        Delete multiple runs.
+
+        Args:
+            run_ids: List of run ids to delete.
+
+        """
+        pass
+
     @abstractmethod
     def restore_run(self, run_id):
         """

mlflow/store/tracking/file_store.py

@@ -1,4 +1,4 @@
 import json
 import logging
 import os
 import shutil
@@ -531,6 +531,19 @@
         new_info = run_info._copy_with_overrides(lifecycle_stage=LifecycleStage.DELETED)
         self._overwrite_run_info(new_info, deleted_time=get_current_time_millis())
 
+    def delete_runs(self, run_ids):
+        for run_id in run_ids:
+            try:
+                run_info = self._get_run_info(run_id)
+                if run_info is None:
+                    raise MlflowException(
+                        f"Run '{run_id}' metadata is in invalid state.", databricks_pb2.INVALID_STATE
+                    )
+                new_info = run_info._copy_with_overrides(lifecycle_stage=LifecycleStage.DELETED)
+                self._overwrite_run_info(new_info, deleted_time=get_current_time_millis())
+            except Exception as e:
+                logging.error(f"Failed to delete run '{run_id}': {str(e)!s}")
+
     def _hard_delete_run(self, run_id):
         """
         Permanently delete a run (metadata and metrics, tags, parameters).

mlflow/store/tracking/rest_store.py

@@ -8,6 +8,7 @@
     CreateRun,
     DeleteExperiment,
     DeleteRun,
+    DeleteRuns,
     DeleteTag,
     GetExperiment,
     GetExperimentByName,
@@ -313,6 +314,10 @@ def delete_run(self, run_id):
         req_body = message_to_json(DeleteRun(run_id=run_id))
         self._call_endpoint(DeleteRun, req_body)
 
+    def delete_runs(self, run_ids):
+        req_body = message_to_json(DeleteRuns(run_ids=run_ids))
+        self._call_endpoint(DeleteRuns, req_body)
+
     def restore_run(self, run_id):
         req_body = message_to_json(RestoreRun(run_id=run_id))
         self._call_endpoint(RestoreRun, req_body)

mlflow/store/tracking/sqlalchemy_store.py

@@ -10,7 +10,7 @@
 
 import sqlalchemy
 import sqlalchemy.sql.expression as sql
-from sqlalchemy import and_, func, sql, text
+from sqlalchemy import and_, func, or_, text
 from sqlalchemy.future import select
 
 import mlflow.store.db.utils
@@ -655,6 +655,15 @@
             run.deleted_time = get_current_time_millis()
             session.add(run)
 
+    def delete_runs(self, run_ids):
+        with self.ManagedSessionMaker() as session:
+            conditions = SqlRun.run_uuid.in_(run_ids)
+            runs = session.query(SqlRun).filter(conditions).all()
+            for run in runs:
+                run.lifecycle_stage = LifecycleStage.DELETED
+                run.deleted_time = get_current_time_millis()
+                session.add(run)
+
     def _hard_delete_run(self, run_id):
         """
         Permanently delete a run (metadata and metrics, tags, parameters).

mlflow/tracking/_tracking_service/client.py

@@ -624,6 +624,12 @@ def delete_run(self, run_id):
         """
         self.store.delete_run(run_id)
 
+    def delete_runs(self, run_ids):
+        """
+        Deletes runs with the given IDs.
+        """
+        self.store.delete_runs(run_ids)
+
     def restore_run(self, run_id):
         """
         Restores a deleted run with the given ID.

mlflow/tracking/client.py

@@ -2099,9 +2099,42 @@
             run_id: a61c7a1851324f7094e8d5014c58c8c8; lifecycle_stage: active
             run_id: a61c7a1851324f7094e8d5014c58c8c8; lifecycle_stage: deleted
 
         """
         self._tracking_client.delete_run(run_id)
 
+    def delete_runs(self, run_ids: List[str]) -> None:
+        """Deletes runs with the given IDs.
+
+        Args:
+            run_ids: List of unique run ids to delete.
+
+        .. code-block:: python
+            :caption: Example
+
+            from mlflow import MlflowClient
+
+            # Create runs under the default experiment (whose id is '0').
+            client = MlflowClient()
+            experiment_id = "0"
+            run1 = client.create_run(experiment_id)
+            run2 = client.create_run(experiment_id)
+            run_ids = [run1.info.run_id, run2.info.run_id]
+            print(f"run_ids: {run_ids}; lifecycle_stage: {[run1.info.lifecycle_stage, run2.info.lifecycle_stage]}")
+            print("--")
+            client.delete_runs(run_ids)
+            del_run1 = client.get_run(run_ids[0])
+            del_run2 = client.get_run(run_ids[1])
+            print(f"run_ids: {run_ids}; lifecycle_stage: {[del_run1.info.lifecycle_stage, del_run2.info.lifecycle_stage]}")
+
+        .. code-block:: text
+            :caption: Output
+
+            run_ids: ['a61c7a1851324f7094e8d5014c58c8c8', 'b71d7a1851324f7094e8d5014c58c8c9']; lifecycle_stage: ['active', 'active']
+            run_ids: ['a61c7a1851324f7094e8d5014c58c8c8', 'b71d7a1851324f7094e8d5014c58c8c9']; lifecycle_stage: ['deleted', 'deleted']
+
+        """
+        self._tracking_client.delete_runs(run_ids)
+
     def restore_run(self, run_id: str) -> None:
         """Restores a deleted run with the given ID.
 

mlflow/tracking/fluent.py

@@ -1755,6 +1755,37 @@ def delete_run(run_id: str) -> None:
     """
     MlflowClient().delete_run(run_id)
 
+def delete_runs(run_ids: List[str]) -> None:
+    """
+    Deletes runs with the given IDs.
+
+    Args:
+        run_ids: List of unique run ids to delete.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+
+        run_ids = []
+        for _ in range(2):
+            with mlflow.start_run() as run:
+                mlflow.log_param("p", 0)
+                run_ids.append(run.info.run_id)
+
+        mlflow.delete_runs(run_ids)
+
+        lifecycle_stages = [mlflow.get_run(run_id).info.lifecycle_stage for run_id in run_ids]
+        print(f"run_ids: {run_ids}; lifecycle_stages: {lifecycle_stages}")
+
+    .. code-block:: text
+        :caption: Output
+
+        run_ids = ['45f4af3e6fd349e58579b27fcb0b8277', 'b71d7a1851324f7094e8d5014c58c8c9']
+        lifecycle_stages = ['deleted', 'deleted']
+
+    """
+    MlflowClient().delete_runs(run_ids)
 
 def get_artifact_uri(artifact_path: Optional[str] = None) -> str:
     """

tests/store/tracking/test_sqlalchemy_store.py

@@ -892,6 +892,24 @@ def test_delete_run(store: SqlAlchemyStore):
         assert actual.run_uuid == deleted_run.info.run_id
 
 
+def test_delete_runs(store: SqlAlchemyStore):
+    runs = [_run_factory(store) for _ in range(2)]
+    run_ids = [run.info.run_id for run in runs]
+
+    store.delete_runs(run_ids)
+
+    with store.ManagedSessionMaker() as session:
+        for run_id in run_ids:
+            actual = session.query(models.SqlRun).filter_by(run_uuid=run_id).first()
+            assert actual.lifecycle_stage == entities.LifecycleStage.DELETED
+            assert (
+                actual.deleted_time is not None
+            )  # deleted time should be updated and thus not None anymore
+
+            deleted_run = store.get_run(run_id)
+            assert actual.run_uuid == deleted_run.info.run_id
+
+
 def test_hard_delete_run(store: SqlAlchemyStore):
     run = _run_factory(store)
     metric = entities.Metric("blahmetric", 100.0, get_current_time_millis(), 0)
@@ -924,6 +942,17 @@ def test_get_deleted_runs(store: SqlAlchemyStore):
     assert deleted_run_ids == [run.info.run_uuid]
 
 
+def test_get_deleted_runs_with_delete_runs(store: SqlAlchemyStore):
+    runs = [_run_factory(store) for _ in range(2)]
+    run_ids = [run.info.run_id for run in runs]
+    deleted_run_ids = store._get_deleted_runs()
+    assert deleted_run_ids == []
+
+    store.delete_runs(run_ids)
+    deleted_run_ids = store._get_deleted_runs()
+    assert set(deleted_run_ids) == set(run_ids)
+
+
 def test_log_metric(store: SqlAlchemyStore):
     run = _run_factory(store)
 

tests/tracking/test_tracking.py

@@ -684,7 +684,18 @@
             pass
     assert mlflow.active_run() is None
 
-
+def test_start_deleted_runs():
+    run_ids = []
+    for _ in range(2):
+        with mlflow.start_run() as active_run:
+            run_ids.append(active_run.info.run_id)
+    tracking.MlflowClient().delete_runs(run_ids)
+    for run_id in run_ids:
+        with pytest.raises(MlflowException, match="because it is in the deleted state."):
+            with mlflow.start_run(run_id=run_id):
+                pass
+    assert mlflow.active_run() is None
+
 @pytest.mark.usefixtures("reset_active_experiment")
 def test_start_run_exp_id_0():
     mlflow.set_experiment("some-experiment")