Check reservation status for a list of processors (#140)

* Add reservation status check function

* Check reservation wip

* Add tests for get_available_processors

* Break lines in comments

* Fix tests

* Use the correct processor id and convert QuantumTimeSlot to EngineTimeSlot

* address comments

* Add test for other project id

* Add TODO for time range parsing problem

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 (

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

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']) == []

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,