Skip to content

Modify run_batch_async #6387

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
merged 16 commits into from
Dec 20, 2023
Merged

Modify run_batch_async #6387

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
eb10318
Created using Colaboratory
senecameeks Nov 27, 2023
21f6d1c
Merge branch 'quantumlib:master' into master
senecameeks Nov 27, 2023
f4e0d88
call run_sweep_async in parallel
senecameeks Dec 15, 2023
ba7de29
call circuits asynchronously
senecameeks Dec 18, 2023
d6c4237
Revert "Created using Colaboratory"
senecameeks Dec 18, 2023
5f8a3ca
Revert "call run_sweep_async in parallel"
senecameeks Dec 18, 2023
ef9b324
revert colab
senecameeks Dec 18, 2023
7173392
lint
senecameeks Dec 18, 2023
3c63bed
use pmap
senecameeks Dec 18, 2023
5d2e52d
add test
senecameeks Dec 19, 2023
2f055f7
lint
senecameeks Dec 19, 2023
cee3241
lint
senecameeks Dec 19, 2023
26fbebd
Merge branch 'main' into parallelize_run_batch
senecameeks Dec 19, 2023
78025e2
nits
senecameeks Dec 20, 2023
7b7dac9
Merge branch 'parallelize_run_batch' of github.com:senecameeks/Cirq i…
senecameeks Dec 20, 2023
db41d92
Merge branch 'main' into parallelize_run_batch
senecameeks Dec 20, 2023
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: 3 additions & 4 deletions cirq-core/cirq/work/sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -294,10 +294,9 @@ async def run_batch_async(
See docs for `cirq.Sampler.run_batch`.
"""
params_list, repetitions = self._normalize_batch_args(programs, params_list, repetitions)
return [
await self.run_sweep_async(circuit, params=params, repetitions=repetitions)
for circuit, params, repetitions in zip(programs, params_list, repetitions)
]
Comment on lines -297 to -300
Copy link
Contributor

Choose a reason for hiding this comment

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

You probably want to use a variant of duet.pmap ("parallel map") to make multiple calls in parallel. In this case, the zip produces tuples of args and we want to call the underlying function like self.run_sweep_async(*args), so we need a "starmap", something like:

return await duet.pstarmap_async(self.run_sweep_async, zip(programs, params_list, repetitions))

Copy link
Contributor

Choose a reason for hiding this comment

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

For reference, here's the definition of duet.pstarmap_async: https://github.com/google/duet/blob/main/duet/api.py#L171

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Thank you for the code pointer!

return await duet.pstarmap_async(
self.run_sweep_async, zip(programs, params_list, repetitions)
)

def _normalize_batch_args(
self,
Expand Down
25 changes: 25 additions & 0 deletions cirq-core/cirq/work/sampler_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,31 @@ def test_sampler_run_batch():
assert np.array_equal(result.measurements['m'], np.array([[0], [0]], dtype='uint8'))


@duet.sync
async def test_run_batch_async_calls_run_sweep_asynchronously():
"""Test run_batch_async calls run_sweep_async without waiting."""
finished = []
a = cirq.LineQubit(0)
circuit1 = cirq.Circuit(cirq.X(a) ** sympy.Symbol('t'), cirq.measure(a, key='m'))
circuit2 = cirq.Circuit(cirq.Y(a) ** sympy.Symbol('t'), cirq.measure(a, key='m'))
params1 = cirq.Points('t', [0.3, 0.7])
params2 = cirq.Points('t', [0.4, 0.6])
params_list = [params1, params2]

class AsyncSampler(cirq.Sampler):
async def run_sweep_async(self, program, params, repetitions: int = 1):
if params == params1:
await duet.sleep(0.001)

finished.append(params)

await AsyncSampler().run_batch_async(
[circuit1, circuit2], params_list=params_list, repetitions=[1, 2]
)

assert finished == list(reversed(params_list))


def test_sampler_run_batch_default_params_and_repetitions():
sampler = cirq.ZerosSampler()
a = cirq.LineQubit(0)
Expand Down