fix: fix Jupyter event loop conflict

src/braket/aws/aws_quantum_task.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
 import contextlib
 import time
 import warnings
@@ -451,13 +452,38 @@ def result(self) -> TaskResult:
         if status in self.NO_RESULT_TERMINAL_STATES:
             return self._result
         try:
+            if self._in_running_loop():
+                return self._result_in_worker_thread()
             async_result = self.async_result()
             return async_result.get_loop().run_until_complete(async_result)
         except asyncio.CancelledError:
             # Future was cancelled, return whatever is in self._result if anything
             self._logger.warning("Task future was cancelled")
             return self._result
 
+    def _result_in_worker_thread(self) -> TaskResult:
+        """Poll on a worker thread with its own loop, cancelling it if the caller is interrupted."""
+        loop = asyncio.new_event_loop()
+        task = loop.create_task(self._wait_for_completion())
+        executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+        try:
+            return executor.submit(loop.run_until_complete, task).result()
+        except BaseException:
+            loop.call_soon_threadsafe(task.cancel)
+            raise
+        finally:
+            executor.shutdown(wait=True)
+            loop.close()
+
+    @staticmethod
+    def _in_running_loop() -> bool:
+        """Return True if an event loop is running in the current thread (e.g. Jupyter)."""
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return False
+        return True
+
     def _get_future(self) -> asyncio.Future:
         try:
             asyncio.get_event_loop()
@@ -472,7 +498,7 @@ def _get_future(self) -> asyncio.Future:
             # timed out and no result
             and self._update_status_if_nonterminal() not in self.NO_RESULT_TERMINAL_STATES
         ):
-            self._future = asyncio.get_event_loop().run_until_complete(self._create_future())
+            self._future = asyncio.get_event_loop().create_task(self._wait_for_completion())
         return self._future
 
     def async_result(self) -> asyncio.Task:
@@ -481,16 +507,6 @@ def async_result(self) -> asyncio.Task:
         """
         return self._get_future()
 
-    async def _create_future(self) -> asyncio.Task:
-        """Wrap the `_wait_for_completion` coroutine inside a future-like object.
-        Invoking this method starts the coroutine and returns back the future-like object
-        that contains it. Note that this does not block on the coroutine to finish.
-
-        Returns:
-            asyncio.Task: An asyncio Task that contains the `_wait_for_completion()` coroutine.
-        """
-        return asyncio.create_task(self._wait_for_completion())
-
     async def _wait_for_completion(
         self,
     ) -> TaskResult:

test/unit_tests/braket/aws/test_aws_quantum_task.py

@@ -12,11 +12,13 @@
 # language governing permissions and limitations under the License.
 
 import asyncio
+import concurrent.futures
+import itertools
 import json
 import threading
 import time
 import warnings
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import DEFAULT, MagicMock, Mock, patch
 
 import pytest
 from botocore.exceptions import ClientError
@@ -376,6 +378,7 @@ def asyncio_get_event_loop_side_effect(*args, **kwargs):
         yield mock
 
 
+@patch.object(AwsQuantumTask, "_wait_for_completion", new=MagicMock())
 @patch("braket.aws.aws_quantum_task.asyncio")
 def test_initialize_asyncio_event_loop_if_required(mock_asyncio, quantum_task):
     mock_asyncio.get_event_loop.side_effect = asyncio_get_event_loop_side_effect()
@@ -403,6 +406,55 @@ def test_result_circuit(circuit_task):
     )
 
 
+def test_result_circuit_from_running_event_loop(circuit_task):
+    _mock_metadata(circuit_task._aws_session, "RUNNING")
+    running = circuit_task._aws_session.get_quantum_task.return_value
+    _mock_metadata(circuit_task._aws_session, "COMPLETED")
+    circuit_task._aws_session.get_quantum_task.side_effect = itertools.chain(
+        [running], itertools.repeat(DEFAULT)
+    )
+    _mock_s3(circuit_task._aws_session, MockS3.MOCK_S3_RESULT_GATE_MODEL)
+    circuit_task._poll_interval_seconds = 0
+    expected = GateModelQuantumTaskResult.from_string(MockS3.MOCK_S3_RESULT_GATE_MODEL)
+
+    # Mimic nested event loop (like running from a Jupyter notebook)
+    async def call_result_from_coroutine():
+        return circuit_task.result()
+
+    assert asyncio.run(call_result_from_coroutine()) == expected
+
+
+def test_result_from_running_event_loop_interrupt_stops_polling(circuit_task):
+    _mock_metadata(circuit_task._aws_session, "RUNNING")
+    circuit_task._poll_interval_seconds = 0.01
+    circuit_task._poll_timeout_seconds = 5
+
+    async def call_result_from_coroutine():
+        return circuit_task.result()
+
+    with (
+        patch.object(concurrent.futures.Future, "result", side_effect=KeyboardInterrupt),
+        pytest.raises(KeyboardInterrupt),
+    ):
+        asyncio.run(call_result_from_coroutine())
+
+    # Check that the loop was interrupted by the user's KeyboardInterrupt, rather than
+    # waiting for a timeout (which would be closer to 500 calls)
+    assert circuit_task._aws_session.get_quantum_task.call_count < 10
+
+
+def test_async_result_from_running_event_loop(circuit_task):
+    _mock_metadata(circuit_task._aws_session, "COMPLETED")
+    _mock_s3(circuit_task._aws_session, MockS3.MOCK_S3_RESULT_GATE_MODEL)
+    expected = GateModelQuantumTaskResult.from_string(MockS3.MOCK_S3_RESULT_GATE_MODEL)
+
+    # Mimic nested event loop when the user calls async_result directly.
+    async def await_async_result():
+        return await circuit_task.async_result()
+
+    assert asyncio.run(await_async_result()) == expected
+
+
 def test_result_program_set(circuit_task):
     _mock_metadata(circuit_task._aws_session, "COMPLETED")
     _mock_s3(circuit_task._aws_session, MockS3.MOCK_S3_RESULT_PROGRAM_SET)
@@ -574,6 +626,7 @@ def set_result_from_callback(future):
     _mock_s3(circuit_task._aws_session, MockS3.MOCK_S3_RESULT_GATE_MODEL)
 
     future = circuit_task.async_result()
+    asyncio.get_event_loop().run_until_complete(asyncio.sleep(0))
 
     # test the different ways to get the result from async