diff --git a/src/bloqade/builder/backend/quera.py b/src/bloqade/builder/backend/quera.py index 208789797..9ed1f7cb8 100644 --- a/src/bloqade/builder/backend/quera.py +++ b/src/bloqade/builder/backend/quera.py @@ -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() diff --git a/src/bloqade/ir/routine/quera.py b/src/bloqade/ir/routine/quera.py index c818caef5..1f3f35f7c 100644 --- a/src/bloqade/ir/routine/quera.py +++ b/src/bloqade/ir/routine/quera.py @@ -1,4 +1,5 @@ -from collections import OrderedDict +from collections import OrderedDict, namedtuple +import time from pydantic.v1.dataclasses import dataclass import json @@ -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__) @@ -41,6 +43,121 @@ def mock( 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( + "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): diff --git a/tests/test_costum_submission.py b/tests/test_costum_submission.py new file mode 100644 index 000000000..e01e0d60c --- /dev/null +++ b/tests/test_costum_submission.py @@ -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)