Support the new quantum engine processor_selector in engine_client

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,25 @@ 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: 'processor_ids' in kwargs or len(args) > 3,
+    )
     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,33 +400,68 @@ 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.
+                `processor_id` should be used instead.
             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 (given by `processor_id`). An Automation Run contains a
+                collection of device configurations for a processor.
+            device_config_name: Configuration identifier used to identify a processor configuration
+                within the automation run.
 
         Returns:
             Tuple of created job id and job.
 
         Raises:
             ValueError: If the priority is not between 0 and 1000.
+            ValueError: If `processor_ids` and `processor_id` are both set.
+            ValueError: If neither `processor_ids` or `processor_id` are set.
+            ValueError: If either `run_name` and `device_config_name` are set but
+                `processor_id` is empty.
+            ValueError: If `run_name` is set but `device_config_name` 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 processor_id and processor_ids:
+            raise ValueError(
+                '`processor_ids` and `processor_id` cannot both be set.'
+                'Please just use `processor_id` instead.'
+            )
+        if not processor_id and not processor_ids:
+            raise ValueError('`processor_id` must be set.')
+        if not processor_id and (run_name or device_config_name):
+            raise ValueError(
+                'Cannot specify `run_name` or `device_config_name` if `processor_id` is empty.'
+            )
+        if bool(run_name) ^ bool(device_config_name):
+            raise ValueError('Cannot specify only one of `run_name` and `device_config_name`')
 
         # Create job.
+        processor_selector = (
+            quantum.SchedulingConfig.ProcessorSelector(
+                processor_names=[
+                    _processor_name_from_ids(project_id, processor_id)
+                    for processor_id in processor_ids
+                ]
+            )
+            if (processor_ids is not None)
+            else quantum.SchedulingConfig.ProcessorSelector(
+                processor=_processor_name_from_ids(project_id, processor_id),
+                device_config_key=quantum.DeviceConfigKey(
+                    run_name=run_name, 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_names=[
-                        _processor_name_from_ids(project_id, processor_id)
-                        for processor_id in processor_ids
-                    ]
-                )
-            ),
+            scheduling_config=quantum.SchedulingConfig(processor_selector=processor_selector),
             run_context=run_context,
         )
         if priority:
@@ -431,7 +476,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,

cirq-google/cirq_google/engine/engine_client_test.py

@@ -14,6 +14,7 @@
 """Tests for EngineClient."""
 import asyncio
 import datetime
+import os
 from unittest import mock
 
 import duet
@@ -341,8 +342,9 @@ def test_delete_program(client_constructor):
     )
 
 
+@mock.patch.dict(os.environ, clear='CIRQ_TESTING')
 @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
-def test_create_job(client_constructor):
+def test_create_job_with_legacy_processor_ids(client_constructor):
     grpc_client = setup_mock_(client_constructor)
 
     result = quantum.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
@@ -467,6 +469,133 @@ def test_create_job(client_constructor):
         )
 
 
+@mock.patch.dict(os.environ, clear='CIRQ_TESTING')
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+def test_create_job_with_run_name_and_device_config_name_and_no_processor_id_throws(
+    client_constructor,
+):
+    grpc_client = setup_mock_(client_constructor)
+    result = quantum.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
+    grpc_client.create_quantum_job.return_value = result
+    run_context = any_pb2.Any()
+    client = EngineClient()
+
+    with pytest.raises(
+        ValueError,
+        match="Cannot specify `run_name` or `device_config_name` if `processor_id` is empty",
+    ):
+        client.create_job(
+            project_id='proj',
+            program_id='prog',
+            job_id=None,
+            processor_ids=['processor0'],
+            run_name="RUN_NAME",
+            device_config_name="device_config_name",
+            run_context=run_context,
+        )
+
+
+@mock.patch.dict(os.environ, clear='CIRQ_TESTING')
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+def test_create_job_with_processor_id_and_processor_ids_throws(client_constructor):
+    grpc_client = setup_mock_(client_constructor)
+    result = quantum.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
+    grpc_client.create_quantum_job.return_value = result
+    run_context = any_pb2.Any()
+    client = EngineClient()
+
+    with pytest.raises(ValueError, match="`processor_ids` and `processor_id` cannot both be set"):
+        client.create_job(
+            project_id='proj',
+            program_id='prog',
+            job_id=None,
+            processor_id='processor0',
+            processor_ids=['processor0'],
+            run_context=run_context,
+        )
+
+
+@mock.patch.dict(os.environ, clear='CIRQ_TESTING')
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+def test_create_job_with_no_processor_id_throws(client_constructor):
+    grpc_client = setup_mock_(client_constructor)
+    result = quantum.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
+    grpc_client.create_quantum_job.return_value = result
+    run_context = any_pb2.Any()
+    client = EngineClient()
+
+    with pytest.raises(ValueError, match="`processor_id` must be set."):
+        client.create_job(
+            project_id='proj', program_id='prog', job_id=None, run_context=run_context
+        )
+
+
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+@pytest.mark.parametrize('run_name, device_config_name', [('RUN_NAME', ''), ('', 'CONFIG_NAME')])
+def test_create_job_with_incomplete_device_config_throws(
+    client_constructor, run_name, device_config_name
+):
+    grpc_client = setup_mock_(client_constructor)
+    result = quantum.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
+    grpc_client.create_quantum_job.return_value = result
+    run_context = any_pb2.Any()
+    client = EngineClient()
+
+    with pytest.raises(
+        ValueError, match="Cannot specify only one of `run_name` and `device_config_name`"
+    ):
+        client.create_job(
+            project_id='proj',
+            program_id='prog',
+            job_id=None,
+            processor_id='processor0',
+            run_name=run_name,
+            device_config_name=device_config_name,
+            run_context=run_context,
+        )
+
+
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+@pytest.mark.parametrize('run_name, device_config_name', [('RUN_NAME', 'CONFIG_NAME'), ('', '')])
+def test_create_job_with_run_name_and_device_config_name(
+    client_constructor, run_name, device_config_name
+):
+    grpc_client = setup_mock_(client_constructor)
+    result = quantum.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
+    grpc_client.create_quantum_job.return_value = result
+    run_context = any_pb2.Any()
+    client = EngineClient()
+
+    assert client.create_job(
+        project_id='proj',
+        program_id='prog',
+        job_id='job0',
+        processor_id='processor0',
+        run_name=run_name,
+        device_config_name=device_config_name,
+        run_context=run_context,
+        priority=10,
+    ) == ('job0', result)
+    grpc_client.create_quantum_job.assert_called_with(
+        quantum.CreateQuantumJobRequest(
+            parent='projects/proj/programs/prog',
+            quantum_job=quantum.QuantumJob(
+                name='projects/proj/programs/prog/jobs/job0',
+                run_context=run_context,
+                scheduling_config=quantum.SchedulingConfig(
+                    priority=10,
+                    processor_selector=quantum.SchedulingConfig.ProcessorSelector(
+                        processor='projects/proj/processors/processor0',
+                        device_config_key=quantum.DeviceConfigKey(
+                            run_name=run_name, config_alias=device_config_name
+                        ),
+                    ),
+                ),
+            ),
+        )
+    )
+
+
 @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
 def test_get_job(client_constructor):
     grpc_client = setup_mock_(client_constructor)