mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-01 23:08:43 +01:00
eeshsaxena
f74c1c0e70
- repo/experiments/queue/base.py: `scm` -> `repo` in BaseStashQueue.__init__ (signature takes a Repo, not an SCM instance) - repo/experiments/queue/tasks.py: stale `tmp_dir`/`entry_dict` args in cleanup_exp replaced with the actual `executor`/`infofile` parameters
124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
from typing import TYPE_CHECKING, Any, Optional
|
|
|
|
from celery import shared_task
|
|
from celery.utils.log import get_task_logger
|
|
|
|
from dvc.repo.experiments.executor.base import ExecutorInfo
|
|
from dvc.repo.experiments.executor.local import TempDirExecutor
|
|
|
|
from .base import BaseStashQueue, QueueEntry
|
|
|
|
if TYPE_CHECKING:
|
|
from dvc.repo.experiments.executor.base import BaseExecutor
|
|
|
|
|
|
logger = get_task_logger(__name__)
|
|
|
|
|
|
@shared_task
|
|
def setup_exp(entry_dict: dict[str, Any]) -> "BaseExecutor":
|
|
"""Setup an experiment.
|
|
|
|
Arguments:
|
|
entry_dict: Serialized QueueEntry for this experiment.
|
|
|
|
Returns:
|
|
Root executor (temp) directory for this experiment.
|
|
"""
|
|
from dvc.repo import Repo
|
|
|
|
entry = QueueEntry.from_dict(entry_dict)
|
|
with Repo(entry.dvc_root) as repo:
|
|
# TODO: split executor.init_cache into separate subtask - we can release
|
|
# exp.scm_lock before DVC push
|
|
executor = BaseStashQueue.init_executor(
|
|
repo.experiments,
|
|
entry,
|
|
TempDirExecutor,
|
|
location="dvc-task",
|
|
)
|
|
infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev)
|
|
executor.info.dump_json(infofile)
|
|
return executor
|
|
|
|
|
|
@shared_task
|
|
def collect_exp(
|
|
proc_dict: dict[str, Any], # noqa: ARG001
|
|
entry_dict: dict[str, Any],
|
|
) -> str:
|
|
"""Collect results for an experiment.
|
|
|
|
Arguments:
|
|
proc_dict: Serialized ProcessInfo for experiment executor process.
|
|
entry_dict: Serialized QueueEntry for this experiment.
|
|
|
|
Returns:
|
|
Directory to be cleaned up after this experiment.
|
|
"""
|
|
from dvc.repo import Repo
|
|
|
|
entry = QueueEntry.from_dict(entry_dict)
|
|
with Repo(entry.dvc_root) as repo:
|
|
celery_queue = repo.experiments.celery_queue
|
|
infofile = celery_queue.get_infofile_path(entry.stash_rev)
|
|
executor_info = ExecutorInfo.load_json(infofile)
|
|
logger.debug("Collecting experiment info '%s'", str(executor_info))
|
|
executor = TempDirExecutor.from_info(executor_info)
|
|
exec_result = executor_info.result
|
|
try:
|
|
if exec_result is not None:
|
|
BaseStashQueue.collect_executor(repo.experiments, executor, exec_result)
|
|
else:
|
|
logger.debug("Experiment failed (Exec result was None)")
|
|
celery_queue.stash_failed(entry)
|
|
except Exception:
|
|
# Log exceptions but do not re-raise so that task chain execution
|
|
# continues
|
|
logger.exception("Failed to collect experiment")
|
|
return executor.root_dir
|
|
|
|
|
|
@shared_task
|
|
def cleanup_exp(executor: TempDirExecutor, infofile: str) -> None:
|
|
"""Cleanup after an experiment.
|
|
|
|
Arguments:
|
|
executor: Executor to clean up.
|
|
infofile: Path to the executor info file.
|
|
"""
|
|
executor.cleanup(infofile)
|
|
|
|
|
|
@shared_task
|
|
def run_exp(
|
|
entry_dict: dict[str, Any],
|
|
copy_paths: Optional[list[str]] = None,
|
|
message: Optional[str] = None,
|
|
) -> None:
|
|
"""Run a full experiment.
|
|
|
|
Experiment subtasks are executed inline as one atomic operation.
|
|
|
|
Arguments:
|
|
entry_dict: Serialized QueueEntry for this experiment.
|
|
"""
|
|
from dvc.repo import Repo
|
|
|
|
entry = QueueEntry.from_dict(entry_dict)
|
|
with Repo(entry.dvc_root) as repo:
|
|
queue = repo.experiments.celery_queue
|
|
infofile = queue.get_infofile_path(entry.stash_rev)
|
|
executor = setup_exp.s(entry_dict)()
|
|
try:
|
|
cmd = ["dvc", "exp", "exec-run", "--infofile", infofile]
|
|
if copy_paths:
|
|
for path in copy_paths:
|
|
cmd.extend(["--copy-paths", path])
|
|
if message:
|
|
cmd.extend(["--message", message])
|
|
proc_dict = queue.proc.run_signature(cmd, name=entry.stash_rev)()
|
|
collect_exp.s(proc_dict, entry_dict)()
|
|
finally:
|
|
cleanup_exp.s(executor, infofile)()
|