Additional debugging for mp tests in CI

Author: timsaucer

python/tests/_pickle_multiprocessing_helpers.py

@@ -122,6 +122,71 @@ def diag_init():
     _diag("worker init: ready for tasks")
 
 
+def _read_text(path: str) -> str:
+    """Read a /proc file; return ``"<unreadable>"`` if not accessible."""
+    try:
+        return Path(path).read_text(encoding="utf-8", errors="replace").strip()
+    except OSError:
+        return "<unreadable>"
+
+
+def _descendants(root_pid: int) -> list[int]:
+    """Return ``root_pid`` plus all descendant pids via /proc/<pid>/task/.../children.
+
+    Linux-only; returns ``[root_pid]`` on platforms without ``/proc``.
+    """
+    out: list[int] = [root_pid]
+    if not Path("/proc").is_dir():
+        return out
+    queue = [root_pid]
+    while queue:
+        pid = queue.pop()
+        try:
+            task_dir = Path(f"/proc/{pid}/task")
+            if not task_dir.is_dir():
+                continue
+            for tdir in task_dir.iterdir():
+                children_file = tdir / "children"
+                try:
+                    children = children_file.read_text(encoding="utf-8").split()
+                except OSError:
+                    continue
+                for child in children:
+                    try:
+                        cpid = int(child)
+                    except ValueError:
+                        continue
+                    out.append(cpid)
+                    queue.append(cpid)
+        except OSError:
+            continue
+    return out
+
+
+def snapshot_processes(label: str, root_pid: int | None = None) -> None:
+    """Dump a process-state snapshot to the diagnostic log.
+
+    For each descendant of ``root_pid`` (default: current process), record
+    cmdline, status (``R``/``S``/``D``), wchan (kernel function the task
+    is blocked in), and kernel stack. Use this to localize a worker hang:
+    a wchan of ``do_futex`` points at a lock; ``poll_schedule_timeout``
+    points at a blocking I/O wait; ``do_select`` at multiprocessing's
+    pipe read.
+    """
+    pid = root_pid if root_pid is not None else os.getpid()
+    _diag(f"snapshot[{label}] root_pid={pid}")
+    for cpid in _descendants(pid):
+        cmd = _read_text(f"/proc/{cpid}/cmdline").replace("\x00", " ").strip()
+        stat = _read_text(f"/proc/{cpid}/status").splitlines()
+        state_line = next((s for s in stat if s.startswith("State:")), "State: ?")
+        wchan = _read_text(f"/proc/{cpid}/wchan")
+        stack = _read_text(f"/proc/{cpid}/stack")
+        _diag(f"snapshot[{label}] pid={cpid} {state_line} wchan={wchan} cmd={cmd!r}")
+        if stack and stack != "<unreadable>":
+            for line in stack.splitlines()[:10]:
+                _diag(f"snapshot[{label}] pid={cpid} stack: {line}")
+
+
 def unpickle_and_describe(blob: bytes) -> str:
     """Unpickle a proto-bytes blob and return its canonical name."""
     import pickle

python/tests/test_pickle_multiprocessing.py

@@ -27,17 +27,40 @@
 
 from __future__ import annotations
 
+import contextlib
 import functools
 import multiprocessing as mp
 import pickle
 import sys
+import threading
 
 import pytest
 from datafusion import col, lit
 
 from . import _pickle_multiprocessing_helpers as helpers
 
 
+@contextlib.contextmanager
+def _snapshot_on_hang(label: str, fire_after_seconds: float = 30.0):
+    """Schedule a process-state snapshot ``fire_after_seconds`` from now.
+
+    Cancelled if the ``with`` block exits before then. Used to capture
+    worker state mid-hang — fork tests return in well under the delay,
+    so the timer only fires when something is actually stuck.
+    """
+    timer = threading.Timer(
+        fire_after_seconds,
+        helpers.snapshot_processes,
+        args=(label,),
+    )
+    timer.daemon = True
+    timer.start()
+    try:
+        yield
+    finally:
+        timer.cancel()
+
+
 @functools.cache
 def _multiprocessing_available() -> tuple[bool, str]:
     """Return (available, reason). Some sandboxed environments deny semaphore
@@ -86,7 +109,10 @@ def test_builtin_pickle_via_pool(start_method):
 
     ctx = mp.get_context(start_method)
     helpers._diag(f"test_builtin_pickle_via_pool[{start_method}]: creating Pool")
-    with ctx.Pool(processes=2, initializer=helpers.diag_init) as pool:
+    with (
+        ctx.Pool(processes=2, initializer=helpers.diag_init) as pool,
+        _snapshot_on_hang(f"builtin[{start_method}]"),
+    ):
         helpers._diag(f"test_builtin_pickle_via_pool[{start_method}]: pool ready, map")
         results = pool.map(helpers.unpickle_and_describe, [blob, blob, blob])
     helpers._diag(f"test_builtin_pickle_via_pool[{start_method}]: pool closed")
@@ -109,7 +135,10 @@ def test_udf_pickle_self_contained(start_method):
 
     ctx = mp.get_context(start_method)
     helpers._diag(f"test_udf_pickle_self_contained[{start_method}]: creating Pool")
-    with ctx.Pool(processes=2, initializer=helpers.diag_init) as pool:
+    with (
+        ctx.Pool(processes=2, initializer=helpers.diag_init) as pool,
+        _snapshot_on_hang(f"udf[{start_method}]"),
+    ):
         helpers._diag(
             f"test_udf_pickle_self_contained[{start_method}]: pool ready, starmap"
         )
@@ -134,7 +163,10 @@ def test_closure_capturing_udf_via_pool(start_method):
 
     ctx = mp.get_context(start_method)
     helpers._diag(f"test_closure_capturing_udf_via_pool[{start_method}]: creating Pool")
-    with ctx.Pool(processes=2, initializer=helpers.diag_init) as pool:
+    with (
+        ctx.Pool(processes=2, initializer=helpers.diag_init) as pool,
+        _snapshot_on_hang(f"closure[{start_method}]"),
+    ):
         helpers._diag(
             f"test_closure_capturing_udf_via_pool[{start_method}]: pool ready, apply"
         )