Add max_concurrent_jobs to Engine.get_sampler(). (quantumlib#7442)

Adds `max_concurrent_jobs` to the utility function
[Engine.get_sampler](https://github.com/dyates/Cirq/blob/b8ca4844ef80102aa83ace1093552f25c02b101a/cirq-google/cirq_google/engine/engine.py#L580)
to allow setting the [same field in the
`ProcessSampler`]((https://github.com/quantumlib/Cirq/blob/4c7e393eb9fcefcba0c3c509e32cc354197cf014/cirq-google/cirq_google/engine/processor_sampler.py#L37)).
This field throttles the number concurrent jobs that are sent to the
backend (via `run_batch`).

---------

Co-authored-by: Pavol Juhas <juhas@google.com>

Author: 2 people

cirq-google/cirq_google/engine/engine.py

@@ -583,6 +583,7 @@ def get_sampler(
         run_name: str = "",
         device_config_name: str = "",
         snapshot_id: str = "",
+        max_concurrent_jobs: int = 10,
     ) -> cirq_google.ProcessorSampler:
         """Returns a sampler backed by the engine.
 
@@ -596,6 +597,10 @@ def get_sampler(
                 available qubits, couplers, and supported gates in the processor.
             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
+                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.
 
         Returns:
             A `cirq.Sampler` instance (specifically a `engine_sampler.ProcessorSampler`
@@ -612,7 +617,10 @@ def get_sampler(
                 'you need to specify a list.'
             )
         return self.get_processor(processor_id).get_sampler(
-            run_name=run_name, device_config_name=device_config_name, snapshot_id=snapshot_id
+            run_name=run_name,
+            device_config_name=device_config_name,
+            snapshot_id=snapshot_id,
+            max_concurrent_jobs=max_concurrent_jobs,
         )
 
 

cirq-google/cirq_google/engine/engine_test.py

@@ -800,6 +800,14 @@ def test_get_processor():
     assert cg.Engine(project_id='proj').get_processor('xmonsim').processor_id == 'xmonsim'
 
 
+def test_get_sampler_initializes_max_concurrent_jobs():
+    max_concurrent_jobs = 5
+    engine = cg.Engine(project_id='proj')
+    sampler = engine.get_sampler(processor_id='tmp', max_concurrent_jobs=max_concurrent_jobs)
+
+    assert sampler.max_concurrent_jobs == max_concurrent_jobs
+
+
 @mock.patch('cirq_google.engine.engine_client.EngineClient', autospec=True)
 def test_sampler_with_unary_rpcs(client):
     setup_run_circuit_with_result_(client, _RESULTS)

cirq-google/cirq_google/engine/processor_sampler.py

@@ -113,3 +113,8 @@ def snapshot_id(self) -> str:
     @property
     def device_config_name(self) -> str:
         return self._device_config_name
+
+    @property
+    def max_concurrent_jobs(self) -> int:
+        assert self._concurrent_job_limiter.capacity is not None
+        return self._concurrent_job_limiter.capacity