Skip to content

feat: Automatically split/merge large program sets #331

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
speller26 wants to merge 2 commits into
base: main
Choose a base branch
from
+133 −95
Open

feat: Automatically split/merge large program sets #331

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5fbf1db
feat: Automatically split/merge large program sets
speller26 Jun 19, 2026
37faa62
Merge branch 'main' into split
speller26 Jul 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 40 additions & 30 deletions src/braket/pennylane_plugin/braket_device.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,11 @@
)
from braket.program_sets import ProgramSet
from braket.simulator import BraketSimulator
from braket.tasks import GateModelQuantumTaskResult, QuantumTask
from braket.tasks import (
GateModelQuantumTaskResult,
ProgramSetQuantumTaskResult,
QuantumTask,
)
from braket.tasks.local_quantum_task_batch import LocalQuantumTaskBatch

from ._version import __version__
Expand Down Expand Up @@ -222,12 +226,6 @@ def batch_execute(self, circuits, **run_kwargs):
if not self._parallel and not self._supports_program_sets:
return super().batch_execute(circuits)

if self._supports_program_sets and (
len(circuits)
> self._device.properties.action["braket.ir.openqasm.program_set"].maximumExecutables
):
return super().batch_execute(circuits)

for circuit in circuits:
self.check_validity(circuit.operations, circuit.observables)
all_trainable = []
Expand Down Expand Up @@ -735,22 +733,41 @@ def use_grouping(self) -> bool:
caps = self.capabilities()
return not (caps.get("provides_jacobian"))

def _run_program_set(
self, program_set: ProgramSet, shots_per_executable: int
) -> ProgramSetQuantumTaskResult:
program_sets, index_map = program_set.split(
self._device.properties.action[DeviceActionType.OPENQASM_PROGRAM_SET].maximumExecutables
)
return ProgramSetQuantumTaskResult.merge(
[
self._device.run(
sub_program_set,
s3_destination_folder=self._s3_folder,
shots=sub_program_set.total_executables * shots_per_executable,
poll_timeout_seconds=self._poll_timeout_seconds,
poll_interval_seconds=self._poll_interval_seconds,
**self._run_kwargs,
).result()
for sub_program_set in program_sets
],
Comment on lines +743 to +753

@rmshaffer rmshaffer Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is going to wait for each program set task to run to completion before submitting the next one. Could we first call run() on each task, and then accumulate the results here? e.g.

tasks = [self._device.run(...) for sub_program_set in program_sets]
return ProgramSetQuantumTaskResult.merge([task.result() for task in tasks], ...)

program_set,
index_map,
)

def _run_task_batch(self, braket_circuits, pl_circuits, batch_shots: int, inputs):
if self._supports_program_sets:
program_set = (
ProgramSet.zip(braket_circuits, input_sets=inputs)
if inputs
else ProgramSet(braket_circuits)
)
task = self._device.run(
program_set,
s3_destination_folder=self._s3_folder,
shots=len(program_set) * batch_shots,
poll_timeout_seconds=self._poll_timeout_seconds,
poll_interval_seconds=self._poll_interval_seconds,
**self._run_kwargs,
return self._braket_program_set_to_pl_result(
self._run_program_set(
(
ProgramSet.zip(braket_circuits, input_sets=inputs)
if inputs
else ProgramSet(braket_circuits)
Comment on lines +763 to +765

@rmshaffer rmshaffer Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

readability nit: these could be separate variables instead of nested calls

program_set = (
    ProgramSet.zip(braket_circuits, input_sets=inputs)
    if inputs
    else ProgramSet(braket_circuits)
)
result = self._run_program_set(program_set, batch_shots)
return self._braket_program_set_to_pl_result(result, pl_circuits)

),
batch_shots,
),
pl_circuits,
)
return self._braket_program_set_to_pl_result(task.result(), pl_circuits)
task_batch = self._device.run_batch(
braket_circuits,
s3_destination_folder=self._s3_folder,
Expand Down Expand Up @@ -794,16 +811,9 @@ def _run_snapshots(self, snapshot_circuits, n_qubits, mapped_wires):
n_snapshots = len(snapshot_circuits)
outcomes = np.zeros((n_snapshots, n_qubits))
if self._supports_program_sets:
program_set = ProgramSet(snapshot_circuits)
task = self._device.run(
program_set,
s3_destination_folder=self._s3_folder,
shots=len(program_set),
poll_timeout_seconds=self._poll_timeout_seconds,
poll_interval_seconds=self._poll_interval_seconds,
**self._run_kwargs,
)
for t, result in enumerate(task.result()):
for t, result in enumerate(
self._run_program_set(ProgramSet(snapshot_circuits), shots_per_executable=1)
):
outcomes[t] = np.array(result[0].measurements[0])[mapped_wires]
elif self._parallel:
task_batch = self._device.run_batch(
Expand Down
154 changes: 89 additions & 65 deletions test/unit_tests/test_braket_device.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,37 @@
)


def _make_program_set_result(num_programs):
"""Builds a ProgramSetQuantumTaskResult with ``num_programs`` single-executable programs,
each returning the measurements in PROGRAM_RESULT."""
return ProgramSetQuantumTaskResult.from_object(
ProgramSetTaskResult(
**{
"braketSchemaHeader": {
"name": "braket.task_result.program_set_task_result",
"version": "1",
},
"programResults": [PROGRAM_RESULT] * num_programs,
"taskMetadata": {
"braketSchemaHeader": {
"name": "braket.task_result.program_set_task_metadata",
"version": "1",
},
"id": "arn:aws:braket:us-west-2:667256736152:quantum-task/bfebc86f-e4ed-4d6f-8131-addd1a49d6dc", # noqa
"deviceId": "arn:aws:braket:::device/quantum-simulator/amazon/sv1",
"requestedShots": 20 * num_programs,
"successfulShots": 20 * num_programs,
"programMetadata": [{"executables": [{}]} for _ in range(num_programs)],
"createdAt": "2024-10-15T19:06:58.986Z",
"endedAt": "2024-10-15T19:07:00.382Z",
"status": "COMPLETED",
"totalFailedExecutables": 0,
},
}
)
)


DEVICE_ARN = "baz"


Expand Down Expand Up @@ -1161,71 +1192,28 @@ def test_batch_execute_program_set_noncommuting():

@patch.object(AwsDevice, "run")
def test_batch_execute_program_set_exceeds_max_executables(mock_run):
"""Test batch_execute falls back to individual programs when exceeding maximumExecutables"""
custom_result = GateModelQuantumTaskResult.from_string(
json.dumps(
{
"braketSchemaHeader": {
"name": "braket.task_result.gate_model_task_result",
"version": "1",
},
"measurements": [[0, 0], [1, 1], [0, 1], [1, 0]],
"resultTypes": [
{
"type": {
"observable": ["x", "y"],
"targets": [0, 1],
"type": "expectation",
},
"value": 0.5,
},
{
"type": {
"observable": ["z", "z"],
"targets": [0, 1],
"type": "expectation",
},
"value": 0.5,
},
],
"measuredQubits": [0, 1],
"taskMetadata": {
"braketSchemaHeader": {
"name": "braket.task_result.task_metadata",
"version": "1",
},
"id": "task_arn",
"shots": 10000,
"deviceId": "default",
},
"additionalMetadata": {
"action": {
"braketSchemaHeader": {
"name": "braket.ir.openqasm.program",
"version": "1",
},
"source": "qubit[2] q; h q[0]; cnot q[0], q[1]; measure q;",
},
},
}
)
)
"""Test batch_execute splits the program set and merges the results when the number of
executables exceeds the device's maximumExecutables."""

# Mock the task to return our custom result
task = Mock()
task.result.return_value = custom_result
type(task).id = PropertyMock(return_value="task_arn")
task.state.return_value = "COMPLETED"
mock_run.return_value = task
# The program set is split into one task per chunk of maximumExecutables circuits, so each
# call to run returns a result sized to the chunk it was given.
def run_side_effect(program_set, **kwargs):
task = Mock()
task.result.return_value = _make_program_set_result(program_set.total_executables)
return task

dev = _aws_device(wires=4, foo="bar", parallel=True, supports_program_sets=True)
mock_run.side_effect = run_side_effect

dev = _aws_device(wires=4, foo="bar", parallel=False, supports_program_sets=True)

# Verify the device properties are set correctly to ensure we hit the second condition
assert dev._parallel == True
assert dev._supports_program_sets == True
assert dev._device.properties.action["braket.ir.openqasm.program_set"].maximumExecutables == 100
max_executables = dev._device.properties.action[
"braket.ir.openqasm.program_set"
].maximumExecutables
assert max_executables == 100

# Create 101 circuits (exceeds maximumExecutables of 100, defined in ACTION_PROPERTIES_PROGRAMSET)
# Create 101 circuits, exceeding maximumExecutables of 100 (defined in
# ACTION_PROPERTIES_PROGRAMSET), so the program set is split into two tasks.
circuits = []
for _ in range(101):
with QuantumTape() as circuit:
Expand All @@ -1235,16 +1223,52 @@ def test_batch_execute_program_set_exceeds_max_executables(mock_run):
circuits.append(circuit)

assert len(circuits) == 101
assert (
len(circuits)
> dev._device.properties.action["braket.ir.openqasm.program_set"].maximumExecutables
)
assert len(circuits) > max_executables

result = dev.batch_execute(circuits)
assert mock_run.call_count == 101

# Two tasks: 100 executables in the first, 1 in the second.
assert mock_run.call_count == 2
run_sizes = sorted(call.args[0].total_executables for call in mock_run.call_args_list)
assert run_sizes == [1, 100]

# The merged result preserves the shape of the original (unsplit) batch.
assert len(result) == 101


@patch.object(AwsDevice, "run")
def test_run_snapshots_program_set_exceeds_max_executables(mock_run):
"""Test _run_snapshots splits the program set and merges the results when the number of
snapshots exceeds the device's maximumExecutables."""

def run_side_effect(program_set, **kwargs):
task = Mock()
task.result.return_value = _make_program_set_result(program_set.total_executables)
return task

mock_run.side_effect = run_side_effect

dev = _aws_device(wires=2, foo="bar", parallel=False, supports_program_sets=True)
max_executables = dev._device.properties.action[
"braket.ir.openqasm.program_set"
].maximumExecutables
assert max_executables == 100

# 101 snapshots exceeds maximumExecutables of 100, so the program set is split into two tasks.
n_snapshots = 101
snapshot_circuits = [Circuit().h(0).cnot(0, 1) for _ in range(n_snapshots)]
mapped_wires = np.array([0, 1])

outcomes = dev._run_snapshots(snapshot_circuits, n_qubits=2, mapped_wires=mapped_wires)

assert mock_run.call_count == 2
run_sizes = sorted(call.args[0].total_executables for call in mock_run.call_args_list)
assert run_sizes == [1, 100]

# One outcome per snapshot, each shots=1 measurement projected onto the mapped wires.
assert outcomes.shape == (n_snapshots, 2)

@rmshaffer rmshaffer Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could we do some basic verification of the content of outcomes to ensure that the merge/split happened correctly?



@patch.object(AwsDevice, "properties", new_callable=mock.PropertyMock)
@patch.object(AwsDevice, "run_batch")
def test_aws_device_batch_execute_parallel(mock_run_batch, mock_properties):
Expand Down
4 changes: 4 additions & 0 deletions test/unit_tests/test_shadow_expval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -374,6 +374,10 @@ def test_shadow_expval_aws_device(
):
mock_action = Mock()
mock_action.action = {"braket.ir.openqasm.program": None}
if supports_program_sets:
program_set_action = Mock()
program_set_action.maximumExecutables = 100
mock_action.action[DeviceActionType.OPENQASM_PROGRAM_SET] = program_set_action
mock_properties.return_value = mock_action
dev = _aws_device(
wires=2,
Expand Down
Loading