From 0ce4e769f11e988f1981ca17381cc3c49d03d243 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Wed, 3 Aug 2022 20:28:24 +0300 Subject: [PATCH] repo: make lock/rwlock kinda work on non local fs --- dvc/repo/__init__.py | 25 ++++++++++++++-------- dvc/rwlock.py | 15 +++++++------ dvc/stage/decorators.py | 4 +++- tests/unit/test_rwlock.py | 45 ++++++++++++++++++++------------------- 4 files changed, 50 insertions(+), 39 deletions(-) diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index d1954314e..eb63e41bd 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -155,7 +155,7 @@ class Repo: ): from dvc.config import Config from dvc.data_cloud import DataCloud - from dvc.fs import GitFileSystem, localfs + from dvc.fs import GitFileSystem, LocalFileSystem, localfs from dvc.lock import LockNoop, make_lock from dvc.odbmgr import ODBManager from dvc.repo.metrics import Metrics @@ -199,15 +199,22 @@ class Repo: self.tmp_dir = None else: self.fs.makedirs(self.tmp_dir, exist_ok=True) - self.lock = make_lock( - os.path.join(self.tmp_dir, "lock"), - tmp_dir=self.tmp_dir, - hardlink_lock=self.config["core"].get("hardlink_lock", False), - friendly=True, - ) - state_db_dir = self._get_database_dir("state") - self.state = State(self.root_dir, state_db_dir, self.dvcignore) + if isinstance(self.fs, LocalFileSystem): + self.lock = make_lock( + self.fs.path.join(self.tmp_dir, "lock"), + tmp_dir=self.tmp_dir, + hardlink_lock=self.config["core"].get( + "hardlink_lock", False + ), + friendly=True, + ) + state_db_dir = self._get_database_dir("state") + self.state = State(self.root_dir, state_db_dir, self.dvcignore) + else: + self.lock = LockNoop() + self.state = StateNoop() + self.odb = ODBManager(self) self.stage_cache = StageCache(self) diff --git a/dvc/rwlock.py b/dvc/rwlock.py index 35a0685e3..e1125b050 100644 --- a/dvc/rwlock.py +++ b/dvc/rwlock.py @@ -33,10 +33,10 @@ class RWLockFileFormatError(DvcException): @contextmanager -def _edit_rwlock(lock_dir): - path = os.path.join(lock_dir, "rwlock") +def _edit_rwlock(lock_dir, fs): + path = fs.path.join(lock_dir, "rwlock") try: - with open(path, encoding="utf-8") as fobj: + with fs.open(path, encoding="utf-8") as fobj: lock = SCHEMA(json.load(fobj)) except FileNotFoundError: lock = SCHEMA({}) @@ -47,7 +47,7 @@ def _edit_rwlock(lock_dir): lock["read"] = defaultdict(list, lock["read"]) lock["write"] = defaultdict(dict, lock["write"]) yield lock - with open(path, "w+", encoding="utf-8") as fobj: + with fs.open(path, "w", encoding="utf-8") as fobj: json.dump(lock, fobj) @@ -133,11 +133,12 @@ def _release_read(lock, info, changes): @contextmanager -def rwlock(tmp_dir, cmd, read, write): +def rwlock(tmp_dir, fs, cmd, read, write): """Create non-thread-safe RWLock for file paths. Args: tmp_dir (str): existing directory where to create the rwlock file. + fs (FileSystem): fs instance that tmp_dir belongs to. cmd (str): command that will be working on these file path. read ([str]): file paths that are going to be read. write ([str]): file paths that are going to be written. @@ -152,7 +153,7 @@ def rwlock(tmp_dir, cmd, read, write): """ info = {"pid": os.getpid(), "cmd": cmd} - with _edit_rwlock(tmp_dir) as lock: + with _edit_rwlock(tmp_dir, fs) as lock: _check_blockers(lock, info, mode="write", waiters=read + write) _check_blockers(lock, info, mode="read", waiters=write) @@ -163,6 +164,6 @@ def rwlock(tmp_dir, cmd, read, write): try: yield finally: - with _edit_rwlock(tmp_dir) as lock: + with _edit_rwlock(tmp_dir, fs) as lock: _release_write(lock, info, wchanges) _release_read(lock, info, rchanges) diff --git a/dvc/stage/decorators.py b/dvc/stage/decorators.py index bdece2310..056d55407 100644 --- a/dvc/stage/decorators.py +++ b/dvc/stage/decorators.py @@ -32,7 +32,9 @@ def rwlocked(call, read=None, write=None): cmd = " ".join(sys.argv) - with rwlock(stage.repo.tmp_dir, cmd, _chain(read), _chain(write)): + with rwlock( + stage.repo.tmp_dir, stage.repo.fs, cmd, _chain(read), _chain(write) + ): return call() diff --git a/tests/unit/test_rwlock.py b/tests/unit/test_rwlock.py index bbd6541a2..4e932901c 100644 --- a/tests/unit/test_rwlock.py +++ b/tests/unit/test_rwlock.py @@ -2,6 +2,7 @@ import os import pytest +from dvc.fs import localfs from dvc.lock import LockError from dvc.rwlock import ( RWLockFileCorruptedError, @@ -15,19 +16,19 @@ def test_rwlock(tmp_path): path = os.fspath(tmp_path) foo = "foo" - with rwlock(path, "cmd1", [foo], []): + with rwlock(path, localfs, "cmd1", [foo], []): with pytest.raises(LockError): - with rwlock(path, "cmd2", [], [foo]): + with rwlock(path, localfs, "cmd2", [], [foo]): pass - with rwlock(path, "cmd1", [], [foo]): + with rwlock(path, localfs, "cmd1", [], [foo]): with pytest.raises(LockError): - with rwlock(path, "cmd2", [foo], []): + with rwlock(path, localfs, "cmd2", [foo], []): pass - with rwlock(path, "cmd1", [], [foo]): + with rwlock(path, localfs, "cmd1", [], [foo]): with pytest.raises(LockError): - with rwlock(path, "cmd2", [], [foo]): + with rwlock(path, localfs, "cmd2", [], [foo]): pass @@ -35,19 +36,19 @@ def test_rwlock_reentrant(tmp_path): path = os.fspath(tmp_path) foo = "foo" - with rwlock(path, "cmd1", [], [foo]): - with rwlock(path, "cmd1", [], [foo]): + with rwlock(path, localfs, "cmd1", [], [foo]): + with rwlock(path, localfs, "cmd1", [], [foo]): pass - with _edit_rwlock(path) as lock: + with _edit_rwlock(path, localfs) as lock: assert lock == { "read": {}, "write": {"foo": {"cmd": "cmd1", "pid": os.getpid()}}, } - with rwlock(path, "cmd", [foo], []): - with rwlock(path, "cmd", [foo], []): + with rwlock(path, localfs, "cmd", [foo], []): + with rwlock(path, localfs, "cmd", [foo], []): pass - with _edit_rwlock(path) as lock: + with _edit_rwlock(path, localfs) as lock: assert lock == { "read": {"foo": [{"cmd": "cmd", "pid": os.getpid()}]}, "write": {}, @@ -59,23 +60,23 @@ def test_rwlock_subdirs(tmp_path): foo = "foo" subfoo = os.path.join("foo", "subfoo") - with rwlock(path, "cmd1", [foo], []): + with rwlock(path, localfs, "cmd1", [foo], []): with pytest.raises(LockError, match=r"subfoo(.|\n)*cmd1"): - with rwlock(path, "cmd2", [], [subfoo]): + with rwlock(path, localfs, "cmd2", [], [subfoo]): pass - with rwlock(path, "cmd1", [], [subfoo]): + with rwlock(path, localfs, "cmd1", [], [subfoo]): with pytest.raises(LockError, match=r"'foo'(.|\n)*cmd1"): - with rwlock(path, "cmd2", [foo], []): + with rwlock(path, localfs, "cmd2", [foo], []): pass - with rwlock(path, "cmd1", [], [subfoo]): + with rwlock(path, localfs, "cmd1", [], [subfoo]): with pytest.raises(LockError): - with rwlock(path, "cmd2", [], [foo]): + with rwlock(path, localfs, "cmd2", [], [foo]): pass - with rwlock(path, "cmd1", [subfoo], []): - with rwlock(path, "cmd2", [foo], []): + with rwlock(path, localfs, "cmd1", [subfoo], []): + with rwlock(path, localfs, "cmd2", [foo], []): pass @@ -85,10 +86,10 @@ def test_broken_rwlock(tmp_path): path.write_text('{"broken": "format"}', encoding="utf-8") with pytest.raises(RWLockFileFormatError): - with _edit_rwlock(dir_path): + with _edit_rwlock(dir_path, localfs): pass path.write_text("{broken json", encoding="utf-8") with pytest.raises(RWLockFileCorruptedError): - with _edit_rwlock(dir_path): + with _edit_rwlock(dir_path, localfs): pass