diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index 43aa61304..f13e12c3a 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -157,11 +157,15 @@ class WorkspaceQueue(BaseStashQueue): raise NotImplementedError def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]: + from dvc.rwlock import check_rwlock from dvc.utils.serialize import load_json assert self._EXEC_NAME - result: Dict[str, Dict] = {} + + if not check_rwlock(self.repo.tmp_dir, self.repo.fs, autocorrect=True): + return result + infofile = self.get_infofile_path(self._EXEC_NAME) try: diff --git a/dvc/rwlock.py b/dvc/rwlock.py index b6937c155..a1e5a9af5 100644 --- a/dvc/rwlock.py +++ b/dvc/rwlock.py @@ -1,14 +1,25 @@ import json +import logging import os from collections import defaultdict from contextlib import contextmanager +from typing import TYPE_CHECKING, Dict, List +import psutil +from funcy import first from voluptuous import Invalid, Optional, Required, Schema from .exceptions import DvcException from .fs import localfs from .lock import make_lock from .utils import relpath +from .utils.fs import remove + +if TYPE_CHECKING: + from .fs import FileSystem + +logger = logging.getLogger(__name__) + INFO_SCHEMA = {Required("pid"): int, Required("cmd"): str} @@ -19,6 +30,9 @@ SCHEMA = Schema( } ) +RWLOCK_FILE = "rwlock" +RWLOCK_LOCK = "rwlock.lock" + class RWLockFileCorruptedError(DvcException): def __init__(self, path): @@ -35,10 +49,10 @@ class RWLockFileFormatError(DvcException): @contextmanager def _edit_rwlock(lock_dir, fs, hardlink): - path = fs.path.join(lock_dir, "rwlock") + path = fs.path.join(lock_dir, RWLOCK_FILE) rwlock_guard = make_lock( - fs.path.join(lock_dir, "rwlock.lock"), + fs.path.join(lock_dir, RWLOCK_LOCK), tmp_dir=lock_dir, hardlink_lock=hardlink, ) @@ -176,3 +190,62 @@ def rwlock(tmp_dir, fs, cmd, read, write, hardlink): with _edit_rwlock(tmp_dir, fs, hardlink) as lock: _release_write(lock, info, wchanges) _release_read(lock, info, rchanges) + + +def check_rwlock( + tmp_dir: str, + fs: "FileSystem", + hardlink: bool = False, + autocorrect: bool = False, +) -> bool: + """Check and autocorrect the RWLock status for file paths. + + Args: + tmp_dir (str): existing directory where to create the rwlock file. + fs (FileSystem): fs instance that tmp_dir belongs to. + hardlink (bool): use hardlink lock to guard rwlock file when on edit. + autocorrect (bool): autocorrect corrupted rwlock file. + + Return: + (bool): if the pid alive. + """ + path = fs.path.join(tmp_dir, RWLOCK_FILE) + + rwlock_guard = make_lock( + fs.path.join(tmp_dir, RWLOCK_LOCK), + tmp_dir=tmp_dir, + hardlink_lock=hardlink, + ) + with rwlock_guard: + try: + with fs.open(path, encoding="utf-8") as fobj: + lock: Dict[str, List[Dict]] = SCHEMA(json.load(fobj)) + file_path = first(lock["read"]) + if not file_path: + return False + lock_info = first(lock["read"][file_path]) + pid = int(lock_info["pid"]) + if psutil.pid_exists(pid): + return True + cmd = lock_info["cmd"] + logger.warning( + "Process '%s' with (Pid %s), in RWLock-file '%s'" + " had been killed.", + cmd, + pid, + relpath(path), + ) + except FileNotFoundError: + return False + except json.JSONDecodeError: + logger.warning( + "Unable to read RWLock-file '%s'. JSON structure is" + " corrupted", + relpath(path), + ) + except Invalid: + logger.warning("RWLock-file '%s' format error.", relpath(path)) + if autocorrect: + logger.warning("Delete corrupted RWLock-file '%s'", relpath(path)) + remove(path) + return False diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py index 6851ff9ff..08ee4c8b7 100644 --- a/tests/func/experiments/test_show.py +++ b/tests/func/experiments/test_show.py @@ -17,10 +17,20 @@ from dvc.repo.experiments.executor.base import ( from dvc.repo.experiments.queue.base import QueueEntry from dvc.repo.experiments.refs import CELERY_STASH, ExpRefInfo from dvc.repo.experiments.utils import exp_refs_by_rev +from dvc.utils import relpath from dvc.utils.fs import makedirs from dvc.utils.serialize import YAMLFileCorruptedError from tests.func.test_repro_multistage import COPY_SCRIPT +LOCK_CONTENTS = { + "read": { + "data/MNIST": [{"pid": 54062, "cmd": "dvc exp run"}], + }, + "write": { + "data/MNIST": {"pid": 54062, "cmd": "dvc exp run"}, + }, +} + def make_executor_info(**kwargs): # set default values for required info fields @@ -410,9 +420,21 @@ def test_show_sort(tmp_dir, scm, dvc, exp_stage, caplog): @pytest.mark.vscode -@pytest.mark.parametrize("status", [TaskStatus.RUNNING, TaskStatus.FAILED]) -def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys, status): +@pytest.mark.parametrize( + "status, pid_exists", + [ + (TaskStatus.RUNNING, True), + (TaskStatus.RUNNING, False), + (TaskStatus.FAILED, False), + ], +) +def test_show_running_workspace( + tmp_dir, scm, dvc, exp_stage, capsys, caplog, status, pid_exists, mocker +): + from dvc.rwlock import RWLOCK_FILE + pid_dir = os.path.join(dvc.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR) + lock_file = relpath(os.path.join(dvc.tmp_dir, RWLOCK_FILE), str(tmp_dir)) info = make_executor_info( location=BaseExecutor.DEFAULT_LOCATION, status=status ) @@ -423,6 +445,9 @@ def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys, status): ) makedirs(os.path.dirname(pidfile), True) (tmp_dir / pidfile).dump_json(info.asdict()) + (tmp_dir / lock_file).dump_json(LOCK_CONTENTS) + + mocker.patch("psutil.pid_exists", return_value=pid_exists) assert dvc.experiments.show().get("workspace") == { "baseline": { @@ -438,10 +463,10 @@ def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys, status): "params": {"params.yaml": {"data": {"foo": 1}}}, "outs": {}, "status": "Running" - if status == TaskStatus.RUNNING + if status == TaskStatus.RUNNING and pid_exists else "Success", "executor": info.location - if status == TaskStatus.RUNNING + if status == TaskStatus.RUNNING and pid_exists else None, "timestamp": None, } @@ -451,8 +476,19 @@ def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys, status): assert main(["exp", "show", "--csv"]) == 0 cap = capsys.readouterr() if status == TaskStatus.RUNNING: - assert "Running" in cap.out - assert info.location in cap.out + if pid_exists: + assert "Running" in cap.out + assert info.location in cap.out + else: + cmd = LOCK_CONTENTS["read"]["data/MNIST"][0]["cmd"] + pid = LOCK_CONTENTS["read"]["data/MNIST"][0]["pid"] + assert ( + f"Process '{cmd}' with (Pid {pid}), in RWLock-file " + f"'{lock_file}' had been killed." + ) in caplog.text + assert ( + f"Delete corrupted RWLock-file '{lock_file}'" + ) in caplog.text def test_show_running_tempdir(tmp_dir, scm, dvc, exp_stage, mocker):