Add device configuration parameters to processor sampler class

cirq-google/cirq_google/engine/processor_sampler.py

@@ -24,19 +24,45 @@
 class ProcessorSampler(cirq.Sampler):
     """A wrapper around AbstractProcessor to implement the cirq.Sampler interface."""
 
-    def __init__(self, *, processor: 'cg.engine.AbstractProcessor'):
-        """Inits ProcessorSampler.
+    def __init__(
+        self,
+        *,
+        processor: 'cg.engine.AbstractProcessor',
+        run_name: str = "",
+        device_config_name: str = "",
+    ):
+        """Inits ProcessorSampler. Either both `run_name` and `device_config_name`
+        must be set, or neither of them must be set. If none of them are set, a
+        default internal device configuration will be used.
 
         Args:
             processor: AbstractProcessor instance to use.
+            run_name: A unique identifier representing an automation run for the
+                specified processor. An Automation Run contains a collection of
+                device configurations for a processor.
+            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.
+
+        Raises:
+            ValueError: If  only one of `run_name` and `device_config_name` are specified.
         """
+        if bool(run_name) ^ bool(device_config_name):
+            raise ValueError('Cannot specify only one of `run_name` and `device_config_name`')
+
         self._processor = processor
+        self._run_name = run_name
+        self._device_config_name = device_config_name
 
     async def run_sweep_async(
         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
+            program=program,
+            params=params,
+            repetitions=repetitions,
+            run_name=self._run_name,
+            device_config_name=self._device_config_name,
         )
         return await job.results_async()
 
@@ -61,7 +87,11 @@ async def run_batch_async(
         if len(set(repetitions)) == 1:
             # All repetitions are the same so batching can be done efficiently
             job = await self._processor.run_batch_async(
-                programs=programs, params_list=params_list, repetitions=repetitions[0]
+                programs=programs,
+                params_list=params_list,
+                repetitions=repetitions[0],
+                run_name=self._run_name,
+                device_config_name=self._device_config_name,
             )
             return await job.batched_results_async()
         # Varying number of repetitions so no speedup

cirq-google/cirq_google/engine/processor_sampler_test.py

@@ -24,15 +24,29 @@
 @pytest.mark.parametrize('circuit', [cirq.Circuit(), cirq.FrozenCircuit()])
 def test_run_circuit(circuit):
     processor = mock.create_autospec(AbstractProcessor)
-    sampler = cg.ProcessorSampler(processor=processor)
+    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
+    )
     params = [cirq.ParamResolver({'a': 1})]
     sampler.run_sweep(circuit, params, 5)
-    processor.run_sweep_async.assert_called_with(params=params, program=circuit, repetitions=5)
+    processor.run_sweep_async.assert_called_with(
+        params=params,
+        program=circuit,
+        repetitions=5,
+        run_name=run_name,
+        device_config_name=device_config_name,
+    )
 
 
 def test_run_batch():
     processor = mock.create_autospec(AbstractProcessor)
-    sampler = cg.ProcessorSampler(processor=processor)
+    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
+    )
     a = cirq.LineQubit(0)
     circuit1 = cirq.Circuit(cirq.X(a))
     circuit2 = cirq.Circuit(cirq.Y(a))
@@ -42,13 +56,21 @@ def test_run_batch():
     params_list = [params1, params2]
     sampler.run_batch(circuits, params_list, 5)
     processor.run_batch_async.assert_called_with(
-        params_list=params_list, programs=circuits, repetitions=5
+        params_list=params_list,
+        programs=circuits,
+        repetitions=5,
+        run_name=run_name,
+        device_config_name=device_config_name,
     )
 
 
 def test_run_batch_identical_repetitions():
     processor = mock.create_autospec(AbstractProcessor)
-    sampler = cg.ProcessorSampler(processor=processor)
+    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
+    )
     a = cirq.LineQubit(0)
     circuit1 = cirq.Circuit(cirq.X(a))
     circuit2 = cirq.Circuit(cirq.Y(a))
@@ -58,7 +80,11 @@ def test_run_batch_identical_repetitions():
     params_list = [params1, params2]
     sampler.run_batch(circuits, params_list, [5, 5])
     processor.run_batch_async.assert_called_with(
-        params_list=params_list, programs=circuits, repetitions=5
+        params_list=params_list,
+        programs=circuits,
+        repetitions=5,
+        run_name=run_name,
+        device_config_name=device_config_name,
     )
 
 
@@ -78,10 +104,14 @@ def test_run_batch_bad_number_of_repetitions():
 
 def test_run_batch_differing_repetitions():
     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
+    )
     job = mock.Mock()
     job.results.return_value = []
     processor.run_sweep.return_value = job
-    sampler = cg.ProcessorSampler(processor=processor)
     a = cirq.LineQubit(0)
     circuit1 = cirq.Circuit(cirq.X(a))
     circuit2 = cirq.Circuit(cirq.Y(a))
@@ -91,7 +121,13 @@ def test_run_batch_differing_repetitions():
     params_list = [params1, params2]
     repetitions = [1, 2]
     sampler.run_batch(circuits, params_list, repetitions)
-    processor.run_sweep_async.assert_called_with(params=params2, program=circuit2, repetitions=2)
+    processor.run_sweep_async.assert_called_with(
+        params=params2,
+        program=circuit2,
+        repetitions=2,
+        run_name=run_name,
+        device_config_name=device_config_name,
+    )
     processor.run_batch_async.assert_not_called()