Store job identifiers in calibration results (quantumlib#4053)

* Store job identifiers in calibration results

* Fix documentation

* Try to mock EngineClient

* Fix import

* Use normal caching for engine job and calibration

* Revert inward files changes

* Fixes after merge

* Fix lint error

* Fix tests

cirq-google/cirq_google/calibration/engine_simulator_test.py

@@ -1,4 +1,4 @@
-from typing import Iterable, Tuple
+from typing import Iterable, Optional, Tuple
 
 import collections
 from unittest import mock
@@ -27,7 +27,9 @@ class DummyPhasedFSimCalibrationRequest(PhasedFSimCalibrationRequest):
     def to_calibration_layer(self) -> cirq_google.CalibrationLayer:
         return NotImplemented
 
-    def parse_result(self, result: cirq_google.CalibrationResult) -> PhasedFSimCalibrationResult:
+    def parse_result(
+        self, result: cirq_google.CalibrationResult, job: Optional[cirq_google.EngineJob] = None
+    ) -> PhasedFSimCalibrationResult:
         return NotImplemented
 
 

cirq-google/cirq_google/calibration/phased_fsim.py

@@ -53,7 +53,7 @@
     rz,
 )
 from cirq_google.api import v2
-from cirq_google.engine import CalibrationLayer, CalibrationResult
+from cirq_google.engine import Calibration, CalibrationLayer, CalibrationResult, Engine, EngineJob
 
 if TYPE_CHECKING:
     import cirq
@@ -238,7 +238,7 @@ def create_phased_fsim_request(
         """
 
 
-@dataclasses.dataclass(frozen=True)
+@dataclasses.dataclass
 class PhasedFSimCalibrationResult:
     """The PhasedFSimGate characterization result.
 
@@ -248,11 +248,19 @@ class PhasedFSimCalibrationResult:
         gate: Characterized gate for each qubit pair. This is copied from the matching
             PhasedFSimCalibrationRequest and is included to preserve execution context.
         options: The options used to gather this result.
+        project_id: Google's job project id.
+        program_id: Google's job program id.
+        job_id: Google's job job id.
     """
 
     parameters: Dict[Tuple[Qid, Qid], PhasedFSimCharacterization]
     gate: Gate
     options: PhasedFSimCalibrationOptions
+    project_id: Optional[str] = None
+    program_id: Optional[str] = None
+    job_id: Optional[str] = None
+    _engine_job: Optional[EngineJob] = None
+    _calibration: Optional[Calibration] = None
 
     def override(self, parameters: PhasedFSimCharacterization) -> 'PhasedFSimCalibrationResult':
         """Creates the new results with certain parameters overridden for all characterizations.
@@ -285,6 +293,27 @@ def get_parameters(self, a: Qid, b: Qid) -> Optional['PhasedFSimCharacterization
         else:
             return None
 
+    @property
+    def engine_job(self) -> Optional[EngineJob]:
+        """The cirq_google.EngineJob associated with this calibration request.
+
+        Available only when project_id, program_id and job_id attributes are present.
+        """
+        if self._engine_job is None and self.project_id and self.program_id and self.job_id:
+            engine = Engine(project_id=self.project_id)
+            self._engine_job = engine.get_program(self.program_id).get_job(self.job_id)
+        return self._engine_job
+
+    @property
+    def engine_calibration(self) -> Optional[Calibration]:
+        """The underlying device calibration that was used for this user-specific calibration.
+
+        This is a cached property that triggers a network call at the first use.
+        """
+        if self._calibration is None and self.engine_job is not None:
+            self._calibration = self.engine_job.get_calibration()
+        return self._calibration
+
     @classmethod
     def _create_parameters_dict(
         cls,
@@ -316,6 +345,9 @@ def _json_dict_(self) -> Dict[str, Any]:
             'gate': self.gate,
             'parameters': [(q_a, q_b, params) for (q_a, q_b), params in self.parameters.items()],
             'options': self.options,
+            'project_id': self.project_id,
+            'program_id': self.program_id,
+            'job_id': self.job_id,
         }
 
 
@@ -398,7 +430,9 @@ def to_calibration_layer(self) -> CalibrationLayer:
         """Encodes this characterization request in a CalibrationLayer object."""
 
     @abc.abstractmethod
-    def parse_result(self, result: CalibrationResult) -> PhasedFSimCalibrationResult:
+    def parse_result(
+        self, result: CalibrationResult, job: Optional[EngineJob] = None
+    ) -> PhasedFSimCalibrationResult:
         """Decodes the characterization result issued for this request."""
 
 
@@ -657,7 +691,9 @@ def to_calibration_layer(self) -> CalibrationLayer:
             args=args,
         )
 
-    def parse_result(self, result: CalibrationResult) -> PhasedFSimCalibrationResult:
+    def parse_result(
+        self, result: CalibrationResult, job: Optional[EngineJob] = None
+    ) -> PhasedFSimCalibrationResult:
         if result.code != v2.calibration_pb2.SUCCESS:
             raise PhasedFSimCalibrationError(result.error_message)
 
@@ -683,7 +719,14 @@ def parse_result(self, result: CalibrationResult) -> PhasedFSimCalibrationResult
                 phi=data.get('phi_est', None),
             )
 
-        return PhasedFSimCalibrationResult(parameters=parsed, gate=self.gate, options=self.options)
+        return PhasedFSimCalibrationResult(
+            parameters=parsed,
+            gate=self.gate,
+            options=self.options,
+            project_id=None if job is None else job.project_id,
+            program_id=None if job is None else job.program_id,
+            job_id=None if job is None else job.job_id,
+        )
 
     @classmethod
     def _from_json_dict_(
@@ -794,7 +837,9 @@ class LocalXEBPhasedFSimCalibrationRequest(PhasedFSimCalibrationRequest):
 
     options: LocalXEBPhasedFSimCalibrationOptions
 
-    def parse_result(self, result: CalibrationResult) -> PhasedFSimCalibrationResult:
+    def parse_result(
+        self, result: CalibrationResult, job: Optional[EngineJob] = None
+    ) -> PhasedFSimCalibrationResult:
         raise NotImplementedError('Not applicable for local calibrations')
 
     def to_calibration_layer(self) -> CalibrationLayer:
@@ -819,7 +864,9 @@ def to_calibration_layer(self) -> CalibrationLayer:
             args=self.options.to_args(),
         )
 
-    def parse_result(self, result: CalibrationResult) -> PhasedFSimCalibrationResult:
+    def parse_result(
+        self, result: CalibrationResult, job: Optional[EngineJob] = None
+    ) -> PhasedFSimCalibrationResult:
         if result.code != v2.calibration_pb2.SUCCESS:
             raise PhasedFSimCalibrationError(result.error_message)
 
@@ -837,7 +884,12 @@ def parse_result(self, result: CalibrationResult) -> PhasedFSimCalibrationResult
 
         # TODO: Return initial_fids, final_fids somehow.
         return PhasedFSimCalibrationResult(
-            parameters=final_params, gate=self.gate, options=self.options
+            parameters=final_params,
+            gate=self.gate,
+            options=self.options,
+            project_id=None if job is None else job.project_id,
+            program_id=None if job is None else job.program_id,
+            job_id=None if job is None else job.job_id,
         )
 
     @classmethod

cirq-google/cirq_google/calibration/phased_fsim_test.py

@@ -14,6 +14,7 @@
 import os
 import re
 
+from unittest import mock
 import numpy as np
 import pandas as pd
 import pytest
@@ -809,6 +810,59 @@ def test_result_override():
     )
 
 
+@mock.patch('cirq_google.engine.engine_client.EngineClient')
+def test_result_engine_job(_client):
+    result = PhasedFSimCalibrationResult(
+        parameters={},
+        gate=cirq.FSimGate(theta=np.pi / 4, phi=0.0),
+        options=WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
+        project_id='project_id',
+        program_id='program_id',
+        job_id='job_id',
+    )
+
+    assert result.engine_job.project_id == 'project_id'
+    assert result.engine_job.program_id == 'program_id'
+    assert result.engine_job.job_id == 'job_id'
+
+
+def test_result_engine_job_none():
+    result = PhasedFSimCalibrationResult(
+        parameters={},
+        gate=cirq.FSimGate(theta=np.pi / 4, phi=0.0),
+        options=WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
+    )
+
+    assert result.engine_job is None
+
+
+@mock.patch('cirq_google.engine.engine_client.EngineClient')
+def test_result_engine_calibration(_client):
+    result = PhasedFSimCalibrationResult(
+        parameters={},
+        gate=cirq.FSimGate(theta=np.pi / 4, phi=0.0),
+        options=WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
+        project_id='project_id',
+        program_id='program_id',
+        job_id='job_id',
+    )
+
+    test_calibration = cirq_google.Calibration()
+    result.engine_job.get_calibration = lambda: test_calibration
+
+    assert result.engine_calibration == test_calibration
+
+
+def test_result_engine_calibration_none():
+    result = PhasedFSimCalibrationResult(
+        parameters={},
+        gate=cirq.FSimGate(theta=np.pi / 4, phi=0.0),
+        options=WITHOUT_CHI_FLOQUET_PHASED_FSIM_CHARACTERIZATION,
+    )
+
+    assert result.engine_calibration is None
+
+
 def test_options_phase_corrected_override():
     assert (
         ALL_ANGLES_FLOQUET_PHASED_FSIM_CHARACTERIZATION.zeta_chi_gamma_correction_override()

cirq-google/cirq_google/calibration/workflow.py

@@ -632,7 +632,7 @@ def _run_calibrations_via_engine(
         job = engine.run_calibration(cal_layers, processor_id=processor_id, gate_set=gate_set)
         request_results = job.calibration_results()
         results += [
-            calibration.parse_result(result)
+            calibration.parse_result(result, job)
             for calibration, result in zip(calibration_requests, request_results)
         ]
         if progress_func:

cirq-google/cirq_google/calibration/workflow_test.py

@@ -780,7 +780,7 @@ def test_make_zeta_chi_gamma_compensation_for_operations_permit_mixed_moments():
         )
 
 
-def test_run_characterization():
+def test_run_calibrations():
     q_00, q_01, q_02, q_03 = [cirq.GridQubit(0, index) for index in range(4)]
     gate = cirq.FSimGate(theta=np.pi / 4, phi=0.0)
 
@@ -836,7 +836,7 @@ def test_run_characterization():
         ),
     )
 
-    job = cirq_google.engine.EngineJob('', '', '', None)
+    job = cirq_google.engine.EngineJob('project_id', 'program_id', 'job_id', None)
     job._calibration_results = [result]
 
     engine = mock.MagicMock(spec=cirq_google.Engine)
@@ -871,6 +871,9 @@ def progress(step: int, steps: int) -> None:
                 characterize_gamma=False,
                 characterize_phi=True,
             ),
+            project_id='project_id',
+            program_id='program_id',
+            job_id='job_id',
         )
     ]
 
@@ -934,7 +937,7 @@ def test_run_characterization_with_engine():
         ),
     )
 
-    job = cirq_google.engine.EngineJob('', '', '', None)
+    job = cirq_google.engine.EngineJob('project_id', 'program_id', 'job_id', None)
     job._calibration_results = [result]
 
     engine = mock.MagicMock(spec=cirq_google.Engine)
@@ -967,18 +970,21 @@ def progress(step: int, steps: int) -> None:
                 characterize_gamma=False,
                 characterize_phi=True,
             ),
+            project_id='project_id',
+            program_id='program_id',
+            job_id='job_id',
         )
     ]
 
     assert actual == expected
     assert progress_calls == [(1, 1)]
 
 
-def test_run_characterization_empty():
+def test_run_calibrations_empty():
     assert workflow.run_calibrations([], None, 'qproc', cirq_google.FSIM_GATESET) == []
 
 
-def test_run_characterization_fails_when_invalid_arguments():
+def test_run_calibrations_fails_when_invalid_arguments():
     with pytest.raises(ValueError):
         assert workflow.run_calibrations(
             [], None, 'qproc', cirq_google.FSIM_GATESET, max_layers_per_request=0
@@ -1001,7 +1007,7 @@ def test_run_characterization_fails_when_invalid_arguments():
         assert workflow.run_calibrations([request], 0, 'qproc', cirq_google.FSIM_GATESET)
 
 
-def test_run_characterization_with_simulator():
+def test_run_calibrations_with_simulator():
     q_00, q_01, q_02, q_03 = [cirq.GridQubit(0, index) for index in range(4)]
     gate = SQRT_ISWAP_INV_GATE
 
@@ -1057,7 +1063,7 @@ def test_run_floquet_characterization_for_moments():
         characterize_phi=True,
     )
 
-    job = cirq_google.engine.EngineJob('', '', '', None)
+    job = cirq_google.engine.EngineJob('project_id', 'program_id', 'job_id', None)
     job._calibration_results = [
         cirq_google.CalibrationResult(
             code=cirq_google.api.v2.calibration_pb2.SUCCESS,
@@ -1119,6 +1125,9 @@ def test_run_floquet_characterization_for_moments():
             },
             gate=gate,
             options=options,
+            project_id='project_id',
+            program_id='program_id',
+            job_id='job_id',
         )
     ]
     assert circuit_with_calibration.circuit == circuit

cirq-google/cirq_google/json_test_data/PhasedFSimCalibrationResult.json

@@ -55,5 +55,8 @@
     "characterize_gamma": false,
     "characterize_phi": true,
     "readout_error_tolerance": 0.4
-  }
+  },
+  "project_id": "project_id",
+  "program_id": "program_id",
+  "job_id": "job_id"
 }

cirq-google/cirq_google/json_test_data/PhasedFSimCalibrationResult.repr

@@ -16,4 +16,7 @@ cirq_google.PhasedFSimCalibrationResult(
                 characterize_phi=True,
                 readout_error_tolerance=0.4,
             ),
+            project_id='project_id',
+            program_id='program_id',
+            job_id='job_id',
         )