continue debugging

ciaranra · ciaranra · commit b5fb7ec5bb1b · 2026-02-16T20:52:23.000-07:00
diff --git a/.github/workflows/debug-multiprocessing.yml b/.github/workflows/debug-multiprocessing.yml
@@ -133,6 +133,9 @@ jobs:
         export LIBRARY_PATH=/usr/lib
         echo "LIBRARY_PATH=/usr/lib" >> $GITHUB_ENV
 
+    - name: Install PECOS CLI
+      run: cargo install --path crates/pecos --features cli
+
     - name: Build PECOS (debug)
       run: just build-debug
 
diff --git a/python/quantum-pecos/src/pecos/_mp_workers.py b/python/quantum-pecos/src/pecos/_mp_workers.py
@@ -0,0 +1,86 @@
+# Copyright 2026 The PECOS Developers
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+# or implied. See the License for the specific language governing permissions and limitations under
+# the License.
+
+"""Worker functions for multiprocessing with spawn context.
+
+When using ``multiprocessing`` with the ``spawn`` start method (default on
+macOS/Windows), worker functions must be importable by the child process.
+Functions defined in test files are NOT importable because test directories
+are not installed packages.  This module provides generic workers that
+live inside the installed ``pecos`` package so child processes can always
+find them.
+"""
+
+from __future__ import annotations
+
+import pickle
+
+
+def deserialize_and_call(args: tuple[bytes, str, tuple, str, tuple]) -> object:
+    """Unpickle an object, call a method on it, and return a result.
+
+    Args:
+        args: A tuple of ``(data, method, method_args, result_method, result_args)``
+            where *data* is the pickled object, *method* is the method to invoke
+            on the deserialized object, *method_args* are its positional arguments,
+            *result_method* is the method/property to call for the return value,
+            and *result_args* are its positional arguments (empty tuple for properties).
+
+    Returns:
+        The value of ``getattr(obj, result_method)(*result_args)``.
+    """
+    data, method, method_args, result_method, result_args = args
+    obj = pickle.loads(data)  # noqa: S301
+    getattr(obj, method)(*method_args)
+    result = getattr(obj, result_method)
+    return result(*result_args) if callable(result) else result
+
+
+def run_callable_worker(args: tuple[object, dict]) -> dict:
+    """Worker that calls a callable with kwargs and returns the result dict.
+
+    Mirrors the production pattern in
+    ``pecos.engines.hybrid_engine_multiprocessing.worker_wrapper``
+    where a callable and its keyword arguments are passed to each pool worker.
+
+    Args:
+        args: A tuple of ``(callable, kwargs)`` where *callable* is the function
+            to invoke and *kwargs* is a dict of keyword arguments.
+
+    Returns:
+        The dict returned by ``callable(**kwargs)``.
+    """
+    fn, kwargs = args
+    return fn(**kwargs)
+
+
+def sim_run_from_bytes(**kwargs: object) -> dict:
+    """Deserialize a simulator from bytes, run a gate, and return results.
+
+    This is a test helper that mimics the pattern of ``HybridEngine.run``:
+    a callable that receives kwargs, operates on a simulator, and returns
+    a results dict.
+
+    Expected kwargs:
+        sim_bytes: Pickled simulator bytes.
+        method: Method name to call on the deserialized simulator.
+        method_args: Positional arguments for the method.
+        result_attr: Attribute name to read from the simulator as the result.
+
+    Returns:
+        A dict with a ``"measurements"`` key containing a list with the result.
+    """
+    sim = pickle.loads(kwargs["sim_bytes"])  # noqa: S301
+    getattr(sim, kwargs["method"])(*kwargs["method_args"])
+    result = getattr(sim, kwargs["result_attr"])
+    value = result() if callable(result) else result
+    return {"measurements": [value]}
diff --git a/python/quantum-pecos/tests/pecos/integration/test_pickle_multiprocessing.py b/python/quantum-pecos/tests/pecos/integration/test_pickle_multiprocessing.py
@@ -10,46 +10,43 @@
 # or implied. See the License for the specific language governing permissions and limitations under
 # the License.
 
-"""Integration tests for pickle-based multiprocessing of simulators."""
+"""Integration tests for pickle-based multiprocessing of simulators.
+
+Worker functions are defined in ``pecos._mp_workers`` (an installed module)
+rather than in this test file, because ``multiprocessing`` with the ``spawn``
+start method (default on macOS/Windows) requires workers to be importable
+by the child process.  Test files live outside the installed package and
+cannot be imported by spawned children.
+"""
 
 import multiprocessing
 import pickle
 import sys
 
 import pytest
+from pecos._mp_workers import (
+    deserialize_and_call,
+    run_callable_worker,
+    sim_run_from_bytes,
+)
+from pecos.engines.hybrid_engine_multiprocessing import worker_wrapper
 from pecos_rslib import CoinToss, PauliProp, SparseSim, StateVec
 
-
-def _statevec_worker(sim_bytes: bytes) -> int:
-    sim = pickle.loads(sim_bytes)
-    sim.run_1q_gate("H", 0)
-    return sim.num_qubits
-
-
-def _sparsesim_worker(sim_bytes: bytes) -> int:
-    sim = pickle.loads(sim_bytes)
-    sim.run_1q_gate("H", 0)
-    return sim.num_qubits
-
-
-def _cointoss_worker(sim_bytes: bytes) -> int:
-    sim = pickle.loads(sim_bytes)
-    sim.run_measure(0)
-    return sim.num_qubits
-
-
-def _pauliprop_worker(sim_bytes: bytes) -> int:
-    sim = pickle.loads(sim_bytes)
-    sim.h(0)
-    return sim.weight()
-
-
 # Use fork context on Linux (fast, avoids spawn serialization issues with test files).
 # On macOS/Windows where fork is unavailable or unsafe, use spawn.
 _MP_CONTEXT = "fork" if sys.platform == "linux" else "spawn"
 _POOL_TIMEOUT = 60  # seconds -- fail fast instead of hanging CI
 
 
+def _get_pool_context() -> multiprocessing.context.BaseContext:
+    return multiprocessing.get_context(_MP_CONTEXT)
+
+
+# ---------------------------------------------------------------------------
+# Basic pickle round-trip tests via deserialize_and_call
+# ---------------------------------------------------------------------------
+
+
 @pytest.mark.timeout(120)
 class TestMultiprocessingStateVec:
     """Tests for multiprocessing StateVec simulators via pickle."""
@@ -59,9 +56,10 @@ def test_pool_map(self) -> None:
         sim = StateVec(3, seed=42)
         sim.run_1q_gate("H", 0)
         sim_bytes = pickle.dumps(sim)
-        ctx = multiprocessing.get_context(_MP_CONTEXT)
+        args = [(sim_bytes, "run_1q_gate", ("H", 0), "num_qubits", ())] * 2
+        ctx = _get_pool_context()
         with ctx.Pool(processes=2) as pool:
-            results = pool.map_async(_statevec_worker, [sim_bytes, sim_bytes]).get(
+            results = pool.map_async(deserialize_and_call, args).get(
                 timeout=_POOL_TIMEOUT,
             )
         assert results == [3, 3]
@@ -77,9 +75,10 @@ def test_pool_map(self) -> None:
         sim.run_1q_gate("H", 0)
         sim.run_2q_gate("CX", (0, 1), None)
         sim_bytes = pickle.dumps(sim)
-        ctx = multiprocessing.get_context(_MP_CONTEXT)
+        args = [(sim_bytes, "run_1q_gate", ("H", 0), "num_qubits", ())] * 2
+        ctx = _get_pool_context()
         with ctx.Pool(processes=2) as pool:
-            results = pool.map_async(_sparsesim_worker, [sim_bytes, sim_bytes]).get(
+            results = pool.map_async(deserialize_and_call, args).get(
                 timeout=_POOL_TIMEOUT,
             )
         assert results == [4, 4]
@@ -93,9 +92,10 @@ def test_pool_map(self) -> None:
         """Test CoinToss serialization works with multiprocessing Pool.map."""
         sim = CoinToss(5, prob=0.3)
         sim_bytes = pickle.dumps(sim)
-        ctx = multiprocessing.get_context(_MP_CONTEXT)
+        args = [(sim_bytes, "run_measure", (0,), "num_qubits", ())] * 2
+        ctx = _get_pool_context()
         with ctx.Pool(processes=2) as pool:
-            results = pool.map_async(_cointoss_worker, [sim_bytes, sim_bytes]).get(
+            results = pool.map_async(deserialize_and_call, args).get(
                 timeout=_POOL_TIMEOUT,
             )
         assert results == [5, 5]
@@ -110,10 +110,163 @@ def test_pool_map(self) -> None:
         sim = PauliProp(num_qubits=3, track_sign=True)
         sim.add_x(0)
         sim_bytes = pickle.dumps(sim)
-        ctx = multiprocessing.get_context(_MP_CONTEXT)
+        args = [(sim_bytes, "h", (0,), "weight", ())] * 2
+        ctx = _get_pool_context()
         with ctx.Pool(processes=2) as pool:
-            results = pool.map_async(_pauliprop_worker, [sim_bytes, sim_bytes]).get(
+            results = pool.map_async(deserialize_and_call, args).get(
                 timeout=_POOL_TIMEOUT,
             )
         # After H on qubit 0: X->Z, so weight should still be 1
         assert all(r == 1 for r in results)
+
+
+# ---------------------------------------------------------------------------
+# Production-pattern tests: Manager queue + worker_wrapper
+#
+# These mirror the pattern used in hybrid_engine_multiprocessing.run_multisim:
+#   1. Create a Manager().Queue() for inter-process messaging
+#   2. Pass (queue, callable, kwargs, index) to worker_wrapper via pool.map
+#   3. worker_wrapper redirects stdout/stderr to WriteStream on the queue
+#   4. worker_wrapper calls the callable and returns (result_dict, run_info)
+#   5. Parent drains the queue and aggregates results
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.timeout(120)
+class TestWorkerWrapperPattern:
+    """Tests that mirror the production worker_wrapper + Manager queue pattern."""
+
+    def test_worker_wrapper_with_statevec(self) -> None:
+        """Test the production worker_wrapper pattern with StateVec."""
+        sim = StateVec(3, seed=42)
+        sim.run_1q_gate("H", 0)
+        sim_bytes = pickle.dumps(sim)
+
+        ctx = _get_pool_context()
+        manager = ctx.Manager()
+        queue = manager.Queue()
+
+        kwargs = {
+            "sim_bytes": sim_bytes,
+            "method": "run_1q_gate",
+            "method_args": ("H", 0),
+            "result_attr": "num_qubits",
+            "seed": 1,
+            "shots": 1,
+            "foreign_object": None,
+        }
+        worker_args = [
+            (queue, sim_run_from_bytes, {**kwargs, "seed": 1}, 0),
+            (queue, sim_run_from_bytes, {**kwargs, "seed": 2}, 1),
+        ]
+
+        with ctx.Pool(processes=2) as pool:
+            presults = pool.map_async(worker_wrapper, worker_args).get(
+                timeout=_POOL_TIMEOUT,
+            )
+
+        for result_dict, run_info in presults:
+            assert result_dict == {"measurements": [3]}
+            assert "pid" in run_info
+            assert "i" in run_info
+
+    def test_worker_wrapper_with_sparsesim(self) -> None:
+        """Test the production worker_wrapper pattern with SparseSim."""
+        sim = SparseSim(4)
+        sim.run_1q_gate("H", 0)
+        sim.run_2q_gate("CX", (0, 1), None)
+        sim_bytes = pickle.dumps(sim)
+
+        ctx = _get_pool_context()
+        manager = ctx.Manager()
+        queue = manager.Queue()
+
+        kwargs = {
+            "sim_bytes": sim_bytes,
+            "method": "run_1q_gate",
+            "method_args": ("H", 0),
+            "result_attr": "num_qubits",
+            "seed": 1,
+            "shots": 1,
+            "foreign_object": None,
+        }
+        worker_args = [
+            (queue, sim_run_from_bytes, {**kwargs, "seed": 1}, 0),
+            (queue, sim_run_from_bytes, {**kwargs, "seed": 2}, 1),
+        ]
+
+        with ctx.Pool(processes=2) as pool:
+            presults = pool.map_async(worker_wrapper, worker_args).get(
+                timeout=_POOL_TIMEOUT,
+            )
+
+        for result_dict, run_info in presults:
+            assert result_dict == {"measurements": [4]}
+            assert "pid" in run_info
+
+    def test_queue_message_passing(self) -> None:
+        """Test that stdout/stderr from workers is captured on the queue."""
+        sim = StateVec(2, seed=0)
+        sim_bytes = pickle.dumps(sim)
+
+        ctx = _get_pool_context()
+        manager = ctx.Manager()
+        queue = manager.Queue()
+
+        kwargs = {
+            "sim_bytes": sim_bytes,
+            "method": "run_1q_gate",
+            "method_args": ("H", 0),
+            "result_attr": "num_qubits",
+            "seed": 1,
+            "shots": 1,
+            "foreign_object": None,
+        }
+        worker_args = [(queue, sim_run_from_bytes, kwargs, 0)]
+
+        with ctx.Pool(processes=1) as pool:
+            pool.map_async(worker_wrapper, worker_args).get(timeout=_POOL_TIMEOUT)
+
+        # The queue may contain stdout/stderr messages captured by WriteStream.
+        # We just verify the queue is accessible and drainable (no deadlock).
+        messages = []
+        while not queue.empty():
+            messages.append(queue.get())
+        # Messages are (pid, stream_type, data) tuples if any output occurred.
+        for msg in messages:
+            assert len(msg) == 3
+
+
+# ---------------------------------------------------------------------------
+# Callable-with-kwargs pattern tests (run_callable_worker)
+#
+# This tests the simpler pattern where a callable + kwargs dict are passed
+# to the pool, similar to how run_multisim passes eng.run + kwargs to workers.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.timeout(120)
+class TestRunCallableWorker:
+    """Tests for the callable+kwargs worker pattern used in production."""
+
+    def test_callable_worker_statevec(self) -> None:
+        """Test passing a callable + kwargs through the pool."""
+        sim = StateVec(3, seed=42)
+        sim.run_1q_gate("H", 0)
+        sim_bytes = pickle.dumps(sim)
+
+        kwargs = {
+            "sim_bytes": sim_bytes,
+            "method": "run_1q_gate",
+            "method_args": ("H", 0),
+            "result_attr": "num_qubits",
+        }
+        args = [(sim_run_from_bytes, kwargs)] * 2
+
+        ctx = _get_pool_context()
+        with ctx.Pool(processes=2) as pool:
+            results = pool.map_async(run_callable_worker, args).get(
+                timeout=_POOL_TIMEOUT,
+            )
+
+        assert results == [{"measurements": [3]}, {"measurements": [3]}]
diff --git a/scripts/debug_multiprocessing_spawn.py b/scripts/debug_multiprocessing_spawn.py
