EngineJob with methods to cancel and get state and results.

Author: lindmarkm

cirq/google/engine/engine.py

@@ -20,6 +20,7 @@
 import re
 import string
 import time
+import urllib.parse
 from collections import Iterable
 from typing import Dict, List, Optional, Union, cast
 
@@ -40,6 +41,7 @@
 from cirq.study.sweeps import Points, Unit, Zip
 
 gcs_prefix_pattern = re.compile('gs://[a-z0-9._/-]+')
+TERMINAL_STATES = ['SUCCESS', 'FAILURE', 'CANCELLED']
 
 
 class EngineTrialResult(TrialResult):
@@ -74,7 +76,6 @@ class EngineOptions:
     """
 
     def __init__(self, project_id: str,
-                 credentials: Optional[oauth2client.client.Credentials] = None,
                  program_id: Optional[str] = None,
                  job_id: Optional[str] = None,
                  gcs_prefix: Optional[str] = None,
@@ -84,7 +85,6 @@ def __init__(self, project_id: str,
         project_id and either gcs_prefix or gcs_program and gcs_results.
 
         Args:
-            credentials: Credentials to use.
             project_id: The project id string of the Google Cloud Project to
                 use.
             program_id: Id of the program to create, defaults to a random
@@ -97,38 +97,46 @@ def __init__(self, project_id: str,
             gcs_results: Explicit override for the results storage location.
         """
         self.project_id = project_id
-        self.credentials = credentials
-        self.program_id = program_id or 'prog-%s' % ''.join(
-            random.choice(string.ascii_uppercase + string.digits) for _ in
-            range(6))
-        self.job_id = job_id or 'job-0'
-        if not gcs_prefix and (not gcs_program or not gcs_results):
-            raise TypeError('Either gcs_prefix must be provided or both'
-                            ' gcs_program and gcs_results are required.')
-        if gcs_prefix and not gcs_prefix_pattern.match(gcs_prefix):
-            raise TypeError('gcs_prefix must be of the form "gs://'
-                            '<bucket name and optional object prefix>/"')
-        self.gcs_prefix = gcs_prefix if not gcs_prefix or gcs_prefix.endswith(
-            '/') else gcs_prefix + '/'
-        self.gcs_program = gcs_program or '%sprograms/%s/%s' % (
-            self.gcs_prefix, self.program_id, self.program_id)
-        self.gcs_results = gcs_results or '%sprograms/%s/jobs/%s' % (
-            self.gcs_prefix, self.program_id, self.job_id)
+        self.program_id = program_id
+        self.job_id = job_id
+        self.gcs_prefix = gcs_prefix
+        self.gcs_program = gcs_program
+        self.gcs_results = gcs_results
 
 
 class Engine:
-    """Executor for Google Quantum Engine
+    """Executor for Google Quantum Engine.
     """
 
     def __init__(self, api_key: str, api: str = 'quantum',
                  version: str = 'v1alpha1',
-                 discovery_url: Optional[str] = None) -> None:
+                 discovery_url: Optional[str] = None,
+                 credentials: Optional[oauth2client.client.Credentials] = None,
+                 gcs_prefix: Optional[str] = None
+    ) -> None:
+        """Engine service client.
+
+        Args:
+            api_key: API key to use to retrieve discovery doc.
+            api: API name.
+            version: API version.
+            discovery_url: Discovery url to use.
+            credentials: Credentials to use.
+        """
         self.api_key = api_key
         self.api = api
         self.version = version
         self.discovery_url = discovery_url or ('https://{api}.googleapis.com/'
                                                '$discovery/rest'
                                                '?version={apiVersion}&key=%s')
+        self.credentials = credentials
+        self.gcs_prefix = gcs_prefix
+        self.service = discovery.build(
+            self.api,
+            self.version,
+            discoveryServiceUrl=self.discovery_url % urllib.parse.quote_plus(
+                self.api_key),
+            credentials=credentials)
 
     def run(self,
             options: EngineOptions,
@@ -153,8 +161,8 @@ def run(self,
         Returns:
             Results for this run.
         """
-        return self.run_sweep(options, circuit, device, [param_resolver],
-                              repetitions, priority, target_route)[0]
+        return list(self.run_sweep(options, circuit, device, [param_resolver],
+                                   repetitions, priority, target_route))[0]
 
     def run_sweep(self,
                   options: EngineOptions,
@@ -164,7 +172,7 @@ def run_sweep(self,
                   repetitions: int = 1,
                   priority: int = 500,
                   target_route: str = '/xmonsim',
-    ) -> List[EngineTrialResult]:
+    ) -> 'EngineJob':
         """Runs the entire supplied Circuit or Schedule via Google Quantum
          Engine.
 
@@ -181,6 +189,29 @@ def run_sweep(self,
         Returns:
             Results for this run.
         """
+        # Check and compute engine options.
+        gcs_prefix = options.gcs_prefix or self.gcs_prefix
+        gcs_prefix = gcs_prefix if not gcs_prefix or gcs_prefix.endswith(
+            '/') else gcs_prefix + '/'
+        if gcs_prefix and not gcs_prefix_pattern.match(gcs_prefix):
+            raise TypeError('gcs_prefix must be of the form "gs://'
+                            '<bucket name and optional object prefix>/"')
+        if not gcs_prefix and (not options.gcs_program or
+                               not options.gcs_results):
+            raise TypeError('Either gcs_prefix must be provided or both'
+                            ' gcs_program and gcs_results are required.')
+
+        project_id = options.project_id
+        program_id = options.program_id or 'prog-%s' % ''.join(
+            random.choice(string.ascii_uppercase + string.digits) for _ in
+            range(6))
+        job_id = options.job_id or 'job-0'
+        gcs_program = options.gcs_program or '%sprograms/%s/%s' % (
+            gcs_prefix, program_id, program_id)
+        gcs_results = options.gcs_results or '%sprograms/%s/jobs/%s' % (
+            gcs_prefix, program_id, job_id)
+
+        # Check program to run and program parameters.
         if not 0 <= priority < 1000:
             raise TypeError('priority must be between 0 and 1000')
 
@@ -201,64 +232,53 @@ def run_sweep(self,
         else:
             raise TypeError('Unexpected execution type')
 
+        # Create program.
         sweeps = _sweepable_to_sweeps(params or ParamResolver({}))
-
-        service = discovery.build(self.api, self.version,
-                                  discoveryServiceUrl=self.discovery_url % (
-                                      self.api_key,),
-                                  credentials=options.credentials)
-
         proto_program = program_pb2.Program()
         for sweep in sweeps:
             sweep_proto = proto_program.parameter_sweeps.add()
             sweep_to_proto(sweep, sweep_proto)
             sweep_proto.repetitions = repetitions
         proto_program.operations.extend(list(schedule_to_proto(schedule)))
-
         code = {
             '@type': 'type.googleapis.com/cirq.api.google.v1.Program'}
         code.update(MessageToDict(proto_program))
-
         request = {
-            'name': 'projects/%s/programs/%s' % (options.project_id,
-                                                 options.program_id,),
-            'gcs_code_location': {'uri': options.gcs_program, },
+            'name': 'projects/%s/programs/%s' % (project_id,
+                                                 program_id,),
+            'gcs_code_location': {'uri': gcs_program, },
             'code': code,
         }
+        response = self.service.projects().programs().create(
+            parent='projects/%s' % project_id, body=request).execute()
 
-        response = service.projects().programs().create(
-            parent='projects/%s' % options.project_id, body=request).execute()
-
+        # Create job.
         request = {
-            'name': '%s/jobs/%s' % (response['name'], options.job_id),
+            'name': '%s/jobs/%s' % (response['name'], job_id),
             'output_config': {
                 'gcs_results_location': {
-                    'uri': options.gcs_results
+                    'uri': gcs_results
                 }
             },
             'scheduling_config': {
                 'priority': priority,
                 'target_route': target_route
             },
         }
-        response = service.projects().programs().jobs().create(
+        response = self.service.projects().programs().jobs().create(
             parent=response['name'], body=request).execute()
 
-        for _ in range(1000):
-            if response['executionStatus']['state'] in ['SUCCESS', 'FAILURE',
-                                                        'CANCELLED']:
-                break
-            time.sleep(0.5)
-            response = service.projects().programs().jobs().get(
-                name=response['name']).execute()
-
-        if response['executionStatus']['state'] != 'SUCCESS':
-            raise RuntimeError('Job %s did not succeed. It is in state %s.' % (
-                response['name'], response['executionStatus']['state']))
+        return EngineJob(
+            EngineOptions(project_id, program_id, job_id, gcs_prefix,
+                          gcs_program, gcs_results), response, self)
 
-        response = service.projects().programs().jobs().getResult(
-            parent=response['name']).execute()
+    def get_job(self, job_resource_name) -> Dict:
+        return self.service.projects().programs().jobs().get(
+            name=job_resource_name).execute()
 
+    def get_job_results(self, job_resource_name) -> List[EngineTrialResult]:
+        response = self.service.projects().programs().jobs().getResult(
+            parent=job_resource_name).execute()
         trial_results = []
         for sweep_result in response['result']['sweepResults']:
             sweep_repetitions = sweep_result['repetitions']
@@ -276,6 +296,66 @@ def run_sweep(self,
                     measurements=measurements))
         return trial_results
 
+    def cancel_job(self, job_resource_name):
+        self.service.projects().programs().jobs().cancel(
+            name=job_resource_name, body={}).execute()
+
+
+class EngineJob:
+    """A job running on the engine.
+
+    Attributes:
+      engine_options: The engine options used for the job.
+      job_resource_name: The full resource name of the engine job.
+    """
+
+    def __init__(self,
+        engine_options: EngineOptions,
+        job: Dict,
+        engine: Engine) -> None:
+        """A job submitted to the engine.
+
+        Args:
+            engine_options: The EngineOptions used to create the job.
+            job: A full Job Dict.
+            engine: Engine connected to the job.
+        """
+        self.engine_options = engine_options
+        self._job = job
+        self._engine = engine
+        self.job_resource_name = job['name']
+        self._results = None  # type: List[EngineTrialResult]
+
+    def _update_job(self):
+        if self._job['executionStatus']['state'] not in TERMINAL_STATES:
+            self._job = self._engine.get_job(self.job_resource_name)
+        return self._job
+
+    def state(self):
+        return self._update_job()['executionStatus']['state']
+
+    def cancel(self):
+        self._engine.cancel_job(self.job_resource_name)
+
+    def results(self) -> List[EngineTrialResult]:
+        if not self._results:
+            job = self._update_job()
+            for _ in range(1000):
+                if job['executionStatus']['state'] in TERMINAL_STATES:
+                    break
+                time.sleep(0.5)
+                job = self._update_job()
+            if job['executionStatus']['state'] != 'SUCCESS':
+                raise RuntimeError(
+                    'Job %s did not succeed. It is in state %s.' % (
+                        job['name'], job['executionStatus']['state']))
+            self._results = self._engine.get_job_results(
+                self.job_resource_name)
+        return self._results
+
+    def __iter__(self):
+        return self.results().__iter__()
+
 
 def _sweepable_to_sweeps(sweepable: Sweepable) -> List[Sweep]:
     if isinstance(sweepable, ParamResolver):

cirq/google/engine/engine_test.py

@@ -125,10 +125,11 @@ def test_run_sweep_params(build):
     jobs.getResult().execute.return_value = {
         'result': MessageToDict(_RESULTS)}
 
-    results = Engine(api_key="key").run_sweep(
+    job = Engine(api_key="key").run_sweep(
         EngineOptions('project-id', gcs_prefix='gs://bucket/folder'),
         moment_by_moment_schedule(UnconstrainedDevice, Circuit()),
         params=[ParamResolver({'a': 1}), ParamResolver({'a': 2})])
+    results = job.results()
     assert len(results) == 2
     for i, v in enumerate([1, 2]):
         assert results[i].repetitions == 1
@@ -168,10 +169,11 @@ def test_run_sweep_sweeps(build):
     jobs.getResult().execute.return_value = {
         'result': MessageToDict(_RESULTS)}
 
-    results = Engine(api_key="key").run_sweep(
+    job = Engine(api_key="key").run_sweep(
         EngineOptions('project-id', gcs_prefix='gs://bucket/folder'),
         moment_by_moment_schedule(UnconstrainedDevice, Circuit()),
         params=Points('a', [1, 2]))
+    results = job.results()
     assert len(results) == 2
     for i, v in enumerate([1, 2]):
         assert results[i].repetitions == 1
@@ -193,11 +195,37 @@ def test_run_sweep_sweeps(build):
     assert jobs.getResult().execute.call_count == 1
 
 
-def test_bad_priority():
+@python3_mock_test(discovery, 'build')
+def test_bad_priority(build):
     with pytest.raises(TypeError, match='priority must be between 0 and 1000'):
         Engine(api_key="key").run(
             EngineOptions('project-id', gcs_prefix='gs://bucket/folder'),
             Circuit(),
             UnconstrainedDevice,
             priority=1001)
 
+
+@python3_mock_test(discovery, 'build')
+def test_cancel(build):
+    service = mock.Mock()
+    build.return_value = service
+    programs = service.projects().programs()
+    jobs = programs.jobs()
+    programs.create().execute.return_value = {
+        'name': 'projects/project-id/programs/test'}
+    jobs.create().execute.return_value = {
+        'name': 'projects/project-id/programs/test/jobs/test',
+        'executionStatus': {'state': 'READY'}}
+    jobs.get().execute.return_value = {
+        'name': 'projects/project-id/programs/test/jobs/test',
+        'executionStatus': {'state': 'CANCELLED'}}
+
+    job = Engine(api_key="key").run_sweep(
+        EngineOptions('project-id', gcs_prefix='gs://bucket/folder'),
+        Circuit(), device=UnconstrainedDevice)
+    job.cancel()
+    assert job.job_resource_name == ('projects/project-id/programs/test/'
+                                     'jobs/test')
+    assert job.state() == 'CANCELLED'
+    assert jobs.cancel.call_args[1][
+               'name'] == 'projects/project-id/programs/test/jobs/test'