Skip to content

Commit

Permalink
Add jobs filtering by executed and scheduled processor IDs (#4589)
Browse files Browse the repository at this point in the history
* Add jobs filtering by executed and schedule processor IDs

* Filter by any of provided scheduled processor IDs

* Make executed processor IDs to be passed as list
  • Loading branch information
aleksandrpak authored Oct 26, 2021
1 parent 009b3da commit ec9fc45
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 8 deletions.
17 changes: 17 additions & 0 deletions cirq-google/cirq_google/engine/engine_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -367,6 +367,8 @@ def list_jobs(
created_after: Optional[Union[datetime.datetime, datetime.date]] = None,
has_labels: Optional[Dict[str, str]] = None,
execution_states: Optional[Set[quantum.enums.ExecutionStatus.State]] = None,
executed_processor_ids: Optional[List[str]] = None,
scheduled_processor_ids: Optional[List[str]] = None,
):
"""Returns the list of jobs for a given program.
Expand All @@ -390,6 +392,11 @@ def list_jobs(
execution_states: retrieve jobs that have an execution state that
is contained in `execution_states`. See
`quantum.enums.ExecutionStatus.State` enum for accepted values.
executed_processor_ids: filters jobs by processor ID used for
execution. Matches any of provided IDs.
scheduled_processor_ids: filters jobs by any of provided
scheduled processor IDs.
"""
filters = []

Expand All @@ -407,6 +414,16 @@ def list_jobs(
for execution_state in execution_states:
state_filter.append(f"execution_status.state = {execution_state.name}")
filters.append(f"({' OR '.join(state_filter)})")
if executed_processor_ids is not None:
ids_filter = []
for processor_id in executed_processor_ids:
ids_filter.append(f"executed_processor_id = {processor_id}")
filters.append(f"({' OR '.join(ids_filter)})")
if scheduled_processor_ids is not None:
ids_filter = []
for processor_id in scheduled_processor_ids:
ids_filter.append(f"scheduled_processor_ids: {processor_id}")
filters.append(f"({' OR '.join(ids_filter)})")

if program_id is None:
program_id = "-"
Expand Down
62 changes: 54 additions & 8 deletions cirq-google/cirq_google/engine/engine_client_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -664,9 +664,13 @@ def test_list_jobs(client_constructor):
'created_after, '
'created_before, '
'labels, '
'execution_states',
'execution_states, '
'executed_processor_ids, '
'scheduled_processor_ids, ',
[
('',
None,
None,
None,
None,
None,
Expand All @@ -675,23 +679,31 @@ def test_list_jobs(client_constructor):
datetime.date(2020, 9, 1),
None,
None,
None,
None,
None),
('create_time >= 1598918400',
datetime.datetime(2020, 9, 1, 0, 0, 0,
tzinfo=datetime.timezone.utc),
None,
None,
None,
None,
None),
('create_time <= 2020-10-01',
None,
datetime.date(2020, 10, 1),
None,
None,
None,
None),
('create_time >= 2020-09-01 AND create_time <= 1598918410',
datetime.date(2020, 9, 1),
datetime.datetime(2020, 9, 1, 0, 0, 10,
tzinfo=datetime.timezone.utc),
None,
None,
None,
None),
('labels.color:red AND labels.shape:*',
None,
Expand All @@ -700,16 +712,18 @@ def test_list_jobs(client_constructor):
'color': 'red',
'shape': '*'
},
None
),
None,
None,
None),
('(execution_status.state = FAILURE OR '
'execution_status.state = CANCELLED)',
None,
None,
None,
[quantum.enums.ExecutionStatus.State.FAILURE,
quantum.enums.ExecutionStatus.State.CANCELLED,]
),
quantum.enums.ExecutionStatus.State.CANCELLED,],
None,
None),
('create_time >= 2020-08-01 AND '
'create_time <= 1598918400 AND '
'labels.color:red AND labels.shape:* AND '
Expand All @@ -721,20 +735,52 @@ def test_list_jobs(client_constructor):
'shape': '*'
},
[quantum.enums.ExecutionStatus.State.SUCCESS,],
),
None,
None),
('(executed_processor_id = proc1)',
None,
None,
None,
None,
['proc1'],
None),
('(executed_processor_id = proc1 OR executed_processor_id = proc2)',
None,
None,
None,
None,
['proc1', 'proc2'],
None),
('(scheduled_processor_ids: proc1)',
None,
None,
None,
None,
None,
['proc1']),
('(scheduled_processor_ids: proc1 OR scheduled_processor_ids: proc2)',
None,
None,
None,
None,
None,
['proc1', 'proc2']),
])
# yapf: enable
@mock.patch.object(quantum, 'QuantumEngineServiceClient', autospec=True)
def test_list_jobs_filters(client_constructor, expected_filter, created_before,
created_after, labels, execution_states):
created_after, labels, execution_states,
executed_processor_ids, scheduled_processor_ids):
grpc_client = setup_mock_(client_constructor)
client = EngineClient()
client.list_jobs(project_id='proj',
program_id='prog',
created_before=created_before,
created_after=created_after,
has_labels=labels,
execution_states=execution_states)
execution_states=execution_states,
executed_processor_ids=executed_processor_ids,
scheduled_processor_ids=scheduled_processor_ids)
assert grpc_client.list_quantum_jobs.call_args[1] == {
'filter_': expected_filter,
}
Expand Down

0 comments on commit ec9fc45

Please sign in to comment.