Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding Generic submission option. #961

Merged
merged 6 commits into from
May 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions src/bloqade/builder/backend/quera.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,3 +126,14 @@ def mock(self, state_file: str = ".mock_state.txt", submission_error: bool = Fal
return self.parse().quera.mock(
state_file=state_file, submission_error=submission_error
)

def custom(self):
"""
Specify custom backend

Return:
CustomSubmissionRoutine

"""

return self.parse().quera.custom()
121 changes: 119 additions & 2 deletions src/bloqade/ir/routine/quera.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from collections import OrderedDict
from collections import OrderedDict, namedtuple
import time
from pydantic.v1.dataclasses import dataclass
import json

Expand All @@ -10,8 +11,9 @@
from bloqade.task.batch import RemoteBatch
from bloqade.task.quera import QuEraTask

from beartype.typing import Tuple, Union, Optional
from beartype.typing import Tuple, Union, Optional, NamedTuple, List, Dict, Any
from beartype import beartype
from requests import Response, request


@dataclass(frozen=True, config=__pydantic_dataclass_config__)
Expand Down Expand Up @@ -41,6 +43,121 @@
backend = MockBackend(state_file=state_file, submission_error=submission_error)
return QuEraHardwareRoutine(self.source, self.circuit, self.params, backend)

def custom(self) -> "CustomSubmissionRoutine":
return CustomSubmissionRoutine(self.source, self.circuit, self.params)


@dataclass(frozen=True, config=__pydantic_dataclass_config__)
class CustomSubmissionRoutine(RoutineBase):
def _compile(
self,
shots: int,
args: Tuple[LiteralType, ...] = (),
):
from bloqade.compiler.passes.hardware import (
analyze_channels,
canonicalize_circuit,
assign_circuit,
validate_waveforms,
generate_ahs_code,
generate_quera_ir,
)
from bloqade.submission.capabilities import get_capabilities

circuit, params = self.circuit, self.params
capabilities = get_capabilities()

for batch_params in params.batch_assignments(*args):
assignments = {**batch_params, **params.static_params}
final_circuit, metadata = assign_circuit(circuit, assignments)

level_couplings = analyze_channels(final_circuit)
final_circuit = canonicalize_circuit(final_circuit, level_couplings)

validate_waveforms(level_couplings, final_circuit)
ahs_components = generate_ahs_code(
capabilities, level_couplings, final_circuit
)

task_ir = generate_quera_ir(ahs_components, shots).discretize(capabilities)
MetaData = namedtuple("MetaData", metadata.keys())

yield MetaData(**metadata), task_ir

def submit(
self,
shots: int,
url: str,
json_body_template: str,
method: str = "POST",
args: Tuple[LiteralType] = (),
request_options: Dict[str, Any] = {},
sleep_time: float = 0.1,
) -> List[Tuple[NamedTuple, Response]]:
"""Compile to QuEraTaskSpecification and submit to a custom service.

Args:
shots (int): number of shots
url (str): url of the custom service
json_body_template (str): json body template, must contain '{task_ir}'
which is a placeholder for a string representation of the task ir.
The task ir string will be inserted into the template with
`json_body_template.format(task_ir=task_ir_string)`.
to be replaced by QuEraTaskSpecification
method (str): http method to be used. Defaults to "POST".
args (Tuple[LiteralType]): additional arguments to be passed into the
compiler coming from `args` option of the build. Defaults to ().
request_options: additional options to be passed into the request method,
Note the `data` option will be overwritten by the
`json_body_template.format(task_ir=task_ir_string)`.
sleep_time (float): time to sleep between each request. Defaults to 0.1.

Returns:
List[Tuple[NamedTuple, Response]]: List of parameters for each batch in
the task and the response from the post request.

Examples:
Here is a simple example of how to use this method. Note the body_template
has double curly braces on the outside to escape the string formatting.

```python
>>> body_template = "{{"token": "my_token", "task": {task_ir}}}"
>>> responses = (
program.quera.custom.submit(
100,
"http://my_custom_service.com",
body_template
)
)
```
"""

if r"{task_ir}" not in json_body_template:
raise ValueError(r"body_template must contain '{task_ir}'")

partial_eval = json_body_template.format(task_ir='"task_ir"')
try:
_ = json.loads(partial_eval)
except json.JSONDecodeError as e:
raise ValueError(

Check warning on line 142 in src/bloqade/ir/routine/quera.py

View check run for this annotation

Codecov / codecov/patch

src/bloqade/ir/routine/quera.py#L141-L142

Added lines #L141 - L142 were not covered by tests
"body_template must be a valid json template. "
'When evaluating template with task_ir="task_ir", '
f"the template evaluated to: {partial_eval!r}.\n"
f"JSONDecodeError: {e}"
)

out = []
for metadata, task_ir in self._compile(shots, args):
json_request_body = json_body_template.format(
task_ir=task_ir.json(exclude_none=True, exclude_unset=True)
)
request_options.update(data=json_request_body)
response = request(method, url, **request_options)
out.append((metadata, response))
time.sleep(sleep_time)

return out


@dataclass(frozen=True, config=__pydantic_dataclass_config__)
class QuEraHardwareRoutine(RoutineBase):
Expand Down
90 changes: 90 additions & 0 deletions tests/test_costum_submission.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
from bloqade import start
from requests import Response
from typing import Dict, List, Union
import simplejson as json

import bloqade.ir.routine.quera # noqa: F401
from unittest.mock import patch
import pytest


def create_response(
status_code: int, content: Union[Dict, List[Dict]], content_type="application/json"
) -> Response:
response = Response()
response.status_code = status_code
response.headers["Content-Type"] = content_type
response._content = bytes(json.dumps(content, use_decimal=True), "utf-8")
return response


@patch("bloqade.ir.routine.quera.request")
def test_custom_submission(request_mock):
body_template = '{{"token": "token", "body":{task_ir}}}'

task_ids = ["1", "2", "3"]
request_mock.side_effect = [
create_response(200, {"task_id": task_id}) for task_id in task_ids
]

# build bloqade program
program = (
start.add_position((0, 0))
.rydberg.rabi.amplitude.uniform.piecewise_linear(
[0.1, "time", 0.1], [0, 15, 15, 0]
)
.batch_assign(time=[0.0, 0.1, 0.5])
)

# submit and get responses and meta data associated with each task in the batch
responses = program.quera.custom().submit(
100, "https://my_service.test", body_template
)

for task_id, (metadata, response) in zip(task_ids, responses):
response_json = response.json()
assert response_json["task_id"] == task_id


@patch("bloqade.ir.routine.quera.request")
def test_custom_submission_error_missing_task_ir_key(request_mock):
body_template = '{{"token": "token", "body":}}'

task_ids = ["1", "2", "3"]
request_mock.side_effect = [
create_response(200, {"task_id": task_id}) for task_id in task_ids
]

# build bloqade program
program = (
start.add_position((0, 0))
.rydberg.rabi.amplitude.uniform.piecewise_linear(
[0.1, "time", 0.1], [0, 15, 15, 0]
)
.batch_assign(time=[0.0, 0.1, 0.5])
)
with pytest.raises(ValueError):
# submit and get responses and meta data associated with each task in the batch
program.quera.custom().submit(100, "https://my_service.test", body_template)


@patch("bloqade.ir.routine.quera.request")
def test_custom_submission_error_malformed_template(request_mock):
body_template = '{{"token": token", "body":}}'

task_ids = ["1", "2", "3"]
request_mock.side_effect = [
create_response(200, {"task_id": task_id}) for task_id in task_ids
]

# build bloqade program
program = (
start.add_position((0, 0))
.rydberg.rabi.amplitude.uniform.piecewise_linear(
[0.1, "time", 0.1], [0, 15, 15, 0]
)
.batch_assign(time=[0.0, 0.1, 0.5])
)
with pytest.raises(ValueError):
# submit and get responses and meta data associated with each task in the batch
program.quera.custom().submit(100, "https://my_service.test", body_template)