Allow external users to use snapshot_id (#6739)

* add snapshot_id * test * update * fixes * add changes to generated client * coverage * review comments * clean u tests * format * review comments * review comment * nit
quantumlib · Oct 14, 2024 · 54f9e8c · 54f9e8c
1 parent 5718365
commit 54f9e8c
Show file tree

Hide file tree

Showing 12 changed files with 390 additions and 94 deletions.
diff --git a/cirq-google/cirq_google/cloud/quantum_v1alpha1/types/quantum.py b/cirq-google/cirq_google/cloud/quantum_v1alpha1/types/quantum.py
@@ -156,15 +156,48 @@ class QuantumJob(proto.Message):
 
 class DeviceConfigSelector(proto.Message):
     r"""-
+    This message has `oneof`_ fields (mutually exclusive fields).
+    For each oneof, at most one member field can be set at the same time.
+    Setting any member of the oneof automatically clears all other
+    members.
+    .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields
     Attributes:
         run_name (str):
             -
+            This field is a member of `oneof`_ ``top_level_identifier``.
+        snapshot_id (str):
+            -
+            This field is a member of `oneof`_ ``top_level_identifier``.
         config_alias (str):
             -
     """
 
-    run_name = proto.Field(proto.STRING, number=1)
-    config_alias = proto.Field(proto.STRING, number=2)
+    run_name: str = proto.Field(proto.STRING, number=1, oneof='top_level_identifier')
+    snapshot_id: str = proto.Field(proto.STRING, number=3, oneof='top_level_identifier')
+    config_alias: str = proto.Field(proto.STRING, number=2)
+
+
+class DeviceConfigKey(proto.Message):
+    r"""-
+    This message has `oneof`_ fields (mutually exclusive fields).
+    For each oneof, at most one member field can be set at the same time.
+    Setting any member of the oneof automatically clears all other
+    members.
+    .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields
+    Attributes:
+        run (str):
+            -
+            This field is a member of `oneof`_ ``top_level_identifier``.
+        snapshot_id (str):
+            -
+            This field is a member of `oneof`_ ``top_level_identifier``.
+        config_alias (str):
+            -
+    """
+
+    run: str = proto.Field(proto.STRING, number=1, oneof='top_level_identifier')
+    snapshot_id: str = proto.Field(proto.STRING, number=3, oneof='top_level_identifier')
+    config_alias: str = proto.Field(proto.STRING, number=2)
 
 
 class SchedulingConfig(proto.Message):
@@ -611,18 +644,4 @@ class QuantumReservation(proto.Message):
     whitelisted_users = proto.RepeatedField(proto.STRING, number=5)
 
 
-class DeviceConfigKey(proto.Message):
-    r"""-
-
-    Attributes:
-        run (str):
-            -
-        config_alias (str):
-            -
-    """
-
-    run = proto.Field(proto.STRING, number=1)
-    config_alias = proto.Field(proto.STRING, number=2)
-
-
 __all__ = tuple(sorted(__protobuf__.manifest))
diff --git a/cirq-google/cirq_google/engine/abstract_processor.py b/cirq-google/cirq_google/engine/abstract_processor.py
@@ -57,8 +57,9 @@ async def run_async(
         self,
         program: cirq.Circuit,
         *,
-        run_name: str,
         device_config_name: str,
+        run_name: str = "",
+        snapshot_id: str = "",
         program_id: Optional[str] = None,
         job_id: Optional[str] = None,
         param_resolver: Optional[cirq.ParamResolver] = None,
@@ -75,7 +76,12 @@ async def run_async(
                 provided, a moment by moment schedule will be used.
             run_name: A unique identifier representing an automation run for the
                 processor. An Automation Run contains a collection of device
-                configurations for the processor.
+                configurations for the processor. `snapshot_id` and `run_name`
+                should not both be set. Choose one.
+            snapshot_id: A unique identifier for an immutable snapshot reference.
+                A snapshot contains a collection of device configurations for the
+                processor. `snapshot_id` and `run_name` should not both be set.
+                Choose one.
             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.
@@ -109,6 +115,7 @@ async def run_async(
             job_description=job_description,
             job_labels=job_labels,
             run_name=run_name,
+            snapshot_id=snapshot_id,
             device_config_name=device_config_name,
         )
         return job.results()[0]
@@ -120,8 +127,9 @@ async def run_sweep_async(
         self,
         program: cirq.AbstractCircuit,
         *,
-        run_name: str,
         device_config_name: str,
+        run_name: str = "",
+        snapshot_id: str = "",
         program_id: Optional[str] = None,
         job_id: Optional[str] = None,
         params: cirq.Sweepable = None,
@@ -144,6 +152,9 @@ async def run_sweep_async(
             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. Both `snapshot_id` and `run_name` should not be set.
             program_id: A user-provided identifier for the program. This must
                 be unique within the Google Cloud project being used. If this
                 parameter is not provided, a random id of the format

diff --git a/cirq-google/cirq_google/engine/engine.py b/cirq-google/cirq_google/engine/engine.py
@@ -223,6 +223,7 @@ def run(
         job_labels: Optional[Dict[str, str]] = None,
         *,
         run_name: str = "",
+        snapshot_id: str = "",
         device_config_name: str = "",
     ) -> cirq.Result:
         """Runs the supplied Circuit via Quantum Engine.
@@ -250,6 +251,9 @@ def run(
                 specified processor. An Automation Run contains a collection of
                 device configurations for a processor. If specified, `processor_id`
                 is required to be set.
+            snapshot_id: A unique identifier for an immutable snapshot reference.
+                A snapshot contains a collection of device configurations for the
+                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.
@@ -277,6 +281,7 @@ def run(
                 job_description=job_description,
                 job_labels=job_labels,
                 run_name=run_name,
+                snapshot_id=snapshot_id,
                 device_config_name=device_config_name,
             )
         )[0]
@@ -295,6 +300,7 @@ async def run_sweep_async(
         job_labels: Optional[Dict[str, str]] = None,
         *,
         run_name: str = "",
+        snapshot_id: str = "",
         device_config_name: str = "",
     ) -> engine_job.EngineJob:
         """Runs the supplied Circuit via Quantum Engine.
@@ -325,6 +331,9 @@ async def run_sweep_async(
                 specified processor. An Automation Run contains a collection of
                 device configurations for a processor. If specified, `processor_id`
                 is required to be set.
+            snapshot_id: A unique identifier for an immutable snapshot reference.
+                A snapshot contains a collection of device configurations for the
+                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.
@@ -360,6 +369,7 @@ async def run_sweep_async(
                 job_labels=job_labels,
                 processor_id=processor_id,
                 run_name=run_name,
+                snapshot_id=snapshot_id,
                 device_config_name=device_config_name,
             )
             return engine_job.EngineJob(
@@ -381,6 +391,7 @@ async def run_sweep_async(
             description=job_description,
             labels=job_labels,
             run_name=run_name,
+            snapshot_id=snapshot_id,
             device_config_name=device_config_name,
         )
 

diff --git a/cirq-google/cirq_google/engine/engine_client.py b/cirq-google/cirq_google/engine/engine_client.py
@@ -388,6 +388,7 @@ async def create_job_async(
         labels: Optional[Dict[str, str]] = None,
         *,
         run_name: str = "",
+        snapshot_id: str = "",
         device_config_name: str = "",
     ) -> Tuple[str, quantum.QuantumJob]:
         """Creates and runs a job on Quantum Engine.
@@ -409,6 +410,9 @@ async def create_job_async(
                 specified processor. An Automation Run contains a collection of
                 device configurations for a processor. If specified, `processor_id`
                 is required to be set.
+            snapshot_id: A unique identifier for an immutable snapshot reference.
+                A snapshot contains a collection of device configurations for the
+                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.
@@ -421,25 +425,37 @@ async def create_job_async(
             ValueError: If  only one of `run_name` and `device_config_name` are specified.
             ValueError: If either `run_name` and `device_config_name` are set but
                 `processor_id` is empty.
+            ValueError: If both `run_name` and `snapshot_id` are specified.
         """
         # Check program to run and program parameters.
         if priority and not 0 <= priority < 1000:
             raise ValueError('priority must be between 0 and 1000')
         if not processor_id:
             raise ValueError('Must specify a processor id when creating a job.')
-        if bool(run_name) ^ bool(device_config_name):
-            raise ValueError('Cannot specify only one of `run_name` and `device_config_name`')
+        if run_name and snapshot_id:
+            raise ValueError('Cannot specify both `run_name` and `snapshot_id`')
+        if (bool(run_name) or bool(snapshot_id)) ^ bool(device_config_name):
+            raise ValueError(
+                'Cannot specify only one of top level identifier (e.g `run_name`, `snapshot_id`)'
+                ' and `device_config_name`'
+            )
 
         # Create job.
+        if snapshot_id:
+            selector = quantum.DeviceConfigSelector(
+                snapshot_id=snapshot_id or None, config_alias=device_config_name
+            )
+        else:
+            selector = quantum.DeviceConfigSelector(
+                run_name=run_name or None, config_alias=device_config_name
+            )
         job_name = _job_name_from_ids(project_id, program_id, job_id) if job_id else ''
         job = quantum.QuantumJob(
             name=job_name,
             scheduling_config=quantum.SchedulingConfig(
                 processor_selector=quantum.SchedulingConfig.ProcessorSelector(
                     processor=_processor_name_from_ids(project_id, processor_id),
-                    device_config_selector=quantum.DeviceConfigSelector(
-                        run_name=run_name, config_alias=device_config_name
-                    ),
+                    device_config_selector=selector,
                 )
             ),
             run_context=run_context,
@@ -737,6 +753,7 @@ def run_job_over_stream(
         job_labels: Optional[Dict[str, str]] = None,
         processor_id: str = "",
         run_name: str = "",
+        snapshot_id: str = "",
         device_config_name: str = "",
     ) -> duet.AwaitableFuture[Union[quantum.QuantumResult, quantum.QuantumJob]]:
         """Runs a job with the given program and job information over a stream.
@@ -761,6 +778,9 @@ def run_job_over_stream(
                 specified processor. An Automation Run contains a collection of
                 device configurations for a processor. If specified, `processor_id`
                 is required to be set.
+            snapshot_id: A unique identifier for an immutable snapshot reference.
+                A snapshot contains a collection of device configurations for the
+                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.
@@ -773,14 +793,19 @@ def run_job_over_stream(
             ValueError: If the priority is not between 0 and 1000.
             ValueError: If `processor_id` is not set.
             ValueError: If only one of `run_name` and `device_config_name` are specified.
+            ValueError: If both `run_name` and `snapshot_id` are specified.
         """
         # Check program to run and program parameters.
         if priority and not 0 <= priority < 1000:
             raise ValueError('priority must be between 0 and 1000')
         if not processor_id:
             raise ValueError('Must specify a processor id when creating a job.')
-        if bool(run_name) ^ bool(device_config_name):
-            raise ValueError('Cannot specify only one of `run_name` and `device_config_name`')
+        if run_name and snapshot_id:
+            raise ValueError('Cannot specify both `run_name` and `snapshot_id`')
+        if (bool(run_name) or bool(snapshot_id)) ^ bool(device_config_name):
+            raise ValueError(
+                'Cannot specify only one of top level identifier and `device_config_name`'
+            )
 
         project_name = _project_name(project_id)
 
@@ -791,14 +816,21 @@ def run_job_over_stream(
         if program_labels:
             program.labels.update(program_labels)
 
+        if snapshot_id:
+            selector = quantum.DeviceConfigSelector(
+                snapshot_id=snapshot_id or None, config_alias=device_config_name
+            )
+        else:
+            selector = quantum.DeviceConfigSelector(
+                run_name=run_name or None, config_alias=device_config_name
+            )
+
         job = quantum.QuantumJob(
             name=_job_name_from_ids(project_id, program_id, job_id),
             scheduling_config=quantum.SchedulingConfig(
                 processor_selector=quantum.SchedulingConfig.ProcessorSelector(
                     processor=_processor_name_from_ids(project_id, processor_id),
-                    device_config_selector=quantum.DeviceConfigSelector(
-                        run_name=run_name, config_alias=device_config_name
-                    ),
+                    device_config_selector=selector,
                 )
             ),
             run_context=run_context,