diff --git a/src/openfermion/linalg/linear_qubit_operator.py b/src/openfermion/linalg/linear_qubit_operator.py index 6395b8972..86de9b40b 100644 --- a/src/openfermion/linalg/linear_qubit_operator.py +++ b/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): diff --git a/src/openfermion/linalg/linear_qubit_operator_test.py b/src/openfermion/linalg/linear_qubit_operator_test.py index 898b85746..0d044cb7a 100644 --- a/src/openfermion/linalg/linear_qubit_operator_test.py +++ b/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."""