Support the new quantum engine processor_selector in engine_client (#6254)

* Add DeviceConfigKey to Quantum Engine API types

* Add new processor_selector arguments in engine_client

* fix lint

* change processor_id type

* Add test coverage

* Address cxing comments

* fix formatting

* fix raises docstrings

* make new parameters to be mandatory keyword arguments

* fix typo

* Fix create_job deprecated arguments

* fix format

* small refactor

* Address wcourtney comments

---------

Co-authored-by: Jose Urruticoechea <urruti@google.com>

Author: jurruti

cirq-google/cirq_google/cloud/quantum/__init__.py

@@ -57,6 +57,7 @@
 from cirq_google.cloud.quantum_v1alpha1.types.engine import UpdateQuantumJobRequest
 from cirq_google.cloud.quantum_v1alpha1.types.engine import UpdateQuantumProgramRequest
 from cirq_google.cloud.quantum_v1alpha1.types.engine import UpdateQuantumReservationRequest
+from cirq_google.cloud.quantum_v1alpha1.types.quantum import DeviceConfigKey
 from cirq_google.cloud.quantum_v1alpha1.types.quantum import ExecutionStatus
 from cirq_google.cloud.quantum_v1alpha1.types.quantum import GcsLocation
 from cirq_google.cloud.quantum_v1alpha1.types.quantum import InlineData
@@ -115,6 +116,7 @@
     'UpdateQuantumJobRequest',
     'UpdateQuantumProgramRequest',
     'UpdateQuantumReservationRequest',
+    'DeviceConfigKey',
     'ExecutionStatus',
     'GcsLocation',
     'InlineData',

cirq-google/cirq_google/cloud/quantum_v1alpha1/__init__.py

@@ -57,6 +57,7 @@
 from .types.engine import UpdateQuantumJobRequest
 from .types.engine import UpdateQuantumProgramRequest
 from .types.engine import UpdateQuantumReservationRequest
+from .types.quantum import DeviceConfigKey
 from .types.quantum import ExecutionStatus
 from .types.quantum import GcsLocation
 from .types.quantum import InlineData
@@ -84,6 +85,7 @@
 'DeleteQuantumJobRequest',
 'DeleteQuantumProgramRequest',
 'DeleteQuantumReservationRequest',
+'DeviceConfigKey',
 'ExecutionStatus',
 'GcsLocation',
 'GetQuantumCalibrationRequest',

cirq-google/cirq_google/cloud/quantum_v1alpha1/types/__init__.py

@@ -56,6 +56,7 @@
     UpdateQuantumReservationRequest,
 )
 from .quantum import (
+    DeviceConfigKey,
     ExecutionStatus,
     GcsLocation,
     InlineData,
@@ -114,6 +115,7 @@
     'UpdateQuantumJobRequest',
     'UpdateQuantumProgramRequest',
     'UpdateQuantumReservationRequest',
+    'DeviceConfigKey',
     'ExecutionStatus',
     'GcsLocation',
     'InlineData',

cirq-google/cirq_google/cloud/quantum_v1alpha1/types/quantum.py

@@ -226,6 +226,25 @@ class QuantumJob(proto.Message):
     )
 
 
+class DeviceConfigKey(proto.Message):
+    r"""-
+    Attributes:
+        run_name (str):
+            -
+        config_alias (str):
+            -
+    """
+
+    run_name = proto.Field(
+        proto.STRING,
+        number=1,
+    )
+    config_alias = proto.Field(
+        proto.STRING,
+        number=2,
+    )
+
+
 class SchedulingConfig(proto.Message):
     r"""-
 
@@ -244,12 +263,25 @@ class ProcessorSelector(proto.Message):
         Attributes:
             processor_names (Sequence[str]):
                 -
+            processor (str):
+                -
+            device_config_key ((google.cloud.quantum_v1alpha1.types.DeviceConfigKey):
+                -
         """
 
         processor_names = proto.RepeatedField(
             proto.STRING,
             number=1,
         )
+        processor = proto.Field(
+            proto.STRING,
+            number=2,
+        )
+        device_config_key = proto.Field(
+            proto.MESSAGE,
+            number=3,
+            message=DeviceConfigKey
+        )
 
     target_route = proto.Field(
         proto.STRING,

cirq-google/cirq_google/engine/engine_client.py

@@ -36,6 +36,7 @@
 from google.protobuf.timestamp_pb2 import Timestamp
 
 from cirq._compat import cached_property
+from cirq._compat import deprecated_parameter
 from cirq_google.cloud import quantum
 from cirq_google.engine.asyncio_executor import AsyncioExecutor
 
@@ -372,16 +373,27 @@ async def delete_program_async(
 
     delete_program = duet.sync(delete_program_async)
 
+    @deprecated_parameter(
+        deadline='v1.4',
+        fix='Use `processor_id` instead of `processor_ids`.',
+        parameter_desc='processor_ids',
+        match=lambda args, kwargs: _match_deprecated_processor_ids(args, kwargs),
+        rewrite=lambda args, kwargs: rewrite_processor_ids_to_processor_id(args, kwargs),
+    )
     async def create_job_async(
         self,
         project_id: str,
         program_id: str,
         job_id: Optional[str],
-        processor_ids: Sequence[str],
-        run_context: any_pb2.Any,
+        processor_ids: Optional[Sequence[str]] = None,
+        run_context: any_pb2.Any = any_pb2.Any(),
         priority: Optional[int] = None,
         description: Optional[str] = None,
         labels: Optional[Dict[str, str]] = None,
+        *,
+        processor_id: str = "",
+        run_name: str = "",
+        device_config_name: str = "",
     ) -> Tuple[str, quantum.QuantumJob]:
         """Creates and runs a job on Quantum Engine.
 
@@ -390,31 +402,53 @@ async def create_job_async(
             program_id: Unique ID of the program within the parent project.
             job_id: Unique ID of the job within the parent program.
             run_context: Properly serialized run context.
-            processor_ids: List of processor id for running the program.
+            processor_ids: Deprecated list of processor ids for running the program.
+                Only allowed to contain one processor_id. If the argument `processor_id`
+                is non-empty, `processor_ids` will be ignored. Otherwise the deprecated
+                decorator will fix the arguments and call create_job_async using
+                `processor_id` instead of `processor_ids`.
             priority: Optional priority to run at, 0-1000.
             description: Optional description to set on the job.
             labels: Optional set of labels to set on the job.
-
+            processor_id: Processor id for running the program. If not set,
+                `processor_ids` will be used.
+            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. If specified, `processor_id`
+                is required to be set.
+            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.
+                If specified, `processor_id` is required to be set.
         Returns:
             Tuple of created job id and job.
 
         Raises:
             ValueError: If the priority is not between 0 and 1000.
+            ValueError: If neither `processor_id` or `processor_ids` are set.
+            ValueError: If  only one of `run_name` and `device_config_name` are specified.
+            ValueError: If `processor_ids` has more than one processor id.
+            ValueError: If either `run_name` and `device_config_name` are set but
+                `processor_id` is empty.
         """
         # 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`')
 
         # Create job.
         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_names=[
-                        _processor_name_from_ids(project_id, processor_id)
-                        for processor_id in processor_ids
-                    ]
+                    processor=_processor_name_from_ids(project_id, processor_id),
+                    device_config_key=quantum.DeviceConfigKey(
+                        run_name=run_name, config_alias=device_config_name
+                    ),
                 )
             ),
             run_context=run_context,
@@ -431,7 +465,8 @@ async def create_job_async(
         job = await self._send_request_async(self.grpc_client.create_quantum_job, request)
         return _ids_from_job_name(job.name)[2], job
 
-    create_job = duet.sync(create_job_async)
+    # TODO(cxing): Remove type ignore once @deprecated_parameter decorator is removed
+    create_job = duet.sync(create_job_async)  # type: ignore
 
     async def list_jobs_async(
         self,
@@ -1090,3 +1125,40 @@ def _date_or_time_to_filter_expr(param_name: str, param: Union[datetime.datetime
         f"type {type(param)}. Supported types: datetime.datetime and"
         f"datetime.date"
     )
+
+
+def rewrite_processor_ids_to_processor_id(args, kwargs):
+    """Rewrites the create_job parameters so that `processor_id` is used instead of the deprecated
+    `processor_ids`.
+
+        Raises:
+            ValueError: If `processor_ids` has more than one processor id.
+            ValueError: If `run_name` or `device_config_name` are set but `processor_id` is not.
+    """
+
+    # Use `processor_id` keyword argument instead of `processor_ids`
+    processor_ids = args[4] if len(args) > 4 else kwargs['processor_ids']
+    if len(processor_ids) > 1:
+        raise ValueError("The use of multiple processors is no longer supported.")
+    if 'processor_id' not in kwargs or not kwargs['processor_id']:
+        if ('run_name' in kwargs and kwargs['run_name']) or (
+            'device_config_name' in kwargs and kwargs['device_config_name']
+        ):
+            raise ValueError(
+                'Cannot specify `run_name` or `device_config_name` if `processor_id` is empty.'
+            )
+        kwargs['processor_id'] = processor_ids[0]
+
+    # Erase `processor_ids` from args and kwargs
+    if len(args) > 4:
+        args_list = list(args)
+        args_list[4] = None
+        args = tuple(args_list)
+    else:
+        kwargs.pop('processor_ids')
+
+    return args, kwargs
+
+
+def _match_deprecated_processor_ids(args, kwargs):
+    return ('processor_ids' in kwargs and kwargs['processor_ids']) or len(args) > 4