async def test_async_failure_in_main_task(self):
after_await = False
with pytest.raises(Fail):
async with duet.new_scope() as scope:
scope.spawn(mul, 2, 3)
await duet.failed_future(Fail())
after_await = True # This should not run.
assert not after_await
async def test_sync_failure_in_main_task(self):
# pylint: disable=unreachable
after_await = False
with pytest.raises(Fail):
async with duet.new_scope() as scope:
scope.spawn(mul, 2, 3)
raise Fail()
after_await = True # This should not run.
assert not after_await
async def test_failure_in_spawned_task(self, fail_func):
after_fail = False
with pytest.raises(Fail):
async with duet.new_scope() as scope:
for a in range(10):
scope.spawn(mul, a, a)
scope.spawn(fail_func)
after_fail = True # This should still run.
assert after_fail
async def test_run_all(self):
results = {}
async def func(a, b):
results[a, b] = await mul(a, b)
async with duet.new_scope() as scope:
for a in range(10):
for b in range(10):
scope.spawn(func, a, b)
assert results == {(a, b): a * b for a in range(10) for b in range(10)}
async def test_ordering(self):
"""Check that waiting coroutines acquire limiter in order."""
limiter = duet.Limiter(1)
acquired = []
async def func(i):
async with limiter:
acquired.append(i)
await duet.completed_future(None)
async with duet.new_scope() as scope:
for i in range(10):
scope.spawn(func, i)
assert acquired == sorted(acquired)
async def test_multiple_calls_to_future_set_result():
"""This checks a scenario that caused deadlocks in earlier versions."""
async def set_results(*fs):
for f in fs:
await duet.completed_future(None)
f.set_result(None)
async with duet.new_scope() as scope:
f0 = duet.AwaitableFuture()
f1 = duet.AwaitableFuture()
scope.spawn(set_results, f0)
await f0
# Calling f0.set_result again should not mark this main task as ready.
# If it does, then the duet scheduler will try to advance the task and
# will block on getting the result of f1. This prevents the background
# `set_results` task from advancing and actually calling f1.set_result,
# so we would deadlock.
scope.spawn(set_results, f0, f1)
await f1
async def func():
async with duet.new_scope() as scope:
f = duet.AwaitableFuture()
scope.spawn(lambda: f)
f.set_exception(ValueError("oops!"))
await duet.AwaitableFuture()
async def collect_async(
self,
sampler: 'cirq.Sampler',
*,
concurrency: int = 2,
max_total_samples: Optional[int] = None,
) -> None:
"""Asynchronously collects needed samples from a sampler.
Examples:
```
collector = cirq.PauliStringCollector(...)
await sampler.collect_async(collector, concurrency=3)
print(collector.estimated_energy())
```
Args:
sampler: The simulator or service to collect samples from.
concurrency: Desired number of sampling jobs to have in flight at
any given time.
max_total_samples: Optional limit on the maximum number of samples
to collect.
Returns:
The collector's result after all desired samples have been
collected.
"""
results: duet.AsyncCollector[Tuple[CircuitSampleJob, 'cirq.Result']] = duet.AsyncCollector()
job_error = None
running_jobs = 0
queued_jobs: List[CircuitSampleJob] = []
remaining_samples = np.infty if max_total_samples is None else max_total_samples
async def run_job(job):
nonlocal job_error
try:
result = await sampler.run_async(job.circuit, repetitions=job.repetitions)
except Exception as error:
if not job_error:
results.error(error)
job_error = error
else:
if not job_error:
results.add((job, result))
# Keep dispatching and processing work.
async with duet.new_scope() as scope:
while True:
# Fill up the work pool.
while remaining_samples > 0 and running_jobs < concurrency:
if not queued_jobs:
queued_jobs.extend(_flatten_jobs(self.next_job()))
# If no jobs were given, stop asking until something completes.
if not queued_jobs:
break
# Start new sampling job.
new_job = queued_jobs.pop(0)
remaining_samples -= new_job.repetitions
running_jobs += 1
scope.spawn(run_job, new_job)
# If no jobs are running, we're in a steady state. Halt.
if not running_jobs:
break
# Get result from next completed job and call on_job_result.
job, result = await results.__anext__()
running_jobs -= 1
self.on_job_result(job, result)