Adding Generic submission option. (#961)

* adding custom service backend for submitting quera tasks.

* allowing more options into the request.

* fix docs.

* enabling routing from builder and adding some better error messages and docs. adding happy path test.

* fixing tests.

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()

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):

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)