diff --git a/recirq/__init__.py b/recirq/__init__.py index cd8944ce..10a128c6 100644 --- a/recirq/__init__.py +++ b/recirq/__init__.py @@ -36,6 +36,7 @@ get_processor_id_by_device_name, get_sampler_by_name, execute_in_queue, + get_available_processors ) from recirq.documentation_utils import ( diff --git a/recirq/engine_utils.py b/recirq/engine_utils.py index e4c833a9..f5206f17 100644 --- a/recirq/engine_utils.py +++ b/recirq/engine_utils.py @@ -16,6 +16,7 @@ import math import os import uuid +import datetime from dataclasses import dataclass, field from typing import List, Any, Optional, Callable, Dict, Union @@ -25,6 +26,8 @@ from cirq import work, study, circuits, ops from cirq.google import devices as cg_devices, gate_sets, engine as cg_engine from cirq.google.engine.engine_job import TERMINAL_STATES +from cirq.google.engine.client.quantum_v1alpha1.gapic import enums +from cirq.google import EngineTimeSlot def _get_program_id(program: Any): @@ -372,3 +375,46 @@ async def worker(): await queue.join() for wjob in worker_jobs: wjob.cancel() + + +def _get_current_time(): + return datetime.datetime.now() + + +def get_available_processors(processor_names: List[str]): + """Returns a list of available processors. + + Checks the reservation status of the processors and returns a list of + processors that are available to run on at the present time. + + Args: + processor_names: A list of processor names which are keys from QUANTUM_PROCESSORS. + """ + project_id = os.environ['GOOGLE_CLOUD_PROJECT'] + engine = cirq.google.get_engine() + available_processors = [] + current_time = _get_current_time() + for processor_name in processor_names: + processor_id = get_processor_id_by_device_name(processor_name) + if processor_id is None: + # Skip the check if this is a simulator. + continue + processor = engine.get_processor(processor_id) + for time_slot in processor.get_schedule(): + try: + time_slot = EngineTimeSlot.from_proto(time_slot) + # Parsing the end_time of the last time range in the schedule might give throw an error. + # TODO: remove this check once it is fixed in cirq. + except ValueError as e: + continue + # Ignore time slots that do not contain the current time. + if time_slot.start_time < current_time < time_slot.end_time: + # Time slots need to be either in OPEN_SWIM or reserved by the + # current project to be considered available. + if (time_slot.slot_type + == enums.QuantumTimeSlot.TimeSlotType.OPEN_SWIM + ) or (time_slot.slot_type + == enums.QuantumTimeSlot.TimeSlotType.RESERVATION and + time_slot.project_id == project_id): + available_processors.append(processor_name) + return available_processors diff --git a/recirq/engine_utils_test.py b/recirq/engine_utils_test.py index 4e7f4230..6b430f98 100644 --- a/recirq/engine_utils_test.py +++ b/recirq/engine_utils_test.py @@ -14,14 +14,21 @@ import os import uuid +from datetime import datetime import pytest +from unittest.mock import patch +from google.protobuf.timestamp_pb2 import Timestamp import cirq +from cirq.google.engine import EngineTimeSlot +from cirq.google.engine.client.quantum_v1alpha1.gapic import enums +from cirq.google.engine.client.quantum_v1alpha1 import types as qtypes import recirq from recirq.engine_utils import _get_program_id + def test_get_program_id(): circuit = cirq.Circuit(cirq.H(cirq.LineQubit(0))) prog_id = _get_program_id(circuit) @@ -108,3 +115,75 @@ def test_sampler_by_name(): assert isinstance(recirq.get_sampler_by_name('Syc23-zeros', gateset='sqrt-iswap'), recirq.ZerosSampler) + + +@patch('cirq.google.engine.engine_client.quantum.QuantumEngineServiceClient') +@patch('recirq.engine_utils._get_current_time') +@patch('cirq.google.engine.EngineProcessor.get_schedule') +def test_get_available_processors_open_swim_in_time_window(schedule_mock, time_mock, engine_mock): + os.environ['GOOGLE_CLOUD_PROJECT'] = 'some_project' + schedule_mock.return_value = [ + qtypes.QuantumTimeSlot( + processor_name='Sycamore23', + start_time=Timestamp(seconds=100), + end_time=Timestamp(seconds=500), + slot_type=enums.QuantumTimeSlot.TimeSlotType.OPEN_SWIM) + ] + time_mock.return_value = datetime.fromtimestamp(300) + assert 'Sycamore23' in recirq.get_available_processors(['Sycamore23']) + + +@patch('cirq.google.engine.engine_client.quantum.QuantumEngineServiceClient') +@patch('recirq.engine_utils._get_current_time') +@patch('cirq.google.engine.EngineProcessor.get_schedule') +def test_get_available_processors_open_swim_outside_window(schedule_mock, time_mock, engine_mock): + os.environ['GOOGLE_CLOUD_PROJECT'] = 'some_project' + schedule_mock.return_value = [ + qtypes.QuantumTimeSlot( + processor_name='Sycamore23', + start_time=Timestamp(seconds=300), + end_time=Timestamp(seconds=500), + slot_type=enums.QuantumTimeSlot.TimeSlotType.OPEN_SWIM) + ] + time_mock.return_value = datetime.fromtimestamp(700) + assert recirq.get_available_processors(['Sycamore23']) == [] + + +@patch('cirq.google.engine.engine_client.quantum.QuantumEngineServiceClient') +@patch('recirq.engine_utils._get_current_time') +@patch('cirq.google.engine.EngineProcessor.get_schedule') +def test_get_available_processors_current_project_reservation( + schedule_mock, time_mock, engine_mock): + os.environ['GOOGLE_CLOUD_PROJECT'] = 'some_project' + schedule_mock.return_value = [ + qtypes.QuantumTimeSlot( + processor_name='Sycamore23', + start_time=Timestamp(seconds=100), + end_time=Timestamp(seconds=500), + slot_type=enums.QuantumTimeSlot.TimeSlotType.RESERVATION, + reservation_config=qtypes.QuantumTimeSlot.ReservationConfig( + project_id='some_project'), + ) + ] + time_mock.return_value = datetime.fromtimestamp(300) + assert 'Sycamore23' in recirq.get_available_processors(['Sycamore23']) + + +@patch('cirq.google.engine.engine_client.quantum.QuantumEngineServiceClient') +@patch('recirq.engine_utils._get_current_time') +@patch('cirq.google.engine.EngineProcessor.get_schedule') +def test_get_available_processors_other_project_reservation( + schedule_mock, time_mock, engine_mock): + os.environ['GOOGLE_CLOUD_PROJECT'] = 'some_project' + schedule_mock.return_value = [ + qtypes.QuantumTimeSlot( + processor_name='Sycamore23', + start_time=Timestamp(seconds=100), + end_time=Timestamp(seconds=500), + slot_type=enums.QuantumTimeSlot.TimeSlotType.RESERVATION, + reservation_config=qtypes.QuantumTimeSlot.ReservationConfig( + project_id='other_project'), + ) + ] + time_mock.return_value = datetime.fromtimestamp(300) + assert recirq.get_available_processors(['Sycamore23']) == [] diff --git a/recirq/quantum_chess/experiments/batch_moves.py b/recirq/quantum_chess/experiments/batch_moves.py index 00935e99..f22bb089 100644 --- a/recirq/quantum_chess/experiments/batch_moves.py +++ b/recirq/quantum_chess/experiments/batch_moves.py @@ -11,8 +11,9 @@ by recirq.quantum_chess.move (see also interactive_board for examples). -PROCESSOR_NAME is a processor name from engine_utils.py. -Defaults to a 54 qubit sycamore noiseless simulator. +PROCESSOR_NAME is a processor name from engine_utils.py. If empty, it will find +an available quantum processor to run on. Default to 54 qubit sycamore +noiseless simulator if none are available. FEN is a initial position in chess FEN notation. Optional. Default is the normal classical chess starting position. @@ -52,7 +53,17 @@ def apply_moves(b: ab.AsciiBoard, moves: List[str]) -> bool: def main_loop(args): f = open(args.filename, 'r') moves = [line.strip() for line in f] - board = create_board(processor_name=args.processor_name, + if args.processor_name: + processor_name = args.processor_name + else: + # Execute on a quantum processor if it is available. + available_processors = utils.get_available_processors(utils.QUANTUM_PROCESSORS.keys()) + if available_processors: + processor_name = available_processors[0] + else: + processor_name = 'Syc54-noiseless' + print(f'Using processor {processor_name}') + board = create_board(processor_name=processor_name, noise_mitigation=0.1) b = ab.AsciiBoard(board=board) if args.position: @@ -72,7 +83,6 @@ def parse(): help='path to file that contains one move per line') parser.add_argument('--processor_name', type=str, - default='Syc54-noiseless', help='name of the QuantumProcessor object to use') parser.add_argument('--position', type=str,