Integrate StreamManager with run_sweep() (#6285)

* Integrate StreamManager with run_sweep()

* Added a feature flag defaulting to True.
* Added logic to temporarily rewrite `processor_ids` to `processor_id` prior to full deprecation of `processor_ids` in `Engine.run_sweep()`.

* [WIP] Addressed maffoo's comments

* Addressed wcourtney's comments

* Modified engine_processor_test to take into account the enable_streaming feature flag

* Address wcourtney's comments

* Rename job_response_future to job_result_future

Author: verult

cirq-google/cirq_google/engine/engine.py

@@ -88,6 +88,8 @@ def __init__(
         client: 'Optional[engine_client.EngineClient]' = None,
         timeout: Optional[int] = None,
         serializer: Serializer = CIRCUIT_SERIALIZER,
+        # TODO(#5996) Remove enable_streaming once the feature is stable.
+        enable_streaming: bool = True,
     ) -> None:
         """Context and client for using Quantum Engine.
 
@@ -103,6 +105,9 @@ def __init__(
             timeout: Timeout for polling for results, in seconds.  Default is
                 to never timeout.
             serializer: Used to serialize circuits when running jobs.
+            enable_streaming: Feature gate for making Quantum Engine requests using the stream RPC.
+                If True, the Quantum Engine streaming RPC is used for creating jobs
+                and getting results. Otherwise, unary RPCs are used.
 
         Raises:
             ValueError: If either `service_args` and `verbose` were supplied
@@ -115,6 +120,7 @@ def __init__(
         if self.proto_version == ProtoVersion.V1:
             raise ValueError('ProtoVersion V1 no longer supported')
         self.serializer = serializer
+        self.enable_streaming = enable_streaming
 
         if not client:
             client = engine_client.EngineClient(service_args=service_args, verbose=verbose)
@@ -306,7 +312,7 @@ async def run_sweep_async(
         run_name: str = "",
         device_config_name: str = "",
     ) -> engine_job.EngineJob:
-        """Runs the supplied Circuit via Quantum Engine.Creates
+        """Runs the supplied Circuit via Quantum Engine.
 
         In contrast to run, this runs across multiple parameter sweeps, and
         does not block until a result is returned.
@@ -355,6 +361,44 @@ async def run_sweep_async(
             ValueError: If either `run_name` and `device_config_name` are set but
                 `processor_id` is empty.
         """
+
+        if self.context.enable_streaming:
+            # This logic is temporary prior to deprecating the processor_ids parameter.
+            # TODO(#6271) Remove after deprecating processor_ids elsewhere prior to v1.4.
+            if processor_ids:
+                if len(processor_ids) > 1:
+                    raise ValueError("The use of multiple processors is no longer supported.")
+                if len(processor_ids) == 1 and not processor_id:
+                    processor_id = processor_ids[0]
+
+            if not program_id:
+                program_id = _make_random_id('prog-')
+            if not job_id:
+                job_id = _make_random_id('job-')
+            run_context = self.context._serialize_run_context(params, repetitions)
+
+            job_result_future = self.context.client.run_job_over_stream(
+                project_id=self.project_id,
+                program_id=str(program_id),
+                program_description=program_description,
+                program_labels=program_labels,
+                code=self.context._serialize_program(program),
+                job_id=str(job_id),
+                run_context=run_context,
+                job_description=job_description,
+                job_labels=job_labels,
+                processor_id=processor_id,
+                run_name=run_name,
+                device_config_name=device_config_name,
+            )
+            return engine_job.EngineJob(
+                self.project_id,
+                str(program_id),
+                str(job_id),
+                self.context,
+                job_result_future=job_result_future,
+            )
+
         engine_program = await self.create_program_async(
             program, program_id, description=program_description, labels=program_labels
         )
@@ -372,6 +416,7 @@ async def run_sweep_async(
 
     run_sweep = duet.sync(run_sweep_async)
 
+    # TODO(#5996) Migrate to stream client
     # TODO(#6271): Deprecate and remove processor_ids before v1.4
     async def run_batch_async(
         self,
@@ -475,6 +520,7 @@ async def run_batch_async(
 
     run_batch = duet.sync(run_batch_async)
 
+    # TODO(#5996) Migrate to stream client
     async def run_calibration_async(
         self,
         layers: List['cirq_google.CalibrationLayer'],

cirq-google/cirq_google/engine/engine_client.py

@@ -39,6 +39,7 @@
 from cirq._compat import deprecated_parameter
 from cirq_google.cloud import quantum
 from cirq_google.engine.asyncio_executor import AsyncioExecutor
+from cirq_google.engine import stream_manager
 
 _M = TypeVar('_M', bound=proto.Message)
 _R = TypeVar('_R')
@@ -106,6 +107,10 @@ async def make_client():
 
         return self._executor.submit(make_client).result()
 
+    @cached_property
+    def _stream_manager(self) -> stream_manager.StreamManager:
+        return stream_manager.StreamManager(self.grpc_client)
+
     async def _send_request_async(self, func: Callable[[_M], Awaitable[_R]], request: _M) -> _R:
         """Sends a request by invoking an asyncio callable."""
         return await self._run_retry_async(func, request)
@@ -736,6 +741,97 @@ async def get_job_results_async(
 
     get_job_results = duet.sync(get_job_results_async)
 
+    def run_job_over_stream(
+        self,
+        *,
+        project_id: str,
+        program_id: str,
+        code: any_pb2.Any,
+        run_context: any_pb2.Any,
+        program_description: Optional[str] = None,
+        program_labels: Optional[Dict[str, str]] = None,
+        job_id: str,
+        priority: Optional[int] = None,
+        job_description: Optional[str] = None,
+        job_labels: Optional[Dict[str, str]] = None,
+        processor_id: str = "",
+        run_name: 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.
+
+        Sends the request over the Quantum Engine QuantumRunStream bidirectional stream, and returns
+        a future for the stream response. The future will be completed with a `QuantumResult` if
+        the job is successful; otherwise, it will be completed with a QuantumJob.
+
+        Args:
+            project_id: A project_id of the parent Google Cloud Project.
+            program_id: Unique ID of the program within the parent project.
+            code: Properly serialized program code.
+            run_context: Properly serialized run context.
+            program_description: An optional description to set on the program.
+            program_labels: Optional set of labels to set on the program.
+            job_id: Unique ID of the job within the parent program.
+            priority: Optional priority to run at, 0-1000.
+            job_description: Optional description to set on the job.
+            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:
+            A future for the job result, or the job if the job has failed.
+
+        Raises:
+            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.
+        """
+        # 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`')
+
+        project_name = _project_name(project_id)
+
+        program_name = _program_name_from_ids(project_id, program_id)
+        program = quantum.QuantumProgram(name=program_name, code=code)
+        if program_description:
+            program.description = program_description
+        if program_labels:
+            program.labels.update(program_labels)
+
+        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_key=quantum.DeviceConfigKey(
+                        run_name=run_name, config_alias=device_config_name
+                    ),
+                )
+            ),
+            run_context=run_context,
+        )
+        if priority:
+            job.scheduling_config.priority = priority
+        if job_description:
+            job.description = job_description
+        if job_labels:
+            job.labels.update(job_labels)
+
+        return self._stream_manager.submit(project_name, program, job)
+
     async def list_processors_async(self, project_id: str) -> List[quantum.QuantumProcessor]:
         """Returns a list of Processors that the user has visibility to in the
         current Engine project. The names of these processors are used to