Skip to content

All multi-circuit programs to be batched in the same request in Engine. #8036

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

Merged
dstrain115 merged 7 commits into from
Apr 22, 2026
Merged

All multi-circuit programs to be batched in the same request in Engine. #8036

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b32a750
All multi-circuit programs to be batched in the same request in Engine.
dstrain115 Apr 13, 2026
72c9b87
Fix coverage and tests. Add support for mapping.
dstrain115 Apr 13, 2026
42883be
Fix type check
dstrain115 Apr 14, 2026
f9c66c7
Fix issue with cross-product of sweeps.
dstrain115 Apr 20, 2026
b752186
formatting
dstrain115 Apr 20, 2026
f18fc5c
Address maffoo comments.
dstrain115 Apr 22, 2026
2904720
return "Oops"
dstrain115 Apr 22, 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
7 changes: 6 additions & 1 deletion cirq-google/cirq_google/engine/abstract_processor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@

import abc
import datetime
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING

import duet
Expand Down Expand Up @@ -126,7 +127,11 @@ async def run_async(
@abc.abstractmethod
async def run_sweep_async(
self,
program: cirq.AbstractCircuit,
program: (
cirq.AbstractCircuit
| Sequence[cirq.AbstractCircuit]
| Mapping[str, cirq.AbstractCircuit]
),
*,
device_config_name: str,
run_name: str = "",
Expand Down
53 changes: 43 additions & 10 deletions cirq-google/cirq_google/engine/engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import enum
import random
import string
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, TypeVar

import duet
Expand All @@ -46,7 +47,7 @@
processor_config,
util,
)
from cirq_google.serialization import CIRCUIT_SERIALIZER, Serializer
from cirq_google.serialization import CIRCUIT_SERIALIZER, CircuitSerializer

if TYPE_CHECKING:
from google.protobuf import any_pb2
Expand Down Expand Up @@ -87,7 +88,7 @@ def __init__(
verbose: bool | None = None,
client: engine_client.EngineClient | None = None,
timeout: int | None = None,
serializer: Serializer = CIRCUIT_SERIALIZER,
serializer: CircuitSerializer = CIRCUIT_SERIALIZER,
# TODO(#5996) Remove enable_streaming once the feature is stable.
enable_streaming: bool = True,
compress_run_context: bool = False,
Expand Down Expand Up @@ -145,6 +146,22 @@ def _serialize_program(self, program: cirq.AbstractCircuit) -> any_pb2.Any:
raise ValueError(f'invalid program proto version: {self.proto_version}')
return util.pack_any(self.serializer.serialize(program))

def _serialize_multi_program(
self, programs: Sequence[cirq.AbstractCircuit] | Mapping[str, cirq.AbstractCircuit]
) -> any_pb2.Any:
if isinstance(programs, Mapping):
if any(not isinstance(value, cirq.AbstractCircuit) for value in programs.values()):
raise TypeError(f'Unrecognized program type: {type(programs)}')
elif isinstance(programs, Sequence):
if any(not isinstance(value, cirq.AbstractCircuit) for value in programs):
raise TypeError(f'Unrecognized program type: {type(programs)}')
else:
raise TypeError(f'Unrecognized program type: {type(programs)}')

if self.proto_version != ProtoVersion.V2:
raise ValueError(f'invalid program proto version: {self.proto_version}')
return util.pack_any(self.serializer.serialize_multi_program(programs))

def _serialize_run_context(self, sweeps: cirq.Sweepable, repetitions: int) -> any_pb2.Any:
if self.proto_version != ProtoVersion.V2:
raise ValueError(f'invalid run context proto version: {self.proto_version}')
Expand Down Expand Up @@ -306,7 +323,11 @@ def run(

async def run_sweep_async(
self,
program: cirq.AbstractCircuit,
program: (
cirq.AbstractCircuit
| Sequence[cirq.AbstractCircuit]
| Mapping[str, cirq.AbstractCircuit]
),
processor_id: str,
program_id: str | None = None,
job_id: str | None = None,
Expand All @@ -329,6 +350,8 @@ async def run_sweep_async(
Args:
program: The Circuit to execute. If a circuit is
provided, a moment by moment schedule will be used.
A list or mapping of programs can also be provided.
These will be executed as KeyedCircuits.
program_id: A user-provided identifier for the program. This must
be unique within the Google Cloud project being used. If this
parameter is not provided, a random id of the format
Expand Down Expand Up @@ -375,12 +398,17 @@ async def run_sweep_async(
job_id = _make_random_id('job-')
run_context = self.context._serialize_run_context(params, repetitions)

if isinstance(program, cirq.AbstractCircuit):
code = self.context._serialize_program(program)
else:
code = self.context._serialize_multi_program(program)

job_result_future = self.context.client.run_job_over_stream(
project_id=self.project_id,
program_id=str(program_id),
program_description=program_description,
program_labels=program_labels,
code=self.context._serialize_program(program),
code=code,
job_id=str(job_id),
run_context=run_context,
job_description=job_description,
Expand Down Expand Up @@ -417,7 +445,11 @@ async def run_sweep_async(

async def create_program_async(
self,
program: cirq.AbstractCircuit,
program: (
cirq.AbstractCircuit
| Sequence[cirq.AbstractCircuit]
| Mapping[str, cirq.AbstractCircuit]
),
program_id: str | None = None,
description: str | None = None,
labels: dict[str, str] | None = None,
Expand All @@ -443,12 +475,13 @@ async def create_program_async(
if not program_id:
program_id = _make_random_id('prog-')

if isinstance(program, cirq.AbstractCircuit):
code = self.context._serialize_program(program)
else:
code = self.context._serialize_multi_program(program)

new_program_id, new_program = await self.context.client.create_program_async(
self.project_id,
program_id,
code=self.context._serialize_program(program),
description=description,
labels=labels,
self.project_id, program_id, code=code, description=description, labels=labels
)

return engine_program.EngineProgram(
Expand Down
9 changes: 8 additions & 1 deletion cirq-google/cirq_google/engine/engine_processor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from __future__ import annotations

import datetime
from collections.abc import Mapping, Sequence
from typing import Any, TYPE_CHECKING

from cirq import _compat
Expand Down Expand Up @@ -155,7 +156,11 @@ def get_sampler(

async def run_sweep_async(
self,
program: cirq.AbstractCircuit,
program: (
cirq.AbstractCircuit
| Sequence[cirq.AbstractCircuit]
| Mapping[str, cirq.AbstractCircuit]
),
*,
device_config_name: str,
run_name: str = "",
Expand All @@ -177,6 +182,8 @@ async def run_sweep_async(
Args:
program: The Circuit to execute. If a circuit is
provided, a moment by moment schedule will be used.
A list or mapping of circuits can also be provided. If so,
it will be executed as a KeyedCircuit.
run_name: A unique identifier representing an automation run for the
processor. An Automation Run contains a collection of device
configurations for the processor.
Expand Down
40 changes: 40 additions & 0 deletions cirq-google/cirq_google/engine/engine_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1151,3 +1151,43 @@ def test_list_processor_configs_from_snapshot(list_processor_configs_async):
('test_config_1', processor_id, '', snapshot.id),
('test_config_2', processor_id, '', snapshot.id),
]


def test_engine_context_serialize_multi_program_errors():
context = EngineContext()

# Mapping with non-circuits
with pytest.raises(TypeError, match='Unrecognized program type'):
context._serialize_multi_program({'a': 'not a circuit'})

# Sequence with non-circuits
with pytest.raises(TypeError, match='Unrecognized program type'):
context._serialize_multi_program(['not a circuit'])

# Invalid type
with pytest.raises(TypeError, match='Unrecognized program type'):
context._serialize_multi_program(123)

# invalid proto version
# Note: ProtoVersion.V1 is blocked in __init__, but UNDEFINED is not.
context.proto_version = cg.engine.engine.ProtoVersion.UNDEFINED
with pytest.raises(ValueError, match='invalid program proto version'):
context._serialize_multi_program([_CIRCUIT])


def test_engine_context_serialize_multi_program_success():
context = EngineContext()
# success
# This should not raise
context._serialize_multi_program([_CIRCUIT])
context._serialize_multi_program({'a': _CIRCUIT})


@mock.patch('cirq_google.engine.engine_client.EngineClient', autospec=True)
def test_engine_create_program_multi(client_mock):
client_mock().create_program_async.return_value = ('prog', quantum.QuantumProgram())
engine = cg.Engine(project_id='proj')

# program is not AbstractCircuit (it's a list)
engine.create_program([_CIRCUIT], 'prog')
client_mock().create_program_async.assert_called_once()
61 changes: 57 additions & 4 deletions cirq-google/cirq_google/engine/processor_sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

from __future__ import annotations

from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import cast, TYPE_CHECKING

import duet
Expand All @@ -36,6 +36,7 @@ def __init__(
snapshot_id: str = "",
device_config_name: str = "",
max_concurrent_jobs: int = 100,
jobs_per_batch: int = 1,
):
"""Inits ProcessorSampler.

Expand All @@ -56,6 +57,9 @@ def __init__(
concurrently to the Engine. This client-side throttle can be
used to proactively reduce load to the backends and avoid quota
violations when pipelining circuit executions.
jobs_per_batch: If set to greater than 1, this will batch multiple
circuits within the same API call when calling run_batch() or
run_batch_async() up to a maximum of `jobs_per_batch`.

Raises:
ValueError: If only one of `run_name` and `device_config_name` are specified.
Expand All @@ -68,9 +72,17 @@ def __init__(
self._snapshot_id = snapshot_id
self._device_config_name = device_config_name
self._concurrent_job_limiter = duet.Limiter(max_concurrent_jobs)
self._jobs_per_batch = jobs_per_batch

async def run_sweep_async(
self, program: cirq.AbstractCircuit, params: cirq.Sweepable, repetitions: int = 1
self,
program: (
cirq.AbstractCircuit
| Sequence[cirq.AbstractCircuit]
| Mapping[str, cirq.AbstractCircuit]

@maffoo maffoo Apr 16, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

In our internal code when we do the equivalent of passing a Sequence or Mapping here, we iterate over those in the "inner loop", e.g. the execution order would be something like:

results = [
    self.run(circuit, params)
    for params in sweep
    for circuit in circuits
]

where the results would have to be flattened into a single list to conform to the specified return type.

On the other hand, for run_batch as defined in the Sampler base class, we expect to return a nested list where the loop over circuits is the outer loop, e.g.

results = [
    [self.run(circuit, params) for params in sweep]
    for circuit, sweep in zip(circuits, sweeps)
]

I think we need to define what the intended execution order is, in particular for the first case since that is not supported at all on the Sampler base class. If for the first case we align with internal code where the loop over circuits is the "inner" loop, then we would either need a way to communicate the ordering to the server since the run_batch case instead runs over servers on the outer loop. Or if we don't actually care about the order things execute on hardware then we might need to "reshape" the output from quantum engine so that we return results in the appropriate order for both multi-circuit case and run_batch case.

@dstrain115 dstrain115 Apr 20, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

The current behavior is that the engine splits all the calls into separate requests, so I don't think it is currently possible to do the first proposal (circuits in the inner loop) externally. I believe that, even after this change, the internal version will still execute in circuit order, but this mode of batching at least opens the door for the other ordering.

What if we put a disclaimer in run_batch saying that the execution order is not guaranteed if you have jobs_per_batch > 1? That way, if you depend on the circuits executing in a certain order, you can disable proto-level batching. If you just want them to execute as fast as possible, then you knowingly forgo the strict ordering of the circuits.

),
params: cirq.Sweepable | Sequence[cirq.Sweepable],
repetitions: int = 1,
) -> Sequence[cg.EngineResult]:
async with self._concurrent_job_limiter:
job = await self._processor.run_sweep_async(
Expand All @@ -88,13 +100,54 @@ async def run_sweep_async(

async def run_batch_async(
self,
programs: Sequence[cirq.AbstractCircuit],
programs: Sequence[cirq.AbstractCircuit] | Mapping[str, cirq.AbstractCircuit],
params_list: Sequence[cirq.Sweepable] | None = None,
repetitions: int | Sequence[int] = 1,
) -> Sequence[Sequence[cg.EngineResult]]:
if self._jobs_per_batch > 1:
# Treat programs as a sequence for iteration, but keep keys if it's a mapping
prog_keys = list(programs.keys()) if isinstance(programs, Mapping) else []
prog_values = (
list(programs.values()) if isinstance(programs, Mapping) else list(programs)
)

params_list, repetitions = self._normalize_batch_args(
prog_values, params_list, repetitions
)
# Batch programs that have the same number of repetitions.
program_batches = []
params_list_batches = []
repetition_batches = []

i = 0
while i < len(prog_values):
batch_reps = repetitions[i]
batch_programs = {prog_keys[i]: prog_values[i]} if prog_keys else [prog_values[i]]
batch_params = [params_list[i]]
i += 1
while (
i < len(prog_values)
and len(batch_programs) < self._jobs_per_batch
and repetitions[i] == batch_reps
):
if isinstance(batch_programs, dict):
batch_programs[prog_keys[i]] = prog_values[i]
else:
batch_programs.append(prog_values[i])
batch_params.append(params_list[i])
i += 1
program_batches.append(batch_programs)
params_list_batches.append(batch_params)
repetition_batches.append(batch_reps)

return await duet.pstarmap_async(
self.run_sweep_async, zip(program_batches, params_list_batches, repetition_batches)
)

prog_values = list(programs.values()) if isinstance(programs, Mapping) else list(programs)
return cast(
Sequence[Sequence['cg.EngineResult']],
await super().run_batch_async(programs, params_list, repetitions),
await super().run_batch_async(prog_values, params_list, repetitions),
)

run_batch = duet.sync(run_batch_async)
Expand Down
Loading
Loading