Fix issue with oversized batches in offline simulators (#210)

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+* Fix issue with oversized batches on offline simulators (#210)
+
 ## qiskit-aqt-provider v1.9.0
 
 * Fix source installation problem caused by removed `debugpy` 1.8.3 package (#188)

qiskit_aqt_provider/aqt_resource.py

@@ -33,7 +33,7 @@
 from qiskit.providers.models import BackendConfiguration
 from qiskit.transpiler import Target
 from qiskit_aer import AerJob, AerSimulator, noise
-from typing_extensions import override
+from typing_extensions import TypeAlias, override
 
 from qiskit_aqt_provider import api_client
 from qiskit_aqt_provider.api_client import models as api_models
@@ -452,7 +452,7 @@ def __init__(
             resource_id=resource_id,
         )
 
-        self.job: Optional[SimulatorJob] = None
+        self.jobs: dict[UUID, SimulatorJob] = {}
 
         if not with_noise_model:
             noise_model = None
@@ -491,12 +491,17 @@ def submit(self, job: AQTJob) -> UUID:
             for circuit in job.api_submit_payload.payload.circuits
         ]
 
-        self.job = SimulatorJob(
+        sim_job = SimulatorJob(
             job=self.simulator.run(circuits, shots=job.options.shots),
             circuits=job.circuits,
             shots=job.options.shots,
         )
-        return self.job.job_id
+
+        # The Aer job is freshly created above, so its ID is unique
+        # among the keys in self.jobs.
+        self.jobs[sim_job.job_id] = sim_job
+
+        return sim_job.job_id
 
     @override
     def result(self, job_id: UUID) -> api_models.JobResponse:
@@ -514,13 +519,13 @@ def result(self, job_id: UUID) -> api_models.JobResponse:
         Raises:
             UnknownJobError: ``job_id`` doesn't correspond to a simulator job on this resource.
         """
-        if self.job is None or job_id != self.job.job_id:
+        if (job := self.jobs.get(job_id)) is None:
             raise api_models.UnknownJobError(str(job_id))
 
-        qiskit_result = self.job.job.result()
+        qiskit_result = job.job.result()
 
         results: dict[str, list[list[int]]] = {}
-        for circuit_index, circuit in enumerate(self.job.circuits):
+        for circuit_index, circuit in enumerate(job.circuits):
             samples: list[list[int]] = []
 
             # Use data()["counts"] instead of get_counts() to access the raw counts
@@ -543,3 +548,7 @@ def result(self, job_id: UUID) -> api_models.JobResponse:
             resource_id=self.resource_id.resource_id,
             results=results,
         )
+
+
+AnyAQTResource: TypeAlias = Union[AQTResource, AQTDirectAccessResource]
+"""Type of any remote or direct access resource."""

qiskit_aqt_provider/primitives/estimator.py

@@ -16,17 +16,17 @@
 from qiskit.primitives import BackendEstimator
 
 from qiskit_aqt_provider import transpiler_plugin
-from qiskit_aqt_provider.aqt_resource import AQTResource, make_transpiler_target
+from qiskit_aqt_provider.aqt_resource import AnyAQTResource, make_transpiler_target
 
 
 class AQTEstimator(BackendEstimator):
     """:class:`BaseEstimatorV1 <qiskit.primitives.BaseEstimatorV1>` primitive for AQT backends."""
 
-    _backend: AQTResource
+    _backend: AnyAQTResource
 
     def __init__(
         self,
-        backend: AQTResource,
+        backend: AnyAQTResource,
         options: Optional[dict[str, Any]] = None,
         abelian_grouping: bool = True,
         skip_transpilation: bool = False,
@@ -66,6 +66,6 @@ def __init__(
         )
 
     @property
-    def backend(self) -> AQTResource:
+    def backend(self) -> AnyAQTResource:
         """Computing resource used for circuit evaluation."""
         return self._backend

qiskit_aqt_provider/primitives/sampler.py

@@ -16,17 +16,17 @@
 from qiskit.primitives import BackendSampler
 
 from qiskit_aqt_provider import transpiler_plugin
-from qiskit_aqt_provider.aqt_resource import AQTResource, make_transpiler_target
+from qiskit_aqt_provider.aqt_resource import AnyAQTResource, make_transpiler_target
 
 
 class AQTSampler(BackendSampler):
     """:class:`BaseSamplerV1 <qiskit.primitives.BaseSamplerV1>` primitive for AQT backends."""
 
-    _backend: AQTResource
+    _backend: AnyAQTResource
 
     def __init__(
         self,
-        backend: AQTResource,
+        backend: AnyAQTResource,
         options: Optional[dict[str, Any]] = None,
         skip_transpilation: bool = False,
     ) -> None:
@@ -101,6 +101,6 @@ def __init__(
         )
 
     @property
-    def backend(self) -> AQTResource:
+    def backend(self) -> AnyAQTResource:
         """Computing resource used for circuit evaluation."""
         return self._backend

qiskit_aqt_provider/test/fixtures.py

@@ -17,13 +17,13 @@
 
 import json
 import re
+import typing
 import uuid
 
 import httpx
 import pytest
 from pytest_httpx import HTTPXMock
 from qiskit.circuit import QuantumCircuit
-from qiskit.providers import BackendV2
 from qiskit_aer import AerSimulator
 from typing_extensions import override
 
@@ -33,6 +33,7 @@
 from qiskit_aqt_provider.aqt_job import AQTJob
 from qiskit_aqt_provider.aqt_provider import AQTProvider
 from qiskit_aqt_provider.aqt_resource import (
+    AnyAQTResource,
     AQTDirectAccessResource,
     OfflineSimulatorResource,
     qubit_states_from_int,
@@ -142,10 +143,11 @@ def handle_result(request: httpx.Request) -> httpx.Response:
     name="any_offline_simulator_no_noise",
     params=["offline_simulator_no_noise", "offline_simulator_no_noise_direct_access"],
 )
-def fixture_any_offline_simulator_no_noise(request: pytest.FixtureRequest) -> BackendV2:
+def fixture_any_offline_simulator_no_noise(request: pytest.FixtureRequest) -> AnyAQTResource:
     """Noiseless, offline simulator backend.
 
     The fixture is parametrized to successively run the dependent tests
     with a regular cloud-bound backend, and a direct-access one.
     """
-    return request.getfixturevalue(request.param)
+    # cast: all fixture parameters have types compatible with this function's return type.
+    return typing.cast(AnyAQTResource, request.getfixturevalue(request.param))

test/test_execution.py

@@ -31,15 +31,15 @@
 from qiskit_aer import AerProvider, AerSimulator
 
 from qiskit_aqt_provider import AQTProvider
-from qiskit_aqt_provider.aqt_resource import AQTResource
+from qiskit_aqt_provider.aqt_resource import AnyAQTResource, AQTResource
 from qiskit_aqt_provider.test.circuits import assert_circuits_equivalent
 from qiskit_aqt_provider.test.fixtures import MockSimulator
 from qiskit_aqt_provider.test.resources import TestResource
 from qiskit_aqt_provider.test.timeout import timeout
 
 
 @pytest.mark.parametrize("shots", [200])
-def test_empty_circuit(shots: int, any_offline_simulator_no_noise: BackendV2) -> None:
+def test_empty_circuit(shots: int, any_offline_simulator_no_noise: AnyAQTResource) -> None:
     """Run an empty circuit."""
     qc = QuantumCircuit(1)
     qc.measure_all()
@@ -102,7 +102,7 @@ def test_cancelled_circuit() -> None:
 
 
 @pytest.mark.parametrize("shots", [1, 100, 200])
-def test_simple_backend_run(shots: int, any_offline_simulator_no_noise: BackendV2) -> None:
+def test_simple_backend_run(shots: int, any_offline_simulator_no_noise: AnyAQTResource) -> None:
     """Run a simple circuit with `backend.run`."""
     qc = QuantumCircuit(1)
     qc.rx(pi, 0)
@@ -144,7 +144,7 @@ def test_simple_backend_execute_noisy(resource: MockSimulator) -> None:
 
 
 @pytest.mark.parametrize("shots", [100])
-def test_ancilla_qubits_mapping(shots: int, any_offline_simulator_no_noise: BackendV2) -> None:
+def test_ancilla_qubits_mapping(shots: int, any_offline_simulator_no_noise: AnyAQTResource) -> None:
     """Run a circuit with two quantum registers, with only one mapped to the classical memory."""
     qr = QuantumRegister(2)
     qr_aux = QuantumRegister(3)
@@ -164,7 +164,7 @@ def test_ancilla_qubits_mapping(shots: int, any_offline_simulator_no_noise: Back
 
 @pytest.mark.parametrize("shots", [100])
 def test_multiple_classical_registers(
-    shots: int, any_offline_simulator_no_noise: BackendV2
+    shots: int, any_offline_simulator_no_noise: AnyAQTResource
 ) -> None:
     """Run a circuit with the final state mapped to multiple classical registers."""
     qr = QuantumRegister(5)
@@ -187,7 +187,7 @@ def test_multiple_classical_registers(
 @pytest.mark.parametrize("shots", [123])
 @pytest.mark.parametrize("memory_opt", [True, False])
 def test_get_memory_simple(
-    shots: int, memory_opt: bool, any_offline_simulator_no_noise: BackendV2
+    shots: int, memory_opt: bool, any_offline_simulator_no_noise: AnyAQTResource
 ) -> None:
     """Check that the raw bitstrings can be accessed for each shot via the
     get_memory() method in Qiskit's Result.
@@ -214,7 +214,9 @@ def test_get_memory_simple(
 
 
 @pytest.mark.parametrize("shots", [123])
-def test_get_memory_ancilla_qubits(shots: int, any_offline_simulator_no_noise: BackendV2) -> None:
+def test_get_memory_ancilla_qubits(
+    shots: int, any_offline_simulator_no_noise: AnyAQTResource
+) -> None:
     """Check that the raw bistrings returned by get_memory() in Qiskit's Result only
     contain the mapped classical bits.
     """
@@ -238,7 +240,9 @@ def test_get_memory_ancilla_qubits(shots: int, any_offline_simulator_no_noise: B
 
 
 @pytest.mark.parametrize("shots", [123])
-def test_get_memory_bit_ordering(shots: int, any_offline_simulator_no_noise: BackendV2) -> None:
+def test_get_memory_bit_ordering(
+    shots: int, any_offline_simulator_no_noise: AnyAQTResource
+) -> None:
     """Check that the bitstrings returned by the results produced by AQT jobs have the same
     bit order as the Qiskit Aer simulators.
     """
@@ -304,7 +308,9 @@ def test_regression_issue_85(backend: BackendV2) -> None:
 
 
 @pytest.mark.parametrize(("shots", "qubits"), [(100, 5), (100, 8)])
-def test_bell_states(shots: int, qubits: int, any_offline_simulator_no_noise: BackendV2) -> None:
+def test_bell_states(
+    shots: int, qubits: int, any_offline_simulator_no_noise: AnyAQTResource
+) -> None:
     """Create a N qubits Bell state."""
     qc = QuantumCircuit(qubits)
     qc.h(0)
@@ -334,7 +340,7 @@ def test_bell_states(shots: int, qubits: int, any_offline_simulator_no_noise: Ba
 def test_state_preparation(
     target_state: Union[int, str, quantum_info.Statevector, list[complex]],
     optimization_level: int,
-    any_offline_simulator_no_noise: BackendV2,
+    any_offline_simulator_no_noise: AnyAQTResource,
 ) -> None:
     """Test the state preparation unitary factory.
 
@@ -357,7 +363,7 @@ def test_state_preparation(
 
 @pytest.mark.parametrize("optimization_level", range(4))
 def test_state_preparation_single_qubit(
-    optimization_level: int, any_offline_simulator_no_noise: BackendV2
+    optimization_level: int, any_offline_simulator_no_noise: AnyAQTResource
 ) -> None:
     """Test the state preparation unitary factory, targeting a single qubit in the register."""
     qreg = QuantumRegister(4)
@@ -396,7 +402,7 @@ def test_initialize_not_supported(offline_simulator_no_noise: AQTResource) -> No
 
 
 @pytest.mark.parametrize("optimization_level", range(4))
-def test_cswap(optimization_level: int, any_offline_simulator_no_noise: BackendV2) -> None:
+def test_cswap(optimization_level: int, any_offline_simulator_no_noise: AnyAQTResource) -> None:
     """Verify that CSWAP (Fredkin) gates can be transpiled and executed (in a trivial case)."""
     qc = QuantumCircuit(3)
     qc.prepare_state("101")

test/test_primitives.py

@@ -23,13 +23,14 @@
     BaseSamplerV1,
     Sampler,
 )
-from qiskit.providers import Backend, BackendV2
+from qiskit.providers import Backend
 from qiskit.quantum_info import SparsePauliOp
 from qiskit.transpiler.exceptions import TranspilerError
 
+from qiskit_aqt_provider.aqt_resource import AnyAQTResource
 from qiskit_aqt_provider.primitives import AQTSampler
 from qiskit_aqt_provider.primitives.estimator import AQTEstimator
-from qiskit_aqt_provider.test.circuits import assert_circuits_equal
+from qiskit_aqt_provider.test.circuits import assert_circuits_equal, random_circuit
 from qiskit_aqt_provider.test.fixtures import MockSimulator
 
 
@@ -66,7 +67,8 @@ def test_backend_primitives_are_v1() -> None:
 )
 @pytest.mark.httpx_mock(assert_all_responses_were_requested=False)
 def test_circuit_sampling_primitive(
-    get_sampler: Callable[[Backend], BaseSamplerV1], any_offline_simulator_no_noise: BackendV2
+    get_sampler: Callable[[Backend], BaseSamplerV1],
+    any_offline_simulator_no_noise: AnyAQTResource,
 ) -> None:
     """Check that a `Sampler` primitive using an AQT backend can sample parametric circuits."""
     theta = Parameter("θ")
@@ -183,3 +185,22 @@ def test_aqt_sampler_transpilation(theta: float, offline_simulator_no_noise: Moc
     tr_expected = qiskit.transpile(expected, offline_simulator_no_noise)
 
     assert_circuits_equal(transpiled_circuit, tr_expected)
+
+
+@pytest.mark.httpx_mock(can_send_already_matched_responses=True)  # for the direct-access mocks
+def test_sampler_circuit_batching(any_offline_simulator_no_noise: AnyAQTResource) -> None:
+    """Check that a Sampler primitive on an offline simulator can split oversized job batches.
+
+    Regression test for #203.
+    """
+    # Arbitrary circuit.
+    qc = random_circuit(2)
+
+    sampler = AQTSampler(any_offline_simulator_no_noise)
+
+    # Use a Sampler batch larger than the maximum number of circuits
+    # per batch in the API.
+    batch_size = 3 * any_offline_simulator_no_noise.max_circuits
+    result = sampler.run([qc] * batch_size).result()
+
+    assert len(result.quasi_dists) == batch_size