feature: Program sets (#214)

Author: speller26

qiskit_braket_provider/providers/braket_backend.py

@@ -8,11 +8,12 @@
 from collections.abc import Iterable
 from typing import Optional, Union
 
-from braket.aws import AwsDevice, AwsQuantumTask, AwsQuantumTaskBatch
+from braket.aws import AwsDevice, AwsQuantumTask
 from braket.aws.queue_information import QueueDepthInfo
 from braket.circuits import Circuit
 from braket.device_schema import DeviceActionType
 from braket.devices import Device, LocalSimulator
+from braket.program_sets import ProgramSet
 from braket.tasks.local_quantum_task import LocalQuantumTask
 from qiskit import QuantumCircuit
 from qiskit.providers import BackendV2, Options, Provider, QubitProperties
@@ -175,7 +176,7 @@ def run(
             task_id=task_id,
             tasks=tasks,
             backend=self,
-            shots=shots,  # type: ignore[arg-type]
+            shots=shots,
         )
 
 
@@ -230,6 +231,9 @@ def __init__(  # pylint: disable=too-many-positional-arguments
             f"QiskitBraketProvider/{version.__version__}"
         )
         self._target = aws_device_to_target(device=self._aws_device)
+        self._supports_program_sets = (
+            DeviceActionType.OPENQASM_PROGRAM_SET in self._aws_device.properties.action
+        )
 
     def retrieve_job(self, task_id: str) -> BraketQuantumTask:
         """Return a single job submitted to AWS backend.
@@ -355,15 +359,28 @@ def run(self, run_input, verbatim: bool = False, native: bool = False, **options
             )
             for circ in circuits
         ]
+        shots = options.pop("shots", None)
+        return (
+            self._run_program_set(braket_circuits, shots, **options)
+            if self._supports_program_sets and shots != 0
+            else self._run_batch(braket_circuits, shots, **options)
+        )
 
-        batch_task: AwsQuantumTaskBatch = self._device.run_batch(
-            braket_circuits, **options
+    def _run_program_set(
+        self, braket_circuits: list[Circuit], shots: Optional[int], **options
+    ):
+        program_set = ProgramSet(braket_circuits, shots_per_executable=shots)
+        task = self._aws_device.run(program_set, **options)
+        return BraketQuantumTask(
+            task_id=task.id, tasks=task, backend=self, shots=program_set.total_shots
         )
+
+    def _run_batch(self, braket_circuits: list[Circuit], shots: int, **options):
+        batch_task = self._aws_device.run_batch(braket_circuits, shots=shots, **options)
         tasks: list[AwsQuantumTask] = batch_task.tasks
         task_id = _TASK_ID_DIVIDER.join(task.id for task in tasks)
-
         return BraketQuantumTask(
-            task_id=task_id, tasks=tasks, backend=self, shots=options.get("shots")
+            task_id=task_id, tasks=tasks, backend=self, shots=shots
         )
 
 

qiskit_braket_provider/providers/braket_quantum_task.py

@@ -1,80 +1,69 @@
 """Amazon Braket task."""
 
 from datetime import datetime
-from typing import List, Optional, Union
+from typing import List, Union
 
 from braket.aws import AwsQuantumTask, AwsQuantumTaskBatch
 from braket.aws.queue_information import QuantumTaskQueueInfo
+from braket.tasks import GateModelQuantumTaskResult, QuantumTask
 from braket.tasks.local_quantum_task import LocalQuantumTask
 from qiskit.providers import BackendV2, JobStatus, JobV1
 from qiskit.quantum_info import Statevector
 from qiskit.result import Result
 from qiskit.result.models import ExperimentResult, ExperimentResultData
 
+_TASK_STATUS_MAP = {
+    "INITIALIZED": JobStatus.INITIALIZING,
+    "QUEUED": JobStatus.INITIALIZING,
+    "FAILED": JobStatus.ERROR,
+    "CANCELLING": JobStatus.CANCELLED,
+    "CANCELLED": JobStatus.CANCELLED,
+    "COMPLETED": JobStatus.DONE,
+    "RUNNING": JobStatus.RUNNING,
+}
+
 
 def retry_if_result_none(result):
     """Retry on result function."""
     return result is None
 
 
-def _get_result_from_tasks(
-    tasks: Union[List[LocalQuantumTask], List[AwsQuantumTask]],
-) -> Optional[List[ExperimentResult]]:
-    """Returns experiment results of AWS tasks.
-
-    Args:
-        tasks: AWS Quantum tasks
-        shots: number of shots
-
-    Returns:
-        List of experiment results.
-    """
-    experiment_results: List[ExperimentResult] = []
-
-    results = AwsQuantumTaskBatch._retrieve_results(
-        tasks, AwsQuantumTaskBatch.MAX_CONNECTIONS_DEFAULT
-    )
-
-    # For each task we create an ExperimentResult object with the downloaded results.
-    for task, result in zip(tasks, results):
-        if not result:
-            return None
-
-        if result.task_metadata.shots == 0:
-            braket_statevector = result.values[
-                result._result_types_indices[
-                    "{'type': <Type.statevector: 'statevector'>}"
-                ]
-            ]
-            data = ExperimentResultData(
-                statevector=Statevector(braket_statevector).reverse_qargs().data,
-            )
-        else:
-            counts = {
-                k[::-1]: v for k, v in dict(result.measurement_counts).items()
-            }  # convert to little-endian
-
-            data = ExperimentResultData(
-                counts=counts,
-                memory=[
-                    "".join(shot_result[::-1].astype(str))
-                    for shot_result in result.measurements
-                ],
-            )
+def _result_from_circuit_task(
+    task: Union[LocalQuantumTask, AwsQuantumTask], result: GateModelQuantumTaskResult
+) -> ExperimentResult:
+    if not result:
+        return None
 
-        experiment_result = ExperimentResult(
-            shots=result.task_metadata.shots,
-            success=True,
-            status=(
-                task.state()
-                if isinstance(task, LocalQuantumTask)
-                else result.task_metadata.status
-            ),
-            data=data,
+    if result.task_metadata.shots == 0:
+        braket_statevector = result.values[
+            result._result_types_indices["{'type': <Type.statevector: 'statevector'>}"]
+        ]
+        data = ExperimentResultData(
+            statevector=Statevector(braket_statevector).reverse_qargs().data,
+        )
+    else:
+        counts = {
+            k[::-1]: v for k, v in dict(result.measurement_counts).items()
+        }  # convert to little-endian
+
+        data = ExperimentResultData(
+            counts=counts,
+            memory=[
+                "".join(shot_result[::-1].astype(str))
+                for shot_result in result.measurements
+            ],
         )
-        experiment_results.append(experiment_result)
 
-    return experiment_results
+    return ExperimentResult(
+        shots=result.task_metadata.shots,
+        success=True,
+        status=(
+            task.state()
+            if isinstance(task, LocalQuantumTask)
+            else result.task_metadata.status
+        ),
+        data=data,
+    )
 
 
 class BraketQuantumTask(JobV1):
@@ -84,8 +73,8 @@ def __init__(
         self,
         task_id: str,
         backend: BackendV2,
-        tasks: Union[List[LocalQuantumTask], List[AwsQuantumTask]],
-        **metadata: Optional[dict],
+        tasks: Union[List[LocalQuantumTask], List[AwsQuantumTask], AwsQuantumTask],
+        **metadata,
     ):
         """BraketQuantumTask for execution of circuits on Amazon Braket or locally.
 
@@ -147,6 +136,8 @@ def queue_position(self) -> QuantumTaskQueueInfo:
             queue_position=None, message='Task is in COMPLETED status. AmazonBraket does
             not show queue position for this status.')
         """
+        if isinstance(self._tasks, QuantumTask):
+            return self._tasks.queue_position()
         for task in self._tasks:
             if isinstance(task, LocalQuantumTask):
                 raise NotImplementedError(
@@ -159,7 +150,43 @@ def task_id(self) -> str:
         return self._task_id
 
     def result(self) -> Result:
-        experiment_results = _get_result_from_tasks(tasks=self._tasks)
+        tasks = self._tasks
+        if isinstance(tasks, QuantumTask):
+            # Guaranteed to be program set result
+            experiment_results = [
+                ExperimentResult(
+                    shots=len((executable_result := program_result[0]).measurements),
+                    success=True,
+                    data=ExperimentResultData(
+                        counts=executable_result.counts,
+                        memory=[
+                            "".join(shot_result[::-1].astype(str))
+                            for shot_result in executable_result.measurements
+                        ],
+                    ),
+                )
+                for program_result in tasks.result()
+            ]
+            status = tasks.state()
+            return Result(
+                backend_name=self._backend.name,
+                backend_version=self._backend.version,
+                job_id=self._task_id,
+                qobj_id=0,
+                success=status not in AwsQuantumTask.NO_RESULT_TERMINAL_STATES,
+                results=experiment_results,
+                status=status,
+            )
+
+        experiment_results = [
+            _result_from_circuit_task(task, result)
+            for task, result in zip(
+                tasks,
+                AwsQuantumTaskBatch._retrieve_results(
+                    tasks, AwsQuantumTaskBatch.MAX_CONNECTIONS_DEFAULT
+                ),
+            )
+        ]
         status = self.status(use_cached_value=True)
 
         return Result(
@@ -173,10 +200,15 @@ def result(self) -> Result:
         )
 
     def cancel(self):
-        for task in self._tasks:
-            task.cancel()
+        if isinstance(self._tasks, QuantumTask):
+            self._tasks.cancel()
+        else:
+            for task in self._tasks:
+                task.cancel()
 
     def status(self, use_cached_value: bool = False):
+        if isinstance(self._tasks, QuantumTask):
+            return _TASK_STATUS_MAP[self._tasks.state()]
         braket_tasks_states = [
             (
                 task.state()

qiskit_braket_provider/version.py

@@ -1,3 +1,3 @@
 """Qiskit-Braket provider version."""
 
-__version__ = "0.4.6"
+__version__ = "0.5.0"

requirements.txt

@@ -1,7 +1,7 @@
 certifi>=2021.5.30
 qiskit>=0.34.2, <2.0
 qiskit-ionq>=0.5.2
-amazon-braket-sdk>=1.76.0
+amazon-braket-sdk>=1.97.0
 
 setuptools>=40.1.0
 numpy>=1.3