Fix #1405: potential pool.join deadlock in ParallelLinearQubitOperator

src/openfermion/linalg/linear_qubit_operator.py

@@ -194,8 +194,13 @@ def _matvec(self, x):
             apply_operator, [(operator, x) for operator in self.linear_operators]
         )
         pool.close()
-        pool.join()
-        return functools.reduce(numpy.add, vecs)
+        # Consume results before join(): imap_unordered uses a bounded pipe and
+        # workers block on write if the main process has not read them yet.
+        try:
+            result = functools.reduce(numpy.add, vecs)
+        finally:
+            pool.join()
+        return result
 
 
 def apply_operator(args):

src/openfermion/linalg/linear_qubit_operator_test.py

@@ -11,6 +11,7 @@
 #   limitations under the License.
 """Tests for linear_qubit_operator.py."""
 
+import sys
 import unittest
 
 import numpy
@@ -270,6 +271,18 @@ def test_matvec_single_process(self):
         )
         self.assertTrue(numpy.allclose(parallel_qubit_op * self.vec, self.expected_matvec))
 
+    @unittest.skipIf(sys.platform == 'win32', 'forkserver multiprocessing is Unix-only')
+    def test_matvec_large_vector_multiprocess(self):
+        """Regression test for imap_unordered/pool.join() pipe deadlock (#1405)."""
+        n_qubits = 14
+        qubit_operator = QubitOperator('Z0') + QubitOperator('Z1') + QubitOperator('Z2')
+        options = LinearQubitOperatorOptions(processes=2)
+        parallel_op = ParallelLinearQubitOperator(qubit_operator, n_qubits, options=options)
+        serial_op = LinearQubitOperator(qubit_operator, n_qubits)
+
+        vec = numpy.ones(2**n_qubits, dtype=complex)
+        self.assertTrue(numpy.allclose(parallel_op * vec, serial_op * vec))
+
 
 class UtilityFunctionTest(unittest.TestCase):
     """Tests for utility functions."""