add snapshot_id to get_sampler interface on engine object (#7005)

* add snapshot to get_sampler interface on engine object

* backwards compat

Author: senecameeks

cirq-google/cirq_google/engine/engine.py

@@ -575,7 +575,11 @@ def get_processor(self, processor_id: str) -> engine_processor.EngineProcessor:
         return engine_processor.EngineProcessor(self.project_id, processor_id, self.context)
 
     def get_sampler(
-        self, processor_id: Union[str, List[str]], run_name: str = "", device_config_name: str = ""
+        self,
+        processor_id: Union[str, List[str]],
+        run_name: str = "",
+        device_config_name: str = "",
+        snapshot_id: str = "",
     ) -> 'cirq_google.ProcessorSampler':
         """Returns a sampler backed by the engine.
 
@@ -587,6 +591,8 @@ def get_sampler(
             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.
+            snapshot_id: A unique identifier for an immutable snapshot reference. A
+                snapshot contains a collection of device configurations for the processor.
 
         Returns:
             A `cirq.Sampler` instance (specifically a `engine_sampler.ProcessorSampler`
@@ -603,7 +609,7 @@ 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
+            run_name=run_name, device_config_name=device_config_name, snapshot_id=snapshot_id
         )
 
 

cirq-google/cirq_google/engine/engine_test.py

@@ -281,6 +281,24 @@ def test_run_circuit_with_unary_rpcs(client):
     client().get_job_results_async.assert_called_once_with('proj', 'prog', 'job-id')
 
 
+@mock.patch('cirq_google.engine.engine_client.EngineClient', autospec=True)
+def test_engine_get_sampler_with_snapshot_id_passes_to_unary_rpc(client):
+    setup_run_circuit_with_result_(client, _A_RESULT)
+    engine = cg.Engine(
+        project_id='proj',
+        context=EngineContext(service_args={'client_info': 1}, enable_streaming=False),
+    )
+    sampler = engine.get_sampler('mysim', device_config_name="config", snapshot_id="123")
+    _ = sampler.run_sweep(_CIRCUIT, params=[cirq.ParamResolver({'a': 1})])
+
+    kwargs = client().create_job_async.call_args_list[0].kwargs
+
+    # We care about asserting that the snapshot_id is correctly passed.
+    assert kwargs["snapshot_id"] == "123"
+    assert kwargs["run_name"] == ""
+    assert kwargs["device_config_name"] == "config"
+
+
 @mock.patch('cirq_google.engine.engine_client.EngineClient', autospec=True)
 def test_run_circuit_with_stream_rpcs_passes(client):
     setup_run_circuit_with_result_(client, _A_RESULT)