mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-06 07:18:28 +01:00
exp: migrate workspace repro to queues
This commit is contained in:
Peter Rowlands
committed by
Peter Rowlands (변기호)
Peter Rowlands (변기호)
parent
e74f8f9e11
commit
a0ae9d3f45
@@ -1,12 +1,10 @@
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import Dict, Iterable, List, Mapping, Optional, Type
|
||||
from typing import Dict, Iterable, Mapping, Optional
|
||||
|
||||
from funcy import cached_property, first
|
||||
|
||||
from dvc.dependency.param import MissingParamsError
|
||||
from dvc.env import DVCLIVE_RESUME
|
||||
from dvc.exceptions import DvcException
|
||||
from dvc.utils import relpath
|
||||
|
||||
@@ -22,22 +20,19 @@ from .executor.base import (
|
||||
BaseExecutor,
|
||||
ExecutorInfo,
|
||||
)
|
||||
from .executor.manager.base import BaseExecutorManager
|
||||
from .executor.manager.local import (
|
||||
TempDirExecutorManager,
|
||||
WorkspaceExecutorManager,
|
||||
)
|
||||
from .executor.manager.ssh import SSHExecutorManager
|
||||
from .queue.base import BaseStashQueue, QueueEntry
|
||||
from .queue.local import LocalCeleryQueue, WorkspaceQueue
|
||||
from .refs import (
|
||||
CELERY_STASH,
|
||||
EXEC_APPLY,
|
||||
EXEC_BRANCH,
|
||||
EXEC_CHECKPOINT,
|
||||
EXEC_NAMESPACE,
|
||||
EXPS_NAMESPACE,
|
||||
EXPS_STASH,
|
||||
WORKSPACE_STASH,
|
||||
ExpRefInfo,
|
||||
)
|
||||
from .stash import ExpStash, ExpStashEntry
|
||||
from .stash import ExpStashEntry
|
||||
from .utils import exp_refs_by_rev, scm_locked, unlocked_repo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -82,255 +77,19 @@ class Experiments:
|
||||
return os.path.join(self.repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE)
|
||||
|
||||
@cached_property
|
||||
def stash(self) -> ExpStash:
|
||||
return ExpStash(self.scm, EXPS_STASH)
|
||||
def workspace_queue(self) -> WorkspaceQueue:
|
||||
return WorkspaceQueue(self.repo, WORKSPACE_STASH)
|
||||
|
||||
@cached_property
|
||||
def celery_queue(self) -> LocalCeleryQueue:
|
||||
return LocalCeleryQueue(self.repo, CELERY_STASH)
|
||||
|
||||
@property
|
||||
def stash_revs(self) -> Dict[str, ExpStashEntry]:
|
||||
return self.stash.stash_revs
|
||||
|
||||
def _stash_exp(
|
||||
self,
|
||||
*args,
|
||||
params: Optional[dict] = None,
|
||||
resume_rev: Optional[str] = None,
|
||||
baseline_rev: Optional[str] = None,
|
||||
branch: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Stash changes from the workspace as an experiment.
|
||||
|
||||
Args:
|
||||
params: Optional dictionary of parameter values to be used.
|
||||
Values take priority over any parameters specified in the
|
||||
user's workspace.
|
||||
resume_rev: Optional checkpoint resume rev.
|
||||
baseline_rev: Optional baseline rev for this experiment, defaults
|
||||
to the current SCM rev.
|
||||
branch: Optional experiment branch name. If specified, the
|
||||
experiment will be added to `branch` instead of creating
|
||||
a new branch.
|
||||
name: Optional experiment name. If specified this will be used as
|
||||
the human-readable name in the experiment branch ref. Has no
|
||||
effect of branch is specified.
|
||||
"""
|
||||
with self.scm.detach_head(client="dvc") as orig_head:
|
||||
stash_head = orig_head
|
||||
if baseline_rev is None:
|
||||
baseline_rev = orig_head
|
||||
|
||||
with self.scm.stash_workspace() as workspace:
|
||||
try:
|
||||
if workspace:
|
||||
self.stash.apply(workspace)
|
||||
|
||||
if resume_rev:
|
||||
# move HEAD to the resume rev so that the stashed diff
|
||||
# only contains changes relative to resume rev
|
||||
stash_head = resume_rev
|
||||
self.scm.set_ref(
|
||||
"HEAD",
|
||||
resume_rev,
|
||||
message=f"dvc: resume from HEAD {resume_rev[:7]}",
|
||||
)
|
||||
self.scm.reset()
|
||||
|
||||
# update experiment params from command line
|
||||
if params:
|
||||
self._update_params(params)
|
||||
|
||||
# DVC commit data deps to preserve state across workspace
|
||||
# & tempdir runs
|
||||
self._stash_commit_deps(*args, **kwargs)
|
||||
|
||||
if resume_rev:
|
||||
if branch:
|
||||
branch_name = ExpRefInfo.from_ref(branch).name
|
||||
else:
|
||||
branch_name = f"{resume_rev[:7]}"
|
||||
if self.scm.is_dirty(untracked_files=False):
|
||||
logger.info(
|
||||
"Modified checkpoint experiment based on "
|
||||
"'%s' will be created",
|
||||
branch_name,
|
||||
)
|
||||
branch = None
|
||||
elif (
|
||||
not branch
|
||||
or self.scm.get_ref(branch) != resume_rev
|
||||
):
|
||||
msg = [
|
||||
(
|
||||
"Nothing to do for unchanged checkpoint "
|
||||
f"'{resume_rev[:7]}'. "
|
||||
)
|
||||
]
|
||||
if branch:
|
||||
msg.append(
|
||||
"To resume from the head of this "
|
||||
"experiment, use "
|
||||
f"'dvc exp apply {branch_name}'."
|
||||
)
|
||||
else:
|
||||
names = [
|
||||
ref_info.name
|
||||
for ref_info in exp_refs_by_rev(
|
||||
self.scm, resume_rev
|
||||
)
|
||||
]
|
||||
if len(names) > 3:
|
||||
names[3:] = [
|
||||
f"... ({len(names) - 3} more)"
|
||||
]
|
||||
msg.append(
|
||||
"To resume an experiment containing this "
|
||||
"checkpoint, apply one of these heads:\n"
|
||||
"\t{}".format(", ".join(names))
|
||||
)
|
||||
raise DvcException("".join(msg))
|
||||
else:
|
||||
logger.info(
|
||||
"Existing checkpoint experiment '%s' will be "
|
||||
"resumed",
|
||||
branch_name,
|
||||
)
|
||||
if name:
|
||||
logger.warning(
|
||||
"Ignoring option '--name %s' for resumed "
|
||||
"experiment. Existing experiment name will"
|
||||
"be preserved instead.",
|
||||
name,
|
||||
)
|
||||
|
||||
# save additional repro command line arguments
|
||||
run_env = {DVCLIVE_RESUME: "1"} if resume_rev else {}
|
||||
self._pack_args(*args, run_env=run_env, **kwargs)
|
||||
|
||||
# save experiment as a stash commit
|
||||
msg = self._stash_msg(
|
||||
stash_head,
|
||||
baseline_rev=baseline_rev,
|
||||
branch=branch,
|
||||
name=name,
|
||||
)
|
||||
stash_rev = self.stash.push(message=msg)
|
||||
logger.debug(
|
||||
(
|
||||
"Stashed experiment '%s' with baseline '%s' "
|
||||
"for future execution."
|
||||
),
|
||||
stash_rev[:7],
|
||||
baseline_rev[:7],
|
||||
)
|
||||
finally:
|
||||
if resume_rev:
|
||||
# NOTE: this set_ref + reset() is equivalent to
|
||||
# `git reset orig_head` (our SCM reset() only operates
|
||||
# on HEAD rather than any arbitrary commit)
|
||||
self.scm.set_ref(
|
||||
"HEAD", orig_head, message="dvc: restore HEAD"
|
||||
)
|
||||
self.scm.reset()
|
||||
# Revert any of our changes before prior unstashing
|
||||
self.scm.reset(hard=True)
|
||||
|
||||
return stash_rev
|
||||
|
||||
def _stash_commit_deps(self, *args, **kwargs):
|
||||
if len(args):
|
||||
targets = args[0]
|
||||
else:
|
||||
targets = kwargs.get("targets")
|
||||
if isinstance(targets, str):
|
||||
targets = [targets]
|
||||
elif not targets:
|
||||
targets = [None]
|
||||
for target in targets:
|
||||
self.repo.commit(
|
||||
target,
|
||||
with_deps=True,
|
||||
recursive=kwargs.get("recursive", False),
|
||||
force=True,
|
||||
allow_missing=True,
|
||||
data_only=True,
|
||||
)
|
||||
|
||||
def _stash_msg(
|
||||
self,
|
||||
rev: str,
|
||||
baseline_rev: str,
|
||||
branch: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
):
|
||||
if not baseline_rev:
|
||||
baseline_rev = rev
|
||||
msg = ExpStash.format_message(rev, baseline_rev, name)
|
||||
if branch:
|
||||
return f"{msg}:{branch}"
|
||||
return msg
|
||||
|
||||
def _pack_args(self, *args, **kwargs):
|
||||
import pickle
|
||||
|
||||
if os.path.exists(self.args_file) and self.scm.is_tracked(
|
||||
self.args_file
|
||||
):
|
||||
logger.warning(
|
||||
(
|
||||
"Temporary DVC file '.dvc/tmp/%s' exists and was "
|
||||
"likely committed to Git by mistake. It should be removed "
|
||||
"with:\n"
|
||||
"\tgit rm .dvc/tmp/%s"
|
||||
),
|
||||
BaseExecutor.PACKED_ARGS_FILE,
|
||||
BaseExecutor.PACKED_ARGS_FILE,
|
||||
)
|
||||
with open(self.args_file, "rb") as fobj:
|
||||
try:
|
||||
data = pickle.load(fobj)
|
||||
except Exception: # pylint: disable=broad-except
|
||||
data = {}
|
||||
extra = int(data.get("extra", 0)) + 1
|
||||
else:
|
||||
extra = None
|
||||
BaseExecutor.pack_repro_args(
|
||||
self.args_file, *args, extra=extra, **kwargs
|
||||
)
|
||||
self.scm.add(self.args_file)
|
||||
|
||||
def _format_new_params_msg(self, new_params, config_path):
|
||||
"""Format an error message for when new parameters are identified"""
|
||||
new_param_count = len(new_params)
|
||||
pluralise = "s are" if new_param_count > 1 else " is"
|
||||
param_list = ", ".join(new_params)
|
||||
return (
|
||||
f"{new_param_count} parameter{pluralise} missing "
|
||||
f"from '{config_path}': {param_list}"
|
||||
)
|
||||
|
||||
def _update_params(self, params: dict):
|
||||
"""Update experiment params files with the specified values."""
|
||||
from dvc.utils.collections import NewParamsFound, merge_params
|
||||
from dvc.utils.serialize import MODIFIERS
|
||||
|
||||
logger.debug("Using experiment params '%s'", params)
|
||||
|
||||
for path in params:
|
||||
suffix = self.repo.fs.path.suffix(path).lower()
|
||||
modify_data = MODIFIERS[suffix]
|
||||
with modify_data(path, fs=self.repo.fs) as data:
|
||||
try:
|
||||
merge_params(data, params[path], allow_new=False)
|
||||
except NewParamsFound as e:
|
||||
msg = self._format_new_params_msg(e.new_params, path)
|
||||
raise MissingParamsError(msg)
|
||||
|
||||
# Force params file changes to be staged in git
|
||||
# Otherwise in certain situations the changes to params file may be
|
||||
# ignored when we `git stash` them since mtime is used to determine
|
||||
# whether the file is dirty
|
||||
self.scm.add(list(params.keys()))
|
||||
revs = {}
|
||||
for queue in (self.workspace_queue, self.celery_queue):
|
||||
revs.update(queue.stash.stash_revs)
|
||||
return revs
|
||||
|
||||
def reproduce_one(
|
||||
self,
|
||||
@@ -372,26 +131,28 @@ class Experiments:
|
||||
else:
|
||||
checkpoint_resume = self._workspace_resume_rev()
|
||||
|
||||
stash_rev = self.new(
|
||||
checkpoint_resume=checkpoint_resume, reset=reset, **kwargs
|
||||
exp_queue = (
|
||||
self.celery_queue if (tmp_dir or queue) else self.workspace_queue
|
||||
)
|
||||
entry = self.new(
|
||||
exp_queue,
|
||||
checkpoint_resume=checkpoint_resume,
|
||||
reset=reset,
|
||||
**kwargs,
|
||||
)
|
||||
if queue:
|
||||
logger.info(
|
||||
"Queued experiment '%s' for future execution.", stash_rev[:7]
|
||||
"Queued experiment '%s' for future execution.",
|
||||
entry.stash_rev[:7],
|
||||
)
|
||||
return [stash_rev]
|
||||
if tmp_dir or queue:
|
||||
manager_cls: Type = TempDirExecutorManager
|
||||
elif machine:
|
||||
manager_cls = SSHExecutorManager
|
||||
else:
|
||||
manager_cls = WorkspaceExecutorManager
|
||||
results = self._reproduce_revs(
|
||||
revs=[stash_rev],
|
||||
keep_stash=False,
|
||||
manager_cls=manager_cls,
|
||||
machine=machine,
|
||||
)
|
||||
return [entry.stash_rev]
|
||||
if tmp_dir:
|
||||
# TODO: start celery worker
|
||||
raise NotImplementedError
|
||||
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)
|
||||
@@ -405,10 +166,12 @@ class Experiments:
|
||||
return None
|
||||
|
||||
def reproduce_queued(self, **kwargs):
|
||||
results = self._reproduce_revs(**kwargs)
|
||||
if results:
|
||||
self._log_reproduced(results, tmp_dir=True)
|
||||
return results
|
||||
# TODO: start celery worker
|
||||
raise NotImplementedError
|
||||
# results = self._reproduce_revs(**kwargs)
|
||||
# if results:
|
||||
# self._log_reproduced(results, tmp_dir=True)
|
||||
# return results
|
||||
|
||||
def _log_reproduced(self, revs: Iterable[str], tmp_dir: bool = False):
|
||||
names = []
|
||||
@@ -443,15 +206,20 @@ class Experiments:
|
||||
raise ExperimentExistsError(exp_ref.name)
|
||||
|
||||
@scm_locked
|
||||
def new(self, *args, checkpoint_resume: Optional[str] = None, **kwargs):
|
||||
"""Create a new experiment.
|
||||
def new(
|
||||
self,
|
||||
queue: BaseStashQueue,
|
||||
*args,
|
||||
checkpoint_resume: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> QueueEntry:
|
||||
"""Create and enqueue a new experiment.
|
||||
|
||||
Experiment will be reproduced and checked out into the user's
|
||||
workspace.
|
||||
Experiment will be derived from the current workspace.
|
||||
"""
|
||||
if checkpoint_resume is not None:
|
||||
return self._resume_checkpoint(
|
||||
*args, resume_rev=checkpoint_resume, **kwargs
|
||||
queue, *args, resume_rev=checkpoint_resume, **kwargs
|
||||
)
|
||||
|
||||
name = kwargs.get("name", None)
|
||||
@@ -464,16 +232,16 @@ class Experiments:
|
||||
if not (kwargs.get("force", False) or kwargs.get("reset", False)):
|
||||
raise err
|
||||
|
||||
return self._stash_exp(*args, **kwargs)
|
||||
return queue.put(*args, **kwargs)
|
||||
|
||||
def _resume_checkpoint(
|
||||
self, *args, resume_rev: Optional[str] = None, **kwargs
|
||||
):
|
||||
"""Resume an existing (checkpoint) experiment.
|
||||
|
||||
Experiment will be reproduced and checked out into the user's
|
||||
workspace.
|
||||
"""
|
||||
self,
|
||||
queue: BaseStashQueue,
|
||||
*args,
|
||||
resume_rev: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> QueueEntry:
|
||||
"""Create and queue a resumed checkpoint experiment."""
|
||||
assert resume_rev
|
||||
|
||||
branch: Optional[str] = None
|
||||
@@ -504,7 +272,7 @@ class Experiments:
|
||||
resume_rev[:7],
|
||||
baseline_rev,
|
||||
)
|
||||
return self._stash_exp(
|
||||
return queue.put(
|
||||
*args,
|
||||
resume_rev=resume_rev,
|
||||
baseline_rev=baseline_rev,
|
||||
@@ -540,90 +308,25 @@ class Experiments:
|
||||
self.scm.remove_ref(EXEC_CHECKPOINT)
|
||||
self.scm.remove_ref(EXEC_APPLY)
|
||||
|
||||
@scm_locked
|
||||
def _reproduce_revs(
|
||||
self,
|
||||
revs: Optional[Iterable] = None,
|
||||
keep_stash: Optional[bool] = True,
|
||||
manager_cls: Type = TempDirExecutorManager,
|
||||
machine: Optional[str] = None,
|
||||
**kwargs,
|
||||
@unlocked_repo
|
||||
def _reproduce_queue(
|
||||
self, queue: BaseStashQueue, **kwargs
|
||||
) -> Mapping[str, str]:
|
||||
"""Reproduce the specified experiments.
|
||||
"""Reproduce queued experiments.
|
||||
|
||||
Args:
|
||||
revs: If revs is not specified, all stashed experiments will be
|
||||
reproduced.
|
||||
keep_stash: If True, stashed experiments will be preserved if they
|
||||
fail to reproduce successfully.
|
||||
Arguments:
|
||||
queue: Experiment queue.
|
||||
|
||||
Returns:
|
||||
dict mapping successfully reproduced experiment revs to their
|
||||
hashes.
|
||||
results.
|
||||
"""
|
||||
stash_revs = self.stash_revs
|
||||
exec_results = queue.reproduce()
|
||||
|
||||
# to_run contains mapping of:
|
||||
# input_rev: (stash_index, rev, baseline_rev)
|
||||
# where input_rev contains the changes to execute (usually a stash
|
||||
# commit), rev is the original SCM commit to be checked out, and
|
||||
# baseline_rev is the experiment baseline.
|
||||
if revs is None:
|
||||
to_run = dict(stash_revs)
|
||||
else:
|
||||
to_run = {
|
||||
rev: stash_revs[rev]
|
||||
if rev in stash_revs
|
||||
else ExpStashEntry(None, rev, rev, None, None)
|
||||
for rev in revs
|
||||
}
|
||||
|
||||
logger.debug(
|
||||
"Reproducing experiment revs '%s'",
|
||||
", ".join(rev[:7] for rev in to_run),
|
||||
)
|
||||
|
||||
manager = manager_cls.from_stash_entries(
|
||||
self.scm,
|
||||
os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR),
|
||||
self.repo,
|
||||
to_run,
|
||||
machine_name=machine,
|
||||
)
|
||||
try:
|
||||
exec_results = {}
|
||||
exec_results.update(self._executors_repro(manager, **kwargs))
|
||||
finally:
|
||||
# only drop successfully run stashed experiments
|
||||
to_drop: List[int] = [
|
||||
entry.stash_index
|
||||
for rev, entry in to_run.items()
|
||||
if (
|
||||
entry.stash_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] = {}
|
||||
results: Dict[str, str] = {}
|
||||
for _, exp_result in exec_results.items():
|
||||
result.update(exp_result)
|
||||
return result
|
||||
|
||||
@unlocked_repo
|
||||
def _executors_repro(
|
||||
self,
|
||||
manager: "BaseExecutorManager",
|
||||
**kwargs,
|
||||
) -> Dict[str, Dict[str, str]]:
|
||||
"""Run dvc repro for the specified BaseExecutors in parallel.
|
||||
|
||||
Returns:
|
||||
dict mapping stash revs to the successfully executed experiments
|
||||
for each stash rev.
|
||||
"""
|
||||
return manager.exec_queue(self.repo, **kwargs)
|
||||
results.update(exp_result)
|
||||
return results
|
||||
|
||||
def check_baseline(self, exp_rev):
|
||||
baseline_sha = self.repo.scm.get_rev()
|
||||
|
||||
@@ -42,11 +42,11 @@ def gc(
|
||||
removed = len(to_remove)
|
||||
|
||||
delete_stashes = []
|
||||
for _, entry in repo.experiments.stash_revs.items():
|
||||
for _, entry in repo.experiments.workspace_queue.stash.stash_revs.items():
|
||||
if not queued or entry.baseline_rev not in keep_revs:
|
||||
delete_stashes.append(entry.stash_index)
|
||||
for index in sorted(delete_stashes, reverse=True):
|
||||
repo.experiments.stash.drop(index)
|
||||
repo.experiments.workspace_queue.stash.drop(index)
|
||||
removed += len(delete_stashes)
|
||||
|
||||
return removed
|
||||
|
||||
@@ -165,7 +165,7 @@ class BaseStashQueue(ABC):
|
||||
branch_name = ExpRefInfo.from_ref(branch).name
|
||||
else:
|
||||
branch_name = f"{resume_rev[:7]}"
|
||||
if self.scm.is_dirty():
|
||||
if self.scm.is_dirty(untracked_files=False):
|
||||
logger.info(
|
||||
"Modified checkpoint experiment based on "
|
||||
"'%s' will be created",
|
||||
|
||||
@@ -6,6 +6,8 @@ from .exceptions import InvalidExpRefError
|
||||
# refs/exps/01/234abcd.../<exp_name>
|
||||
EXPS_NAMESPACE = "refs/exps"
|
||||
EXPS_STASH = f"{EXPS_NAMESPACE}/stash"
|
||||
WORKSPACE_STASH = EXPS_STASH
|
||||
CELERY_STASH = f"{EXPS_NAMESPACE}/celery/stash"
|
||||
EXEC_NAMESPACE = f"{EXPS_NAMESPACE}/exec"
|
||||
EXEC_APPLY = f"{EXEC_NAMESPACE}/EXEC_APPLY"
|
||||
EXEC_CHECKPOINT = f"{EXEC_NAMESPACE}/EXEC_CHECKPOINT"
|
||||
|
||||
@@ -6,6 +6,7 @@ from dvc.repo.scm_context import scm_context
|
||||
from dvc.scm import iter_revs
|
||||
|
||||
from .exceptions import UnresolvedExpNamesError
|
||||
from .queue.base import QueueEntry
|
||||
from .refs import ExpRefInfo
|
||||
from .utils import (
|
||||
exp_refs,
|
||||
@@ -39,16 +40,16 @@ def remove(
|
||||
return removed
|
||||
|
||||
if queue:
|
||||
removed.extend(_clear_stash(repo))
|
||||
removed.extend(_clear_queue(repo))
|
||||
if all_commits:
|
||||
removed.extend(_clear_all_commits(repo, git_remote))
|
||||
return removed
|
||||
|
||||
commit_ref_dict: Dict[ExpRefInfo, str] = {}
|
||||
queued_ref_dict: Dict[int, str] = {}
|
||||
queue_entry_dict: Dict[str, QueueEntry] = {}
|
||||
if exp_names:
|
||||
_resolve_exp_by_name(
|
||||
repo, exp_names, commit_ref_dict, queued_ref_dict, git_remote
|
||||
repo, exp_names, commit_ref_dict, queue_entry_dict, git_remote
|
||||
)
|
||||
|
||||
if rev:
|
||||
@@ -59,8 +60,8 @@ def remove(
|
||||
_remove_commited_exps(repo.scm, commit_ref_dict, git_remote)
|
||||
)
|
||||
|
||||
if queued_ref_dict:
|
||||
removed.extend(_remove_queued_exps(repo, queued_ref_dict))
|
||||
if queue_entry_dict:
|
||||
removed.extend(_remove_queued_exps(repo, queue_entry_dict))
|
||||
|
||||
return removed
|
||||
|
||||
@@ -68,8 +69,8 @@ def remove(
|
||||
def _resolve_exp_by_name(
|
||||
repo: "Repo",
|
||||
exp_names: Union[str, List[str]],
|
||||
commit_ref_dict: Dict[ExpRefInfo, str],
|
||||
queued_ref_dict: Dict[int, str],
|
||||
commit_ref_dict: Dict["ExpRefInfo", str],
|
||||
queue_entry_dict: Dict[str, QueueEntry],
|
||||
git_remote: Optional[str],
|
||||
):
|
||||
remained = set()
|
||||
@@ -84,10 +85,10 @@ def _resolve_exp_by_name(
|
||||
commit_ref_dict[exp_ref] = exp_name
|
||||
|
||||
if not git_remote:
|
||||
stash_index_dict = _get_queued_index_by_names(repo, remained)
|
||||
for exp_name, stash_index in stash_index_dict.items():
|
||||
if stash_index is not None:
|
||||
queued_ref_dict[stash_index] = exp_name
|
||||
_named_entries = _get_queue_entry_by_names(repo, remained)
|
||||
for exp_name, entry in _named_entries.items():
|
||||
if entry is not None:
|
||||
queue_entry_dict[exp_name] = entry
|
||||
remained.remove(exp_name)
|
||||
|
||||
if remained:
|
||||
@@ -111,15 +112,11 @@ def _resolve_exp_by_baseline(
|
||||
commit_ref_dict[ref_info] = ref_info.name
|
||||
|
||||
|
||||
def _clear_stash(repo: "Repo") -> List[str]:
|
||||
def _clear_queue(repo: "Repo") -> List[str]:
|
||||
removed_name_list = []
|
||||
for entry in repo.experiments.stash.list():
|
||||
new_sha = entry.new_sha.decode("utf8")
|
||||
name = repo.experiments.get_exact_name(new_sha)
|
||||
if not name:
|
||||
name = new_sha[:7]
|
||||
removed_name_list.append(name)
|
||||
repo.experiments.stash.clear()
|
||||
for entry in repo.experiments.celery_queue.iter_queued():
|
||||
removed_name_list.append(entry.name or entry.stash_rev[:7])
|
||||
repo.experiments.celery_queue.clear()
|
||||
return removed_name_list
|
||||
|
||||
|
||||
@@ -130,25 +127,25 @@ def _clear_all_commits(repo, git_remote) -> List:
|
||||
return _remove_commited_exps(repo.scm, ref_infos, git_remote)
|
||||
|
||||
|
||||
def _get_queued_index_by_names(
|
||||
repo,
|
||||
def _get_queue_entry_by_names(
|
||||
repo: "Repo",
|
||||
exp_name_set: Set[str],
|
||||
) -> Mapping[str, Optional[int]]:
|
||||
) -> Dict[str, Optional[QueueEntry]]:
|
||||
from scmrepo.exceptions import RevError as InternalRevError
|
||||
|
||||
result = {}
|
||||
stash_revs = repo.experiments.stash_revs
|
||||
for _, entry in stash_revs.items():
|
||||
rev_entries = {}
|
||||
for entry in repo.experiments.celery_queue.iter_queued():
|
||||
if entry.name in exp_name_set:
|
||||
result[entry.name] = entry.stash_index
|
||||
result[entry.name] = entry
|
||||
else:
|
||||
rev_entries[entry.stash_rev] = entry
|
||||
|
||||
for exp_name in exp_name_set:
|
||||
if exp_name in result:
|
||||
continue
|
||||
for exp_name in exp_name_set.difference(result.keys()):
|
||||
try:
|
||||
rev = repo.scm.resolve_rev(exp_name)
|
||||
if rev in stash_revs:
|
||||
result[exp_name] = stash_revs.get(rev).stash_index
|
||||
if rev in rev_entries:
|
||||
result[exp_name] = rev_entries[rev]
|
||||
except InternalRevError:
|
||||
result[exp_name] = None
|
||||
return result
|
||||
@@ -174,9 +171,9 @@ def _remove_commited_exps(
|
||||
return list(exp_ref_dict.values())
|
||||
|
||||
|
||||
def _remove_queued_exps(repo, indexes: Mapping[int, str]) -> List[str]:
|
||||
index_list = list(indexes.keys())
|
||||
index_list.sort(reverse=True)
|
||||
for index in index_list:
|
||||
repo.experiments.stash.drop(index)
|
||||
return list(indexes.values())
|
||||
def _remove_queued_exps(
|
||||
repo: "Repo", named_entries: Mapping[str, QueueEntry]
|
||||
) -> List[str]:
|
||||
for entry in named_entries.values():
|
||||
repo.experiments.celery_queue.remove(entry.stash_rev)
|
||||
return list(named_entries.keys())
|
||||
|
||||
@@ -7,7 +7,7 @@ from scmrepo.git import Stash
|
||||
class ExpStashEntry(NamedTuple):
|
||||
"""Experiment stash entry.
|
||||
|
||||
stash_index: EXPS_STASH index for this entry. Can be None if this commit
|
||||
stash_index: Stash index for this entry. Can be None if this commit
|
||||
is not pushed onto the stash ref.
|
||||
head_rev: HEAD Git commit to be checked out for this experiment.
|
||||
baseline_rev: Experiment baseline commit.
|
||||
|
||||
Reference in New Issue
Block a user