Add timeout to engine object (#2996)

* Add timeout to engine object

- Adds timeout to Engine object and makes it configurable
- This plumbs the timeout down to results object so that the
amount of polling is configurable.

* Moved timeout into EngineContext

cirq/google/engine/engine.py

@@ -71,7 +71,8 @@ def __init__(self,
                  proto_version: Optional[ProtoVersion] = None,
                  service_args: Optional[Dict] = None,
                  verbose: Optional[bool] = None,
-                 client: 'Optional[engine_client.EngineClient]' = None) -> None:
+                 client: 'Optional[engine_client.EngineClient]' = None,
+                 timeout: Optional[int] = None) -> None:
         """Context and client for using Quantum Engine.
 
         Args:
@@ -80,6 +81,8 @@ def __init__(self,
                 configure options on the underlying client.
             verbose: Suppresses stderr messages when set to False. Default is
                 true.
+            timeout: Timeout for polling for results, in seconds.  Default is
+                to never timeout.
         """
         if (service_args or verbose) and client:
             raise ValueError(
@@ -91,6 +94,7 @@ def __init__(self,
             client = engine_client.EngineClient(service_args=service_args,
                                                 verbose=verbose)
         self.client = client
+        self.timeout = timeout
 
     def copy(self) -> 'EngineContext':
         return EngineContext(proto_version=self.proto_version,
@@ -124,6 +128,7 @@ def __init__(
             service_args: Optional[Dict] = None,
             verbose: Optional[bool] = None,
             context: Optional[EngineContext] = None,
+            timeout: Optional[int] = None,
     ) -> None:
         """Supports creating and running programs against the Quantum Engine.
 
@@ -138,6 +143,8 @@ def __init__(
                 configure options on the underlying client.
             verbose: Suppresses stderr messages when set to False. Default is
                 true.
+            timeout: Timeout for polling for results, in seconds.  Default is
+                to never timeout.
         """
         if context and (proto_version or service_args or verbose):
             raise ValueError(
@@ -148,7 +155,8 @@ def __init__(
         if not context:
             context = EngineContext(proto_version=proto_version,
                                     service_args=service_args,
-                                    verbose=verbose)
+                                    verbose=verbose,
+                                    timeout=timeout)
         self.context = context
 
     def run(

cirq/google/engine/engine_job.py

@@ -262,10 +262,15 @@ def results(self) -> List[study.TrialResult]:
         import cirq.google.engine.engine as engine_base
         if not self._results:
             job = self._refresh_job()
-            for _ in range(1000):
+            total_seconds_waited = 0.0
+            timeout = self.context.timeout
+            while True:
+                if timeout and total_seconds_waited >= timeout:
+                    break
                 if job.execution_status.state in TERMINAL_STATES:
                     break
                 time.sleep(0.5)
+                total_seconds_waited += 0.5
                 job = self._refresh_job()
             self._raise_on_failure(job)
             response = self.context.client.get_job_results(

cirq/google/engine/engine_job_test.py

@@ -374,7 +374,7 @@ def test_timeout(patched_time_sleep, get_job):
     qjob = qtypes.QuantumJob(execution_status=qtypes.ExecutionStatus(
         state=qtypes.ExecutionStatus.State.RUNNING))
     get_job.return_value = qjob
-    job = cg.EngineJob('a', 'b', 'steve', EngineContext())
+    job = cg.EngineJob('a', 'b', 'steve', EngineContext(timeout=500))
     with pytest.raises(RuntimeError, match='Timed out'):
         job.results()
 

cirq/google/engine/engine_test.py

@@ -338,7 +338,7 @@ def test_run_circuit_timeout(patched_time_sleep, client):
             'state': 'RUNNING',
         })
 
-    engine = cg.Engine(project_id='project-id')
+    engine = cg.Engine(project_id='project-id', timeout=600)
     with pytest.raises(RuntimeError, match='Timed out'):
         engine.run(program=_CIRCUIT)