Fix #10593 -- add --keep option for dvc experiments remove (#10633)

* Add keep_selected parameter, and corresponding code to keep only the selected exps (and remove all the other ones)

* test keep_selected_by_name

* test keep_selected_by_rev

* test keep_selected multiple, by name

* test keep all by name

* test keep by rev, with num=2

* added option to cli

* refactoring to meet pr needs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fixed test_experiments to add keep_selected=False to remove tests

* rename parameter to match cli option

* follow the normal path, then invert the selection before removing

* fixed tests for list ordering + fixed test with non existent name, it didn't make sense to delete everything if an exp name did not exist

* changed cli option comment

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fixed typing issue

* updated parameter name

* removed handling queued experiments (since --queue would remove them all)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* code simplification, added __eq__ and __hash__ to be able to compare ExpRefs, updated and parametrized tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fixed linting issues

* - --keep and --queue together raise an InvalidArgumentError
    - added a test to check if the error is raised
    - fixed CLI message

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* re-run gh tests. Some tests which did not involve my changes started failing while they were passing fine before.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

tests/func/experiments/test_remove.py

@@ -179,3 +179,85 @@ def test_remove_multi_rev(tmp_dir, scm, dvc, exp_stage):
 
     assert scm.get_ref(str(baseline_exp_ref)) is None
     assert scm.get_ref(str(new_exp_ref)) is None
+
+
+@pytest.mark.parametrize(
+    "keep, expected_removed",
+    [
+        [["exp1"], ["exp2", "exp3"]],
+        [["exp1", "exp2"], ["exp3"]],
+        [["exp1", "exp2", "exp3"], []],
+        [[], []],  # remove does nothing if no experiments are specified
+    ],
+)
+def test_keep_selected_by_name(tmp_dir, scm, dvc, exp_stage, keep, expected_removed):
+    # Setup: Run experiments
+    refs = {}
+    for i in range(1, len(keep) + len(expected_removed) + 1):
+        results = dvc.experiments.run(
+            exp_stage.addressing, params=[f"foo={i}"], name=f"exp{i}"
+        )
+        refs[f"exp{i}"] = first(exp_refs_by_rev(scm, first(results)))
+        assert scm.get_ref(str(refs[f"exp{i}"])) is not None
+
+    removed = dvc.experiments.remove(exp_names=keep, keep=True)
+    assert sorted(removed) == sorted(expected_removed)
+
+    for exp in expected_removed:
+        assert scm.get_ref(str(refs[exp])) is None
+
+    for exp in keep:
+        assert scm.get_ref(str(refs[exp])) is not None
+
+
+def test_keep_selected_by_nonexistent_name(tmp_dir, scm, dvc, exp_stage):
+    # non existent name should raise an error
+    with pytest.raises(UnresolvedExpNamesError):
+        dvc.experiments.remove(exp_names=["nonexistent"], keep=True)
+
+
+@pytest.mark.parametrize(
+    "num_exps, rev, num, expected_removed",
+    [
+        [2, "exp1", 1, ["exp2"]],
+        [3, "exp3", 1, ["exp1", "exp2"]],
+        [3, "exp3", 2, ["exp1"]],
+        [3, "exp3", 3, []],
+        [3, "exp2", 2, ["exp3"]],
+        [4, "exp2", 2, ["exp3", "exp4"]],
+        [4, "exp4", 2, ["exp1", "exp2"]],
+        [1, None, 1, []],  # remove does nothing if no experiments are specified
+    ],
+)
+def test_keep_selected_by_rev(
+    tmp_dir, scm, dvc, exp_stage, num_exps, rev, num, expected_removed
+):
+    refs = {}
+    revs = {}
+    # Setup: Run experiments and commit
+    for i in range(1, num_exps + 1):
+        scm.commit(f"commit{i}")
+        results = dvc.experiments.run(
+            exp_stage.addressing, params=[f"foo={i}"], name=f"exp{i}"
+        )
+        refs[f"exp{i}"] = first(exp_refs_by_rev(scm, first(results)))
+        revs[f"exp{i}"] = scm.get_rev()
+        assert scm.get_ref(str(refs[f"exp{i}"])) is not None
+
+    # Keep the experiment from the new revision
+    removed = dvc.experiments.remove(rev=revs.get(rev), num=num, keep=True)
+    assert sorted(removed) == sorted(expected_removed)
+
+    # Check remaining experiments
+    for exp in expected_removed:
+        assert scm.get_ref(str(refs[exp])) is None
+
+    for exp, ref in refs.items():
+        if exp not in expected_removed:
+            assert scm.get_ref(str(ref)) is not None
+
+
+def test_remove_with_queue_and_keep(tmp_dir, scm, dvc, exp_stage):
+    # This should raise an exception, until decided otherwise
+    with pytest.raises(InvalidArgumentError):
+        dvc.experiments.remove(queue=True, keep=True)

tests/unit/command/test_experiments.py

@@ -384,6 +384,7 @@ def test_experiments_remove_flag(dvc, scm, mocker, capsys, caplog):
         num=2,
         queue=False,
         git_remote="myremote",
+        keep=False,
     )
 
 
@@ -410,6 +411,7 @@ def test_experiments_remove_special(dvc, scm, mocker, capsys, caplog):
         num=1,
         queue=False,
         git_remote="myremote",
+        keep=False,
     )