Add throttle for jobs run with ProcessorSampler (#6786)

* Limit the number of concurrent jobs when running run_batch.

* Add throttle for concurrent jobs run by ProcessorSampler.

* Add tests

* Remove unused imports.

* lint

* format

* moar lint

Author: wcourtney

cirq-core/cirq/work/sampler.py

@@ -291,15 +291,14 @@ async def run_batch_async(
         programs: Sequence['cirq.AbstractCircuit'],
         params_list: Optional[Sequence['cirq.Sweepable']] = None,
         repetitions: Union[int, Sequence[int]] = 1,
-        limiter: duet.Limiter = duet.Limiter(10),
     ) -> Sequence[Sequence['cirq.Result']]:
         """Runs the supplied circuits asynchronously.
 
         See docs for `cirq.Sampler.run_batch`.
         """
         params_list, repetitions = self._normalize_batch_args(programs, params_list, repetitions)
         return await duet.pstarmap_async(
-            self.run_sweep_async, zip(programs, params_list, repetitions, [limiter] * len(programs))
+            self.run_sweep_async, zip(programs, params_list, repetitions)
         )
 
     def _normalize_batch_args(

cirq-google/cirq_google/engine/engine_job.py

@@ -262,14 +262,12 @@ def delete(self) -> None:
         """Deletes the job and result, if any."""
         self.context.client.delete_job(self.project_id, self.program_id, self.job_id)
 
-    async def results_async(
-        self, limiter: duet.Limiter = duet.Limiter(None)
-    ) -> Sequence[EngineResult]:
+    async def results_async(self) -> Sequence[EngineResult]:
         """Returns the job results, blocking until the job is complete."""
         import cirq_google.engine.engine as engine_base
 
         if self._results is None:
-            result_response = await self._await_result_async(limiter)
+            result_response = await self._await_result_async()
             result = result_response.result
             result_type = result.type_url[len(engine_base.TYPE_PREFIX) :]
             if (
@@ -288,9 +286,7 @@ async def results_async(
                 raise ValueError(f'invalid result proto version: {result_type}')
         return self._results
 
-    async def _await_result_async(
-        self, limiter: duet.Limiter = duet.Limiter(None)
-    ) -> quantum.QuantumResult:
+    async def _await_result_async(self) -> quantum.QuantumResult:
         if self._job_result_future is not None:
             response = await self._job_result_future
             if isinstance(response, quantum.QuantumResult):
@@ -303,13 +299,12 @@ async def _await_result_async(
                     'Internal error: The job response type is not recognized.'
                 )  # pragma: no cover
 
-        async with limiter:
-            async with duet.timeout_scope(self.context.timeout):  # type: ignore[arg-type]
-                while True:
-                    job = await self._refresh_job_async()
-                    if job.execution_status.state in TERMINAL_STATES:
-                        break
-                    await duet.sleep(1)
+        async with duet.timeout_scope(self.context.timeout):  # type: ignore[arg-type]
+            while True:
+                job = await self._refresh_job_async()
+                if job.execution_status.state in TERMINAL_STATES:
+                    break
+                await duet.sleep(1)
         _raise_on_failure(job)
         response = await self.context.client.get_job_results_async(
             self.project_id, self.program_id, self.job_id

cirq-google/cirq_google/engine/engine_processor.py

@@ -87,7 +87,11 @@ def engine(self) -> 'engine_base.Engine':
         return engine_base.Engine(self.project_id, context=self.context)
 
     def get_sampler(
-        self, run_name: str = "", device_config_name: str = "", snapshot_id: str = ""
+        self,
+        run_name: str = "",
+        device_config_name: str = "",
+        snapshot_id: str = "",
+        max_concurrent_jobs: int = 10,
     ) -> 'cg.engine.ProcessorSampler':
         """Returns a sampler backed by the engine.
         Args:
@@ -100,6 +104,11 @@ def get_sampler(
             snapshot_id: A unique identifier for an immutable snapshot reference.
                 A snapshot contains a collection of device configurations for the
                 processor.
+            max_concurrent_jobs: The maximum number of jobs to be sent
+                simultaneously 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.
+
         Returns:
             A `cirq.Sampler` instance (specifically a `engine_sampler.ProcessorSampler`
             that will send circuits to the Quantum Computing Service
@@ -127,6 +136,7 @@ def get_sampler(
             run_name=run_name,
             snapshot_id=snapshot_id,
             device_config_name=device_config_name,
+            max_concurrent_jobs=max_concurrent_jobs,
         )
 
     async def run_sweep_async(

cirq-google/cirq_google/engine/processor_sampler.py

@@ -16,7 +16,6 @@
 
 import cirq
 import duet
-from cirq_google.engine.engine_job import EngineJob
 
 if TYPE_CHECKING:
     import cirq_google as cg
@@ -32,6 +31,7 @@ def __init__(
         run_name: str = "",
         snapshot_id: str = "",
         device_config_name: str = "",
+        max_concurrent_jobs: int = 10,
     ):
         """Inits ProcessorSampler.
 
@@ -48,6 +48,10 @@ def __init__(
             device_config_name: An identifier used to select the processor configuration
                 utilized to run the job. A configuration identifies the set of
                 available qubits, couplers, and supported gates in the processor.
+            max_concurrent_jobs: The maximum number of jobs to be sent
+                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.
 
         Raises:
             ValueError: If  only one of `run_name` and `device_config_name` are specified.
@@ -59,28 +63,22 @@ def __init__(
         self._run_name = run_name
         self._snapshot_id = snapshot_id
         self._device_config_name = device_config_name
-        self._result_limiter = duet.Limiter(None)
+        self._concurrent_job_limiter = duet.Limiter(max_concurrent_jobs)
 
     async def run_sweep_async(
-        self,
-        program: 'cirq.AbstractCircuit',
-        params: cirq.Sweepable,
-        repetitions: int = 1,
-        limiter: duet.Limiter = duet.Limiter(None),
+        self, program: 'cirq.AbstractCircuit', params: cirq.Sweepable, repetitions: int = 1
     ) -> Sequence['cg.EngineResult']:
-        job = await self._processor.run_sweep_async(
-            program=program,
-            params=params,
-            repetitions=repetitions,
-            run_name=self._run_name,
-            snapshot_id=self._snapshot_id,
-            device_config_name=self._device_config_name,
-        )
-
-        if isinstance(job, EngineJob):
-            return await job.results_async(limiter)
-
-        return await job.results_async()
+        async with self._concurrent_job_limiter:
+            job = await self._processor.run_sweep_async(
+                program=program,
+                params=params,
+                repetitions=repetitions,
+                run_name=self._run_name,
+                snapshot_id=self._snapshot_id,
+                device_config_name=self._device_config_name,
+            )
+
+            return await job.results_async()
 
     run_sweep = duet.sync(run_sweep_async)
 
@@ -89,12 +87,10 @@ async def run_batch_async(
         programs: Sequence[cirq.AbstractCircuit],
         params_list: Optional[Sequence[cirq.Sweepable]] = None,
         repetitions: Union[int, Sequence[int]] = 1,
-        limiter: duet.Limiter = duet.Limiter(10),
     ) -> Sequence[Sequence['cg.EngineResult']]:
-        self._result_limiter = limiter
         return cast(
             Sequence[Sequence['cg.EngineResult']],
-            await super().run_batch_async(programs, params_list, repetitions, self._result_limiter),
+            await super().run_batch_async(programs, params_list, repetitions),
         )
 
     run_batch = duet.sync(run_batch_async)
@@ -114,7 +110,3 @@ def snapshot_id(self) -> str:
     @property
     def device_config_name(self) -> str:
         return self._device_config_name
-
-    @property
-    def result_limiter(self) -> duet.Limiter:
-        return self._result_limiter

cirq-google/cirq_google/engine/processor_sampler_test.py

@@ -15,6 +15,7 @@
 from unittest import mock
 
 import pytest
+import duet
 
 import cirq
 import cirq_google as cg
@@ -170,29 +171,60 @@ def test_run_batch_differing_repetitions():
     )
 
 
-def test_run_batch_receives_results_using_limiter():
+@duet.sync
+async def test_sampler_with_full_job_queue_blocks():
     processor = mock.create_autospec(AbstractProcessor)
-    run_name = "RUN_NAME"
-    device_config_name = "DEVICE_CONFIG_NAME"
-    sampler = cg.ProcessorSampler(
-        processor=processor, run_name=run_name, device_config_name=device_config_name
-    )
+    sampler = cg.ProcessorSampler(processor=processor, max_concurrent_jobs=2)
 
-    job = mock.AsyncMock(EngineJob)
+    async def wait_forever(**kwargs):
+        await duet.AwaitableFuture[None]()
+
+    processor.run_sweep_async.side_effect = wait_forever
+
+    a = cirq.LineQubit(0)
+    circuit = cirq.Circuit(cirq.X(a))
 
+    with pytest.raises(TimeoutError):
+        async with duet.timeout_scope(0.01):
+            await sampler.run_batch_async([circuit] * 3)
+
+    assert processor.run_sweep_async.call_count == 2
+
+
+@duet.sync
+async def test_sampler_with_job_queue_availability_runs_all():
+    processor = mock.create_autospec(AbstractProcessor)
+    sampler = cg.ProcessorSampler(processor=processor, max_concurrent_jobs=3)
+
+    async def wait_forever(**kwargs):
+        await duet.AwaitableFuture[None]()
+
+    processor.run_sweep_async.side_effect = wait_forever
+
+    a = cirq.LineQubit(0)
+    circuit = cirq.Circuit(cirq.X(a))
+
+    with pytest.raises(TimeoutError):
+        async with duet.timeout_scope(0.01):
+            await sampler.run_batch_async([circuit] * 3)
+
+    assert processor.run_sweep_async.call_count == 3
+
+
+@duet.sync
+async def test_sampler_with_full_job_queue_unblocks_when_available():
+    processor = mock.create_autospec(AbstractProcessor)
+    sampler = cg.ProcessorSampler(processor=processor, max_concurrent_jobs=2)
+
+    job = mock.AsyncMock(EngineJob)
     processor.run_sweep_async.return_value = job
+
     a = cirq.LineQubit(0)
-    circuit1 = cirq.Circuit(cirq.X(a))
-    circuit2 = cirq.Circuit(cirq.Y(a))
-    params1 = [cirq.ParamResolver({'t': 1})]
-    params2 = [cirq.ParamResolver({'t': 2})]
-    circuits = [circuit1, circuit2]
-    params_list = [params1, params2]
-    repetitions = [1, 2]
+    circuit = cirq.Circuit(cirq.X(a))
 
-    sampler.run_batch(circuits, params_list, repetitions)
+    await sampler.run_batch_async([circuit] * 3)
 
-    job.results_async.assert_called_with(sampler.result_limiter)
+    assert processor.run_sweep_async.call_count == 3
 
 
 def test_processor_sampler_processor_property():