mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-01 23:08:43 +01:00
proc: remove dvc.proc
This commit is contained in:
Peter Rowlands
committed by
Peter Rowlands (변기호)
Peter Rowlands (변기호)
parent
3f842a205f
commit
9101fce761
@@ -1,23 +0,0 @@
|
||||
from dvc.exceptions import DvcException
|
||||
|
||||
|
||||
class ProcessNotTerminatedError(DvcException):
|
||||
def __init__(self, name):
|
||||
super().__init__(f"Managed process '{name}' has not been terminated.")
|
||||
|
||||
|
||||
class ProcessNotFoundError(DvcException):
|
||||
def __init__(self, name):
|
||||
super().__init__(f"Managed process '{name}' does not exist.")
|
||||
|
||||
|
||||
class TimeoutExpired(DvcException):
|
||||
def __init__(self, cmd, timeout):
|
||||
super().__init__(f"'{cmd}' did not complete before timeout '{timeout}'")
|
||||
self.cmd = cmd
|
||||
self.timeout = timeout
|
||||
|
||||
|
||||
class UnsupportedSignalError(DvcException):
|
||||
def __init__(self, sig):
|
||||
super().__init__(f"Unsupported signal: {sig}")
|
||||
@@ -1,149 +0,0 @@
|
||||
"""Serverless process manager."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
from typing import Generator, List, Optional, Tuple, Union
|
||||
|
||||
from funcy.flow import reraise
|
||||
from shortuuid import uuid
|
||||
|
||||
from .exceptions import ProcessNotTerminatedError, UnsupportedSignalError
|
||||
from .process import ManagedProcess, ProcessInfo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ProcessManager:
|
||||
"""Manager for controlling background ManagedProcess(es).
|
||||
|
||||
Spawned process entries are kept in the manager directory until they
|
||||
are explicitly removed (with remove() or cleanup()) so that return
|
||||
value and log information can be accessed after a process has completed.
|
||||
"""
|
||||
|
||||
def __init__(self, wdir: Optional[str] = None):
|
||||
self.wdir = wdir or "."
|
||||
|
||||
def __iter__(self) -> Generator[str, None, None]:
|
||||
if not os.path.exists(self.wdir):
|
||||
return
|
||||
yield from os.listdir(self.wdir)
|
||||
|
||||
def __getitem__(self, key: str) -> "ProcessInfo":
|
||||
info_path = os.path.join(self.wdir, key, f"{key}.json")
|
||||
try:
|
||||
with open(info_path, encoding="utf-8") as fobj:
|
||||
return ProcessInfo.from_dict(json.load(fobj))
|
||||
except FileNotFoundError:
|
||||
raise KeyError # noqa: B904
|
||||
|
||||
@reraise(FileNotFoundError, KeyError)
|
||||
def __setitem__(self, key: str, value: "ProcessInfo"):
|
||||
info_path = os.path.join(self.wdir, key, f"{key}.json")
|
||||
with open(info_path, "w", encoding="utf-8") as fobj:
|
||||
return json.dump(value.asdict(), fobj)
|
||||
|
||||
def __delitem__(self, key: str) -> None:
|
||||
from dvc.utils.fs import remove
|
||||
|
||||
path = os.path.join(self.wdir, key)
|
||||
if os.path.exists(path):
|
||||
remove(path)
|
||||
|
||||
def get(self, key: str, default=None):
|
||||
try:
|
||||
return self[key]
|
||||
except KeyError:
|
||||
return default
|
||||
|
||||
def processes(self) -> Generator[Tuple[str, "ProcessInfo"], None, None]:
|
||||
for name in self:
|
||||
try:
|
||||
yield name, self[name]
|
||||
except KeyError:
|
||||
continue
|
||||
|
||||
def spawn(self, args: Union[str, List[str]], name: Optional[str] = None):
|
||||
"""Run the given command in the background."""
|
||||
name = name or uuid()
|
||||
pid = ManagedProcess.spawn(
|
||||
args,
|
||||
wdir=os.path.join(self.wdir, name),
|
||||
name=name,
|
||||
)
|
||||
logger.debug(
|
||||
"Spawned managed process '%s' (PID: '%d')",
|
||||
name,
|
||||
pid,
|
||||
)
|
||||
|
||||
def send_signal(self, name: str, sig: int):
|
||||
"""Send `signal` to the specified named process."""
|
||||
process_info = self[name]
|
||||
if sys.platform == "win32" and sig not in (
|
||||
signal.SIGTERM,
|
||||
signal.CTRL_C_EVENT, # pylint: disable=no-member
|
||||
signal.CTRL_BREAK_EVENT, # pylint: disable=no-member
|
||||
):
|
||||
raise UnsupportedSignalError(sig)
|
||||
|
||||
def handle_closed_process():
|
||||
logging.warning("Process %s had already aborted unexpectedly.", name)
|
||||
process_info.returncode = -1
|
||||
self[name] = process_info
|
||||
|
||||
if process_info.returncode is None:
|
||||
try:
|
||||
os.kill(process_info.pid, sig)
|
||||
except ProcessLookupError:
|
||||
handle_closed_process()
|
||||
raise
|
||||
except OSError as exc:
|
||||
if getattr(exc, "winerror", None) == 87:
|
||||
handle_closed_process()
|
||||
raise ProcessLookupError from exc
|
||||
raise
|
||||
|
||||
def terminate(self, name: str):
|
||||
"""Terminate the specified named process."""
|
||||
self.send_signal(name, signal.SIGTERM)
|
||||
|
||||
def kill(self, name: str):
|
||||
"""Kill the specified named process."""
|
||||
if sys.platform == "win32":
|
||||
self.send_signal(name, signal.SIGTERM)
|
||||
else:
|
||||
self.send_signal(name, signal.SIGKILL) # pylint: disable=no-member
|
||||
|
||||
def remove(self, name: str, force: bool = False):
|
||||
"""Remove the specified named process from this manager.
|
||||
|
||||
If the specified process is still running, it will be forcefully killed
|
||||
if `force` is True`, otherwise an exception will be raised.
|
||||
|
||||
Raises:
|
||||
ProcessNotTerminatedError if the specified process is still
|
||||
running and was not forcefully killed.
|
||||
"""
|
||||
try:
|
||||
process_info = self[name]
|
||||
except KeyError:
|
||||
return
|
||||
if process_info.returncode is None and not force:
|
||||
raise ProcessNotTerminatedError(name)
|
||||
try:
|
||||
self.kill(name)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
del self[name]
|
||||
|
||||
def cleanup(self, force: bool = False):
|
||||
"""Remove stale (terminated) processes from this manager."""
|
||||
for name in self:
|
||||
try:
|
||||
self.remove(name, force)
|
||||
except ProcessNotTerminatedError:
|
||||
continue
|
||||
@@ -1,179 +0,0 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
from contextlib import AbstractContextManager
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import BinaryIO, List, Optional, Union
|
||||
|
||||
from shortuuid import uuid
|
||||
|
||||
from dvc.utils.objects import cached_property
|
||||
|
||||
from .exceptions import TimeoutExpired
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessInfo:
|
||||
pid: int
|
||||
stdin: Optional[str]
|
||||
stdout: Optional[str]
|
||||
stderr: Optional[str]
|
||||
returncode: Optional[int]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d):
|
||||
return cls(**d)
|
||||
|
||||
def asdict(self):
|
||||
return asdict(self)
|
||||
|
||||
|
||||
class ManagedProcess(AbstractContextManager):
|
||||
"""Run the specified command with redirected output.
|
||||
|
||||
stdout and stderr will both be redirected to <name>.out.
|
||||
Interactive processes (requiring stdin input) are currently unsupported.
|
||||
|
||||
Parameters:
|
||||
args: Command to be run.
|
||||
wdir: If specified, redirected output files will be placed in `wdir`.
|
||||
name: Name to use for this process, if not specified a UUID will be
|
||||
generated instead.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
args: Union[str, List[str]],
|
||||
wdir: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
):
|
||||
self.args: List[str] = (
|
||||
shlex.split(args, posix=os.name == "posix")
|
||||
if isinstance(args, str)
|
||||
else list(args)
|
||||
)
|
||||
self.wdir = wdir
|
||||
self.name = name or uuid()
|
||||
self.returncode: Optional[int] = None
|
||||
self._stdout: Optional[BinaryIO] = None
|
||||
self._stderr: Optional[BinaryIO] = None
|
||||
|
||||
self._run()
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.wait()
|
||||
|
||||
def _close_fds(self):
|
||||
if self._stdout:
|
||||
self._stdout.close()
|
||||
self._stdout = None
|
||||
if self._stderr:
|
||||
self._stderr.close()
|
||||
self._stderr = None
|
||||
|
||||
def _make_path(self, path: str) -> str:
|
||||
return os.path.join(self.wdir, path) if self.wdir else path
|
||||
|
||||
@cached_property
|
||||
def stdout_path(self) -> str:
|
||||
return self._make_path(f"{self.name}.out")
|
||||
|
||||
@cached_property
|
||||
def info_path(self) -> str:
|
||||
return self._make_path(f"{self.name}.json")
|
||||
|
||||
@cached_property
|
||||
def pidfile_path(self) -> str:
|
||||
return self._make_path(f"{self.name}.pid")
|
||||
|
||||
@property
|
||||
def info(self) -> "ProcessInfo":
|
||||
return ProcessInfo(
|
||||
pid=self.pid,
|
||||
stdin=None,
|
||||
stdout=self.stdout_path,
|
||||
stderr=None,
|
||||
returncode=self.returncode,
|
||||
)
|
||||
|
||||
def _make_wdir(self):
|
||||
if self.wdir:
|
||||
os.makedirs(self.wdir, exist_ok=True)
|
||||
|
||||
def _dump(self):
|
||||
self._make_wdir()
|
||||
with open(self.info_path, "w", encoding="utf-8") as fobj:
|
||||
json.dump(self.info.asdict(), fobj)
|
||||
with open(self.pidfile_path, "w", encoding="utf-8") as fobj:
|
||||
fobj.write(str(self.pid))
|
||||
|
||||
def _run(self) -> None:
|
||||
self._make_wdir()
|
||||
logger.debug(
|
||||
"Appending output to '%s'",
|
||||
self.stdout_path,
|
||||
)
|
||||
self._stdout = open(self.stdout_path, "ab") # noqa: SIM115
|
||||
try:
|
||||
self._proc = subprocess.Popen(
|
||||
self.args,
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=self._stdout,
|
||||
stderr=subprocess.STDOUT,
|
||||
close_fds=True,
|
||||
shell=False,
|
||||
)
|
||||
self.pid: int = self._proc.pid
|
||||
self._dump()
|
||||
except Exception:
|
||||
if self._proc is not None:
|
||||
self._proc.kill()
|
||||
self._close_fds()
|
||||
raise
|
||||
|
||||
def wait(self, timeout: Optional[int] = None) -> Optional[int]:
|
||||
"""Block until a process started with `run` has completed.
|
||||
|
||||
Raises:
|
||||
TimeoutExpired if `timeout` was set and the process
|
||||
did not terminate after `timeout` seconds.
|
||||
"""
|
||||
if self.returncode is not None:
|
||||
return self.returncode
|
||||
try:
|
||||
self._proc.wait(timeout=timeout)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
raise TimeoutExpired(exc.cmd, exc.timeout) from exc
|
||||
self.returncode = self._proc.returncode
|
||||
self._close_fds()
|
||||
self._dump()
|
||||
return self.returncode
|
||||
|
||||
@classmethod
|
||||
def spawn(cls, *args, **kwargs) -> Optional[int]:
|
||||
"""Spawn a ManagedProcess command in the background.
|
||||
|
||||
Returns: The spawned process PID.
|
||||
"""
|
||||
import multiprocessing as mp
|
||||
|
||||
proc = mp.Process(
|
||||
target=cls._spawn,
|
||||
args=args,
|
||||
kwargs=kwargs,
|
||||
daemon=True,
|
||||
)
|
||||
proc.start()
|
||||
# Do not terminate the child daemon when the main process exits
|
||||
# pylint: disable=protected-access
|
||||
mp.process._children.discard(proc) # type: ignore[attr-defined]
|
||||
return proc.pid
|
||||
|
||||
@classmethod
|
||||
def _spawn(cls, *args, **kwargs):
|
||||
with cls(*args, **kwargs):
|
||||
pass
|
||||
@@ -1,114 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from dvc.proc.exceptions import ProcessNotTerminatedError, UnsupportedSignalError
|
||||
from dvc.proc.manager import ProcessManager
|
||||
from dvc.proc.process import ProcessInfo
|
||||
|
||||
PID_FINISHED = 1234
|
||||
PID_RUNNING = 5678
|
||||
|
||||
|
||||
def create_process(root: str, name: str, pid: int, returncode=None):
|
||||
info_path = os.path.join(root, name, f"{name}.json")
|
||||
os.makedirs(os.path.join(root, name))
|
||||
process_info = ProcessInfo(
|
||||
pid=pid, stdin=None, stdout=None, stderr=None, returncode=returncode
|
||||
)
|
||||
with open(info_path, "w", encoding="utf-8") as fobj:
|
||||
json.dump(process_info.asdict(), fobj)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def finished_process(tmp_dir):
|
||||
key = "finished"
|
||||
create_process(tmp_dir, key, PID_FINISHED, 0)
|
||||
return key
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def running_process(tmp_dir):
|
||||
key = "running"
|
||||
create_process(tmp_dir, key, PID_RUNNING)
|
||||
return key
|
||||
|
||||
|
||||
def test_send_signal(tmp_dir, mocker, finished_process, running_process):
|
||||
m = mocker.patch("os.kill")
|
||||
process_manager = ProcessManager(tmp_dir)
|
||||
process_manager.send_signal(running_process, signal.SIGTERM)
|
||||
m.assert_called_once_with(PID_RUNNING, signal.SIGTERM)
|
||||
|
||||
m.reset_mock()
|
||||
process_manager.send_signal(finished_process, signal.SIGTERM)
|
||||
m.assert_not_called()
|
||||
|
||||
if sys.platform == "win32":
|
||||
with pytest.raises(UnsupportedSignalError):
|
||||
process_manager.send_signal(finished_process, signal.SIGABRT)
|
||||
|
||||
|
||||
def test_dead_process(tmp_dir, mocker, running_process):
|
||||
process_manager = ProcessManager(tmp_dir)
|
||||
|
||||
def side_effect(*args):
|
||||
if sys.platform == "win32":
|
||||
err = OSError()
|
||||
err.winerror = 87
|
||||
raise err
|
||||
raise ProcessLookupError()
|
||||
|
||||
mocker.patch("os.kill", side_effect=side_effect)
|
||||
with pytest.raises(ProcessLookupError):
|
||||
process_manager.send_signal(running_process, signal.SIGTERM)
|
||||
assert process_manager[running_process].returncode == -1
|
||||
|
||||
|
||||
def test_kill(tmp_dir, mocker, finished_process, running_process):
|
||||
m = mocker.patch("os.kill")
|
||||
process_manager = ProcessManager(tmp_dir)
|
||||
process_manager.kill(running_process)
|
||||
if sys.platform == "win32":
|
||||
m.assert_called_once_with(PID_RUNNING, signal.SIGTERM)
|
||||
else:
|
||||
m.assert_called_once_with(PID_RUNNING, signal.SIGKILL)
|
||||
|
||||
m.reset_mock()
|
||||
process_manager.kill(finished_process)
|
||||
m.assert_not_called()
|
||||
|
||||
|
||||
def test_terminate(tmp_dir, mocker, running_process, finished_process):
|
||||
m = mocker.patch("os.kill")
|
||||
process_manager = ProcessManager(tmp_dir)
|
||||
process_manager.terminate(running_process)
|
||||
m.assert_called_once_with(PID_RUNNING, signal.SIGTERM)
|
||||
|
||||
m.reset_mock()
|
||||
process_manager.terminate(finished_process)
|
||||
m.assert_not_called()
|
||||
|
||||
|
||||
def test_remove(mocker, tmp_dir, running_process, finished_process):
|
||||
mocker.patch("os.kill", return_value=None)
|
||||
process_manager = ProcessManager(tmp_dir)
|
||||
process_manager.remove(finished_process)
|
||||
assert not (tmp_dir / finished_process).exists()
|
||||
with pytest.raises(ProcessNotTerminatedError):
|
||||
process_manager.remove(running_process)
|
||||
assert (tmp_dir / running_process).exists()
|
||||
process_manager.remove(running_process, True)
|
||||
assert not (tmp_dir / running_process).exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("force", [True, False])
|
||||
def test_cleanup(mocker, tmp_dir, running_process, finished_process, force):
|
||||
mocker.patch("os.kill", return_value=None)
|
||||
process_manager = ProcessManager(tmp_dir)
|
||||
process_manager.cleanup(force)
|
||||
assert (tmp_dir / running_process).exists() != force
|
||||
assert not (tmp_dir / finished_process).exists()
|
||||
@@ -1,61 +0,0 @@
|
||||
import json
|
||||
import subprocess
|
||||
|
||||
import pytest
|
||||
|
||||
from dvc.proc.process import ManagedProcess, ProcessInfo
|
||||
|
||||
TEST_PID = 1234
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_popen(mocker):
|
||||
mocker.patch(
|
||||
"subprocess.Popen",
|
||||
return_value=mocker.MagicMock(pid=TEST_PID, returncode=None),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args",
|
||||
[
|
||||
"/bin/foo -o option",
|
||||
["/bin/foo", "-o", "option"],
|
||||
],
|
||||
)
|
||||
def test_init_args(tmp_dir, request, args, mocker):
|
||||
expected = ["/bin/foo", "-o", "option"]
|
||||
proc = ManagedProcess(args)
|
||||
request.addfinalizer(proc._close_fds)
|
||||
assert expected == proc.args
|
||||
|
||||
|
||||
def test_run(tmp_dir, request, mocker):
|
||||
proc = ManagedProcess("/bin/foo")
|
||||
|
||||
request.addfinalizer(proc._close_fds)
|
||||
assert proc.pid == TEST_PID
|
||||
|
||||
with open(proc.info_path, encoding="utf-8") as fobj:
|
||||
info = ProcessInfo.from_dict(json.load(fobj))
|
||||
assert info.pid == TEST_PID
|
||||
|
||||
|
||||
def test_wait(tmp_dir, mocker):
|
||||
from dvc.proc.exceptions import TimeoutExpired
|
||||
|
||||
proc = ManagedProcess("/bin/foo")
|
||||
proc._proc.wait = mocker.Mock(side_effect=subprocess.TimeoutExpired("/bin/foo", 5))
|
||||
with pytest.raises(TimeoutExpired):
|
||||
proc.wait(5)
|
||||
|
||||
proc._proc.wait = mocker.Mock(return_value=None)
|
||||
proc._proc.returncode = 0
|
||||
assert proc.wait() == 0
|
||||
proc._proc.wait.assert_called_once()
|
||||
|
||||
# once subprocess return code is set, future ManagedProcess.wait() calls
|
||||
# should not block
|
||||
proc._proc.wait.reset_mock()
|
||||
assert proc.wait() == 0
|
||||
proc._proc.wait.assert_not_called()
|
||||
Reference in New Issue
Block a user