Qiskit session

pennylane_qiskit/qiskit_device2.py

@@ -55,14 +55,21 @@
 
 # pylint: disable=protected-access
 @contextmanager
-def qiskit_session(device):
+def qiskit_session(device, **kwargs):
     """A context manager that creates a Qiskit Session and sets it as a session
     on the device while the context manager is active. Using the context manager
     will ensure the Session closes properly and is removed from the device after
-    completing the tasks.
+    completing the tasks. Any Session that was initialized and passed into the
+    device will be overwritten by the Qiskit Session created by this context
+    manager.
 
     Args:
         device (QiskitDevice2): the device that will create remote tasks using the session
+        **kwargs: session keyword arguments to be used for settings for the Session. At the
+        time of writing, the only relevant keyword argument is "max_time", which lets you
+        set the maximum amount of time the sessin is open. For the most up to date information,
+        please refer to the Qiskit Session
+        `documentation <https://docs.quantum.ibm.com/api/qiskit-ibm-runtime/qiskit_ibm_runtime.Session>`_.
 
     **Example:**
 
@@ -87,18 +94,60 @@ def circuit(x):
 
         angle = 0.1
 
-        with qiskit_session(dev) as session:
-
-            res = circuit(angle)[0] # you queue for the first execution
+        with qiskit_session(dev, max_time=60) as session:
+            # queue for the first execution
+            res = circuit(angle)[0]
 
             # then this loop executes immediately after without queueing again
             while res > 0:
                 angle += 0.3
                 res = circuit(angle)[0]
+
+    Note that if you passed in a session to your device, that session will be overwritten
+    by `qiskit_session`.
+
+    .. code-block:: python
+
+        import pennylane as qml
+        from pennylane_qiskit import qiskit_session
+        from qiskit_ibm_runtime import QiskitRuntimeService, Session
+
+        # get backend
+        service = QiskitRuntimeService(channel="ibm_quantum")
+        backend = service.least_busy(simulator=False, operational=True)
+
+        # initialize device
+        dev = qml.device('qiskit.remote', wires=2, backend=backend, session=Session(backend=backend, max_time=30))
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, 0)
+            qml.CNOT([0, 1])
+            return qml.expval(qml.PauliZ(1))
+
+        angle = 0.1
+
+        # This session will have the Qiskit default settings max_time=900
+        with qiskit_session(dev) as session:
+            res = circuit(angle)[0]
+
+            while res > 0:
+                angle += 0.3
+                res = circuit(angle)[0]
     """
     # Code to acquire session:
     existing_session = device._session
-    session = Session(backend=device.backend)
+
+    session_options = {"backend": device.backend, "service": device.service}
+
+    for k, v in kwargs.items():
+        # Options like service and backend should be tied to the settings set on device
+        if k in session_options:
+            warnings.warn(f"Using '{k}' set in device, {getattr(device, k)}", UserWarning)
+        else:
+            session_options[k] = v
+
+    session = Session(**session_options)
     device._session = session
     try:
         yield session

tests/test_base_device.py

@@ -22,8 +22,7 @@
 
 import pennylane as qml
 from pennylane.tape.qscript import QuantumScript
-
-from qiskit_ibm_runtime import Session, EstimatorV2 as Estimator
+from qiskit_ibm_runtime import EstimatorV2 as Estimator, Session
 from qiskit_ibm_runtime.fake_provider import FakeManila, FakeManilaV2
 from qiskit_aer import AerSimulator
 
@@ -104,8 +103,10 @@ def options(self):
 # pylint: disable=too-few-public-methods
 class MockSession:
     def __init__(self, backend, max_time=None):
-        self.backend = backend
-        self.max_time = max_time
+        self._backend = backend
+        self._max_time = max_time
+        self._args = "random"  # this is to satisfy a mock
+        self._kwargs = "random"  # this is to satisfy a mock
         self.session_id = "123"
 
     def close(self):  # This is just to appease a test
@@ -232,6 +233,94 @@ def test_using_session_context(self, mock_session, initial_session):
 
         assert dev._session == initial_session
 
+    def test_using_session_context_options(self):
+        """Test that you can set session options using qiskit_session"""
+        dev = QiskitDevice2(wires=2, backend=backend)
+
+        assert dev._session is None
+
+        with qiskit_session(dev, max_time=30) as session:
+            assert dev._session == session
+            assert dev._session is not None
+            assert dev._session._max_time == 30
+
+        assert dev._session is None
+
+    def test_error_when_passing_unexpected_kwarg(self):
+        """Test that we accept any keyword argument that the user wants to supply so that if
+        Qiskit allows for more customization we can automatically accomodate those needs. Right
+        now there are no such keyword arguments, so an error on Qiskit's side is raised."""
+
+        dev = QiskitDevice2(wires=2, backend=backend)
+
+        assert dev._session is None
+
+        with pytest.raises(
+            TypeError,  # Type error for wrong keyword argument differs across python versions
+        ):
+            with qiskit_session(dev, any_kwarg=30) as session:
+                assert dev._session == session
+                assert dev._session is not None
+
+        assert dev._session is None
+
+    def test_no_warning_when_using_initial_session_options(self):
+        initial_session = Session(backend=backend, max_time=30)
+        dev = QiskitDevice2(wires=2, backend=backend, session=initial_session)
+
+        assert dev._session == initial_session
+
+        with qiskit_session(dev) as session:
+            assert dev._session == session
+            assert dev._session != initial_session
+            assert dev._session._max_time == session._max_time
+            assert dev._session._max_time != initial_session._max_time
+
+        assert dev._session == initial_session
+        assert dev._session._max_time == initial_session._max_time
+
+    def test_warnings_when_overriding_session_context_options(self, recorder):
+        """Test that warnings are raised when the session options try to override either the
+        device's `backend` or `service`. Also ensures that the session options, even the
+        default options, passed in from the `qiskit_session` take precedence, barring
+        `backend` or `service`"""
+        initial_session = Session(backend=backend)
+        dev = QiskitDevice2(wires=2, backend=backend, session=initial_session)
+
+        assert dev._session == initial_session
+
+        with pytest.warns(
+            UserWarning,
+            match="Using 'backend' set in device",
+        ):
+            with qiskit_session(dev, max_time=30, backend=FakeManilaV2()) as session:
+                assert dev._session == session
+                assert dev._session != initial_session
+                assert dev._session._backend.name == "aer_simulator"
+
+        with pytest.warns(
+            UserWarning,
+            match="Using 'service' set in device",
+        ):
+            with qiskit_session(dev, max_time=30, service="placeholder") as session:
+                assert dev._session == session
+                assert dev._session != initial_session
+                assert dev._session._service != "placeholder"
+
+        # device session should be unchanged by qiskit_session
+        assert dev._session == initial_session
+
+        max_time_session = Session(backend=backend, max_time=60)
+        dev = QiskitDevice2(wires=2, backend=backend, session=max_time_session)
+        with qiskit_session(dev, max_time=30) as session:
+            assert dev._session == session
+            assert dev._session != initial_session
+            assert dev._session._max_time == 30
+            assert dev._session._max_time != 60
+
+        assert dev._session == max_time_session
+        assert dev._session._max_time == 60
+
     @pytest.mark.parametrize("initial_session", [None, MockSession(backend)])
     def test_update_session(self, initial_session):
         """Test that you can update the session stored on the device"""