Add get config functions

cirq-google/cirq_google/engine/__init__.py

@@ -95,3 +95,5 @@
 from cirq_google.engine.engine_result import EngineResult as EngineResult
 
 from cirq_google.engine.processor_sampler import ProcessorSampler as ProcessorSampler
+
+from cirq_google.engine.processor_config import ProcessorConfig as ProcessorConfig

cirq-google/cirq_google/engine/abstract_local_engine_test.py

@@ -81,6 +81,12 @@ def list_programs(
     def get_program(self, program_id: str) -> AbstractProgram:
         return self._programs[program_id]
 
+    def get_config_from_run(self, *args, **kwargs):
+        pass
+
+    def get_config_from_snapshot(self, *args, **kwargs):
+        pass
+
 
 class NothingEngine(AbstractLocalEngine):
     """Engine for Testing."""

cirq-google/cirq_google/engine/engine.py

@@ -43,6 +43,7 @@
     engine_job,
     engine_processor,
     engine_program,
+    processor_config,
     util,
 )
 from cirq_google.serialization import CIRCUIT_SERIALIZER, Serializer
@@ -635,6 +636,62 @@ def get_sampler(
             max_concurrent_jobs=max_concurrent_jobs,
         )
 
+    async def get_processor_config_from_snapshot_async(
+        self, processor_id: str, snapshot_id: str, config_name: str = 'default'
+    ) -> processor_config.ProcessorConfig | None:
+        """Returns a ProcessorConfig from this project and the given processor id.
+
+        Args:
+           processor_id: The processor unique identifier.
+           snapshot_id: The unique identifier for the snapshot.
+           config_name: The identifier for the config.
+
+        Returns:
+           The ProcessorConfig from this project and processor.
+        """
+        client = self.context.client
+        quantum_config = await client.get_quantum_processor_config_from_snapshot_async(
+            project_id=self.project_id,
+            processor_id=processor_id,
+            snapshot_id=snapshot_id,
+            config_name=config_name,
+        )
+        if quantum_config:
+            return processor_config.ProcessorConfig(quantum_processor_config=quantum_config)
+        return None
+
+    get_processor_config_from_snapshot = duet.sync(get_processor_config_from_snapshot_async)
+
+    async def get_processor_config_from_run_async(
+        self, processor_id: str, run_name: str = 'current', config_name: str = 'default'
+    ) -> processor_config.ProcessorConfig | None:
+        """Returns a ProcessorConfig from this project and the given processor id.
+
+        If no `run_name` and `config_name` are specified, the inernally configured default config
+        is returned.
+
+        Args:
+            processor_id: The processor unique identifier.
+            run_name: The unique identifier for the automation run.
+            config_name: The identifier for the config.
+
+        Returns:
+            The ProcessorConfig from this project and processor.
+        """
+        quantum_config = await self.context.client.get_quantum_processor_config_from_run_async(
+            project_id=self.project_id,
+            processor_id=processor_id,
+            run_name=run_name,
+            config_name=config_name,
+        )
+        if quantum_config:
+            return processor_config.ProcessorConfig(
+                quantum_processor_config=quantum_config, run_name=run_name
+            )
+        return None
+
+    get_processor_config_from_run = duet.sync(get_processor_config_from_run_async)
+
 
 def get_engine(project_id: str | None = None) -> Engine:
     """Get an Engine instance assuming some sensible defaults.

cirq-google/cirq_google/engine/engine_client.py

@@ -1181,6 +1181,76 @@ async def list_time_slots_async(
 
     list_time_slots = duet.sync(list_time_slots_async)
 
+    async def _get_quantum_processor_config(
+        self, name: str
+    ) -> quantum.QuantumProcessorConfig | None:
+        """Runs get_quantum_processor_config with the given resource name."""
+        try:
+            request = quantum.GetQuantumProcessorConfigRequest(name=name)
+            return await self._send_request_async(
+                self.grpc_client.get_quantum_processor_config, request
+            )
+        except EngineException as err:
+            if isinstance(err.__cause__, NotFound):
+                return None
+            raise
+
+    async def get_quantum_processor_config_from_snapshot_async(
+        self, project_id: str, processor_id: str, snapshot_id: str, config_name: str
+    ) -> quantum.QuantumProcessorConfig | None:
+        """Returns the QuantumProcessorConfig for the given snapshot id.
+
+        Args:
+            project_id: A project_id of the parent Google Cloud Project.
+            processor_id: The processor unique identifier.
+            snapshot_id: The id of the snapshot that contains the quantum processor config.
+            config_name: The id of the quantum processor config.
+
+        Returns:
+            The quantum procesor config or None if it does not exist.
+
+        Raises:
+            EngineException: If the request to get the config fails.
+        """
+        name = _quantum_processor_name_with_snapshot_id(
+            project_id=project_id,
+            processor_id=processor_id,
+            snapshot_id=snapshot_id,
+            config_name=config_name,
+        )
+        return await self._get_quantum_processor_config(name)
+
+    get_quantum_processor_config_from_snapshot = duet.sync(
+        get_quantum_processor_config_from_snapshot_async
+    )
+
+    async def get_quantum_processor_config_from_run_async(
+        self, project_id: str, processor_id: str, run_name: str, config_name: str
+    ) -> quantum.QuantumProcessorConfig | None:
+        """Returns the QuantumProcessorConfig for the given run_name.
+
+        Args:
+            project_id: A project_id of the parent Google Cloud Project.
+            processor_id: The processor unique identifier.
+            config_name: The id of the quantum processor config.
+            run_name: The run_name that contains the quantum processor config.
+
+        Returns:
+            The quantum procesor config or None if it does not exist.
+
+        Raises:
+            EngineException: If the request to get the config fails.
+        """
+        name = _quantum_processor_name_with_run_name(
+            project_id=project_id,
+            processor_id=processor_id,
+            run_name=run_name,
+            config_name=config_name,
+        )
+        return await self._get_quantum_processor_config(name)
+
+    get_quantum_processor_config_from_run = duet.sync(get_quantum_processor_config_from_run_async)
+
 
 def _project_name(project_id: str) -> str:
     return f'projects/{project_id}'
@@ -1230,6 +1300,28 @@ def _ids_from_calibration_name(calibration_name: str) -> tuple[str, str, int]:
     return parts[1], parts[3], int(parts[5])
 
 
+def _quantum_processor_name_with_snapshot_id(
+    project_id: str, processor_id: str, snapshot_id: str, config_name: str
+) -> str:
+    return (
+        f'projects/{project_id}/'
+        f'processors/{processor_id}/'
+        f'configSnapshots/{snapshot_id}/'
+        f'configs/{config_name}'
+    )
+
+
+def _quantum_processor_name_with_run_name(
+    project_id: str, processor_id: str, run_name: str, config_name: str
+) -> str:
+    return (
+        f'projects/{project_id}/'
+        f'processors/{processor_id}/'
+        f'configAutomationRuns/{run_name}/'
+        f'configs/{config_name}'
+    )
+
+
 def _date_or_time_to_filter_expr(param_name: str, param: datetime.datetime | datetime.date):
     """Formats datetime or date to filter expressions.
 

cirq-google/cirq_google/engine/engine_client_test.py

@@ -1743,3 +1743,144 @@ def test_list_time_slots(client_constructor, default_engine_client):
     grpc_client.list_quantum_time_slots.return_value = Pager(results)
 
     assert default_engine_client.list_time_slots('proj', 'processor0') == results
+
+
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+def test_get_quantum_processor_config_from_snapshot(client_constructor, default_engine_client):
+
+    project_id = "test_project_id"
+    processor_id = "test_processor_id"
+    snapshot_id = "test_snapshot_id"
+    config_name = "test_config_name"
+    resource_name = (
+        f'projects/{project_id}/'
+        f'processors/{processor_id}/'
+        f'configSnapshots/{snapshot_id}/'
+        f'configs/{config_name}'
+    )
+
+    grpc_client = _setup_client_mock(client_constructor)
+    expected_result = quantum.QuantumProcessorConfig(name=resource_name)
+    grpc_client.get_quantum_processor_config.return_value = expected_result
+
+    actual_result = default_engine_client.get_quantum_processor_config_from_snapshot(
+        project_id=project_id,
+        processor_id=processor_id,
+        config_name=config_name,
+        snapshot_id=snapshot_id,
+    )
+    grpc_client.get_quantum_processor_config.assert_called_with(
+        quantum.GetQuantumProcessorConfigRequest(name=resource_name)
+    )
+    assert actual_result == expected_result
+
+
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+def test_get_quantum_processor_config_from_snapshot_not_found(
+    client_constructor, default_engine_client
+):
+    project_id = "test_project_id"
+    processor_id = "test_processor_id"
+    snapshot_id = "test_snapshot_id"
+    config_name = "test_config_name"
+    resource_name = (
+        f'projects/{project_id}/'
+        f'processors/{processor_id}/'
+        f'configSnapshots/{snapshot_id}/'
+        f'configs/{config_name}'
+    )
+
+    grpc_client = _setup_client_mock(client_constructor)
+    grpc_client.get_quantum_processor_config.side_effect = exceptions.NotFound('not found')
+
+    actual_result = default_engine_client.get_quantum_processor_config_from_snapshot(
+        project_id=project_id,
+        processor_id=processor_id,
+        config_name=config_name,
+        snapshot_id=snapshot_id,
+    )
+    grpc_client.get_quantum_processor_config.assert_called_with(
+        quantum.GetQuantumProcessorConfigRequest(name=resource_name)
+    )
+    assert actual_result is None
+
+
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+def test_get_quantum_processor_config_from_snapshot_exception(
+    client_constructor, default_engine_client
+):
+    grpc_client = _setup_client_mock(client_constructor)
+    grpc_client.get_quantum_processor_config.side_effect = exceptions.BadRequest('invalid_reueust')
+
+    with pytest.raises(EngineException, match='invalid_reueust'):
+        _ = default_engine_client.get_quantum_processor_config_from_snapshot(
+            project_id="test_project_id",
+            processor_id="test_processor_id",
+            config_name="test_config_name",
+            snapshot_id="test_snapshot_id",
+        )
+
+
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+def test_get_quantum_processor_config_from_run(client_constructor, default_engine_client):
+
+    project_id = "test_project_id"
+    processor_id = "test_processor_id"
+    run_name = "test_run_name"
+    config_name = "test_config_name"
+    resource_name = (
+        f'projects/{project_id}/'
+        f'processors/{processor_id}/'
+        f'configAutomationRuns/{run_name}/'
+        f'configs/{config_name}'
+    )
+
+    grpc_client = _setup_client_mock(client_constructor)
+    expected_result = quantum.QuantumProcessorConfig(name=resource_name)
+    grpc_client.get_quantum_processor_config.return_value = expected_result
+
+    actual_result = default_engine_client.get_quantum_processor_config_from_run(
+        project_id=project_id, processor_id=processor_id, config_name=config_name, run_name=run_name
+    )
+    grpc_client.get_quantum_processor_config.assert_called_with(
+        quantum.GetQuantumProcessorConfigRequest(name=resource_name)
+    )
+    assert actual_result == expected_result
+
+
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+def test_get_quantum_processor_config_from_run_not_found(client_constructor, default_engine_client):
+    project_id = "test_project_id"
+    processor_id = "test_processor_id"
+    run_name = "test_run_name"
+    config_name = "test_config_name"
+    resource_name = (
+        f'projects/{project_id}/'
+        f'processors/{processor_id}/'
+        f'configAutomationRuns/{run_name}/'
+        f'configs/{config_name}'
+    )
+    grpc_client = _setup_client_mock(client_constructor)
+    grpc_client.get_quantum_processor_config.side_effect = exceptions.NotFound('not found')
+
+    actual_result = default_engine_client.get_quantum_processor_config_from_run(
+        project_id=project_id, processor_id=processor_id, config_name=config_name, run_name=run_name
+    )
+    grpc_client.get_quantum_processor_config.assert_called_with(
+        quantum.GetQuantumProcessorConfigRequest(name=resource_name)
+    )
+    assert actual_result is None
+
+
+@mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True)
+def test_get_quantum_processor_config_from_run_exception(client_constructor, default_engine_client):
+    grpc_client = _setup_client_mock(client_constructor)
+    grpc_client.get_quantum_processor_config.side_effect = exceptions.BadRequest('invalid_reueust')
+
+    with pytest.raises(EngineException, match='invalid_reueust'):
+        _ = default_engine_client.get_quantum_processor_config_from_run(
+            project_id="test_project_id",
+            processor_id="test_processor_id",
+            config_name="test_config_name",
+            run_name="test_run_name",
+        )