diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py index 1fd83e859..c8021f12f 100644 --- a/dvc/command/experiments.py +++ b/dvc/command/experiments.py @@ -151,7 +151,13 @@ def _collect_rows( new_checkpoint = True for i, (rev, exp) in enumerate(experiments.items()): - queued = str(exp.get("queued") or "") + if exp.get("running"): + state = "Running" + elif exp.get("queued"): + state = "Queued" + else: + state = FILL_VALUE + executor = exp.get("executor", FILL_VALUE) is_baseline = rev == "baseline" if is_baseline: @@ -190,10 +196,11 @@ def _collect_rows( row = [ exp_name, name_rev, - queued, typ, _format_time(exp.get("timestamp")), parent, + state, + executor, ] _extend_row( row, metric_names, exp.get("metrics", {}).items(), precision @@ -305,7 +312,15 @@ def experiments_table( from dvc.compare import TabularData - headers = ["Experiment", "rev", "queued", "typ", "Created", "parent"] + headers = [ + "Experiment", + "rev", + "typ", + "Created", + "parent", + "State", + "Executor", + ] td = TabularData( lconcat(headers, metric_headers, param_headers), fill_value=FILL_VALUE ) @@ -341,8 +356,7 @@ def prepare_exp_id(kwargs) -> "Text": text.append(suff) tree = experiment_types[typ] - queued = "*" if kwargs.get("queued") else "" - pref = (f"{tree} " if tree else "") + queued + pref = f"{tree} " if tree else "" return Text(pref) + text @@ -382,9 +396,13 @@ def show_experiments( if no_timestamp: td.drop("Created") + for col in ("State", "Executor"): + if td.is_empty(col): + td.drop(col) + row_styles = lmap(baseline_styler, td.column("typ")) - merge_headers = ["Experiment", "rev", "queued", "typ", "parent"] + merge_headers = ["Experiment", "rev", "typ", "parent"] td.column("Experiment")[:] = map(prepare_exp_id, td.as_dict(merge_headers)) td.drop(*merge_headers[1:]) @@ -392,6 +410,8 @@ def show_experiments( styles = { "Experiment": {"no_wrap": True, "header_style": "black on grey93"}, "Created": {"header_style": "black on grey93"}, + "State": {"header_style": "black on grey93"}, + "Executor": {"header_style": "black on grey93"}, } header_bg_colors = {"metrics": "cornsilk1", "params": "light_cyan1"} styles.update( diff --git a/dvc/compare.py b/dvc/compare.py index 096338344..d185d9b05 100644 --- a/dvc/compare.py +++ b/dvc/compare.py @@ -134,6 +134,10 @@ class TabularData(MutableSequence[Sequence["CellT"]]): self.drop(*(set(self._keys) - set(col_names))) self._keys = list(col_names) + def is_empty(self, col_name: str) -> bool: + col = self.column(col_name) + return not any(item != self._fill_value for item in col) + def to_csv(self) -> str: import csv from io import StringIO diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 64b41623d..8377403e5 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -81,6 +81,7 @@ class Experiments: r"(?P-checkpoint)?$" ) EXEC_TMP_DIR = "exps" + EXEC_PID_DIR = os.path.join(EXEC_TMP_DIR, "run") StashEntry = namedtuple( "StashEntry", ["index", "rev", "baseline_rev", "branch", "name"] @@ -588,26 +589,21 @@ class Experiments: ) executors = self._init_executors(to_run) - exec_results = self._executors_repro(executors, **kwargs) - - if keep_stash: + try: + exec_results = {} + exec_results.update(self._executors_repro(executors, **kwargs)) + finally: # only drop successfully run stashed experiments - to_drop = sorted( - ( - stash_revs[rev][0] - for rev in exec_results - if rev in stash_revs - ), - reverse=True, - ) - else: - # drop all stashed experiments - to_drop = sorted( - (stash_revs[rev][0] for rev in to_run if rev in stash_revs), - reverse=True, - ) - for index in to_drop: - self.stash.drop(index) + to_drop = [ + entry.index + for rev, entry in to_run.items() + if ( + entry.index is not None + and (not keep_stash or rev in exec_results) + ) + ] + for index in sorted(to_drop, reverse=True): + self.stash.drop(index) result: Dict[str, str] = {} for _, exp_result in exec_results.items(): @@ -623,6 +619,9 @@ class Experiments: base_tmp_dir = os.path.join(self.repo.tmp_dir, self.EXEC_TMP_DIR) if not os.path.exists(base_tmp_dir): makedirs(base_tmp_dir) + pid_dir = os.path.join(self.repo.tmp_dir, self.EXEC_PID_DIR) + if not os.path.exists(pid_dir): + makedirs(pid_dir) for stash_rev, item in to_run.items(): self.scm.set_ref(EXEC_HEAD, item.rev) self.scm.set_ref(EXEC_MERGE, stash_rev) @@ -649,9 +648,10 @@ class Experiments: return executors + @unlocked_repo def _executors_repro( self, executors: dict, jobs: Optional[int] = 1 - ) -> Mapping[str, Mapping[str, str]]: + ) -> Dict[str, Dict[str, str]]: """Run dvc repro for the specified BaseExecutors in parallel. Returns: @@ -667,6 +667,11 @@ class Experiments: with ProcessPoolExecutor(max_workers=jobs) as workers: futures = {} for rev, executor in executors.items(): + pidfile = os.path.join( + self.repo.tmp_dir, + self.EXEC_PID_DIR, + f"{rev}{executor.PIDFILE_EXT}", + ) future = workers.submit( executor.reproduce, executor.dvc_dir, @@ -675,6 +680,8 @@ class Experiments: name=executor.name, rel_cwd=rel_cwd, log_level=logger.getEffectiveLevel(), + pidfile=pidfile, + git_url=executor.git_url, ) futures[future] = (rev, executor) @@ -683,12 +690,15 @@ class Experiments: except KeyboardInterrupt: # forward SIGINT to any running executor processes and # cancel any remaining futures + workers.shutdown(wait=False) pids = {} - while not pid_q.empty(): - rev, pid = pid_q.get() - pids[rev] = pid for future, (rev, _) in futures.items(): if future.running(): + # if future has already been started by the scheduler + # we still have to wait until it tells us its PID + while rev not in pids: + rev, pid = pid_q.get() + pids[rev] = pid os.kill(pids[rev], signal.SIGINT) elif not future.done(): future.cancel() @@ -735,7 +745,10 @@ class Experiments: raise ExperimentExistsError(ref_info.name) for ref in executor.fetch_exps( - self.scm, force=exec_result.force, on_diverged=on_diverged + self.scm, + executor.git_url, + force=exec_result.force, + on_diverged=on_diverged, ): exp_rev = self.scm.get_ref(ref) if exp_rev: @@ -747,6 +760,8 @@ class Experiments: @unlocked_repo def _workspace_repro(self) -> Mapping[str, str]: """Run the most recently stashed experiment in the workspace.""" + from dvc.utils.fs import makedirs + from .executor.base import BaseExecutor entry = first(self.stash_revs.values()) @@ -766,12 +781,19 @@ class Experiments: self.scm.remove_ref(EXEC_BRANCH) try: orig_checkpoint = self.scm.get_ref(EXEC_CHECKPOINT) + pid_dir = os.path.join(self.repo.tmp_dir, self.EXEC_PID_DIR) + if not os.path.exists(pid_dir): + makedirs(pid_dir) + pidfile = os.path.join( + pid_dir, f"workspace{BaseExecutor.PIDFILE_EXT}" + ) exec_result = BaseExecutor.reproduce( None, rev, name=entry.name, rel_cwd=relpath(os.getcwd(), self.scm.root_dir), log_errors=False, + pidfile=pidfile, ) if not exec_result.exp_hash: @@ -861,6 +883,51 @@ class Experiments: return ExpRefInfo.from_ref(ref).name return None + def get_running_exps(self) -> Dict[str, int]: + """Return info for running experiments.""" + from dvc.utils.serialize import load_yaml + + from .executor.base import BaseExecutor, ExecutorInfo + + result = {} + for pidfile in self.repo.fs.walk_files( + os.path.join(self.repo.tmp_dir, self.EXEC_PID_DIR) + ): + rev, _ = os.path.splitext(os.path.basename(pidfile)) + + try: + info = ExecutorInfo.from_dict(load_yaml(pidfile)) + if rev == "workspace": + # If we are appending to a checkpoint branch in a workspace + # run, show the latest checkpoint as running. + last_rev = self.scm.get_ref(EXEC_BRANCH) + if last_rev: + result[last_rev] = info.to_dict() + else: + result[rev] = info.to_dict() + else: + result[rev] = info.to_dict() + if info.git_url: + + def on_diverged(_ref: str, _checkpoint: bool): + return False + + for ref in BaseExecutor.fetch_exps( + self.scm, + info.git_url, + on_diverged=on_diverged, + ): + logger.debug( + "Updated running experiment '%s'.", ref + ) + last_rev = self.scm.get_ref(ref) + result[rev]["last"] = last_rev + if last_rev: + result[last_rev] = info.to_dict() + except OSError: + pass + return result + def apply(self, *args, **kwargs): from dvc.repo.experiments.apply import apply diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index a980da0b7..f79772804 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -3,6 +3,7 @@ import os import pickle from abc import ABC, abstractmethod from contextlib import contextmanager +from dataclasses import dataclass from functools import partial from typing import ( TYPE_CHECKING, @@ -52,6 +53,36 @@ class ExecutorResult(NamedTuple): force: bool +@dataclass +class ExecutorInfo: + PARAM_PID = "pid" + PARAM_GIT_URL = "git" + PARAM_BASELINE_REV = "baseline" + PARAM_LOCATION = "location" + + pid: Optional[int] + git_url: Optional[str] + baseline_rev: Optional[str] + location: Optional[str] + + @classmethod + def from_dict(cls, d): + return cls( + d.get(cls.PARAM_PID), + d.get(cls.PARAM_GIT_URL), + d.get(cls.PARAM_BASELINE_REV), + d.get(cls.PARAM_LOCATION), + ) + + def to_dict(self): + return { + self.PARAM_PID: self.pid, + self.PARAM_GIT_URL: self.git_url, + self.PARAM_BASELINE_REV: self.baseline_rev, + self.PARAM_LOCATION: self.location, + } + + class BaseExecutor(ABC): """Base class for executing experiments in parallel. @@ -66,6 +97,8 @@ class BaseExecutor(ABC): PACKED_ARGS_FILE = "repro.dat" WARN_UNTRACKED = False QUIET = False + PIDFILE_EXT = ".run" + DEFAULT_LOCATION: Optional[str] = "workspace" def __init__( self, @@ -161,9 +194,11 @@ class BaseExecutor(ABC): data = pickle.load(fobj) return data["args"], data["kwargs"] + @classmethod def fetch_exps( - self, + cls, dest_scm: "Git", + url: str, force: bool = False, on_diverged: Callable[[str, bool], None] = None, ) -> Iterable[str]: @@ -171,13 +206,17 @@ class BaseExecutor(ABC): Args: dest_scm: Destination Git instance. + url: Git remote URL to fetch from. 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. """ refs = [] - for ref in self.scm.iter_refs(base=EXPS_NAMESPACE): - if not ref.startswith(EXEC_NAMESPACE) and ref != EXPS_STASH: + has_checkpoint = False + for ref in dest_scm.iter_remote_refs(url, base=EXPS_NAMESPACE): + if ref == EXEC_CHECKPOINT: + has_checkpoint = True + elif not ref.startswith(EXEC_NAMESPACE) and ref != EXPS_STASH: refs.append(ref) def on_diverged_ref(orig_ref: str, new_rev: str): @@ -185,24 +224,25 @@ class BaseExecutor(ABC): logger.debug("Replacing existing experiment '%s'", orig_ref) return True - checkpoint = self.scm.get_ref(EXEC_CHECKPOINT) is not None - self._raise_ref_conflict(dest_scm, orig_ref, new_rev, checkpoint) + cls._raise_ref_conflict( + dest_scm, orig_ref, new_rev, has_checkpoint + ) if on_diverged: - on_diverged(orig_ref, checkpoint) + on_diverged(orig_ref, has_checkpoint) logger.debug("Reproduced existing experiment '%s'", orig_ref) return False # fetch experiments dest_scm.fetch_refspecs( - self.git_url, + url, [f"{ref}:{ref}" for ref in refs], on_diverged=on_diverged_ref, force=force, ) # update last run checkpoint (if it exists) - if self.scm.get_ref(EXEC_CHECKPOINT): + if has_checkpoint: dest_scm.fetch_refspecs( - self.git_url, + url, [f"{EXEC_CHECKPOINT}:{EXEC_CHECKPOINT}"], force=True, ) @@ -218,6 +258,7 @@ class BaseExecutor(ABC): name: Optional[str] = None, log_errors: bool = True, log_level: Optional[int] = None, + **kwargs, ) -> "ExecutorResult": """Run dvc repro and return the result. @@ -245,7 +286,12 @@ class BaseExecutor(ABC): exp_ref: Optional["ExpRefInfo"] = None repro_force: bool = False - with cls._repro_dvc(dvc_dir, rel_cwd, log_errors) as dvc: + with cls._repro_dvc( + dvc_dir, + rel_cwd, + log_errors, + **kwargs, + ) as dvc: args, kwargs = cls._repro_args(dvc) if args: targets: Optional[Union[list, str]] = args[0] @@ -339,9 +385,16 @@ class BaseExecutor(ABC): @classmethod @contextmanager def _repro_dvc( - cls, dvc_dir: Optional[str], rel_cwd: Optional[str], log_errors: bool + cls, + dvc_dir: Optional[str], + rel_cwd: Optional[str], + log_errors: bool, + pidfile: Optional[str] = None, + git_url: Optional[str] = None, + **kwargs, ): from dvc.repo import Repo + from dvc.utils.serialize import modify_yaml dvc = Repo(dvc_dir) if cls.QUIET: @@ -354,6 +407,15 @@ class BaseExecutor(ABC): os.chdir(dvc.root_dir) else: old_cwd = None + if pidfile is not None: + info = ExecutorInfo( + os.getpid(), + git_url, + dvc.scm.get_rev(), + cls.DEFAULT_LOCATION, + ) + with modify_yaml(pidfile) as d: + d.update(info.to_dict()) logger.debug("Running repro in '%s'", os.getcwd()) try: @@ -369,6 +431,8 @@ class BaseExecutor(ABC): logger.exception("unexpected error") raise finally: + if pidfile is not None: + remove(pidfile) dvc.close() if old_cwd: os.chdir(old_cwd) diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index fd4e4bf13..f5658e128 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/dvc/repo/experiments/executor/local.py @@ -45,6 +45,7 @@ class TempDirExecutor(BaseLocalExecutor): # suggestions) that are not applicable outside of workspace runs WARN_UNTRACKED = True QUIET = True + DEFAULT_LOCATION: Optional[str] = "temp" def __init__( self, diff --git a/dvc/repo/experiments/show.py b/dvc/repo/experiments/show.py index 3da741ca1..44a73a175 100644 --- a/dvc/repo/experiments/show.py +++ b/dvc/repo/experiments/show.py @@ -5,6 +5,7 @@ from datetime import datetime from dvc.exceptions import InvalidArgumentError from dvc.repo import locked from dvc.repo.experiments.base import ExpRefInfo +from dvc.repo.experiments.executor.base import ExecutorInfo from dvc.repo.experiments.utils import fix_exp_head from dvc.repo.metrics.show import _collect_metrics, _read_metrics from dvc.repo.params.show import _collect_configs, _read_params @@ -14,7 +15,7 @@ logger = logging.getLogger(__name__) def _collect_experiment_commit( - repo, exp_rev, stash=False, sha_only=True, param_deps=False + repo, exp_rev, stash=False, sha_only=True, param_deps=False, running=None ): res = defaultdict(dict) for rev in repo.brancher(revs=[exp_rev]): @@ -34,6 +35,12 @@ def _collect_experiment_commit( res["params"] = params res["queued"] = stash + if running is not None and exp_rev in running: + res["running"] = True + res["executor"] = running[exp_rev].get(ExecutorInfo.PARAM_LOCATION) + else: + res["running"] = False + res["executor"] = None if not stash: metrics = _collect_metrics(repo, None, rev, False) vals = _read_metrics(repo, metrics, rev) @@ -117,9 +124,15 @@ def show( ) ) + running = repo.experiments.get_running_exps() + for rev in revs: res[rev]["baseline"] = _collect_experiment_commit( - repo, rev, sha_only=sha_only, param_deps=param_deps + repo, + rev, + sha_only=sha_only, + param_deps=param_deps, + running=running, ) if rev == "workspace": @@ -142,14 +155,20 @@ def show( rev, sha_only=sha_only, param_deps=param_deps, + running=running, ) # collect queued (not yet reproduced) experiments for stash_rev, entry in repo.experiments.stash_revs.items(): if entry.baseline_rev in revs: - experiment = _collect_experiment_commit( - repo, stash_rev, stash=True, param_deps=param_deps - ) - res[entry.baseline_rev][stash_rev] = experiment + if stash_rev not in running or not running[stash_rev].get("last"): + experiment = _collect_experiment_commit( + repo, + stash_rev, + stash=stash_rev not in running, + param_deps=param_deps, + running=running, + ) + res[entry.baseline_rev][stash_rev] = experiment return res diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py index 030901f02..758b02ebc 100644 --- a/tests/func/experiments/test_show.py +++ b/tests/func/experiments/test_show.py @@ -7,7 +7,9 @@ from funcy import first from dvc.exceptions import InvalidArgumentError from dvc.main import main -from dvc.repo.experiments.base import ExpRefInfo +from dvc.repo.experiments.base import EXPS_STASH, ExpRefInfo +from dvc.repo.experiments.executor.base import BaseExecutor, ExecutorInfo +from dvc.utils.fs import makedirs from dvc.utils.serialize import dump_yaml from tests.func.test_repro_multistage import COPY_SCRIPT @@ -18,6 +20,8 @@ def test_show_simple(tmp_dir, scm, dvc, exp_stage): "metrics": {"metrics.yaml": {"foo": 1}}, "params": {"params.yaml": {"foo": 1}}, "queued": False, + "running": False, + "executor": None, "timestamp": None, } } @@ -39,6 +43,8 @@ def test_show_experiment(tmp_dir, scm, dvc, exp_stage, workspace): "metrics": {"metrics.yaml": {"foo": 1}}, "params": {"params.yaml": {"foo": 1}}, "queued": False, + "running": False, + "executor": None, "timestamp": timestamp, "name": "master", } @@ -56,8 +62,6 @@ def test_show_experiment(tmp_dir, scm, dvc, exp_stage, workspace): def test_show_queued(tmp_dir, scm, dvc, exp_stage): - from dvc.repo.experiments.base import EXPS_STASH - baseline_rev = scm.get_rev() dvc.experiments.run(exp_stage.addressing, params=["foo=2"], queue=True) @@ -346,3 +350,87 @@ def test_show_sort(tmp_dir, scm, dvc, exp_stage, caplog): assert ( main(["exp", "show", "--no-pager", "--sort-by=metrics.yaml:foo"]) == 0 ) + + +def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys): + pid_dir = os.path.join(dvc.tmp_dir, dvc.experiments.EXEC_PID_DIR) + makedirs(pid_dir, True) + info = ExecutorInfo(None, None, None, BaseExecutor.DEFAULT_LOCATION) + pidfile = os.path.join(pid_dir, f"workspace{BaseExecutor.PIDFILE_EXT}") + dump_yaml(pidfile, info.to_dict()) + + assert dvc.experiments.show()["workspace"] == { + "baseline": { + "metrics": {"metrics.yaml": {"foo": 1}}, + "params": {"params.yaml": {"foo": 1}}, + "queued": False, + "running": True, + "executor": info.location, + "timestamp": None, + } + } + + capsys.readouterr() + assert main(["exp", "show", "--no-pager"]) == 0 + cap = capsys.readouterr() + assert "Running" in cap.out + assert info.location in cap.out + + +def test_show_running_executor(tmp_dir, scm, dvc, exp_stage): + baseline_rev = scm.get_rev() + dvc.experiments.run(exp_stage.addressing, params=["foo=2"], queue=True) + exp_rev = dvc.experiments.scm.resolve_rev(f"{EXPS_STASH}@{{0}}") + + pid_dir = os.path.join(dvc.tmp_dir, dvc.experiments.EXEC_PID_DIR) + makedirs(pid_dir, True) + info = ExecutorInfo(None, None, None, BaseExecutor.DEFAULT_LOCATION) + pidfile = os.path.join(pid_dir, f"{exp_rev}{BaseExecutor.PIDFILE_EXT}") + dump_yaml(pidfile, info.to_dict()) + + results = dvc.experiments.show() + assert not results[baseline_rev][exp_rev]["queued"] + assert results[baseline_rev][exp_rev]["running"] + assert results[baseline_rev][exp_rev]["executor"] == info.location + assert not results["workspace"]["baseline"]["running"] + + +@pytest.mark.parametrize("workspace", [True, False]) +def test_show_running_checkpoint( + tmp_dir, scm, dvc, checkpoint_stage, workspace, mocker +): + from dvc.repo.experiments.base import EXEC_BRANCH + from dvc.repo.experiments.executor.local import TempDirExecutor + from dvc.repo.experiments.utils import exp_refs_by_rev + + baseline_rev = scm.get_rev() + dvc.experiments.run( + checkpoint_stage.addressing, params=["foo=2"], queue=True + ) + stash_rev = dvc.experiments.scm.resolve_rev(f"{EXPS_STASH}@{{0}}") + + run_results = dvc.experiments.run(run_all=True) + checkpoint_rev = first(run_results) + exp_ref = first(exp_refs_by_rev(scm, checkpoint_rev)) + + pid_dir = os.path.join(dvc.tmp_dir, dvc.experiments.EXEC_PID_DIR) + makedirs(pid_dir, True) + executor = ( + BaseExecutor.DEFAULT_LOCATION + if workspace + else TempDirExecutor.DEFAULT_LOCATION + ) + info = ExecutorInfo(123, "foo.git", baseline_rev, executor) + rev = "workspace" if workspace else stash_rev + pidfile = os.path.join(pid_dir, f"{rev}{BaseExecutor.PIDFILE_EXT}") + dump_yaml(pidfile, info.to_dict()) + + mocker.patch.object( + BaseExecutor, "fetch_exps", return_value=[str(exp_ref)] + ) + if workspace: + scm.set_ref(EXEC_BRANCH, str(exp_ref), symbolic=True) + results = dvc.experiments.show() + assert results[baseline_rev][checkpoint_rev]["running"] + assert results[baseline_rev][checkpoint_rev]["executor"] == info.location + assert not results["workspace"]["baseline"]["running"]