Fix #10638 : makes remove return correct list when used with both --queue and -A (#10641)

* fixed a bug where using both `--queue` and `-A` would yield a "removed" list without the queued experiments

* added a test to validate the fix

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

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

* clearer name for test

---------

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

.github/workflows/tests.yaml

@@ -19,6 +19,8 @@ jobs:
   lint:
     timeout-minutes: 10
     runs-on: ${{ matrix.os }}
+    env:
+       DVC_NO_ANALYTICS: true
     strategy:
       fail-fast: false
       matrix:

dvc/commands/experiments/__init__.py

@@ -58,6 +58,15 @@ def add_parser(subparsers, parent_parser):
     hide_subparsers_from_help(experiments_subparsers)
 
 
+def add_keep_selection_flag(experiments_subcmd_parser):
+    experiments_subcmd_parser.add_argument(
+        "--keep",
+        action="store_true",
+        default=False,
+        help="Keep the selected experiments instead of removing them.",
+    )
+
+
 def add_rev_selection_flags(
     experiments_subcmd_parser, command: str, default: bool = True
 ):

dvc/commands/experiments/remove.py

@@ -34,6 +34,7 @@ class CmdExperimentsRemove(CmdBase):
             num=self.args.num,
             queue=self.args.queue,
             git_remote=self.args.git_remote,
+            keep=self.args.keep,
         )
         if removed:
             ui.write(f"Removed experiments: {humanize.join(map(repr, removed))}")
@@ -44,7 +45,7 @@ class CmdExperimentsRemove(CmdBase):
 
 
 def add_parser(experiments_subparsers, parent_parser):
-    from . import add_rev_selection_flags
+    from . import add_keep_selection_flag, add_rev_selection_flags
 
     EXPERIMENTS_REMOVE_HELP = "Remove experiments."
     experiments_remove_parser = experiments_subparsers.add_parser(
@@ -57,6 +58,7 @@ def add_parser(experiments_subparsers, parent_parser):
     )
     remove_group = experiments_remove_parser.add_mutually_exclusive_group()
     add_rev_selection_flags(experiments_remove_parser, "Remove", False)
+    add_keep_selection_flag(experiments_remove_parser)
     remove_group.add_argument(
         "--queue", action="store_true", help="Remove all queued experiments."
     )

dvc/repo/experiments/refs.py

@@ -67,3 +67,12 @@ class ExpRefInfo:
         baseline_sha = parts[2] + parts[3]
         name = parts[4] if len(parts) == 5 else None
         return cls(baseline_sha, name)
+
+    def __eq__(self, other):
+        if not isinstance(other, ExpRefInfo):
+            return False
+
+        return self.baseline_sha == other.baseline_sha and self.name == other.name
+
+    def __hash__(self):
+        return hash((self.baseline_sha, self.name))

dvc/repo/experiments/remove.py

@@ -6,7 +6,7 @@ from dvc.repo import locked
 from dvc.repo.scm_context import scm_context
 from dvc.scm import Git, iter_revs
 
-from .exceptions import UnresolvedExpNamesError
+from .exceptions import InvalidArgumentError, UnresolvedExpNamesError
 from .utils import exp_refs, exp_refs_by_baseline, push_refspec
 
 if TYPE_CHECKING:
@@ -30,10 +30,16 @@ def remove(  # noqa: C901, PLR0912
     num: int = 1,
     queue: bool = False,
     git_remote: Optional[str] = None,
+    keep: bool = False,
 ) -> list[str]:
     removed: list[str] = []
+
+    if all([keep, queue]):
+        raise InvalidArgumentError("Cannot use both `--keep` and `--queue`.")
+
     if not any([exp_names, queue, all_commits, rev]):
         return removed
+
     celery_queue: LocalCeleryQueue = repo.experiments.celery_queue
 
     if queue:
@@ -43,6 +49,7 @@ def remove(  # noqa: C901, PLR0912
 
     exp_ref_list: list[ExpRefInfo] = []
     queue_entry_list: list[QueueEntry] = []
+
     if exp_names:
         results: dict[str, ExpRefAndQueueEntry] = (
             celery_queue.get_ref_and_entry_by_names(exp_names, git_remote)
@@ -68,6 +75,10 @@ def remove(  # noqa: C901, PLR0912
         exp_ref_list.extend(exp_ref_dict.values())
     elif all_commits:
         exp_ref_list.extend(exp_refs(repo.scm, git_remote))
+        removed.extend([ref.name for ref in exp_ref_list])
+
+    if keep:
+        exp_ref_list = list(set(exp_refs(repo.scm, git_remote)) - set(exp_ref_list))
         removed = [ref.name for ref in exp_ref_list]
 
     if exp_ref_list:
@@ -83,6 +94,7 @@ def remove(  # noqa: C901, PLR0912
 
         removed_refs = [str(r) for r in exp_ref_list]
         notify_refs_to_studio(repo, git_remote, removed=removed_refs)
+
     return removed
 
 

tests/func/experiments/test_remove.py

@@ -43,6 +43,26 @@ def test_remove_all_queued_experiments(tmp_dir, scm, dvc, exp_stage):
     assert scm.get_ref(str(ref_info)) is not None
 
 
+def test_remove_all_experiments_queued_and_completed(tmp_dir, scm, dvc, exp_stage):
+    queue_length = 3
+    for i in range(queue_length):
+        dvc.experiments.run(
+            exp_stage.addressing, params=[f"foo={i}"], name=f"exp{i}", queue=True
+        )
+
+    results = dvc.experiments.run(
+        exp_stage.addressing, params=[f"foo={queue_length}"], name=f"exp{queue_length}"
+    )
+    ref_info = first(exp_refs_by_rev(scm, first(results)))
+
+    removed = sorted(dvc.experiments.remove(all_commits=True, queue=True))
+
+    assert len(removed) == queue_length + 1
+    assert removed == [f"exp{i}" for i in range(queue_length)] + [ref_info.name]
+    assert len(dvc.experiments.stash_revs) == 0
+    assert scm.get_ref(str(ref_info)) is None
+
+
 def test_remove_special_queued_experiments(tmp_dir, scm, dvc, exp_stage):
     dvc.experiments.run(
         exp_stage.addressing, params=["foo=1"], queue=True, name="queue1"
@@ -179,3 +199,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,
     )