amazon-braket / amazon-braket-algorithm-library

Apache License 2.0
89 stars 33 forks source link

qc_qmc fails with pickling error on lightning.qubit #150

Open licedric opened 1 month ago

licedric commented 1 month ago

Describe the bug The qc_qmc notebook fails if run as-is with the following stack trace:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File <timed exec>:1

File ~/anaconda3/envs/Braket/lib/python3.10/site-packages/braket/experimental/algorithms/qc_qmc/qc_qmc.py:71, in qc_qmc(num_walkers, num_steps, dtau, quantum_evaluations_every_n_steps, trial, prop, trial_state_circuit, dev, max_pool)
     69 # parallelize with multiprocessing
     70 with mp.Pool(max_pool) as pool:
---> 71     results = list(pool.map(q_full_imag_time_evolution_wrapper, inputs))
     73 local_energies, weights, nums, denoms = map(np.array, zip(*results))
     75 energies = np.real(np.average(local_energies, weights=weights, axis=0))

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/pool.py:367, in Pool.map(self, func, iterable, chunksize)
    362 def map(self, func, iterable, chunksize=None):
    363     '''
    364     Apply `func` to each element in `iterable`, collecting the results
    365     in a list that is returned.
    366     '''
--> 367     return self._map_async(func, iterable, mapstar, chunksize).get()

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/pool.py:774, in ApplyResult.get(self, timeout)
    772     return self._value
    773 else:
--> 774     raise self._value

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/pool.py:540, in Pool._handle_tasks(taskqueue, put, outqueue, pool, cache)
    538     break
    539 try:
--> 540     put(task)
    541 except Exception as e:
    542     job, idx = task[:2]

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/connection.py:206, in _ConnectionBase.send(self, obj)
    204 self._check_closed()
    205 self._check_writable()
--> 206 self._send_bytes(_ForkingPickler.dumps(obj))

File ~/anaconda3/envs/Braket/lib/python3.10/multiprocessing/reduction.py:51, in ForkingPickler.dumps(cls, obj, protocol)
     48 @classmethod
     49 def dumps(cls, obj, protocol=None):
     50     buf = io.BytesIO()
---> 51     cls(buf, protocol).dump(obj)
     52     return buf.getbuffer()

TypeError: cannot pickle 'pennylane_lightning.lightning_qubit_ops.StateVectorC128' object

To reproduce I created a new Braket NBI with default settings and ran the notebook.

Expected behavior The code should execute without errors.

Screenshots or logs If applicable, add screenshots or logs to help explain your problem.

System information A description of your system. Please provide:

Additional context

licedric commented 1 month ago

Error comes from sending the lightning.qubit device as an argument to pool.map for multiprocessing: https://github.com/amazon-braket/amazon-braket-algorithm-library/blob/a1895e212cc10bd747156ef491edfecc474561d5/src/braket/experimental/algorithms/qc_qmc/qc_qmc.py#L71

At some point the lightning.qubit device became unpicklable: https://discuss.pennylane.ai/t/qml-device-error-with-qiskit-aer/4556/7

Monkey-patching around this allows the notebook to run, so it's confirmed that this is the issue.