Repository: PennyLaneAI/pennylane
Pull request #3924 (2023-03-20 19:04:09)

[queuing] move the RLock from QuantumTape to AnnotatedQueue
**Context:**
`QuantumTape` has a class-level RLock to ensure parallel tapes don't queue each other's contents. This wasn't ported to `AnnotatedQueue` because we failed to replicate the issue when porting contents over, but I recently confirmed that the issue still exists with `AnnotatedQueue`.

**Description of the Change:**
Move the RLock from `QuantumTape` to `AnnotatedQueue`.

**Benefits:**
Any parallel queuing, regardless of type, will be guaranteed to not overlap with each other.

**Possible Drawbacks:**
There's the lock-acquisition overhead which won't be relevant in the majority of cases when users aren't multi-threading.

**Related GitHub Issues:**
Reoccurrence of #910

A minimal example of the issue (which I add as a unit test in this PR):
```python
import pennylane as qml
from multiprocessing.dummy import Pool as ThreadPool
q1 = qml.queuing.AnnotatedQueue()
q2 = qml.queuing.AnnotatedQueue()
def queue_pauli(arg):
    q, pauli = arg
    with q:
        for _ in range(10000):
            pauli(0)
args = [(q1, qml.PauliX), (q2, qml.PauliY)]
ThreadPool(2).map(queue_pauli, args)
[(all(isinstance(op, expected_op) for op in queue.queue), len(queue)) for queue, expected_op in args]
```

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
josh146: @timmysilv is there somewhere we can keep/store knowledge of the minimal-non-working example you have in the PR comment?

We have 'forgotten' how to reproduce this multiple times in the past, and wondering how we can avoid this in the future 😆 
timmysilv: I'll add it as a test! if someone is poking through the git blame, they should find this PR.
mlxd: Hey folks, I think the problem here is the wrong datastructure design, rather than a need for adding locks to existing code.

The in-built Python queue is already thread-safe, so we should be using that as part of the process, rather than locking our own DS. I'll think more about this, but I'd rather avoid any locking code if we can avoid it at our development level, and especially when using Python --- given the existing implementations can handle this for us.
mlxd: > to be clear on the timeline:
> 
>     1. original QuantumTape didn't have locks, race condition happened
> 
>     2. bug was closed when external user (who found the RC) submitted a fix
> 
>     3. We split QuantumTape into AnnotatedQueue+QuantumScript, but did NOT take the RLock
> 
>     4. I'm realizing that was a mistake (or at least, a regression)
> 
> 
> That said, I agree that we can do better than an RLock here. Just making sure that it's understood that this PR is re-introducing old code (and not proposing new stuff).
> 
> Do we want to hold off on this while investigating more? In my opinion, the change is pretty small in terms of LOC, and I'd most want to have (more) correct queuing in the next release.

Just chatted with @josh146  --- it may make sense to create a bugfix release for this ASAP to avoid the problem appearing in the released 0.29.x branch. We can discuss the details of removing this then for 0.30. @albi3ro that work for your timeline too?
timmysilv: I proposed something different in the PR linked right above this comment. Rather than using an RLock, I essentially make `QueuingManager._active_contexts` get stored in thread-local memory. I'm not sure if this is better, but I've put up a PR to open discussion
timmysilv: hi @PennyLaneAI/performance 👋 

tagging y'all for review because it might affect you most/you'll have the most context on this. tl;dr, a user discovered a while back that if you have multiples tapes queuing in parallel, they will accidentally queue the wrong stuff. So, we added a `threading.RLock` that tapes acquire upon context entry. We did not move that to `AnnotatedQueue` when making that class, but we should have, because it has the same problem. This PR moves the lock from `QuantumTape` (which is now only a user-convenience thing) to `AnnotatedQueue` (which is the underlying queuing object used by everything including QNodes).

I considered something else (like putting queues in thread-local memory) but we decided to not introduce new stuff yet, and first fix what's wrong. We'd actually most prefer to get rid of queuing altogether. One day...

Patch:
diff --git a/doc/code/qml_queuing.rst b/doc/code/qml_queuing.rst
index 9d9d469f0..841e600cb 100644
--- a/doc/code/qml_queuing.rst
+++ b/doc/code/qml_queuing.rst
@@ -15,4 +15,4 @@ Overview
     :no-heading:
     :include-all-objects:
     :no-inherited-members:
-    :skip: contextmanager, OrderedDict
+    :skip: contextmanager, OrderedDict, RLock
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 7b37cb752..bacc51d9a 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -263,6 +263,9 @@
   Now only `DefaultQubitJax` supports the operator, as expected.
   [(#3964)](https://github.com/PennyLaneAI/pennylane/pull/3964)
 
+* Ensure that parallel `AnnotatedQueues` do not queue each other's contents.
+  [(#3924)](https://github.com/PennyLaneAI/pennylane/pull/3924)
+
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/queuing.py b/pennylane/queuing.py
index 23d4014bc..e3f46384b 100644
--- a/pennylane/queuing.py
+++ b/pennylane/queuing.py
@@ -166,6 +166,7 @@ use the :meth:`~.QueuingManager.stop_recording` context or specify `do_queue=Fal
 import copy
 from collections import OrderedDict
 from contextlib import contextmanager
+from threading import RLock
 
 
 class QueuingError(Exception):
@@ -335,12 +336,16 @@ class AnnotatedQueue(OrderedDict):
     """Lightweight class that maintains a basic queue of operations, in addition
     to metadata annotations."""
 
+    _lock = RLock()
+    """threading.RLock: Used to synchronize appending to/popping from global QueueingContext."""
+
     def __enter__(self):
         """Adds this instance to the global list of active contexts.
 
         Returns:
             AnnotatedQueue: this instance
         """
+        AnnotatedQueue._lock.acquire()
         QueuingManager.add_active_queue(self)
 
         return self
@@ -348,6 +353,7 @@ class AnnotatedQueue(OrderedDict):
     def __exit__(self, exception_type, exception_value, traceback):
         """Remove this instance from the global list of active contexts."""
         QueuingManager.remove_active_queue()
+        AnnotatedQueue._lock.release()
 
     def append(self, obj, **kwargs):
         """Append ``obj`` into the queue with ``kwargs`` metadata."""
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index a9fa37381..3872b924d 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -352,25 +352,15 @@ class QuantumTape(QuantumScript, AnnotatedQueue):
 
     def __enter__(self):
         QuantumTape._lock.acquire()
-        try:
-            if self.do_queue:
-                QueuingManager.append(self)
-            return AnnotatedQueue.__enter__(self)
-        except Exception as _:
-            QuantumTape._lock.release()
-            raise
+        if self.do_queue:
+            QueuingManager.append(self)
+        QueuingManager.add_active_queue(self)
+        return self
 
     def __exit__(self, exception_type, exception_value, traceback):
-        try:
-            AnnotatedQueue.__exit__(self, exception_type, exception_value, traceback)
-            # After other optimizations in #2963, #2986 and follow-up work, we should check whether
-            # calling `_process_queue` only if there is no `exception_type` saves time. This would
-            # be done via the following:
-            # if exception_type is None:
-            #    self._process_queue()
-            self._process_queue()
-        finally:
-            QuantumTape._lock.release()
+        QueuingManager.remove_active_queue()
+        QuantumTape._lock.release()
+        self._process_queue()
 
     # ========================================================
     # construction methods


Test patch:
diff --git a/tests/test_queuing.py b/tests/test_queuing.py
index 03dc25acf..2bba3a0c9 100644
--- a/tests/test_queuing.py
+++ b/tests/test_queuing.py
@@ -14,9 +14,10 @@
 """
 Unit tests for the :mod:`pennylane` :class:`QueuingManager` class.
 """
+from multiprocessing.dummy import Pool as ThreadPool
 import pytest
-import pennylane as qml
 import numpy as np
+import pennylane as qml
 
 from pennylane.queuing import (
     AnnotatedQueue,
@@ -225,7 +226,7 @@ class TestAnnotatedQueue:
         for a non-existent object"""
 
         with AnnotatedQueue() as q:
-            A = qml.PauliZ(0)
+            qml.PauliZ(0)
 
         B = qml.PauliY(1)
 
@@ -264,7 +265,7 @@ class TestAnnotatedQueue:
         for a non-existent object."""
 
         with AnnotatedQueue() as q:
-            A = qml.PauliZ(0)
+            qml.PauliZ(0)
 
         B = qml.PauliY(1)
 
@@ -281,6 +282,25 @@ class TestAnnotatedQueue:
 
         assert q.queue == [tensor_op]
 
+    def test_parallel_queues_are_isolated(self):
+        """Tests that parallel queues do not queue each other's constituents."""
+        q1 = AnnotatedQueue()
+        q2 = AnnotatedQueue()
+        n = 10000
+
+        def queue_pauli(arg):
+            q, pauli = arg
+            with q:
+                for _ in range(n):
+                    pauli(0)
+
+        args = [(q1, qml.PauliX), (q2, qml.PauliY)]
+        ThreadPool(2).map(queue_pauli, args)
+        assert len(q1) == n
+        assert len(q2) == n
+        for queue, expected_op in args:
+            assert all(isinstance(op, expected_op) for op in queue.queue)
+
 
 test_observables = [
     qml.PauliZ(0) @ qml.PauliZ(1),
@@ -335,7 +355,7 @@ class TestApplyOp:
         assert tape.measurements == [op]
 
     @pytest.mark.parametrize("obs", test_observables)
-    def test_default_queue_measurements_outside(self, obs):
+    def test_default_queue_measurements_inside(self, obs):
         """Test applying a measurement instantiated inside a queuing context
         to an existing active queue"""
 
@@ -389,7 +409,7 @@ class TestApplyOp:
         assert tape2.measurements == []
 
     @pytest.mark.parametrize("obs", test_observables)
-    def test_different_queue_measurements_outside(self, obs):
+    def test_different_queue_measurements_inside(self, obs):
         """Test applying a measurement instantiated inside a queuing context
         to a specfied queuing context"""
 
diff --git a/tests/tests_passing_pylint.txt b/tests/tests_passing_pylint.txt
index e1c400221..b799ab8c5 100644
--- a/tests/tests_passing_pylint.txt
+++ b/tests/tests_passing_pylint.txt
@@ -30,8 +30,9 @@ tests/transforms/test_qfunc_transform.py
 tests/ops/op_math/test_exp.py
 tests/ops/qubit/test_state_prep.py
 tests/devices/qubit/**
-tests/test_typing.py
 tests/test_operation.py
+tests/test_queuing.py
+tests/test_typing.py
 tests/pulse/**
 tests/math/test_init.py
 tests/math/test_functions.py