[cirqflow] Execution Loop I/O (#4599)

Following #4584 comment.

![ExecutableGroupResult](https://user-images.githubusercontent.com/4967059/138186981-ef4c82fc-6f05-400e-a761-d9a9f0b3257d.png)

- Update the constituent parts only when necessary to avoid re-writing the full `ExecutableGroupResult` when only a part has changed. Re-writing would cause a performance hit and potential corruption risk.
- Add a new dataclass to keep track of filenames to make loading in the data easier; see the test.

Author: mpharrigan

cirq-google/cirq_google/__init__.py

@@ -132,6 +132,7 @@
     RuntimeInfo,
     ExecutableResult,
     ExecutableGroupResult,
+    ExecutableGroupResultFilesystemRecord,
     QuantumRuntimeConfiguration,
     execute,
 )

cirq-google/cirq_google/json_resolver_cache.py

@@ -49,5 +49,9 @@ def _class_resolver_dictionary() -> Dict[str, ObjectFactory]:
         'cirq.google.RuntimeInfo': cirq_google.RuntimeInfo,
         'cirq.google.ExecutableResult': cirq_google.ExecutableResult,
         'cirq.google.ExecutableGroupResult': cirq_google.ExecutableGroupResult,
+        # Pylint fights with the black formatter.
+        # pylint: disable=line-too-long
+        'cirq.google.ExecutableGroupResultFilesystemRecord': cirq_google.ExecutableGroupResultFilesystemRecord,
+        # pylint: enable=line-too-long
         'cirq.google.QuantumRuntimeConfiguration': cirq_google.QuantumRuntimeConfiguration,
     }

cirq-google/cirq_google/json_test_data/cirq.google.ExecutableGroupResultFilesystemRecord.json

@@ -0,0 +1,10 @@
+{
+  "cirq_type": "cirq.google.ExecutableGroupResultFilesystemRecord",
+  "runtime_configuration_path": "RuntimeConfiguration.json.gz",
+  "shared_runtime_info_path": "SharedRuntimeInfo.jzon.gz",
+  "executable_result_paths": [
+    "ExecutableResult.1.json.gz",
+    "ExecutableResult.2.json.gz"
+  ],
+  "run_id": "my-run-id"
+}

cirq-google/cirq_google/json_test_data/cirq.google.ExecutableGroupResultFilesystemRecord.repr

@@ -0,0 +1 @@
+cirq_google.ExecutableGroupResultFilesystemRecord(runtime_configuration_path='RuntimeConfiguration.json.gz', shared_runtime_info_path='SharedRuntimeInfo.jzon.gz', executable_result_paths=['ExecutableResult.1.json.gz', 'ExecutableResult.2.json.gz'], run_id='my-run-id')

cirq-google/cirq_google/json_test_data/spec.py

@@ -65,6 +65,7 @@
             'QuantumRuntimeConfiguration',
             'RuntimeInfo',
             'SharedRuntimeInfo',
+            'ExecutableGroupResultFilesystemRecord',
         ]
     },
     tested_elsewhere=[

cirq-google/cirq_google/workflow/__init__.py

@@ -12,6 +12,7 @@
     RuntimeInfo,
     ExecutableResult,
     ExecutableGroupResult,
+    ExecutableGroupResultFilesystemRecord,
     QuantumRuntimeConfiguration,
     execute,
 )

cirq-google/cirq_google/workflow/quantum_runtime.py

@@ -116,6 +116,51 @@ def __repr__(self) -> str:
         return _compat.dataclass_repr(self, namespace='cirq_google')
 
 
+@dataclasses.dataclass
+class ExecutableGroupResultFilesystemRecord:
+    """Filename references to the constituent parts of a `cg.ExecutableGroupResult`.
+
+    Args:
+        runtime_configuration_path: A filename pointing to the `runtime_configuration` value.
+        shared_runtime_info_path: A filename pointing to the `shared_runtime_info` value.
+        executable_result_paths: A list of filenames pointing to the `executable_results` values.
+        run_id: The unique `str` identifier from this run. This is used to locate the other
+            values on disk.
+    """
+
+    runtime_configuration_path: str
+    shared_runtime_info_path: str
+    executable_result_paths: List[str]
+
+    run_id: str
+
+    def load(self, *, base_data_dir: str = ".") -> ExecutableGroupResult:
+        """Using the filename references in this dataclass, load a `cg.ExecutableGroupResult`
+        from its constituent parts.
+
+        Args:
+            base_data_dir: The base data directory. Files should be found at
+                {base_data_dir}/{run_id}/{this class's paths}
+        """
+        data_dir = f"{base_data_dir}/{self.run_id}"
+        return ExecutableGroupResult(
+            runtime_configuration=cirq.read_json_gzip(
+                f'{data_dir}/{self.runtime_configuration_path}'
+            ),
+            shared_runtime_info=cirq.read_json_gzip(f'{data_dir}/{self.shared_runtime_info_path}'),
+            executable_results=[
+                cirq.read_json_gzip(f'{data_dir}/{exe_path}')
+                for exe_path in self.executable_result_paths
+            ],
+        )
+
+    def _json_dict_(self) -> Dict[str, Any]:
+        return dataclass_json_dict(self, namespace='cirq.google')
+
+    def __repr__(self) -> str:
+        return _compat.dataclass_repr(self, namespace='cirq_google')
+
+
 @dataclasses.dataclass
 class QuantumRuntimeConfiguration:
     """User-requested configuration of how to execute a given `cg.QuantumExecutableGroup`.
@@ -138,22 +183,63 @@ def __repr__(self) -> str:
         return _compat.dataclass_repr(self, namespace='cirq_google')
 
 
+def _safe_to_json(obj: Any, *, part_path: str, nominal_path: str, bak_path: str):
+    """Safely update a json file.
+
+    1. The new value is written to a "part" file
+    2. The previous file atomically replaces the previous backup file, thereby becoming the
+       current backup file.
+    3. The part file is atomically renamed to the desired filename.
+    """
+    cirq.to_json_gzip(obj, part_path)
+    if os.path.exists(nominal_path):
+        os.replace(nominal_path, bak_path)
+    os.replace(part_path, nominal_path)
+
+
+def _update_updatable_files(
+    egr_record: ExecutableGroupResultFilesystemRecord,
+    shared_rt_info: SharedRuntimeInfo,
+    data_dir: str,
+):
+    """Safely update ExecutableGroupResultFilesystemRecord.json.gz and SharedRuntimeInfo.json.gz
+    during an execution run.
+    """
+    _safe_to_json(
+        shared_rt_info,
+        part_path=f'{data_dir}/SharedRuntimeInfo.json.gz.part',
+        nominal_path=f'{data_dir}/SharedRuntimeInfo.json.gz',
+        bak_path=f'{data_dir}/SharedRuntimeInfo.json.gz.bak',
+    )
+    _safe_to_json(
+        egr_record,
+        part_path=f'{data_dir}/ExecutableGroupResultFilesystemRecord.json.gz.part',
+        nominal_path=f'{data_dir}/ExecutableGroupResultFilesystemRecord.json.gz',
+        bak_path=f'{data_dir}/ExecutableGroupResultFilesystemRecord.json.gz.bak',
+    )
+
+
 def execute(
     rt_config: QuantumRuntimeConfiguration,
     executable_group: QuantumExecutableGroup,
     base_data_dir: str = ".",
 ) -> ExecutableGroupResult:
     """Execute a `cg.QuantumExecutableGroup` according to a `cg.QuantumRuntimeConfiguration`.
 
+    The ExecutableGroupResult's constituent parts will be saved to disk as they become
+    available. Within the "{base_data_dir}/{run_id}" directory we save:
+        - The `cg.QuantumRuntimeConfiguration` at the start of the execution as a record
+          of *how* the executable group was run.
+        - A `cg.SharedRuntimeInfo` which is updated throughout the run.
+        - An `cg.ExecutableResult` for each `cg.QuantumExecutable` as they become available.
+        - A `cg.ExecutableGroupResultFilesystemRecord` which is updated throughout the run.
+
     Args:
         rt_config: The `cg.QuantumRuntimeConfiguration` specifying how to execute
             `executable_group`.
         executable_group: The `cg.QuantumExecutableGroup` containing the executables to execute.
-        base_data_dir: A filesystem path to write data. We write
-            "{base_data_dir}/{run_id}/ExecutableGroupResult.json.gz"
-            containing the `cg.ExecutableGroupResult` as well as one file
-            "{base_data_dir}/{run_id}/ExecutableResult.{i}.json.gz" per `cg.ExecutableResult` as
-            each executable result becomes available.
+        base_data_dir: Each data file will be written to the "{base_data_dir}/{run_id}/" directory,
+            which must not already exist.
 
     Returns:
         The `cg.ExecutableGroupResult` containing all data and metadata for an execution.
@@ -174,15 +260,21 @@ def execute(
         # coverage: ignore
         raise ValueError("Please provide a non-empty `base_data_dir`.")
 
-    os.makedirs(f'{base_data_dir}/{run_id}', exist_ok=False)
-
-    # Results object that we will fill in in the main loop.
-    exegroup_result = ExecutableGroupResult(
-        runtime_configuration=rt_config,
-        shared_runtime_info=SharedRuntimeInfo(run_id=run_id),
-        executable_results=list(),
+    # Set up data saving, save runtime configuration.
+    data_dir = f'{base_data_dir}/{run_id}'
+    os.makedirs(data_dir, exist_ok=False)
+    egr_record = ExecutableGroupResultFilesystemRecord(
+        runtime_configuration_path='QuantumRuntimeConfiguration.json.gz',
+        shared_runtime_info_path='SharedRuntimeInfo.json.gz',
+        executable_result_paths=[],
+        run_id=run_id,
     )
-    cirq.to_json_gzip(exegroup_result, f'{base_data_dir}/{run_id}/ExecutableGroupResult.json.gz')
+    cirq.to_json_gzip(rt_config, f'{data_dir}/{egr_record.runtime_configuration_path}')
+
+    # Set up to-be-updated objects.
+    shared_rt_info = SharedRuntimeInfo(run_id=run_id)
+    _update_updatable_files(egr_record, shared_rt_info, data_dir)
+    executable_results = []
 
     # Loop over executables.
     sampler = rt_config.processor.get_sampler()
@@ -206,9 +298,18 @@ def execute(
             runtime_info=runtime_info,
             raw_data=sampler_run_result,
         )
-        cirq.to_json_gzip(exe_result, f'{base_data_dir}/{run_id}/ExecutableResult.{i}.json.gz')
-        exegroup_result.executable_results.append(exe_result)
-        print(f'\r{i+1} / {n_executables}', end='', flush=True)
+        # Do bookkeeping for finished ExecutableResult
+        exe_result_path = f'ExecutableResult.{i}.json.gz'
+        cirq.to_json_gzip(exe_result, f"{data_dir}/{exe_result_path}")
+        executable_results.append(exe_result)
+        egr_record.executable_result_paths.append(exe_result_path)
+
+        _update_updatable_files(egr_record, shared_rt_info, data_dir)
+        print(f'\r{i + 1} / {n_executables}', end='', flush=True)
     print()
 
-    return exegroup_result
+    return ExecutableGroupResult(
+        runtime_configuration=rt_config,
+        shared_runtime_info=shared_rt_info,
+        executable_results=executable_results,
+    )

cirq-google/cirq_google/workflow/quantum_runtime_test.py

@@ -15,13 +15,12 @@
 import re
 import uuid
 from dataclasses import dataclass
-from typing import List
-
-import numpy as np
-import pytest
+from typing import List, cast, Any
 
 import cirq
 import cirq_google as cg
+import numpy as np
+import pytest
 from cirq_google.workflow._abstract_engine_processor_shim import AbstractEngineProcessorShim
 from cirq_google.workflow.quantum_executable_test import _get_quantum_executables, _get_example_spec
 
@@ -69,14 +68,23 @@ def test_executable_result():
     cg_assert_equivalent_repr(er)
 
 
-def _cg_read_json_gzip(fn):
-    def _testing_resolver(cirq_type: str):
-        if cirq_type == 'cirq.google.testing._MockEngineProcessor':
-            return _MockEngineProcessor
+def _testing_resolver(cirq_type: str):
+    if cirq_type == 'cirq.google.testing._MockEngineProcessor':
+        return _MockEngineProcessor
 
+
+def _cg_read_json_gzip(fn):
     return cirq.read_json_gzip(fn, resolvers=[_testing_resolver] + cirq.DEFAULT_RESOLVERS)
 
 
+@pytest.fixture
+def patch_cirq_default_resolvers():
+    backup = cirq.DEFAULT_RESOLVERS.copy()
+    cirq.DEFAULT_RESOLVERS.insert(0, _testing_resolver)
+    yield True
+    cirq.DEFAULT_RESOLVERS = backup
+
+
 def _assert_json_roundtrip(o, tmpdir):
     cirq.to_json_gzip(o, f'{tmpdir}/o.json')
     o2 = _cg_read_json_gzip(f'{tmpdir}/o.json')
@@ -128,28 +136,57 @@ def test_executable_group_result(tmpdir):
     _assert_json_roundtrip(egr, tmpdir)
 
 
-@pytest.mark.parametrize('run_id', ['unit_test_runid', None])
-def test_execute(tmpdir, run_id):
-    rt_config = cg.QuantumRuntimeConfiguration(processor=_MockEngineProcessor(), run_id=run_id)
+def test_egr_filesystem_record_repr():
+    egr_fs_record = cg.ExecutableGroupResultFilesystemRecord(
+        runtime_configuration_path='RuntimeConfiguration.json.gz',
+        shared_runtime_info_path='SharedRuntimeInfo.jzon.gz',
+        executable_result_paths=[
+            'ExecutableResult.1.json.gz',
+            'ExecutableResult.2.json.gz',
+        ],
+        run_id='my-run-id',
+    )
+    cg_assert_equivalent_repr(egr_fs_record)
+
+
+def _load_result_by_hand(tmpdir: str, run_id: str) -> cg.ExecutableGroupResult:
+    """Load `ExecutableGroupResult` "by hand" without using
+    `ExecutableGroupResultFilesystemRecord`."""
+    rt_config = cirq.read_json_gzip(f'{tmpdir}/{run_id}/QuantumRuntimeConfiguration.json.gz')
+    shared_rt_info = cirq.read_json_gzip(f'{tmpdir}/{run_id}/SharedRuntimeInfo.json.gz')
+    fns = glob.glob(f'{tmpdir}/{run_id}/ExecutableResult.*.json.gz')
+    fns = sorted(
+        fns,
+        key=lambda s: int(cast(Any, re.search(r'ExecutableResult\.(\d+)\.json\.gz$', s)).group(1)),
+    )
+    assert len(fns) == 3
+    exe_results: List[cg.ExecutableResult] = [cirq.read_json_gzip(fn) for fn in fns]
+    return cg.ExecutableGroupResult(
+        runtime_configuration=rt_config,
+        shared_runtime_info=shared_rt_info,
+        executable_results=exe_results,
+    )
+
+
+@pytest.mark.parametrize('run_id_in', ['unit_test_runid', None])
+def test_execute(tmpdir, run_id_in, patch_cirq_default_resolvers):
+    assert patch_cirq_default_resolvers
+    rt_config = cg.QuantumRuntimeConfiguration(processor=_MockEngineProcessor(), run_id=run_id_in)
     executable_group = cg.QuantumExecutableGroup(_get_quantum_executables())
     returned_exegroup_result = cg.execute(
         rt_config=rt_config, executable_group=executable_group, base_data_dir=tmpdir
     )
-    actual_run_id = returned_exegroup_result.shared_runtime_info.run_id
-    if run_id is not None:
-        assert run_id == actual_run_id
+    run_id = returned_exegroup_result.shared_runtime_info.run_id
+    if run_id_in is not None:
+        assert run_id_in == run_id
     else:
-        assert isinstance(uuid.UUID(actual_run_id), uuid.UUID)
-    fns = glob.glob(f'{tmpdir}/{actual_run_id}/ExecutableGroupResult.json.gz')
-    assert len(fns) == 1
-    exegroup_result: cg.ExecutableGroupResult = _cg_read_json_gzip(fns[0])
+        assert isinstance(uuid.UUID(run_id), uuid.UUID)
 
-    fns = glob.glob(f'{tmpdir}/{actual_run_id}/ExecutableResult.*.json.gz')
-    fns = sorted(
-        fns, key=lambda s: int(re.search(r'ExecutableResult\.(\d+)\.json\.gz$', s).group(1))
+    manual_exegroup_result = _load_result_by_hand(tmpdir, run_id)
+    egr_record: cg.ExecutableGroupResultFilesystemRecord = cirq.read_json_gzip(
+        f'{tmpdir}/{run_id}/ExecutableGroupResultFilesystemRecord.json.gz'
     )
-    assert len(fns) == 3
-    exe_results: List[cg.ExecutableResult] = [_cg_read_json_gzip(fn) for fn in fns]
+    exegroup_result: cg.ExecutableGroupResult = egr_record.load(base_data_dir=tmpdir)
 
-    exegroup_result.executable_results = exe_results
     assert returned_exegroup_result == exegroup_result
+    assert manual_exegroup_result == exegroup_result