diff --git a/dvc/repo/experiments/apply.py b/dvc/repo/experiments/apply.py index b1b5a4d2e..d7f2d3579 100644 --- a/dvc/repo/experiments/apply.py +++ b/dvc/repo/experiments/apply.py @@ -1,14 +1,14 @@ import logging import os -from contextlib import contextmanager from typing import TYPE_CHECKING, Optional from dvc.repo import locked from dvc.repo.scm_context import scm_context from dvc.scm import Git +from dvc.ui import ui from dvc.utils.fs import remove -from .exceptions import ApplyConflictError, BaselineMismatchError, InvalidExpRevError +from .exceptions import BaselineMismatchError, InvalidExpRevError from .executor.base import BaseExecutor from .refs import EXEC_APPLY @@ -21,11 +21,9 @@ logger = logging.getLogger(__name__) @locked @scm_context -def apply(repo: "Repo", rev: str, force: bool = True, **kwargs): # noqa: C901 - from scmrepo.exceptions import SCMError as _SCMError - +def apply(repo: "Repo", rev: str, **kwargs): # noqa: C901 from dvc.repo.checkout import checkout as dvc_checkout - from dvc.scm import GitMergeError, RevError, resolve_rev + from dvc.scm import RevError, resolve_rev exps: "Experiments" = repo.experiments @@ -34,7 +32,6 @@ def apply(repo: "Repo", rev: str, force: bool = True, **kwargs): # noqa: C901 assert isinstance(repo.scm, Git) try: exp_rev = resolve_rev(repo.scm, rev) - exps.check_baseline(exp_rev) except RevError as exc: ( exp_ref_info, @@ -42,14 +39,9 @@ def apply(repo: "Repo", rev: str, force: bool = True, **kwargs): # noqa: C901 ) = exps.celery_queue.get_ref_and_entry_by_names( rev )[rev] - baseline_sha = repo.scm.get_rev() if exp_ref_info: - if baseline_sha != exp_ref_info.baseline_sha: - raise InvalidExpRevError(rev) # noqa: B904 exp_rev = repo.scm.get_ref(str(exp_ref_info)) elif queue_entry: - if queue_entry.baseline_rev != baseline_sha: - raise InvalidExpRevError(rev) # noqa: B904 exp_rev = queue_entry.stash_rev is_stash = True else: @@ -57,68 +49,22 @@ def apply(repo: "Repo", rev: str, force: bool = True, **kwargs): # noqa: C901 except BaselineMismatchError as exc: raise InvalidExpRevError(rev) from exc - # NOTE: we don't use scmrepo's stash_workspace() here since we need - # finer control over the merge behavior when we unstash everything - with _apply_workspace(repo, rev, force): - try: - repo.scm.merge(exp_rev, commit=False, squash=True) - except _SCMError as exc: - raise GitMergeError(str(exc), scm=repo.scm) # noqa: B904 - - repo.scm.reset() - - if is_stash: - assert repo.tmp_dir is not None - args_path = os.path.join(repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE) - if os.path.exists(args_path): - remove(args_path) - + _apply(repo, exp_rev, name=rev, is_stash=is_stash) dvc_checkout(repo, **kwargs) repo.scm.set_ref(EXEC_APPLY, exp_rev) - logger.info( - "Changes for experiment '%s' have been applied to your current workspace.", - rev, + ui.write( + f"Changes for experiment '{rev}' have been applied to your current workspace.", ) -@contextmanager -def _apply_workspace(repo: "Repo", rev: str, force: bool): - from scmrepo.exceptions import MergeConflictError - from scmrepo.exceptions import SCMError as _SCMError +def _apply(repo: "Repo", rev: str, name: Optional[str] = None, is_stash: bool = False): + exps: "Experiments" = repo.experiments - assert isinstance(repo.scm, Git) - if repo.scm.is_dirty(untracked_files=True): - logger.debug("Stashing workspace") - stash_rev: Optional[str] = repo.scm.stash.push(include_untracked=True) - else: - stash_rev = None - try: - yield - except Exception: # pylint: disable=broad-except - if stash_rev: - _clean_and_pop(repo.scm) - raise - if not stash_rev: - return - - try: - repo.scm.reset() - repo.scm.stash.apply(stash_rev, skip_conflicts=force) - repo.scm.stash.drop() - except (MergeConflictError, _SCMError) as exc: - _clean_and_pop(repo.scm) - raise ApplyConflictError(rev) from exc - except Exception: # pylint: disable=broad-except - _clean_and_pop(repo.scm) - raise - - -def _clean_and_pop(scm: "Git"): - """Revert any changes and pop the last stash entry.""" - scm.reset(hard=True) - if scm.is_dirty(untracked_files=True): - # drop any changes to untracked files before popping stash - scm.stash.push(include_untracked=True) - scm.stash.drop() - scm.stash.pop() + with exps.apply_stash.preserve_workspace(rev, name=name): + with repo.scm.detach_head(rev, force=True): + if is_stash: + assert repo.tmp_dir is not None + args_path = os.path.join(repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE) + if os.path.exists(args_path): + remove(args_path)