diff --git a/pyproject.toml b/pyproject.toml index 4769efb..4befa25 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ classifiers = [ ] keywords = ["verification", "math", "evaluation"] dependencies = [ + "cloudpickle>=3.0.0", "latex2sympy2_extended==1.11.0", ] requires-python = ">=3.10" @@ -80,4 +81,4 @@ line-ending = "auto" [tool.black] line-length = 88 -preview = true \ No newline at end of file +preview = true diff --git a/src/math_verify/timeout_worker.py b/src/math_verify/timeout_worker.py new file mode 100644 index 0000000..fa87b7b --- /dev/null +++ b/src/math_verify/timeout_worker.py @@ -0,0 +1,33 @@ +"""Subprocess entry point for Windows timeout execution.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import cloudpickle + + +def main() -> None: + input_path = Path(sys.argv[1]) + output_path = Path(sys.argv[2]) + try: + func, args, kwargs = cloudpickle.loads(input_path.read_bytes()) + try: + response = (True, func(*args, **kwargs)) + except BaseException as exc: + response = (False, exc) + except BaseException as exc: + response = (False, RuntimeError(f"Unable to start timeout worker: {exc}")) + + try: + payload = cloudpickle.dumps(response) + except BaseException as exc: + payload = cloudpickle.dumps( + (False, RuntimeError(f"Unable to serialize timeout worker response: {exc}")) + ) + output_path.write_bytes(payload) + + +if __name__ == "__main__": + main() diff --git a/src/math_verify/utils.py b/src/math_verify/utils.py index 4261ff9..d00582e 100644 --- a/src/math_verify/utils.py +++ b/src/math_verify/utils.py @@ -20,8 +20,15 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. +import functools import logging import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import cloudpickle from math_verify.errors import TimeoutException @@ -29,6 +36,20 @@ logger = logging.getLogger(__name__) +def _terminate_process(process: subprocess.Popen) -> None: + """Terminate a timeout worker and its descendants.""" + if sys.platform == "win32" and process.poll() is None: + subprocess.run( + ["taskkill", "/F", "/T", "/PID", str(process.pid)], + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + if process.poll() is None: + process.kill() + process.wait() + + def timeout(timeout_seconds: int | None = 10): # noqa: C901 """A decorator that applies a timeout to the decorated function. @@ -38,7 +59,7 @@ def timeout(timeout_seconds: int | None = 10): # noqa: C901 Notes: On Unix systems, uses a signal-based alarm approach which is more efficient as it doesn't require spawning a new process. - On Windows systems, uses a multiprocessing-based approach since signal.alarm is not available. This will incur a huge performance penalty. + On Windows systems, uses a subprocess since signal.alarm is not available. This will incur a performance penalty. """ if timeout_seconds is None or timeout_seconds <= 0: @@ -55,6 +76,7 @@ def decorator(func): def handler(signum, frame): raise TimeoutException("Operation timed out!") + @functools.wraps(func) def wrapper(*args, **kwargs): old_handler = signal.getsignal(signal.SIGALRM) signal.signal(signal.SIGALRM, handler) @@ -71,37 +93,56 @@ def wrapper(*args, **kwargs): return decorator else: - # Windows approach: use multiprocessing - from multiprocessing import Process, Queue def decorator(func): + @functools.wraps(func) def wrapper(*args, **kwargs): - q = Queue() + payload = cloudpickle.dumps((func, args, kwargs)) + with tempfile.TemporaryDirectory(prefix="math-verify-timeout-") as tmp: + input_path = Path(tmp, "input.bin") + output_path = Path(tmp, "output.bin") + input_path.write_bytes(payload) + process = subprocess.Popen( + [ + sys.executable, + "-m", + "math_verify.timeout_worker", + str(input_path), + str(output_path), + ], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + creationflags=( + getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) + if sys.platform == "win32" + else 0 + ), + ) - def run_func(q, args, kwargs): try: - result = func(*args, **kwargs) - q.put((True, result)) - except Exception as e: - q.put((False, e)) - - p = Process(target=run_func, args=(q, args, kwargs)) - p.start() - p.join(timeout_seconds) - - if p.is_alive(): - # Timeout: Terminate the process - p.terminate() - p.join() - raise TimeoutException("Operation timed out!") - - # If we got here, the process completed in time. - success, value = q.get() - if success: - return value - else: - # The child raised an exception; re-raise it here - raise value + process.wait(timeout=timeout_seconds) + except subprocess.TimeoutExpired: + raise TimeoutException("Operation timed out!") from None + finally: + _terminate_process(process) + + if process.returncode: + raise RuntimeError( + f"Timeout worker exited with code {process.returncode}" + ) + try: + success, value = cloudpickle.loads(output_path.read_bytes()) + except Exception as exc: + raise RuntimeError( + "Timeout worker returned an invalid response" + ) from exc + + if success: + return value + if isinstance(value, BaseException): + raise value + raise RuntimeError("Timeout worker returned an invalid exception") return wrapper diff --git a/tests/test_timeout.py b/tests/test_timeout.py index 2b48762..ecbe810 100644 --- a/tests/test_timeout.py +++ b/tests/test_timeout.py @@ -1,8 +1,18 @@ +import multiprocessing +import os +import subprocess +import sys import time +from types import SimpleNamespace from unittest.mock import patch +import pytest + +from math_verify import utils +from math_verify.errors import TimeoutException from math_verify.grader import verify from math_verify.parser import parse +from math_verify.utils import timeout @patch("math_verify.parser.parse_expr") @@ -54,3 +64,69 @@ def delayed_sympy_expr_eq(*args, **kwargs): gold = [parse("1+1")[0]] assert not verify(gold, gold, timeout_seconds=1) + + +def test_windows_timeout_supports_local_callables(monkeypatch): + monkeypatch.setattr(utils, "os", SimpleNamespace(name="nt")) + parent_pid = os.getpid() + + def noisy_local_callable(): + os.write(1, b"worker output must not corrupt the protocol") + return os.getpid(), lambda value: value + 1 + + child_pid, result = timeout(2)(noisy_local_callable)() + + assert child_pid != parent_pid + assert result(2) == 3 + + +def test_windows_timeout_supports_parse_and_verify(monkeypatch): + monkeypatch.setattr(utils, "os", SimpleNamespace(name="nt")) + + assert {str(value) for value in parse(r"\boxed{4}", raise_on_error=True)} == {"4"} + assert verify("4", "4", strict=True, timeout_seconds=2, raise_on_error=True) + + +def test_windows_timeout_terminates_child(monkeypatch): + monkeypatch.setattr(utils, "os", SimpleNamespace(name="nt")) + children_before = {child.pid for child in multiprocessing.active_children()} + + with pytest.raises(TimeoutException, match="Operation timed out"): + timeout(0.1)(lambda: time.sleep(5))() + + children_after = {child.pid for child in multiprocessing.active_children()} + assert children_after == children_before + + +def test_windows_timeout_propagates_child_timeout(monkeypatch): + monkeypatch.setattr(utils, "os", SimpleNamespace(name="nt")) + + def raise_timeout(): + raise TimeoutException("child timeout") + + with pytest.raises(TimeoutException, match="child timeout"): + timeout(2)(raise_timeout)() + + +@pytest.mark.skipif(sys.platform != "win32", reason="requires native taskkill") +def test_windows_timeout_terminates_descendants(tmp_path): + pid_path = tmp_path / "descendant.pid" + + def spawn_descendant(): + child = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], + ) + pid_path.write_text(str(child.pid)) + time.sleep(30) + + with pytest.raises(TimeoutException, match="Operation timed out"): + timeout(1)(spawn_descendant)() + + descendant_pid = pid_path.read_text() + tasklist = subprocess.run( + ["tasklist", "/FI", f"PID eq {descendant_pid}"], + check=True, + capture_output=True, + text=True, + ) + assert descendant_pid not in tasklist.stdout