Refactor get sampler

cirq-google/cirq_google/engine/abstract_processor.py

@@ -180,17 +180,8 @@ async def run_sweep_async(
     run_sweep = duet.sync(run_sweep_async)
 
     @abc.abstractmethod
-    def get_sampler(self, run_name: str = "", device_config_name: str = "") -> cg.ProcessorSampler:
-        """Returns a sampler backed by the processor.
-
-        Args:
-            run_name: A unique identifier representing an automation run for the
-                processor. An Automation Run contains a collection of device
-                configurations for the processor.
-            device_config_name: An identifier used to select the processor configuration
-                utilized to run the job. A configuration identifies the set of
-                available qubits, couplers, and supported gates in the processor.
-        """
+    def get_sampler(self) -> cg.ProcessorSampler:
+        """Returns a sampler backed by the processor."""
 
     @abc.abstractmethod
     def engine(self) -> abstract_engine.AbstractEngine | None:

cirq-google/cirq_google/engine/engine.py

@@ -593,23 +593,18 @@ def get_processor(self, processor_id: str) -> engine_processor.EngineProcessor:
     def get_sampler(
         self,
         processor_id: str | list[str],
-        run_name: str = "",
-        device_config_name: str = "",
-        snapshot_id: str = "",
+        device_config_name: str | None = None,
+        device_version: processor_config.DeviceVersion | None = None,
         max_concurrent_jobs: int = 100,
     ) -> cirq_google.ProcessorSampler:
         """Returns a sampler backed by the engine.
 
         Args:
             processor_id: String identifier of which processor should be used to sample.
-            run_name: A unique identifier representing an automation run for the
-                processor. An Automation Run contains a collection of device
-                configurations for the processor.
             device_config_name: An identifier used to select the processor configuration
                 utilized to run the job. A configuration identifies the set of
                 available qubits, couplers, and supported gates in the processor.
-            snapshot_id: A unique identifier for an immutable snapshot reference. A
-                snapshot contains a collection of device configurations for the processor.
+            device_version: Specifies either the snapshot_id or the run_name.
             max_concurrent_jobs: The maximum number of jobs to be sent
                 concurrently to the Engine. This client-side throttle can be
                 used to proactively reduce load to the backends and avoid quota
@@ -629,41 +624,18 @@ def get_sampler(
                 'to get_sampler() no longer supported. Use Engine.run() instead if '
                 'you need to specify a list.'
             )
+
         return self.get_processor(processor_id).get_sampler(
-            run_name=run_name,
             device_config_name=device_config_name,
-            snapshot_id=snapshot_id,
+            device_version=device_version,
             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'
+    async def get_processor_config_async(
+        self,
+        processor_id: str,
+        device_version: processor_config.DeviceVersion = processor_config.Run(id='current'),
+        config_name: str = 'default',
     ) -> processor_config.ProcessorConfig | None:
         """Returns a ProcessorConfig from this project and the given processor id.
 
@@ -672,25 +644,27 @@ async def get_processor_config_from_run_async(
 
         Args:
             processor_id: The processor unique identifier.
-            run_name: The unique identifier for the automation run.
+            device_version: Specifies either the snapshot_id or the run_name.
             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(
+        quantum_config = await self.context.client.get_quantum_processor_config_async(
             project_id=self.project_id,
             processor_id=processor_id,
-            run_name=run_name,
+            device_version=device_version,
             config_name=config_name,
         )
         if quantum_config:
             return processor_config.ProcessorConfig(
-                quantum_processor_config=quantum_config, run_name=run_name
+                processor=self.get_processor(processor_id),
+                quantum_processor_config=quantum_config,
+                device_version=device_version,
             )
         return None
 
-    get_processor_config_from_run = duet.sync(get_processor_config_from_run_async)
+    get_processor_config = duet.sync(get_processor_config_async)
 
 
 def get_engine(project_id: str | None = None) -> Engine:

cirq-google/cirq_google/engine/engine_client.py

@@ -30,6 +30,7 @@
 from cirq_google.cloud import quantum
 from cirq_google.engine import stream_manager
 from cirq_google.engine.asyncio_executor import AsyncioExecutor
+from cirq_google.engine.processor_config import DeviceVersion, Run, Snapshot
 
 _M = TypeVar('_M', bound=proto.Message)
 _R = TypeVar('_R')
@@ -1181,29 +1182,19 @@ 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
+    async def get_quantum_processor_config_async(
+        self,
+        project_id: str,
+        processor_id: str,
+        config_name: str = 'default',
+        device_version: DeviceVersion = Run(id='current'),
     ) -> 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.
+            device_version: Specifies either the snapshot_id or the run_name.
             config_name: The id of the quantum processor config.
 
         Returns:
@@ -1212,44 +1203,23 @@ async def get_quantum_processor_config_from_snapshot_async(
         Raises:
             EngineException: If the request to get the config fails.
         """
-        name = _quantum_processor_config_name_from_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_config_name_from_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)
+        try:
+            name = _quantum_processor_config_name_from_device_version(
+                project_id=project_id,
+                processor_id=processor_id,
+                config_name=config_name,
+                device_version=device_version,
+            )
+            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
 
-    get_quantum_processor_config_from_run = duet.sync(get_quantum_processor_config_from_run_async)
+    get_quantum_processor_config = duet.sync(get_quantum_processor_config_async)
 
 
 def _project_name(project_id: str) -> str:
@@ -1300,22 +1270,23 @@ def _ids_from_calibration_name(calibration_name: str) -> tuple[str, str, int]:
     return parts[1], parts[3], int(parts[5])
 
 
-def _quantum_processor_config_name_from_snapshot_id(
-    project_id: str, processor_id: str, snapshot_id: str, config_name: str
-) -> str:
-    return (
-        f'{_processor_name_from_ids(project_id, processor_id)}/'
-        f'configSnapshots/{snapshot_id}/'
-        f'configs/{config_name}'
-    )
-
-
-def _quantum_processor_config_name_from_run_name(
-    project_id: str, processor_id: str, run_name: str, config_name: str
+def _quantum_processor_config_name_from_device_version(
+    project_id: str,
+    processor_id: str,
+    config_name: str,
+    device_version: DeviceVersion | None = None,
 ) -> str:
+    processor_resource_name = _processor_name_from_ids(project_id, processor_id)
+    if isinstance(device_version, Snapshot):
+        return (
+            f'{processor_resource_name}/'
+            f'configSnapshots/{device_version.id}/'
+            f'configs/{config_name}'
+        )
+    default_run_name = 'default'
     return (
-        f'{_processor_name_from_ids(project_id, processor_id)}/'
-        f'configAutomationRuns/{run_name}/'
+        f'{processor_resource_name}/'
+        f'configAutomationRuns/{device_version.id if device_version else default_run_name}/'
         f'configs/{config_name}'
     )