Add async methods to AbstractEngine and AbstractJob (quantumlib#5555)

Review: @95-martin-orion

cirq/work/sampler.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 """Abstract base class for things sampling quantum circuits."""
 
-import abc
 import collections
 from typing import Dict, FrozenSet, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
+import duet
 import pandas as pd
 
-from cirq import ops, protocols, study
+from cirq import ops, protocols, study, value
 from cirq.work.observable_measurement import (
     measure_observables,
     RepetitionsStoppingCriteria,
@@ -31,7 +31,7 @@
     import cirq
 
 
-class Sampler(metaclass=abc.ABCMeta):
+class Sampler(metaclass=value.ABCMetaImplementAnyOneOf):
     """Something capable of sampling quantum circuits. Simulator or hardware."""
 
     def run(
@@ -177,7 +177,19 @@ def sample(
 
         return pd.concat(results)
 
-    @abc.abstractmethod
+    def _run_sweep_impl(
+        self, program: 'cirq.AbstractCircuit', params: 'cirq.Sweepable', repetitions: int = 1
+    ) -> Sequence['cirq.Result']:
+        """Implements run_sweep using run_sweep_async"""
+        return duet.run(self.run_sweep_async, program, params, repetitions)
+
+    async def _run_sweep_async_impl(
+        self, program: 'cirq.AbstractCircuit', params: 'cirq.Sweepable', repetitions: int = 1
+    ) -> Sequence['cirq.Result']:
+        """Implements run_sweep_async using run_sweep"""
+        return self.run_sweep(program, params=params, repetitions=repetitions)
+
+    @value.alternative(requires='run_sweep_async', implementation=_run_sweep_impl)
     def run_sweep(
         self, program: 'cirq.AbstractCircuit', params: 'cirq.Sweepable', repetitions: int = 1
     ) -> Sequence['cirq.Result']:
@@ -200,6 +212,7 @@ def run_sweep(
             Result list for this run; one for each possible parameter resolver.
         """
 
+    @value.alternative(requires='run_sweep', implementation=_run_sweep_async_impl)
     async def run_sweep_async(
         self, program: 'cirq.AbstractCircuit', params: 'cirq.Sweepable', repetitions: int = 1
     ) -> Sequence['cirq.Result']:
@@ -217,13 +230,12 @@ async def run_sweep_async(
         Returns:
             Result list for this run; one for each possible parameter resolver.
         """
-        return self.run_sweep(program, params=params, repetitions=repetitions)
 
     def run_batch(
         self,
         programs: Sequence['cirq.AbstractCircuit'],
-        params_list: Optional[List['cirq.Sweepable']] = None,
-        repetitions: Union[int, List[int]] = 1,
+        params_list: Optional[Sequence['cirq.Sweepable']] = None,
+        repetitions: Union[int, Sequence[int]] = 1,
     ) -> Sequence[Sequence['cirq.Result']]:
         """Runs the supplied circuits.
 
@@ -263,6 +275,34 @@ def run_batch(
             ValueError: If length of `programs` is not equal to the length
                 of `params_list` or the length of `repetitions`.
         """
+        params_list, repetitions = self._normalize_batch_args(programs, params_list, repetitions)
+        return [
+            self.run_sweep(circuit, params=params, repetitions=repetitions)
+            for circuit, params, repetitions in zip(programs, params_list, repetitions)
+        ]
+
+    async def run_batch_async(
+        self,
+        programs: Sequence['cirq.AbstractCircuit'],
+        params_list: Optional[Sequence['cirq.Sweepable']] = None,
+        repetitions: Union[int, Sequence[int]] = 1,
+    ) -> Sequence[Sequence['cirq.Result']]:
+        """Runs the supplied circuits.
+
+        This is an asynchronous version of `run_batch`; see full docs there.
+        """
+        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)
+        ]
+
+    def _normalize_batch_args(
+        self,
+        programs: Sequence['cirq.AbstractCircuit'],
+        params_list: Optional[Sequence['cirq.Sweepable']] = None,
+        repetitions: Union[int, Sequence[int]] = 1,
+    ) -> Tuple[Sequence['cirq.Sweepable'], Sequence[int]]:
         if params_list is None:
             params_list = [None] * len(programs)
         if len(programs) != len(params_list):
@@ -277,10 +317,7 @@ def run_batch(
                 'len(programs) and len(repetitions) must match. '
                 f'Got {len(programs)} and {len(repetitions)}.'
             )
-        return [
-            self.run_sweep(circuit, params=params, repetitions=repetitions)
-            for circuit, params, repetitions in zip(programs, params_list, repetitions)
-        ]
+        return params_list, repetitions
 
     def sample_expectation_values(
         self,

cirq/work/sampler_test.py

@@ -59,6 +59,42 @@ def run_sweep(self, program, params, repetitions: int = 1):
         await FailingSampler().run_sweep_async(cirq.Circuit(), repetitions=1, params=None)
 
 
+def test_run_sweep_impl():
+    """Test run_sweep implemented in terms of run_sweep_async."""
+
+    class AsyncSampler(cirq.Sampler):
+        async def run_sweep_async(self, program, params, repetitions: int = 1):
+            await duet.sleep(0.001)
+            return cirq.Simulator().run_sweep(program, params, repetitions)
+
+    results = AsyncSampler().run_sweep(
+        cirq.Circuit(cirq.measure(cirq.GridQubit(0, 0), key='m')),
+        cirq.Linspace('foo', 0, 1, 10),
+        repetitions=10,
+    )
+    assert len(results) == 10
+    for result in results:
+        np.testing.assert_equal(result.records['m'], np.zeros((10, 1, 1)))
+
+
+@duet.sync
+async def test_run_sweep_async_impl():
+    """Test run_sweep_async implemented in terms of run_sweep."""
+
+    class SyncSampler(cirq.Sampler):
+        def run_sweep(self, program, params, repetitions: int = 1):
+            return cirq.Simulator().run_sweep(program, params, repetitions)
+
+    results = await SyncSampler().run_sweep_async(
+        cirq.Circuit(cirq.measure(cirq.GridQubit(0, 0), key='m')),
+        cirq.Linspace('foo', 0, 1, 10),
+        repetitions=10,
+    )
+    assert len(results) == 10
+    for result in results:
+        np.testing.assert_equal(result.records['m'], np.zeros((10, 1, 1)))
+
+
 def test_sampler_sample_multiple_params():
     a, b = cirq.LineQubit.range(2)
     s = sympy.Symbol('s')

requirements.txt

@@ -3,7 +3,7 @@
 # functools.cached_property was introduced in python 3.8
 backports.cached_property~=1.0.1; python_version < '3.8'
 
-duet~=0.2.6
+duet~=0.2.7
 matplotlib~=3.0
 networkx~=2.4
 numpy~=1.16