fix(timeout): support Windows spawn semantics

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
+preview = true

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()

src/math_verify/utils.py

@@ -20,15 +20,36 @@
 # 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
 
 TIMEOUT_WARNING_SHOWN = False
 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
 

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