Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-06 07:18:28 +01:00
Code Issues Packages Projects Releases Wiki Activity

Use a more granular lock on exp operations.

Browse Source 
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.
This commit is contained in:
karajan1001
2022-11-28 17:04:16 +08:00
committed by Gao
parent 5fa33a06bc
commit 6bfce5a5db
12 changed files with 180 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/repo/experiments/__init__.py
+1 -9
View File
@@ -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(
dvc/repo/experiments/executor/base.py
+9 -16
View File
@@ -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:
dvc/repo/experiments/executor/local.py
+46 -39
View File
@@ -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
dvc/repo/experiments/executor/ssh.py
+1
View File
@@ -119,6 +119,7 @@ class SSHExecutor(BaseExecutor):
def init_git(
self,
repo: "Repo",
scm: "Git",
stash_rev: str,
entry: "ExpStashEntry",
dvc/repo/experiments/queue/base.py
+37 -22
View File
@@ -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
dvc/repo/experiments/queue/celery.py
+4 -2
View File
@@ -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
dvc/repo/experiments/queue/tasks.py
+6 -2
View File
@@ -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:
dvc/repo/experiments/queue/utils.py
+36 -10
View File
@@ -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]
dvc/repo/experiments/queue/workspace.py
+4 -2
View File
@@ -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
dvc/repo/experiments/utils.py
+26 -9
View File
@@ -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):
dvc/rwlock.py
+8 -8
View File
@@ -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)
tests/func/experiments/executor/test_ssh.py
+2 -2
View File
@@ -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