Qiskit session

pennylane_qiskit/qiskit_device2.py

@@ -55,7 +55,7 @@
 
 # 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
@@ -98,7 +98,34 @@ def circuit(x):
     """
     # Code to acquire session:
     existing_session = device._session
-    session = Session(backend=device.backend)
+
+    # When an existing session exists, we want to use its settings unless overwritten
+    # by settings in the qiskit_session
+    if existing_session:
+        session_args = inspect.signature(Session).parameters
+        session_options = {arg: getattr(existing_session, "_" + arg) for arg in session_args}
+    else:
+        # when an existing session doesn't exist, we create a session with default settings
+        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 and k != "max_time":
+            warnings.warn(f"Using '{k}' set in device, {getattr(device, k)}", UserWarning)
+
+        # Need "_" since `max_time` attribute on Session is `_max_time`
+        # When there is overlap between the Session options on the device and the session options
+        # passed in via qiskit_session, we prefer the ones passed in via qiskit_session
+        elif existing_session and getattr(existing_session, "_" + k):
+            warnings.warn(
+                f"`{k}` was also set in the Session passed to the device. Using `{k}` '{v}' set in `qiskit_session`.",
+                UserWarning,
+            )
+            session_options[k] = v
+        else:
+            session_options[k] = v
+
+    session = Session(**session_options)
     device._session = session
     try:
         yield session

tests/test_base_device.py

@@ -23,7 +23,7 @@
 
 import pennylane as qml
 from pennylane.tape.qscript import QuantumScript
-from qiskit_ibm_runtime import 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
@@ -105,8 +105,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
@@ -269,6 +271,99 @@ 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, match=r"Session.__init__\(\) got an unexpected keyword argument 'any_kwarg'"
+        ):
+            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 there are is an overlap between the Session options
+        on the device and the session options passed in via qiskit_session. Also ensures that the
+        session options passed in from the qiskit_session take precedence except for `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 pytest.warns(
+            UserWarning,
+            match="`max_time` was set in the Session passed to the device. Using `max_time` '30' set in `qiskit_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"""