repro: parallel execution of stages

Author: sjawhar

dvc/dvcfile.py

@@ -1,5 +1,7 @@
 import contextlib
 import os
+import threading
+from collections import defaultdict
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, TypeVar, Union
 
 from dvc.exceptions import DvcException
@@ -77,6 +79,10 @@ def check_dvcfile_path(repo, path):
         raise FileIsGitIgnored(relpath(path), True)
 
 
+_file_locks: dict[str, threading.Lock] = defaultdict(threading.Lock)
+_file_locks_lock = threading.Lock()
+
+
 class FileMixin:
     SCHEMA: Callable[[_T], _T]
 
@@ -85,6 +91,10 @@ def __init__(self, repo, path, verify=True, **kwargs):
         self.path = path
         self.verify = verify
 
+    def _thread_lock(self) -> threading.Lock:
+        with _file_locks_lock:
+            return _file_locks[self.path]
+
     def __repr__(self):
         return f"{self.__class__.__name__}: {relpath(self.path, self.repo.root_dir)}"
 
@@ -148,15 +158,19 @@ def validate(cls, d: _T, fname: Optional[str] = None) -> _T:
     def _load_yaml(self, **kwargs: Any) -> tuple[Any, str]:
         from dvc.utils import strictyaml
 
-        return strictyaml.load(
-            self.path,
-            self.SCHEMA,  # type: ignore[arg-type]
-            self.repo.fs,
-            **kwargs,
-        )
+        with self._thread_lock():
+            return strictyaml.load(
+                self.path,
+                self.SCHEMA,  # type: ignore[arg-type]
+                self.repo.fs,
+                **kwargs,
+            )
 
     def remove(self, force=False):  # noqa: ARG002
-        with contextlib.suppress(FileNotFoundError):
+        with (
+            self._thread_lock(),
+            contextlib.suppress(FileNotFoundError),
+        ):
             os.unlink(self.path)
 
     def dump(self, stage, **kwargs):
@@ -407,7 +421,10 @@ def _load(self, **kwargs: Any):
             return {}, ""
 
     def dump_dataset(self, dataset: dict):
-        with modify_yaml(self.path, fs=self.repo.fs) as data:
+        with (
+            self._thread_lock(),
+            modify_yaml(self.path, fs=self.repo.fs) as data,
+        ):
             data.update({"schema": "2.0"})
             if not data:
                 logger.info("Generating lock file '%s'", self.relpath)
@@ -430,7 +447,10 @@ def dump_stages(self, stages, **kwargs):
 
         is_modified = False
         log_updated = False
-        with modify_yaml(self.path, fs=self.repo.fs) as data:
+        with (
+            self._thread_lock(),
+            modify_yaml(self.path, fs=self.repo.fs) as data,
+        ):
             if not data:
                 data.update({"schema": "2.0"})
                 # order is important, meta should always be at the top
@@ -468,7 +488,8 @@ def remove_stage(self, stage):
         del data[stage.name]
 
         if data:
-            dump_yaml(self.path, d)
+            with self._thread_lock():
+                dump_yaml(self.path, d)
         else:
             self.remove()
 

dvc/repo/reproduce.py

@@ -1,4 +1,7 @@
+import concurrent.futures
 from collections.abc import Iterable
+from dataclasses import dataclass
+from enum import Enum
 from typing import TYPE_CHECKING, Callable, NoReturn, Optional, TypeVar, Union, cast
 
 from funcy import ldistinct
@@ -120,14 +123,6 @@ def _reproduce_stage(stage: "Stage", **kwargs) -> Optional["Stage"]:
     return ret
 
 
-def _get_upstream_downstream_nodes(
-    graph: Optional["DiGraph"], node: T
-) -> tuple[list[T], list[T]]:
-    succ = list(graph.successors(node)) if graph else []
-    pre = list(graph.predecessors(node)) if graph else []
-    return succ, pre
-
-
 def _repr(stages: Iterable["Stage"]) -> str:
     return humanize.join(repr(stage.addressing) for stage in stages)
 
@@ -155,56 +150,129 @@ def _raise_error(exc: Optional[Exception], *stages: "Stage") -> NoReturn:
     raise ReproductionError(f"failed to reproduce{segment} {names}") from exc
 
 
-def _reproduce(
-    stages: list["Stage"],
-    graph: Optional["DiGraph"] = None,
-    force_downstream: bool = False,
-    on_error: str = "fail",
-    force: bool = False,
+class ReproStatus(Enum):
+    READY = "ready"
+    IN_PROGRESS = "in-progress"
+    COMPLETE = "complete"
+    SKIPPED = "skipped"
+    FAILED = "failed"
+
+
+@dataclass
+class StageInfo:
+    upstream: list["Stage"]
+    upstream_unfinished: set["Stage"]
+    downstream: list["Stage"]
+    force: bool
+    status: ReproStatus
+    result: Optional["Stage"]
+
+
+def _start_ready_stages(
+    to_repro: dict["Stage", StageInfo],
+    executor: concurrent.futures.ThreadPoolExecutor,
     repro_fn: Callable = _reproduce_stage,
     **kwargs,
+) -> dict[concurrent.futures.Future["Stage"], "Stage"]:
+    futures = {
+        executor.submit(
+            repro_fn,
+            stage,
+            upstream=stage_info.upstream,
+            force=stage_info.force,
+            **kwargs,
+        ): stage
+        for stage, stage_info in to_repro.items()
+        if stage_info.status == ReproStatus.READY and not stage_info.upstream_unfinished
+    }
+    for stage in futures.values():
+        to_repro[stage].status = ReproStatus.IN_PROGRESS
+    return futures
+
+
+def _result_or_raise(
+    to_repro: dict["Stage", StageInfo], stages: list["Stage"], on_error: str
 ) -> list["Stage"]:
-    assert on_error in ("fail", "keep-going", "ignore")
-
     result: list[Stage] = []
     failed: list[Stage] = []
-    to_skip: dict[Stage, Stage] = {}
-    ret: Optional[Stage] = None
-
-    force_state = dict.fromkeys(stages, force)
-
+    # Preserve original order
     for stage in stages:
-        if stage in to_skip:
-            continue
-
-        if ret:
-            logger.info("")  # add a newline
-
-        upstream, downstream = _get_upstream_downstream_nodes(graph, stage)
-        force_stage = force_state[stage]
-
-        try:
-            ret = repro_fn(stage, upstream=upstream, force=force_stage, **kwargs)
-        except Exception as exc:  # noqa: BLE001
+        stage_info = to_repro[stage]
+        if stage_info.status == ReproStatus.FAILED:
             failed.append(stage)
-            if on_error == "fail":
-                _raise_error(exc, stage)
-
-            dependents = handle_error(graph, on_error, exc, stage)
-            to_skip.update(dict.fromkeys(dependents, stage))
-            continue
-
-        if force_downstream and (ret or force_stage):
-            force_state.update(dict.fromkeys(downstream, True))
-
-        if ret:
-            result.append(ret)
+        elif stage_info.result:
+            result.append(stage_info.result)
 
     if on_error != "ignore" and failed:
         _raise_error(None, *failed)
+
     return result
 
 
+def _reproduce(
+    stages: list["Stage"],
+    graph: Optional["DiGraph"] = None,
+    force_downstream: bool = False,
+    on_error: str = "fail",
+    force: bool = False,
+    max_workers: int = 10,
+    **kwargs,
+) -> list["Stage"]:
+    assert on_error in ("fail", "keep-going", "ignore")
+
+    ret: Optional[Stage] = None
+    to_repro = {
+        stage: StageInfo(
+            upstream=(upstream := list(graph.successors(stage)) if graph else []),
+            upstream_unfinished=set(upstream).intersection(stages),
+            downstream=list(graph.predecessors(stage)) if graph else [],
+            force=force,
+            status=ReproStatus.READY,
+            result=None,
+        )
+        for stage in stages
+    }
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = _start_ready_stages(to_repro, executor, **kwargs)
+        while futures:
+            done, _ = concurrent.futures.wait(
+                futures, return_when=concurrent.futures.FIRST_COMPLETED
+            )
+            for future in done:
+                stage = futures.pop(future)
+                stage_info = to_repro[stage]
+
+                try:
+                    ret = future.result()
+                    stage_info.status = ReproStatus.COMPLETE
+                except Exception as exc:  # noqa: BLE001
+                    if on_error == "fail":
+                        return _raise_error(exc, stage)
+
+                    stage_info.status = ReproStatus.FAILED
+                    dependents = handle_error(graph, on_error, exc, stage)
+                    for dependent in dependents:
+                        to_repro[dependent].status = ReproStatus.SKIPPED
+
+                success = stage_info.status == ReproStatus.COMPLETE
+                for dependent in stage_info.downstream:
+                    if dependent not in to_repro:
+                        continue
+                    dependent_info = to_repro[dependent]
+                    if stage in dependent_info.upstream_unfinished:
+                        dependent_info.upstream_unfinished.remove(stage)
+                    if success and force_downstream and (ret or stage_info.force):
+                        dependent_info.force = True
+
+                if success and ret:
+                    stage_info.result = ret
+
+            futures.update(_start_ready_stages(to_repro, executor, **kwargs))
+
+    return _result_or_raise(to_repro, stages, on_error)
+
+
 @locked
 @scm_context
 def reproduce(

dvc/rwlock.py

@@ -1,5 +1,6 @@
 import json
 import os
+import threading
 from collections import defaultdict
 from contextlib import contextmanager
 
@@ -25,6 +26,8 @@
     }
 )
 
+RWLOCK_THREAD_LOCK = threading.Lock()
+RWLOCK_THREAD_TIMEOUT = 3
 RWLOCK_FILE = "rwlock"
 RWLOCK_LOCK = "rwlock.lock"
 
@@ -50,21 +53,26 @@ def _edit_rwlock(lock_dir, fs, hardlink):
         tmp_dir=lock_dir,
         hardlink_lock=hardlink,
     )
-    with rwlock_guard:
-        try:
-            with fs.open(path, encoding="utf-8") as fobj:
-                lock = SCHEMA(json.load(fobj))
-        except FileNotFoundError:
-            lock = SCHEMA({})
-        except json.JSONDecodeError as exc:
-            raise RWLockFileCorruptedError(path) from exc
-        except Invalid as exc:
-            raise RWLockFileFormatError(path) from exc
-        lock["read"] = defaultdict(list, lock["read"])
-        lock["write"] = defaultdict(dict, lock["write"])
-        yield lock
-        with fs.open(path, "w", encoding="utf-8") as fobj:
-            json.dump(lock, fobj)
+    RWLOCK_THREAD_LOCK.acquire(timeout=RWLOCK_THREAD_TIMEOUT)
+    try:
+        with rwlock_guard:
+            try:
+                with fs.open(path, encoding="utf-8") as fobj:
+                    lock = SCHEMA(json.load(fobj))
+            except FileNotFoundError:
+                lock = SCHEMA({})
+            except json.JSONDecodeError as exc:
+                raise RWLockFileCorruptedError(path) from exc
+            except Invalid as exc:
+                raise RWLockFileFormatError(path) from exc
+            lock["read"] = defaultdict(list, lock["read"])
+            lock["write"] = defaultdict(dict, lock["write"])
+            yield lock
+            with fs.open(path, "w", encoding="utf-8") as fobj:
+                json.dump(lock, fobj)
+    finally:
+        if RWLOCK_THREAD_LOCK.locked():
+            RWLOCK_THREAD_LOCK.release()
 
 
 def _infos_to_str(infos):

dvc/stage/decorators.py

@@ -1,3 +1,4 @@
+import threading
 from functools import wraps
 
 from funcy import decorator
@@ -47,6 +48,8 @@ def _chain(names):
 def unlocked_repo(f):
     @wraps(f)
     def wrapper(stage, *args, **kwargs):
+        if threading.current_thread() is not threading.main_thread():
+            return f(stage, *args, **kwargs)
         stage.repo.lock.unlock()
         stage.repo._reset()
         try:

tests/unit/test_rwlock.py

@@ -1,5 +1,7 @@
+import concurrent.futures
 import json
 import os
+import time
 
 import pytest
 
@@ -59,6 +61,7 @@ def test_rwlock_reentrant(tmp_path):
 def test_rwlock_edit_is_guarded(tmp_path, mocker):
     # patching to speedup tests
     mocker.patch("dvc.lock.DEFAULT_TIMEOUT", 0.01)
+    mocker.patch("dvc.rwlock.RWLOCK_THREAD_TIMEOUT", 0.01)
 
     path = os.fspath(tmp_path)
 
@@ -68,6 +71,21 @@ def test_rwlock_edit_is_guarded(tmp_path, mocker):
                 pass
 
 
+def test_rwlock_multiple_threads(tmp_path, mocker):
+    # patching to speedup tests
+    mocker.patch("dvc.rwlock.RWLOCK_THREAD_TIMEOUT", 0.01)
+    path = os.fspath(tmp_path)
+    foo = "foo"
+
+    def work():
+        with rwlock(path, localfs, "cmd1", [foo], [], False):
+            time.sleep(1)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+        futures = [executor.submit(work) for _ in range(2)]
+        concurrent.futures.wait(futures)
+
+
 def test_rwlock_subdirs(tmp_path):
     path = os.fspath(tmp_path)
     foo = "foo"