From 6bfce5a5db07fd0eb13451de2c3277c32f8422be Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Wed, 23 Nov 2022 18:48:27 +0800 Subject: [PATCH] Use a more granular lock on exp operations. For now in `experiment.new`, `init_executor` and `collect_exectuor` we lock the whole scm repo. We can make a more granular lock in these operations. --- dvc/repo/experiments/__init__.py | 10 +-- dvc/repo/experiments/executor/base.py | 25 +++--- dvc/repo/experiments/executor/local.py | 85 +++++++++++---------- dvc/repo/experiments/executor/ssh.py | 1 + dvc/repo/experiments/queue/base.py | 59 ++++++++------ dvc/repo/experiments/queue/celery.py | 6 +- dvc/repo/experiments/queue/tasks.py | 8 +- dvc/repo/experiments/queue/utils.py | 46 ++++++++--- dvc/repo/experiments/queue/workspace.py | 6 +- dvc/repo/experiments/utils.py | 35 ++++++--- dvc/rwlock.py | 16 ++-- tests/func/experiments/executor/test_ssh.py | 4 +- 12 files changed, 180 insertions(+), 121 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 287490fcb..b335c2933 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -32,7 +32,7 @@ from .refs import ( ExpRefInfo, ) from .stash import ExpStashEntry -from .utils import exp_refs_by_rev, scm_locked, unlocked_repo +from .utils import exp_refs_by_rev, unlocked_repo logger = logging.getLogger(__name__) @@ -50,18 +50,12 @@ class Experiments: ) def __init__(self, repo): - from dvc.lock import make_lock from dvc.scm import NoSCMError if repo.config["core"].get("no_scm", False): raise NoSCMError self.repo = repo - self.scm_lock = make_lock( - os.path.join(self.repo.tmp_dir, "exp_scm_lock"), - tmp_dir=self.repo.tmp_dir, - hardlink_lock=repo.config["core"].get("hardlink_lock", False), - ) @property def scm(self): @@ -253,7 +247,6 @@ class Experiments: if self.scm.get_ref(str(exp_ref)): raise ExperimentExistsError(exp_ref.name) - @scm_locked def new( self, queue: BaseStashQueue, @@ -279,7 +272,6 @@ class Experiments: except ExperimentExistsError as err: if not (kwargs.get("force", False) or kwargs.get("reset", False)): raise err - return queue.put(*args, **kwargs) def _resume_checkpoint( diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 470f5339c..f85ab4e6a 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -12,6 +12,7 @@ from typing import ( Callable, Dict, Iterable, + List, NamedTuple, Optional, Tuple, @@ -40,9 +41,6 @@ from ..refs import ( EXEC_CHECKPOINT, EXEC_HEAD, EXEC_MERGE, - EXEC_NAMESPACE, - EXPS_NAMESPACE, - EXPS_STASH, ExpRefInfo, ) @@ -168,6 +166,7 @@ class BaseExecutor(ABC): @abstractmethod def init_git( self, + repo: "Repo", scm: "Git", stash_rev: str, entry: "ExpStashEntry", @@ -315,6 +314,7 @@ class BaseExecutor(ABC): def fetch_exps( self, dest_scm: "Git", + refs: List[str], force: bool = False, on_diverged: Callable[[str, bool], None] = None, **kwargs, @@ -323,26 +323,19 @@ class BaseExecutor(ABC): Args: dest_scm: Destination Git instance. + refs: reference names to be fetched from the remotes. force: If True, diverged refs will be overwritten on_diverged: Callback in the form on_diverged(ref, is_checkpoint) to be called when an experiment ref has diverged. Extra kwargs will be passed into the remote git client. """ - from ..utils import iter_remote_refs - refs = [] - has_checkpoint = False - for ref in iter_remote_refs( - dest_scm, - self.git_url, - base=EXPS_NAMESPACE, - **kwargs, - ): - if ref == EXEC_CHECKPOINT: - has_checkpoint = True - elif not ref.startswith(EXEC_NAMESPACE) and ref != EXPS_STASH: - refs.append(ref) + if EXEC_CHECKPOINT in refs: + refs.remove(EXEC_CHECKPOINT) + has_checkpoint = True + else: + has_checkpoint = False def on_diverged_ref(orig_ref: str, new_rev: str): if force: diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index 4d812fa80..9ab2f18f9 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/dvc/repo/experiments/executor/local.py @@ -4,9 +4,10 @@ from contextlib import ExitStack from tempfile import mkdtemp from typing import TYPE_CHECKING, Optional -from funcy import cached_property +from funcy import cached_property, retry from scmrepo.exceptions import SCMError as _SCMError +from dvc.lock import LockError from dvc.scm import SCM, GitMergeError from dvc.utils.fs import makedirs, remove @@ -19,7 +20,7 @@ from ..refs import ( EXEC_MERGE, EXEC_NAMESPACE, ) -from ..utils import EXEC_TMP_DIR +from ..utils import EXEC_TMP_DIR, get_exp_rwlock from .base import BaseExecutor, TaskStatus if TYPE_CHECKING: @@ -68,8 +69,10 @@ class TempDirExecutor(BaseLocalExecutor): QUIET = True DEFAULT_LOCATION = "tempdir" + @retry(180, errors=LockError, timeout=1) def init_git( self, + repo: "Repo", scm: "Git", stash_rev: str, entry: "ExpStashEntry", @@ -86,28 +89,29 @@ class TempDirExecutor(BaseLocalExecutor): if infofile: self.info.dump_json(infofile) - with self.set_exec_refs(scm, stash_rev, entry): - refspec = f"{EXEC_NAMESPACE}/" - push_refspec(scm, self.git_url, refspec, refspec) + with get_exp_rwlock(repo, writes=[EXEC_NAMESPACE]): + with self.set_exec_refs(scm, stash_rev, entry): + refspec = f"{EXEC_NAMESPACE}/" + push_refspec(scm, self.git_url, refspec, refspec) - if branch: - push_refspec(scm, self.git_url, branch, branch) - self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True) - elif self.scm.get_ref(EXEC_BRANCH): - self.scm.remove_ref(EXEC_BRANCH) + if branch: + push_refspec(scm, self.git_url, branch, branch) + self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True) + elif self.scm.get_ref(EXEC_BRANCH): + self.scm.remove_ref(EXEC_BRANCH) - if self.scm.get_ref(EXEC_CHECKPOINT): - self.scm.remove_ref(EXEC_CHECKPOINT) - # checkout EXEC_HEAD and apply EXEC_MERGE on top of it without - # committing - head = EXEC_BRANCH if branch else EXEC_HEAD - self.scm.checkout(head, detach=True) - merge_rev = self.scm.get_ref(EXEC_MERGE) + if self.scm.get_ref(EXEC_CHECKPOINT): + self.scm.remove_ref(EXEC_CHECKPOINT) + # checkout EXEC_HEAD and apply EXEC_MERGE on top of it without + # committing + head = EXEC_BRANCH if branch else EXEC_HEAD + self.scm.checkout(head, detach=True) + merge_rev = self.scm.get_ref(EXEC_MERGE) - try: - self.scm.merge(merge_rev, squash=True, commit=False) - except _SCMError as exc: - raise GitMergeError(str(exc), scm=self.scm) + try: + self.scm.merge(merge_rev, squash=True, commit=False) + except _SCMError as exc: + raise GitMergeError(str(exc), scm=self.scm) def _config(self, cache_dir): local_config = os.path.join( @@ -168,8 +172,10 @@ class WorkspaceExecutor(BaseLocalExecutor): logger.debug("Init workspace executor in '%s'", root_dir) return executor + @retry(180, errors=LockError, timeout=1) def init_git( self, + repo: "Repo", scm: "Git", stash_rev: str, entry: "ExpStashEntry", @@ -180,25 +186,26 @@ class WorkspaceExecutor(BaseLocalExecutor): if infofile: self.info.dump_json(infofile) - scm.set_ref(EXEC_HEAD, entry.head_rev) - scm.set_ref(EXEC_MERGE, stash_rev) - scm.set_ref(EXEC_BASELINE, entry.baseline_rev) - self._detach_stack.enter_context( - self.scm.detach_head( - self.scm.get_ref(EXEC_HEAD), - force=True, - client="dvc", + with get_exp_rwlock(repo, writes=[EXEC_NAMESPACE]): + scm.set_ref(EXEC_HEAD, entry.head_rev) + scm.set_ref(EXEC_MERGE, stash_rev) + scm.set_ref(EXEC_BASELINE, entry.baseline_rev) + self._detach_stack.enter_context( + self.scm.detach_head( + self.scm.get_ref(EXEC_HEAD), + force=True, + client="dvc", + ) ) - ) - merge_rev = self.scm.get_ref(EXEC_MERGE) - try: - self.scm.merge(merge_rev, squash=True, commit=False) - except _SCMError as exc: - raise GitMergeError(str(exc), scm=self.scm) - if branch: - self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True) - elif scm.get_ref(EXEC_BRANCH): - self.scm.remove_ref(EXEC_BRANCH) + merge_rev = self.scm.get_ref(EXEC_MERGE) + try: + self.scm.merge(merge_rev, squash=True, commit=False) + except _SCMError as exc: + raise GitMergeError(str(exc), scm=self.scm) + if branch: + self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True) + elif scm.get_ref(EXEC_BRANCH): + self.scm.remove_ref(EXEC_BRANCH) def init_cache(self, repo: "Repo", rev: str, run_cache: bool = True): pass diff --git a/dvc/repo/experiments/executor/ssh.py b/dvc/repo/experiments/executor/ssh.py index adccf3aad..673a5bc1d 100644 --- a/dvc/repo/experiments/executor/ssh.py +++ b/dvc/repo/experiments/executor/ssh.py @@ -119,6 +119,7 @@ class SSHExecutor(BaseExecutor): def init_git( self, + repo: "Repo", scm: "Git", stash_rev: str, entry: "ExpStashEntry", diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index d160fb7a3..b7e15b53a 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -30,7 +30,8 @@ from ..executor.base import BaseExecutor, ExecutorResult from ..executor.local import WorkspaceExecutor from ..refs import ExpRefInfo from ..stash import ExpStash, ExpStashEntry -from ..utils import EXEC_PID_DIR, EXEC_TMP_DIR, exp_refs_by_rev, scm_locked +from ..utils import EXEC_PID_DIR, EXEC_TMP_DIR, exp_refs_by_rev, get_exp_rwlock +from .utils import get_remote_executor_refs if TYPE_CHECKING: from scmrepo.git import Git @@ -550,28 +551,39 @@ class BaseStashQueue(ABC): @staticmethod @retry(180, errors=LockError, timeout=1) - @scm_locked + def get_stash_entry( + exp: "Experiments", + queue_entry: QueueEntry, + ) -> "ExpStashEntry": + stash = ExpStash(exp.scm, queue_entry.stash_ref) + stash_rev = queue_entry.stash_rev + with get_exp_rwlock(exp.repo, writes=[queue_entry.stash_ref]): + stash_entry = stash.stash_revs.get( + stash_rev, + ExpStashEntry(None, stash_rev, stash_rev, None, None), + ) + if stash_entry.stash_index is not None: + stash.drop(stash_entry.stash_index) + return stash_entry + + @classmethod def init_executor( + cls, exp: "Experiments", queue_entry: QueueEntry, executor_cls: Type[BaseExecutor] = WorkspaceExecutor, **kwargs, ) -> BaseExecutor: - scm = exp.scm - stash = ExpStash(scm, queue_entry.stash_ref) - stash_rev = queue_entry.stash_rev - stash_entry = stash.stash_revs.get( - stash_rev, - ExpStashEntry(None, stash_rev, stash_rev, None, None), - ) - if stash_entry.stash_index is not None: - stash.drop(stash_entry.stash_index) + stash_entry = cls.get_stash_entry(exp, queue_entry) + executor = executor_cls.from_stash_entry( exp.repo, stash_entry, **kwargs ) + stash_rev = queue_entry.stash_rev infofile = exp.celery_queue.get_infofile_path(stash_rev) executor.init_git( + exp.repo, exp.repo.scm, stash_rev, stash_entry, @@ -592,7 +604,6 @@ class BaseStashQueue(ABC): @staticmethod @retry(180, errors=LockError, timeout=1) - @scm_locked def collect_git( exp: "Experiments", executor: BaseExecutor, @@ -606,16 +617,20 @@ class BaseStashQueue(ABC): raise CheckpointExistsError(ref_info.name) raise ExperimentExistsError(ref_info.name) - for ref in executor.fetch_exps( - exp.scm, - force=exec_result.force, - on_diverged=on_diverged, - ): - exp_rev = exp.scm.get_ref(ref) - if exp_rev: - assert exec_result.exp_hash - logger.debug("Collected experiment '%s'.", exp_rev[:7]) - results[exp_rev] = exec_result.exp_hash + refs = get_remote_executor_refs(exp.scm, executor.git_url) + + with get_exp_rwlock(exp.repo, writes=refs): + for ref in executor.fetch_exps( + exp.scm, + refs, + force=exec_result.force, + on_diverged=on_diverged, + ): + exp_rev = exp.scm.get_ref(ref) + if exp_rev: + assert exec_result.exp_hash + logger.debug("Collected experiment '%s'.", exp_rev[:7]) + results[exp_rev] = exec_result.exp_hash return results diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index 04e876d88..b8dca00db 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -24,7 +24,8 @@ from dvc.ui import ui from ..exceptions import UnresolvedQueueExpNamesError from ..executor.base import ExecutorInfo, ExecutorResult -from ..utils import EXEC_TMP_DIR +from ..refs import CELERY_STASH +from ..utils import EXEC_TMP_DIR, get_exp_rwlock from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult from .exceptions import CannotKillTasksError from .tasks import run_exp @@ -165,7 +166,8 @@ class LocalCeleryQueue(BaseStashQueue): def put(self, *args, **kwargs) -> QueueEntry: """Stash an experiment and add it to the queue.""" - entry = self._stash_exp(*args, **kwargs) + with get_exp_rwlock(self.repo, writes=["workspace", CELERY_STASH]): + entry = self._stash_exp(*args, **kwargs) self.celery.signature(run_exp.s(entry.asdict())).delay() return entry diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index 36094d1de..7fcd99a0e 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/dvc/repo/experiments/queue/tasks.py @@ -1,4 +1,4 @@ -from typing import Any, Dict +from typing import TYPE_CHECKING, Any, Dict from celery import shared_task from celery.utils.log import get_task_logger @@ -7,11 +7,15 @@ from ..executor.base import ExecutorInfo from ..executor.local import TempDirExecutor from .base import BaseStashQueue, QueueEntry +if TYPE_CHECKING: + from ..executor.base import BaseExecutor + + logger = get_task_logger(__name__) @shared_task -def setup_exp(entry_dict: Dict[str, Any]) -> TempDirExecutor: +def setup_exp(entry_dict: Dict[str, Any]) -> "BaseExecutor": """Setup an experiment. Arguments: diff --git a/dvc/repo/experiments/queue/utils.py b/dvc/repo/experiments/queue/utils.py index c9acaaf28..3f09108e9 100644 --- a/dvc/repo/experiments/queue/utils.py +++ b/dvc/repo/experiments/queue/utils.py @@ -1,17 +1,40 @@ import logging -from typing import TYPE_CHECKING, Dict +from typing import TYPE_CHECKING, Dict, List from scmrepo.exceptions import SCMError from ..executor.base import ExecutorInfo, TaskStatus +from ..refs import EXEC_CHECKPOINT, EXEC_NAMESPACE, EXPS_NAMESPACE, EXPS_STASH +from ..utils import get_exp_rwlock, iter_remote_refs logger = logging.getLogger(__name__) if TYPE_CHECKING: + from scmrepo.git import Git + from .base import BaseStashQueue +def get_remote_executor_refs(scm: "Git", remote_url: str) -> List[str]: + """Get result list refs from a remote repository + + Args: + remote_url : remote executor's url + """ + refs = [] + for ref in iter_remote_refs( + scm, + remote_url, + base=EXPS_NAMESPACE, + ): + if ref == EXEC_CHECKPOINT or ( + not ref.startswith(EXEC_NAMESPACE) and ref != EXPS_STASH + ): + refs.append(ref) + return refs + + def fetch_running_exp_from_temp_dir( queue: "BaseStashQueue", rev: str, fetch_refs: bool ) -> Dict[str, Dict]: @@ -45,15 +68,18 @@ def fetch_running_exp_from_temp_dir( executor = TempDirExecutor.from_info(info) try: - for ref in executor.fetch_exps( - queue.scm, - on_diverged=on_diverged, - ): - logger.debug("Updated running experiment '%s'.", ref) - last_rev = queue.scm.get_ref(ref) - result[rev]["last"] = last_rev - if last_rev: - result[last_rev] = info.asdict() + refs = get_remote_executor_refs(queue.scm, executor.git_url) + with get_exp_rwlock(queue.repo, writes=refs): + for ref in executor.fetch_exps( + queue.scm, + refs, + on_diverged=on_diverged, + ): + logger.debug("Updated running experiment '%s'.", ref) + last_rev = queue.scm.get_ref(ref) + result[rev]["last"] = last_rev + if last_rev: + result[last_rev] = info.asdict() except (InvalidRemoteSCMRepo, SCMError): # ignore stale info files del result[rev] diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index 68e7bb7d8..4331c4562 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -21,7 +21,8 @@ from ..executor.base import ( TaskStatus, ) from ..executor.local import WorkspaceExecutor -from ..refs import EXEC_BRANCH +from ..refs import EXEC_BRANCH, WORKSPACE_STASH +from ..utils import get_exp_rwlock from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult if TYPE_CHECKING: @@ -34,7 +35,8 @@ class WorkspaceQueue(BaseStashQueue): _EXEC_NAME: Optional[str] = "workspace" def put(self, *args, **kwargs) -> QueueEntry: - return self._stash_exp(*args, **kwargs) + with get_exp_rwlock(self.repo, writes=["workspace", WORKSPACE_STASH]): + return self._stash_exp(*args, **kwargs) def get(self) -> QueueGetResult: revs = self.stash.stash_revs diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index 0397ddc39..54d115f4c 100644 --- a/dvc/repo/experiments/utils.py +++ b/dvc/repo/experiments/utils.py @@ -1,6 +1,9 @@ +import os +import sys from collections import defaultdict from functools import wraps from typing import ( + TYPE_CHECKING, Callable, Dict, Generator, @@ -16,6 +19,7 @@ from scmrepo.git import Git from dvc.exceptions import InvalidArgumentError from dvc.repo.experiments.exceptions import AmbiguousExpRefInfo +from dvc.rwlock import rwlock from .refs import ( EXEC_APPLY, @@ -28,21 +32,34 @@ from .refs import ( ExpRefInfo, ) +if TYPE_CHECKING: + from dvc.repo import Repo + + EXEC_TMP_DIR = "exps" EXEC_PID_DIR = "run" -def scm_locked(f): - # Lock the experiments workspace so that we don't try to perform two - # different sequences of git operations at once - @wraps(f) - def wrapper(exp, *args, **kwargs): - from dvc.scm import map_scm_exception +def get_exp_rwlock( + repo: "Repo", + reads: Optional[List[str]] = None, + writes: Optional[List[str]] = None, +): + reads = reads or [] + writes = writes or [] - with map_scm_exception(), exp.scm_lock: - return f(exp, *args, **kwargs) + cmd = " ".join(sys.argv) + path = os.path.join(repo.tmp_dir, EXEC_TMP_DIR) + repo.fs.makedirs(path, exist_ok=True) - return wrapper + return rwlock( + path, + repo.fs, + cmd, + reads, + writes, + repo.config["core"].get("hardlink_lock", False), + ) def unlocked_repo(f): diff --git a/dvc/rwlock.py b/dvc/rwlock.py index 0bf2a4f91..27f4215b2 100644 --- a/dvc/rwlock.py +++ b/dvc/rwlock.py @@ -72,7 +72,7 @@ def _infos_to_str(infos): ) -def _check_blockers(lock, info, *, mode, waiters): +def _check_blockers(tmp_dir, lock, info, *, mode, waiters): from .lock import LockError for waiter_path in waiters: @@ -88,13 +88,11 @@ def _check_blockers(lock, info, *, mode, waiters): continue raise LockError( - "'{path}' is busy, it is being blocked by:\n" - "{blockers}\n" + f"'{waiter_path}' is busy, it is being blocked by:\n" + f"{_infos_to_str(blockers)}\n" "\n" "If there are no processes with such PIDs, you can manually " - "remove '.dvc/tmp/rwlock' and try again.".format( - path=waiter_path, blockers=_infos_to_str(blockers) - ) + f"remove '{tmp_dir}/rwlock' and try again." ) @@ -171,8 +169,10 @@ def rwlock(tmp_dir, fs, cmd, read, write, hardlink): with _edit_rwlock(tmp_dir, fs, hardlink) as lock: - _check_blockers(lock, info, mode="write", waiters=read + write) - _check_blockers(lock, info, mode="read", waiters=write) + _check_blockers( + tmp_dir, lock, info, mode="write", waiters=read + write + ) + _check_blockers(tmp_dir, lock, info, mode="read", waiters=write) rchanges = _acquire_read(lock, info, read) wchanges = _acquire_write(lock, info, write) diff --git a/tests/func/experiments/executor/test_ssh.py b/tests/func/experiments/executor/test_ssh.py index 6c2091612..4bf724831 100644 --- a/tests/func/experiments/executor/test_ssh.py +++ b/tests/func/experiments/executor/test_ssh.py @@ -38,7 +38,7 @@ def test_init_from_stash(tmp_dir, scm, dvc, machine_instance, mocker): @pytest.mark.needs_internet @pytest.mark.parametrize("cloud", [pytest.lazy_fixture("git_ssh")]) -def test_init_git(tmp_dir, scm, cloud, mocker): +def test_init_git(tmp_dir, dvc, scm, cloud, mocker): tmp_dir.scm_gen({"foo": "foo", "dir": {"bar": "bar"}}, commit="init") baseline_rev = scm.get_rev() tmp_dir.gen("foo", "stashed") @@ -59,7 +59,7 @@ def test_init_git(tmp_dir, scm, cloud, mocker): fs_factory=partial(_ssh_factory, cloud), ) infofile = str((root_url / "foo.run").path) - executor.init_git(scm, rev, mock, infofile=infofile) + executor.init_git(dvc, scm, rev, mock, infofile=infofile) assert root_url.path == executor._repo_abspath fs = cloud._ssh