deps: bumps dvc-azure>=2.21.2

Closes #9506 via https://github.com/iterative/dvc-azure/pull/48

(cherry picked from commit a1b891ce01)

dvc/commands/stage.py

@@ -1,5 +1,6 @@
 import argparse
 import logging
+from contextlib import contextmanager
 from itertools import chain, filterfalse
 from typing import TYPE_CHECKING, Dict, Iterable, List
 
@@ -119,8 +120,22 @@ def parse_cmd(commands: List[str]) -> str:
     return " ".join(map(quote_argument, commands))
 
 
+@contextmanager
+def _disable_logging(highest_level=logging.CRITICAL):
+    previous_level = logging.root.manager.disable
+
+    logging.disable(highest_level)
+
+    try:
+        yield
+    finally:
+        logging.disable(previous_level)
+
+
 class CmdStageAdd(CmdBase):
     def run(self):
+        from dvc.repo import lock_repo
+
         kwargs = vars(self.args)
         kwargs.update(
             {
@@ -128,7 +143,15 @@ class CmdStageAdd(CmdBase):
                 "params": parse_params(self.args.params),
             }
         )
-        self.repo.stage.add(**kwargs)
+
+        with self.repo.scm_context, lock_repo(self.repo):
+            with _disable_logging(logging.INFO):
+                stage = self.repo.stage.add(**kwargs)
+            logger.info("Added stage %r in %r", stage.addressing, stage.relpath)
+            if self.args.run:
+                stage.run()
+                stage.dump(update_pipeline=False)
+
         return 0
 
 
@@ -257,6 +280,12 @@ def _add_common_args(parser):
             "This doesn't affect any DVC operations."
         ),
     )
+    parser.add_argument(
+        "--run",
+        action="store_true",
+        default=False,
+        help="Execute the stage after generating it.",
+    )
     parser.add_argument(
         "command",
         nargs=argparse.REMAINDER,

dvc/repo/experiments/__init__.py

@@ -449,17 +449,6 @@ class Experiments:
             result[rev] = name
         return result
 
-    def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]:
-        """Return info for running experiments."""
-        result = {}
-        for queue in (
-            self.workspace_queue,
-            self.tempdir_queue,
-            self.celery_queue,
-        ):
-            result.update(queue.get_running_exps(fetch_refs))
-        return result
-
     def apply(self, *args, **kwargs):
         from dvc.repo.experiments.apply import apply
 

dvc/repo/experiments/brancher.py

@@ -0,0 +1,47 @@
+from contextlib import ExitStack, contextmanager
+from typing import TYPE_CHECKING, Iterator, Tuple
+
+from dvc.repo.experiments.exceptions import InvalidExpRevError
+from dvc.scm import RevError
+
+if TYPE_CHECKING:
+    from dvc.repo import Repo
+
+
+@contextmanager
+def switch_repo(
+    repo: "Repo",
+    rev: str,
+) -> Iterator[Tuple["Repo", str]]:
+    """Return a repo instance (brancher) switched to rev.
+
+    If rev is the name of a running experiment, the returned instance will be
+    the live repo wherever the experiment is running.
+
+    NOTE: This will not resolve git SHA's that only exist in queued exp workspaces
+    (it will only match queued exp names).
+    """
+    try:
+        with repo.switch(rev):
+            yield repo, rev
+        return
+    except RevError as exc:
+        orig_exc = exc
+    exps = repo.experiments
+
+    if rev == exps.workspace_queue.get_running_exp():
+        yield repo, "workspace"
+        return
+
+    for queue in (exps.tempdir_queue, exps.celery_queue):
+        try:
+            active_repo = queue.active_repo(rev)
+        except InvalidExpRevError:
+            continue
+        stack = ExitStack()
+        stack.enter_context(active_repo)
+        stack.enter_context(active_repo.switch("workspace"))
+        with stack:
+            yield active_repo, rev
+        return
+    raise orig_exc

dvc/repo/experiments/exceptions.py

@@ -114,3 +114,10 @@ class UnresolvedRunningExpNamesError(UnresolvedExpNamesError):
 
 class ExpQueueEmptyError(DvcException):
     pass
+
+
+class ExpNotStartedError(DvcException):
+    def __init__(self, name: str):
+        super().__init__(
+            f"Queued experiment '{name}' exists but has not started running yet"
+        )

dvc/repo/experiments/executor/base.py

@@ -497,6 +497,9 @@ class BaseExecutor(ABC):
         exp_ref: Optional["ExpRefInfo"] = None
         repro_force: bool = False
 
+        if info.name:
+            ui.write(f"Reproducing experiment '{info.name}'")
+
         with cls._repro_dvc(
             info,
             infofile,

dvc/repo/experiments/ls.py

@@ -1,6 +1,6 @@
 import logging
 from collections import defaultdict
-from typing import Optional
+from typing import List, Optional, Union
 
 from dvc.repo import locked
 from dvc.repo.scm_context import scm_context
@@ -15,14 +15,17 @@ logger = logging.getLogger(__name__)
 @scm_context
 def ls(
     repo,
-    rev: Optional[str] = None,
+    rev: Optional[Union[List[str], str]] = None,
     all_commits: bool = False,
     num: int = 1,
     git_remote: Optional[str] = None,
 ):
     rev_set = None
     if not all_commits:
-        revs = iter_revs(repo.scm, [rev or "HEAD"], num)
+        rev = rev or "HEAD"
+        if isinstance(rev, str):
+            rev = [rev]
+        revs = iter_revs(repo.scm, rev, num)
         rev_set = set(revs.keys())
     ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote)
 

dvc/repo/experiments/pull.py

@@ -23,7 +23,7 @@ def pull(  # noqa: C901
     git_remote: str,
     exp_names: Union[Iterable[str], str],
     all_commits=False,
-    rev: Optional[str] = None,
+    rev: Optional[Union[List[str], str]] = None,
     num=1,
     force: bool = False,
     pull_cache: bool = False,
@@ -49,7 +49,9 @@ def pull(  # noqa: C901
                 raise UnresolvedExpNamesError(unresolved_exp_names)
 
         if rev:
-            rev_dict = iter_revs(repo.scm, [rev], num)
+            if isinstance(rev, str):
+                rev = [rev]
+            rev_dict = iter_revs(repo.scm, rev, num)
             rev_set = set(rev_dict.keys())
             ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote)
             for _, ref_info_list in ref_info_dict.items():

dvc/repo/experiments/queue/base.py

@@ -729,14 +729,6 @@ class BaseStashQueue(ABC):
                 message=f"commit: {msg}",
             )
 
-    @abstractmethod
-    def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]:
-        """Get the execution info of the currently running experiments
-
-        Args:
-            fetch_ref (bool): fetch completed checkpoints or not.
-        """
-
     @abstractmethod
     def collect_active_data(
         self,
@@ -790,3 +782,26 @@ class BaseStashQueue(ABC):
         Returns:
             Dict mapping baseline revision to list of queued experiments.
         """
+
+    def active_repo(self, name: str) -> "Repo":
+        """Return a Repo for the specified active experiment if it exists."""
+        from dvc.repo import Repo
+        from dvc.repo.experiments.exceptions import (
+            ExpNotStartedError,
+            InvalidExpRevError,
+        )
+        from dvc.repo.experiments.executor.base import ExecutorInfo, TaskStatus
+
+        for entry in self.iter_active():
+            if entry.name != name:
+                continue
+            infofile = self.get_infofile_path(entry.stash_rev)
+            executor_info = ExecutorInfo.load_json(infofile)
+            if executor_info.status < TaskStatus.RUNNING:
+                raise ExpNotStartedError(name)
+            dvc_root = os.path.join(executor_info.root_dir, executor_info.dvc_dir)
+            try:
+                return Repo(dvc_root)
+            except (FileNotFoundError, DvcException) as exc:
+                raise InvalidExpRevError(name) from exc
+        raise InvalidExpRevError(name)

dvc/repo/experiments/queue/celery.py

@@ -489,19 +489,6 @@ class LocalCeleryQueue(BaseStashQueue):
 
         return celery_remove(self, *args, **kwargs)
 
-    def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]:
-        """Get the execution info of the currently running experiments
-
-        Args:
-            fetch_ref (bool): fetch completed checkpoints or not.
-        """
-        result: Dict[str, Dict] = {}
-        for entry in self.iter_active():
-            result.update(
-                fetch_running_exp_from_temp_dir(self, entry.stash_rev, fetch_refs)
-            )
-        return result
-
     def get_ref_and_entry_by_names(
         self,
         exp_names: Union[str, List[str]],

dvc/repo/experiments/queue/tempdir.py

@@ -144,14 +144,6 @@ class TempDirQueue(WorkspaceQueue):
     ) -> Dict[str, str]:
         return BaseStashQueue.collect_executor(exp, executor, exec_result)
 
-    def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]:
-        result: Dict[str, Dict] = {}
-        for entry in self.iter_active():
-            result.update(
-                fetch_running_exp_from_temp_dir(self, entry.stash_rev, fetch_refs)
-            )
-        return result
-
     def collect_active_data(
         self,
         baseline_revs: Optional[Collection[str]],

dvc/repo/experiments/queue/workspace.py

@@ -191,33 +191,18 @@ class WorkspaceQueue(BaseStashQueue):
     ):
         raise NotImplementedError
 
-    def get_running_exps(
+    def get_running_exp(self) -> Optional[str]:
-        self,
+        """Return the name of the exp running in workspace (if it exists)."""
-        fetch_refs: bool = True,  # noqa: ARG002
-    ) -> Dict[str, Dict]:
         assert self._EXEC_NAME
-        result: Dict[str, Dict] = {}
         if self._active_pid is None:
-            return result
+            return None
 
         infofile = self.get_infofile_path(self._EXEC_NAME)
-
         try:
             info = ExecutorInfo.from_dict(load_json(infofile))
         except OSError:
-            return result
+            return None
-
+        return info.name
-        if info.status < TaskStatus.FAILED:
-            # If we are appending to a checkpoint branch in a workspace
-            # run, show the latest checkpoint as running.
-            if info.status == TaskStatus.SUCCESS:
-                return result
-            last_rev = self.scm.get_ref(EXEC_BRANCH)
-            if last_rev:
-                result[last_rev] = info.asdict()
-            else:
-                result[self._EXEC_NAME] = info.asdict()
-        return result
 
     def collect_active_data(
         self,

dvc/repo/experiments/remove.py

@@ -24,7 +24,7 @@ logger = logging.getLogger(__name__)
 def remove(  # noqa: C901, PLR0912
     repo: "Repo",
     exp_names: Union[None, str, List[str]] = None,
-    rev: Optional[str] = None,
+    rev: Optional[Union[List[str], str]] = None,
     all_commits: bool = False,
     num: int = 1,
     queue: bool = False,
@@ -60,6 +60,8 @@ def remove(  # noqa: C901, PLR0912
         if remained:
             raise UnresolvedExpNamesError(remained, git_remote=git_remote)
     elif rev:
+        if isinstance(rev, str):
+            rev = [rev]
         exp_ref_dict = _resolve_exp_by_baseline(repo, rev, num, git_remote)
         removed.extend(exp_ref_dict.keys())
         exp_ref_list.extend(exp_ref_dict.values())
@@ -85,14 +87,14 @@ def remove(  # noqa: C901, PLR0912
 
 def _resolve_exp_by_baseline(
     repo: "Repo",
-    rev: str,
+    rev: List[str],
     num: int,
     git_remote: Optional[str] = None,
 ) -> Dict[str, "ExpRefInfo"]:
     assert isinstance(repo.scm, Git)
 
     commit_ref_dict: Dict[str, "ExpRefInfo"] = {}
-    rev_dict = iter_revs(repo.scm, [rev], num)
+    rev_dict = iter_revs(repo.scm, rev, num)
     rev_set = set(rev_dict.keys())
     ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote)
     for _, ref_info_list in ref_info_dict.items():

dvc/repo/plots/__init__.py

@@ -121,20 +121,25 @@ class Plots:
 
             }
         """
+        from dvc.repo.experiments.brancher import switch_repo
         from dvc.utils.collections import ensure_list
 
         targets = ensure_list(targets)
         targets = [self.repo.dvcfs.from_os_path(target) for target in targets]
 
-        for rev in self.repo.brancher(revs=revs):
+        if revs is None:
-            # .brancher() adds unwanted workspace
+            revs = ["workspace"]
-            if revs is not None and rev not in revs:
+        else:
-                continue
+            revs = list(revs)
-            rev = rev or "workspace"
+            if "workspace" in revs:
-
+                # reorder revs to match repo.brancher ordering
+                revs.remove("workspace")
+                revs = ["workspace"] + revs
+        for rev in revs:
+            with switch_repo(self.repo, rev) as (repo, _):
                 res: Dict = {}
                 definitions = _collect_definitions(
-                self.repo,
+                    repo,
                     targets=targets,
                     revision=rev,
                     onerror=onerror,
@@ -146,6 +151,7 @@ class Plots:
                     data_targets = _get_data_targets(definitions)
 
                     res[rev]["sources"] = self._collect_data_sources(
+                        repo,
                         targets=data_targets,
                         recursive=recursive,
                         props=props,
@@ -156,16 +162,17 @@ class Plots:
     @error_handler
     def _collect_data_sources(
         self,
+        repo: "Repo",
         targets: Optional[List[str]] = None,
         recursive: bool = False,
         props: Optional[Dict] = None,
         onerror: Optional[Callable] = None,
     ):
-        fs = self.repo.dvcfs
+        fs = repo.dvcfs
 
         props = props or {}
 
-        plots = _collect_plots(self.repo, targets, recursive)
+        plots = _collect_plots(repo, targets, recursive)
         res: Dict[str, Any] = {}
         for fs_path, rev_props in plots.items():
             joined_props = {**rev_props, **props}

dvc/repo/reproduce.py

@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Iterator, List
 
 from dvc.exceptions import DvcException, ReproductionError
 from dvc.repo.scm_context import scm_context
+from dvc.stage.cache import RunCacheNotSupported
 from dvc.stage.exceptions import CheckpointKilledError
 
 from . import locked
@@ -125,9 +126,12 @@ def reproduce(  # noqa: C901, PLR0912
                 )
             )
 
-    if kwargs.get("pull", False):
+    if kwargs.get("pull", False) and kwargs.get("run_cache", True):
         logger.debug("Pulling run cache")
+        try:
             self.stage_cache.pull(None)
+        except RunCacheNotSupported as e:
+            logger.warning("Failed to pull run cache: %s", e)
 
     return _reproduce_stages(self.index.graph, list(stages), **kwargs)
 

dvc/repo/stage.py

@@ -122,7 +122,6 @@ class StageLoad:
             force=force,
             **stage_data,
         )
-        with self.repo.scm_context:
         stage.dump(update_lock=update_lock)
         try:
             stage.ignore_outs()

pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
 ]
 [project.optional-dependencies]
 all = ["dvc[azure,gdrive,gs,hdfs,oss,s3,ssh,webdav,webhdfs]"]
-azure = ["dvc-azure==2.21.1"]
+azure = ["dvc-azure>=2.21.2"]
 dev = ["dvc[azure,gdrive,gs,hdfs,lint,oss,s3,ssh,terraform,tests,webdav,webhdfs]"]
 gdrive = ["dvc-gdrive==2.19.2"]
 gs = ["dvc-gs==2.22.0"]
@@ -303,5 +303,5 @@ ignore-names = ["M", "SCM"]
 max-args = 10
 
 [tool.bandit]
-exclude_dirs = ["tests", "scripts"]
+exclude_dirs = ["tests", "scripts", "dvc/repo/experiments/utils.py"]
 skips = ["B101"]

tests/func/experiments/test_experiments.py

@@ -369,6 +369,12 @@ def test_list(tmp_dir, scm, dvc, exp_stage):
         baseline_a[:7]: {ref_info_a.name, ref_info_b.name}
     }
 
+    exp_list = dvc.experiments.ls(rev=[baseline_a, scm.get_rev()])
+    assert {key: set(val) for key, val in exp_list.items()} == {
+        baseline_a[:7]: {ref_info_a.name, ref_info_b.name},
+        "master": {ref_info_c.name},
+    }
+
     exp_list = dvc.experiments.ls(all_commits=True)
     assert {key: set(val) for key, val in exp_list.items()} == {
         baseline_a[:7]: {ref_info_a.name, ref_info_b.name},

tests/func/experiments/test_remote.py

@@ -251,6 +251,30 @@ def test_pull_args(tmp_dir, scm, dvc, git_downstream, exp_stage, all_, rev, resu
     assert git_downstream.tmp_dir.scm.get_ref(str(ref_info3)) == result3
 
 
+def test_pull_multi_rev(tmp_dir, scm, dvc, git_downstream, exp_stage):
+    baseline = scm.get_rev()
+
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=1"])
+    exp1 = first(results)
+    ref_info1 = first(exp_refs_by_rev(scm, exp1))
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=2"])
+    exp2 = first(results)
+    ref_info2 = first(exp_refs_by_rev(scm, exp2))
+
+    scm.commit("new_baseline")
+
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=3"])
+    exp3 = first(results)
+    ref_info3 = first(exp_refs_by_rev(scm, exp3))
+
+    downstream_exp = git_downstream.tmp_dir.dvc.experiments
+    git_downstream.tmp_dir.scm.fetch_refspecs(str(tmp_dir), ["master:master"])
+    downstream_exp.pull(git_downstream.remote, [], rev=[baseline, scm.get_rev()])
+    assert git_downstream.tmp_dir.scm.get_ref(str(ref_info1)) == exp1
+    assert git_downstream.tmp_dir.scm.get_ref(str(ref_info2)) == exp2
+    assert git_downstream.tmp_dir.scm.get_ref(str(ref_info3)) == exp3
+
+
 def test_pull_diverged(tmp_dir, scm, dvc, git_downstream, exp_stage):
     git_downstream.tmp_dir.scm_gen("foo", "foo", commit="init")
     remote_rev = git_downstream.tmp_dir.scm.get_rev()

tests/func/experiments/test_remove.py

@@ -156,3 +156,26 @@ def test_remove_experiments_by_rev(tmp_dir, scm, dvc, exp_stage):
     assert "queue2" in queue_revs
     assert scm.get_ref(new_exp_ref) is not None
     assert "queue4" in queue_revs
+
+
+def test_remove_multi_rev(tmp_dir, scm, dvc, exp_stage):
+    baseline = scm.get_rev()
+
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=1"])
+    baseline_exp_ref = first(exp_refs_by_rev(scm, first(results)))
+
+    dvc.experiments.run(
+        exp_stage.addressing, params=["foo=2"], queue=True, name="queue2"
+    )
+    scm.commit("new_baseline")
+
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=3"])
+    new_exp_ref = first(exp_refs_by_rev(scm, first(results)))
+
+    assert set(dvc.experiments.remove(rev=[baseline, scm.get_rev()])) == {
+        baseline_exp_ref.name,
+        new_exp_ref.name,
+    }
+
+    assert scm.get_ref(str(baseline_exp_ref)) is None
+    assert scm.get_ref(str(new_exp_ref)) is None

tests/func/experiments/test_utils.py

@@ -2,18 +2,19 @@ from funcy import first
 
 
 def test_generate_random_exp_name(tmp_dir, dvc, scm, exp_stage, mocker):
+    mocked_generator = mocker.MagicMock()
+    mocked_generator.choice.side_effect = [
+        0,
+        0,
+        0,
+        0,
+        1,
+        1,
+        0,
+        0,
+    ]
     mocker.patch(
-        "dvc.repo.experiments.utils.random.choice",
+        "dvc.repo.experiments.utils.random.Random", return_value=mocked_generator
-        side_effect=[
-            0,
-            0,
-            0,
-            0,
-            1,
-            1,
-            0,
-            0,
-        ],
     )
 
     ref = first(dvc.experiments.run(exp_stage.addressing, params=["foo=1"]))

tests/func/test_repro_multistage.py

@@ -9,6 +9,7 @@ from dvc.cli import main
 from dvc.dvcfile import LOCK_FILE, PROJECT_FILE
 from dvc.exceptions import CyclicGraphError, ReproductionError
 from dvc.stage import PipelineStage
+from dvc.stage.cache import RunCacheNotSupported
 from dvc.stage.exceptions import StageNotFound
 from dvc.utils.fs import remove
 
@@ -457,3 +458,30 @@ def test_repro_allow_missing_and_pull(tmp_dir, dvc, mocker, local_remote):
     ret = dvc.reproduce(pull=True, allow_missing=True)
     # create-foo is skipped ; copy-foo pulls missing dep
     assert len(ret) == 1
+
+
+def test_repro_pulls_continue_without_run_cache(tmp_dir, dvc, mocker, local_remote):
+    (foo,) = tmp_dir.dvc_gen("foo", "foo")
+
+    dvc.push()
+    mocker.patch.object(
+        dvc.stage_cache, "pull", side_effect=RunCacheNotSupported("foo")
+    )
+    dvc.stage.add(name="copy-foo", cmd="cp foo bar", deps=["foo"], outs=["bar"])
+    remove("foo")
+    remove(foo.outs[0].cache_path)
+
+    assert dvc.reproduce(pull=True)
+
+
+def test_repro_skip_pull_if_no_run_cache_is_passed(tmp_dir, dvc, mocker, local_remote):
+    (foo,) = tmp_dir.dvc_gen("foo", "foo")
+
+    dvc.push()
+    spy_pull = mocker.spy(dvc.stage_cache, "pull")
+    dvc.stage.add(name="copy-foo", cmd="cp foo bar", deps=["foo"], outs=["bar"])
+    remove("foo")
+    remove(foo.outs[0].cache_path)
+
+    assert dvc.reproduce(pull=True, run_cache=False)
+    assert not spy_pull.called

tests/unit/command/test_stage.py

@@ -86,3 +86,15 @@ def test_stage_add(mocker, dvc, command, parsed_command):
         cmd=parsed_command,
         force=True,
     )
+
+
+def test_stage_add_and_run(mocker, dvc):
+    cli_args = parse_args(["stage", "add", "--run", "-n", "foo", "-o", "foo", "cmd"])
+    cmd = cli_args.func(cli_args)
+    add_mock = mocker.patch.object(cmd.repo.stage, "add")
+
+    assert cmd.run() == 0
+
+    assert called_once_with_subset(add_mock, name="foo", outs=["foo"], cmd="cmd")
+    add_mock.return_value.run.assert_called_once()
+    add_mock.return_value.dump.assert_called_once_with(update_pipeline=False)