Return processor_id and project_id in get_qcs_objects_for_notebooks (#5759)

This adds two fields in the QCSObjectForNotebook

This will allow us some flexibility in google notebooks so as not to specify the processor_id.  This will also be a step in enabling tests for these notebooks.

Author: dstrain115

cirq-google/cirq_google/engine/abstract_local_processor.py

@@ -99,6 +99,11 @@ def __init__(
             if self._schedule[idx].end_time > self._schedule[idx + 1].start_time:
                 raise ValueError('Time slots cannot overlap!')
 
+    @property
+    def project_id(self) -> str:
+        """Project name of the processor."""
+        return self._project_name
+
     @property
     def processor_id(self) -> str:
         """Unique string id of the processor."""

cirq-google/cirq_google/engine/qcs_notebook.py

@@ -13,38 +13,57 @@
 # limitations under the License.
 
 import dataclasses
-from typing import Union, Optional
+from typing import cast, Optional, Sequence, Union
 
 import cirq
 
-from cirq_google import (
-    PhasedFSimEngineSimulator,
-    ProcessorSampler,
-    Sycamore,
-    SQRT_ISWAP_INV_PARAMETERS,
-    PhasedFSimCharacterization,
-    get_engine,
+from cirq_google import ProcessorSampler, get_engine
+from cirq_google.engine import (
+    AbstractEngine,
+    AbstractProcessor,
+    AbstractLocalProcessor,
+    create_noiseless_virtual_engine_from_latest_templates,
+    EngineProcessor,
 )
 
 
 @dataclasses.dataclass
 class QCSObjectsForNotebook:
+    """All the objects you might need to run a notbook with QCS.
+
+    Contains an (Abstract) Engine, Processor, Device, and Sampler,
+    as well as associated meta-data signed_in, processor_id, and project_id.
+
+    This removes the need for boiler plate in notebooks, and provides a
+    central place to handle the various environments (testing vs production),
+    (stand-alone vs colab vs jupyter).
+    """
+
+    engine: AbstractEngine
+    processor: AbstractProcessor
     device: cirq.Device
-    sampler: Union[PhasedFSimEngineSimulator, ProcessorSampler]
+    sampler: ProcessorSampler
     signed_in: bool
-
-    @property
-    def is_simulator(self):
-        return isinstance(self.sampler, PhasedFSimEngineSimulator)
+    processor_id: Optional[str]
+    project_id: Optional[str]
+    is_simulator: bool
 
 
-# Disable missing-raises-doc lint check, since pylint gets confused
-# by exceptions that are raised and caught within this function.
-# pylint: disable=missing-raises-doc
 def get_qcs_objects_for_notebook(
-    project_id: Optional[str] = None, processor_id: Optional[str] = None
-) -> QCSObjectsForNotebook:  # pragma: nocover
-    """Authenticates on Google Cloud, can return a Device and Simulator.
+    project_id: Optional[str] = None, processor_id: Optional[str] = None, virtual=False
+) -> QCSObjectsForNotebook:
+    """Authenticates on Google Cloud and returns Engine related objects.
+
+    This function will authenticate to Google Cloud and attempt to
+    instantiate an Engine object.  If it does not succeed, it will instead
+    return a virtual AbstractEngine that is backed by a noisy simulator.
+    This function is designed for maximum versatility and
+    to work in colab notebooks, as a stand-alone, and in tests.
+
+    Note that, if you are using this to connect to QCS and do not care about
+    the added versatility, you may want to use `cirq_google.get_engine()` or
+    `cirq_google.Engine()` instead to guarantee the use of a production instance
+    and to avoid accidental use of a noisy simulator.
 
     Args:
         project_id: Optional explicit Google Cloud project id. Otherwise,
@@ -53,9 +72,14 @@ def get_qcs_objects_for_notebook(
             personal project IDs in shared code.
         processor_id: Engine processor ID (from Cloud console or
             ``Engine.list_processors``).
+        virtual: If set to True, will create a noisy virtual Engine instead.
+            This is useful for testing and simulation.
 
     Returns:
-        An instance of DeviceSamplerInfo.
+        An instance of QCSObjectsForNotebook which contains all the objects .
+
+    Raises:
+        ValueError: if processor_id is not specified and no processors are available.
     """
 
     # Check for Google Application Default Credentials and run
@@ -80,32 +104,46 @@ def get_qcs_objects_for_notebook(
             print(f"Authentication failed: {exc}")
 
     # Attempt to connect to the Quantum Engine API, and use a simulator if unable to connect.
-    sampler: Union[PhasedFSimEngineSimulator, ProcessorSampler]
-    try:
-        engine = get_engine(project_id)
-        if processor_id:
-            processor = engine.get_processor(processor_id)
-        else:
-            processors = engine.list_processors()
-            if not processors:
-                raise ValueError("No processors available.")
-            processor = processors[0]
-            print(f"Available processors: {[p.processor_id for p in processors]}")
-            print(f"Using processor: {processor.processor_id}")
-        device = processor.get_device()
-        sampler = processor.get_sampler()
-        signed_in = True
-    except Exception as exc:
-        print(f"Unable to connect to quantum engine: {exc}")
-        print("Using a noisy simulator.")
-        sampler = PhasedFSimEngineSimulator.create_with_random_gaussian_sqrt_iswap(
-            mean=SQRT_ISWAP_INV_PARAMETERS,
-            sigma=PhasedFSimCharacterization(theta=0.01, zeta=0.10, chi=0.01, gamma=0.10, phi=0.02),
-        )
-        device = Sycamore
+    if virtual:
+        engine: AbstractEngine = create_noiseless_virtual_engine_from_latest_templates()
         signed_in = False
-
-    return QCSObjectsForNotebook(device=device, sampler=sampler, signed_in=signed_in)
-
-
-# pylint: enable=missing-raises-doc
+        is_simulator = True
+    else:
+        try:
+            engine = get_engine(project_id)
+            signed_in = True
+            is_simulator = False
+        except Exception as exc:
+            print(f"Unable to connect to quantum engine: {exc}")
+            print("Using a noisy simulator.")
+            engine = create_noiseless_virtual_engine_from_latest_templates()
+            signed_in = False
+            is_simulator = True
+    if processor_id:
+        processor = engine.get_processor(processor_id)
+    else:
+        # All of these are either local processors or engine processors
+        # Either way, tell mypy they have a processor_id field.
+        processors = cast(
+            Sequence[Union[EngineProcessor, AbstractLocalProcessor]], engine.list_processors()
+        )
+        if not processors:
+            raise ValueError("No processors available.")
+        processor = processors[0]
+        processor_id = processor.processor_id
+        print(f"Available processors: {[p.processor_id for p in processors]}")
+        print(f"Using processor: {processor_id}")
+    if not project_id:
+        project_id = getattr(processor, 'project_id', None)
+    device = processor.get_device()
+    sampler = processor.get_sampler()
+    return QCSObjectsForNotebook(
+        engine=engine,
+        processor=processor,
+        device=device,
+        sampler=sampler,
+        signed_in=signed_in,
+        project_id=project_id,
+        processor_id=processor_id,
+        is_simulator=is_simulator,
+    )

cirq-google/cirq_google/engine/qcs_notebook_test.py

@@ -12,16 +12,86 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import unittest.mock as mock
+import pytest
+
 import cirq_google as cg
-from cirq_google.engine.qcs_notebook import get_qcs_objects_for_notebook
+from cirq_google.engine.qcs_notebook import get_qcs_objects_for_notebook, QCSObjectsForNotebook
 
 
-def test_get_device_sampler():
-    result = get_qcs_objects_for_notebook()
-    assert result.device is cg.Sycamore
+def _assert_correct_types(result: QCSObjectsForNotebook):
+    assert isinstance(result.device, cg.GridDevice)
+    assert isinstance(result.sampler, cg.ProcessorSampler)
+    assert isinstance(result.engine, cg.engine.AbstractEngine)
+    assert isinstance(result.processor, cg.engine.AbstractProcessor)
+
+
+def _assert_simulated_values(result: QCSObjectsForNotebook):
     assert not result.signed_in
-    assert isinstance(result.sampler, cg.PhasedFSimEngineSimulator)
     assert result.is_simulator
+    assert result.project_id == 'fake_project'
 
-    result = get_qcs_objects_for_notebook("", "")
-    assert not result.signed_in
+
+def test_get_qcs_objects_for_notebook_virtual():
+    result = get_qcs_objects_for_notebook(virtual=True)
+    _assert_correct_types(result)
+    _assert_simulated_values(result)
+    assert result.processor_id == 'rainbow'
+    assert len(result.device.metadata.qubit_set) == 23
+
+    result = get_qcs_objects_for_notebook(processor_id='weber', virtual=True)
+    _assert_correct_types(result)
+    _assert_simulated_values(result)
+    assert result.processor_id == 'weber'
+    assert len(result.device.metadata.qubit_set) == 53
+
+
+@mock.patch('cirq_google.engine.qcs_notebook.get_engine')
+def test_get_qcs_objects_for_notebook_mocked_engine_fails(engine_mock):
+    """Tests creating an engine object which fails."""
+    engine_mock.side_effect = EnvironmentError('This is a mock, not real credentials.')
+    result = get_qcs_objects_for_notebook()
+    _assert_correct_types(result)
+    _assert_simulated_values(result)
+
+
+@mock.patch('cirq_google.engine.qcs_notebook.get_engine')
+def test_get_qcs_objects_for_notebook_mocked_engine_succeeds(engine_mock):
+    """Uses a mocked engine call to test a 'prod' Engine."""
+    fake_processor = cg.engine.SimulatedLocalProcessor(
+        processor_id='tester', project_name='mock_project', device=cg.Sycamore
+    )
+    fake_processor2 = cg.engine.SimulatedLocalProcessor(
+        processor_id='tester23', project_name='mock_project', device=cg.Sycamore23
+    )
+    fake_engine = cg.engine.SimulatedLocalEngine([fake_processor, fake_processor2])
+    engine_mock.return_value = fake_engine
+
+    result = get_qcs_objects_for_notebook()
+    _assert_correct_types(result)
+    assert result.signed_in
+    assert not result.is_simulator
+    assert result.project_id == 'mock_project'
+    assert len(result.device.metadata.qubit_set) == 54
+
+    result = get_qcs_objects_for_notebook(processor_id='tester')
+    _assert_correct_types(result)
+    assert result.signed_in
+    assert not result.is_simulator
+    assert result.project_id == 'mock_project'
+    assert len(result.device.metadata.qubit_set) == 54
+
+    result = get_qcs_objects_for_notebook(processor_id='tester23')
+    _assert_correct_types(result)
+    assert result.signed_in
+    assert not result.is_simulator
+    assert result.project_id == 'mock_project'
+    assert len(result.device.metadata.qubit_set) == 23
+
+
+@mock.patch('cirq_google.engine.qcs_notebook.get_engine')
+def test_get_qcs_objects_for_notebook_no_processors(engine_mock):
+    fake_engine = cg.engine.SimulatedLocalEngine([])
+    engine_mock.return_value = fake_engine
+    with pytest.raises(ValueError, match='processors'):
+        _ = get_qcs_objects_for_notebook()