From 63b710b7c08bc5b74c0526166f9d10a8d21fbc23 Mon Sep 17 00:00:00 2001
From: Maxime Rischard <rischardmaxime@gmail.com>
Date: Mon, 9 Aug 2021 23:49:21 +0100
Subject: [PATCH] More granular read-write lock requests in the run() and
 reproduce() methods of the Stage class (#5815)

* Remove rwlock decorator for the run() method of the Stage class.

Instead of pre-emptively grabbing read locks on the dependencies
and write locks on the outputs when calling the run() method,
delegate the locking mechanism to other methods.
This results in more selective locking that is less likely
to cause a LockError when executing steps in parallel.

* Remove rwlocked decorator from reproduce() method of Stage class.

* Reinstate read-lock (not write) for run method of Stage class

* fix bugs and typos

Co-authored-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/stage/__init__.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py
index 7aa8a121c..325c90704 100644
--- a/dvc/stage/__init__.py
+++ b/dvc/stage/__init__.py
@@ -315,6 +315,7 @@ class Stage(params.StageParams):
 
         return self._changed_deps()
 
+    @rwlocked(read=["deps"])
     def _changed_deps(self):
         for dep in self.deps:
             status = dep.status()
@@ -328,6 +329,7 @@ class Stage(params.StageParams):
                 return True
         return False
 
+    @rwlocked(read=["outs"])
     def changed_outs(self):
         for out in self.outs:
             status = out.status()
@@ -513,7 +515,7 @@ class Stage(params.StageParams):
         if link_failures:
             raise CacheLinkError(link_failures)
 
-    @rwlocked(read=["deps"], write=["outs"])
+    @rwlocked(read=["deps", "outs"])
     def run(
         self,
         dry=False,
@@ -527,16 +529,16 @@ class Stage(params.StageParams):
 
         if not self.frozen and self.is_import:
             jobs = kwargs.get("jobs", None)
-            sync_import(self, dry, force, jobs)
+            self._sync_import(dry, force, jobs)
         elif not self.frozen and self.cmd:
-            run_stage(self, dry, force, **kwargs)
+            self._run_stage(dry, force, **kwargs)
         else:
             args = (
                 ("outputs", "frozen ") if self.frozen else ("data sources", "")
             )
             logger.info("Verifying %s in %s%s", *args, self)
             if not dry:
-                check_missing_outputs(self)
+                self._check_missing_outputs()
 
         if not dry:
             if kwargs.get("checkpoint_func", None):
@@ -545,6 +547,18 @@ class Stage(params.StageParams):
             if not no_commit:
                 self.commit(allow_missing=allow_missing)
 
+    @rwlocked(read=["deps"], write=["outs"])
+    def _run_stage(self, dry, force, **kwargs):
+        return run_stage(self, dry, force, **kwargs)
+
+    @rwlocked(read=["deps"], write=["outs"])
+    def _sync_import(self, dry, force, jobs):
+        sync_import(self, dry, force, jobs)
+
+    @rwlocked(read=["outs"])
+    def _check_missing_outputs(self):
+        check_missing_outputs(self)
+
     def filter_outs(self, path_info):
         def _func(o):
             return path_info.isin_or_eq(o.path_info)