From 9e0788aa3a536e14338f4b971de5e2e9877745b0 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Thu, 16 Jun 2022 15:52:34 +0900 Subject: [PATCH] exp run: preserve --temp standalone behavior --- dvc/repo/experiments/__init__.py | 62 +++++++++++----------- dvc/repo/experiments/executor/local.py | 6 ++- dvc/repo/experiments/run.py | 20 ++++++- tests/func/experiments/test_experiments.py | 52 +++++++++++------- 4 files changed, 85 insertions(+), 55 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 43a1e4c76..a315b4a84 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -2,7 +2,7 @@ import logging import os import re import time -from typing import Any, Dict, Iterable, Mapping, Optional +from typing import Any, Dict, Iterable, Optional from funcy import cached_property, first @@ -99,21 +99,12 @@ class Experiments: def reproduce_one( self, - queue: bool = False, tmp_dir: bool = False, - checkpoint_resume: Optional[str] = None, - reset: bool = False, machine: Optional[str] = None, **kwargs, ): - """Reproduce and checkout a single experiment.""" - if queue and not checkpoint_resume: - reset = True - - if reset: - self.reset_checkpoints() - - if not (queue or tmp_dir or machine): + """Reproduce and checkout a single (standalone) experiment.""" + if not (tmp_dir or machine): staged, _, _ = self.scm.status(untracked_files="no") if staged: logger.warning( @@ -122,6 +113,30 @@ class Experiments: ) self.scm.reset() + exp_queue: BaseStashQueue = ( + self.tempdir_queue if tmp_dir else self.workspace_queue + ) + self.queue_one(exp_queue, **kwargs) + if machine: + # TODO: decide how to handle queued remote execution + raise NotImplementedError + results = self._reproduce_queue(exp_queue) + exp_rev = first(results) + if exp_rev is not None: + self._log_reproduced(results, tmp_dir=tmp_dir) + return results + + def queue_one( + self, + queue: BaseStashQueue, + checkpoint_resume: Optional[str] = None, + reset: bool = False, + **kwargs, + ) -> QueueEntry: + """Queue a single experiment.""" + if reset: + self.reset_checkpoints() + if checkpoint_resume: from dvc.scm import resolve_rev @@ -137,29 +152,12 @@ class Experiments: else: checkpoint_resume = self._workspace_resume_rev() - exp_queue = ( - self.celery_queue if (tmp_dir or queue) else self.workspace_queue - ) - entry = self.new( - exp_queue, + return self.new( + queue, checkpoint_resume=checkpoint_resume, reset=reset, **kwargs, ) - if queue: - name = entry.name or entry.stash_rev[:7] - ui.write(f"Queued experiment '{name}' for future execution.") - return [entry.stash_rev] - if tmp_dir: - return self.reproduce_celery(entries=[entry]) - if machine: - # TODO: decide how to handle queued remote execution - raise NotImplementedError - results = self._reproduce_queue(exp_queue) - exp_rev = first(results) - if exp_rev is not None: - self._log_reproduced(results, tmp_dir=tmp_dir) - return results def _workspace_resume_rev(self) -> Optional[str]: last_checkpoint = self._get_last_checkpoint() @@ -342,7 +340,7 @@ class Experiments: @unlocked_repo def _reproduce_queue( self, queue: BaseStashQueue, **kwargs - ) -> Mapping[str, str]: + ) -> Dict[str, str]: """Reproduce queued experiments. Arguments: diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index 3490c1394..dc41763e4 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/dvc/repo/experiments/executor/local.py @@ -8,7 +8,7 @@ from funcy import cached_property from scmrepo.exceptions import SCMError as _SCMError from dvc.scm import SCM, GitMergeError -from dvc.utils.fs import remove +from dvc.utils.fs import makedirs, remove from ..refs import ( EXEC_APPLY, @@ -123,7 +123,9 @@ class TempDirExecutor(BaseLocalExecutor): wdir: Optional[str] = None, **kwargs, ): - tmp_dir = mkdtemp(dir=wdir or os.path.join(repo.tmp_dir, EXEC_TMP_DIR)) + parent_dir: str = wdir or os.path.join(repo.tmp_dir, EXEC_TMP_DIR) + makedirs(parent_dir, exist_ok=True) + tmp_dir = mkdtemp(dir=parent_dir) try: executor = cls._from_stash_entry(repo, stash_rev, entry, tmp_dir) logger.debug("Init temp dir executor in '%s'", tmp_dir) diff --git a/dvc/repo/experiments/run.py b/dvc/repo/experiments/run.py index 8b63adc6b..b1a601b0f 100644 --- a/dvc/repo/experiments/run.py +++ b/dvc/repo/experiments/run.py @@ -1,7 +1,8 @@ import logging -from typing import Iterable, Optional +from typing import Dict, Iterable, Optional from dvc.repo import locked +from dvc.ui import ui from dvc.utils.cli_parse import loads_param_overrides logger = logging.getLogger(__name__) @@ -15,8 +16,9 @@ def run( run_all: bool = False, jobs: int = 1, tmp_dir: bool = False, + queue: bool = False, **kwargs, -) -> dict: +) -> Dict[str, str]: """Reproduce the specified targets as an experiment. Accepts the same additional kwargs as Repo.reproduce. @@ -30,6 +32,20 @@ def run( if params: params = loads_param_overrides(params) + + if queue: + if not kwargs.get("checkpoint_resume", None): + kwargs["reset"] = True + queue_entry = repo.experiments.queue_one( + repo.experiments.celery_queue, + targets=targets, + params=params, + **kwargs, + ) + name = queue_entry.name or queue_entry.stash_rev[:7] + ui.write(f"Queued experiment '{name}' for future execution.") + return {} + return repo.experiments.reproduce_one( targets=targets, params=params, tmp_dir=tmp_dir, **kwargs ) diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 2cde0feaf..54d6eb1de 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -211,19 +211,17 @@ def test_add_params(tmp_dir, scm, dvc, changes): dvc.experiments.run(stage.addressing, params=changes) -@pytest.mark.parametrize("queue", [True, False]) -def test_apply(tmp_dir, scm, dvc, exp_stage, queue, test_queue): +def test_apply(tmp_dir, scm, dvc, exp_stage): from dvc.exceptions import InvalidArgumentError from dvc.repo.experiments.exceptions import ApplyConflictError - metrics_original = (tmp_dir / "metrics.yaml").read_text().strip() results = dvc.experiments.run( - exp_stage.addressing, params=["foo=2"], queue=queue, tmp_dir=True + exp_stage.addressing, params=["foo=2"], tmp_dir=True ) exp_a = first(results) results = dvc.experiments.run( - exp_stage.addressing, params=["foo=3"], queue=queue, tmp_dir=True + exp_stage.addressing, params=["foo=3"], tmp_dir=True ) exp_b = first(results) @@ -232,29 +230,45 @@ def test_apply(tmp_dir, scm, dvc, exp_stage, queue, test_queue): dvc.experiments.apply(exp_a) assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 2" - assert ( - (tmp_dir / "metrics.yaml").read_text().strip() == metrics_original - if queue - else "foo: 2" - ) + assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 2" with pytest.raises(ApplyConflictError): dvc.experiments.apply(exp_b, force=False) # failed apply should revert everything to prior state assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 2" - assert ( - (tmp_dir / "metrics.yaml").read_text().strip() == metrics_original - if queue - else "foo: 2" - ) + assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 2" dvc.experiments.apply(exp_b) assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 3" - assert ( - (tmp_dir / "metrics.yaml").read_text().strip() == metrics_original - if queue - else "foo: 3" + assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 3" + + +def test_apply_queued(tmp_dir, scm, dvc, exp_stage, test_queue): + # from dvc.exceptions import InvalidArgumentError + from dvc.repo.experiments.exceptions import ApplyConflictError + + metrics_original = (tmp_dir / "metrics.yaml").read_text().strip() + dvc.experiments.run( + exp_stage.addressing, params=["foo=2"], name="exp-a", queue=True ) + dvc.experiments.run( + exp_stage.addressing, params=["foo=3"], name="exp-b", queue=True + ) + queue_revs = { + entry.name: entry.stash_rev + for entry in dvc.experiments.celery_queue.iter_queued() + } + + dvc.experiments.apply(queue_revs["exp-a"]) + assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 2" + assert (tmp_dir / "metrics.yaml").read_text().strip() == metrics_original + + with pytest.raises(ApplyConflictError): + dvc.experiments.apply(queue_revs["exp-b"], force=False) + + dvc.experiments.apply(queue_revs["exp-b"]) + assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 3" + assert (tmp_dir / "metrics.yaml").read_text().strip() == metrics_original def test_apply_untracked(tmp_dir, scm, dvc, exp_stage):