Add ability to use experimental capabilities in program submission (#967)

* Add ability to use experimental capabilities in run and run_async

* immediately return saved experimental capabilities instead of remotely fetching

* Fix old typos

* add experimental flag to __call__ and mock

* lint

* Fix failing tests

* Fix lint

Author: johnzl-777

src/bloqade/ir/routine/braket.py

@@ -34,6 +34,7 @@ class BraketHardwareRoutine(RoutineBase):
     def _compile(
         self,
         shots: int,
+        use_experimental: bool = False,
         args: Tuple[LiteralType, ...] = (),
         name: Optional[str] = None,
     ) -> RemoteBatch:
@@ -47,7 +48,7 @@ def _compile(
             generate_quera_ir,
         )
 
-        capabilities = self.backend.get_capabilities()
+        capabilities = self.backend.get_capabilities(use_experimental)
 
         circuit, params = self.circuit, self.params
 
@@ -85,6 +86,7 @@ def run_async(
         shots: int,
         args: Tuple[LiteralType, ...] = (),
         name: Optional[str] = None,
+        use_experimental: bool = False,
         shuffle: bool = False,
         **kwargs,
     ) -> RemoteBatch:
@@ -99,14 +101,15 @@ def run_async(
             shots (int): number of shots
             args (Tuple): Values of the parameter defined in `args`, defaults to ()
             name (str | None): custom name of the batch, defaults to None
+            use_experimental (bool): Use experimental hardware capabilities
             shuffle (bool): shuffle the order of jobs
 
         Return:
             RemoteBatch
 
         """
 
-        batch = self._compile(shots, args, name)
+        batch = self._compile(shots, use_experimental, args, name)
         batch._submit(shuffle, **kwargs)
         return batch
 
@@ -116,6 +119,7 @@ def run(
         shots: int,
         args: Tuple[LiteralType, ...] = (),
         name: Optional[str] = None,
+        use_experimental: bool = False,
         shuffle: bool = False,
         **kwargs,
     ) -> RemoteBatch:
@@ -139,7 +143,7 @@ def run(
 
         """
 
-        batch = self.run_async(shots, args, name, shuffle, **kwargs)
+        batch = self.run_async(shots, args, name, use_experimental, shuffle, **kwargs)
         batch.pull()
         return batch
 
@@ -149,6 +153,7 @@ def __call__(
         *args: LiteralType,
         shots: int = 1,
         name: Optional[str] = None,
+        use_experimental: bool = False,
         shuffle: bool = False,
         **kwargs,
     ):
@@ -171,7 +176,7 @@ def __call__(
             RemoteBatch
 
         """
-        return self.run(shots, args, name, shuffle, **kwargs)
+        return self.run(shots, args, name, use_experimental, shuffle, **kwargs)
 
 
 @dataclass(frozen=True, config=__pydantic_dataclass_config__)

src/bloqade/ir/routine/quera.py

@@ -52,6 +52,7 @@ class CustomSubmissionRoutine(RoutineBase):
     def _compile(
         self,
         shots: int,
+        use_experimental: bool = False,
         args: Tuple[LiteralType, ...] = (),
     ):
         from bloqade.compiler.passes.hardware import (
@@ -65,7 +66,7 @@ def _compile(
         from bloqade.submission.capabilities import get_capabilities
 
         circuit, params = self.circuit, self.params
-        capabilities = get_capabilities()
+        capabilities = get_capabilities(use_experimental)
 
         for batch_params in params.batch_assignments(*args):
             assignments = {**batch_params, **params.static_params}
@@ -92,6 +93,7 @@ def submit(
         method: str = "POST",
         args: Tuple[LiteralType] = (),
         request_options: Dict[str, Any] = {},
+        use_experimental: bool = False,
         sleep_time: float = 0.1,
     ) -> List[Tuple[NamedTuple, Response]]:
         """Compile to QuEraTaskSpecification and submit to a custom service.
@@ -110,6 +112,7 @@ def submit(
             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)`.
+            use_experimental (bool): Enable experimental hardware capabilities
             sleep_time (float): time to sleep between each request. Defaults to 0.1.
 
         Returns:
@@ -147,7 +150,7 @@ def submit(
             )
 
         out = []
-        for metadata, task_ir in self._compile(shots, args):
+        for metadata, task_ir in self._compile(shots, use_experimental, args):
             json_request_body = json_body_template.format(
                 task_ir=task_ir.json(exclude_none=True, exclude_unset=True)
             )
@@ -166,6 +169,7 @@ class QuEraHardwareRoutine(RoutineBase):
     def _compile(
         self,
         shots: int,
+        use_experimental: bool = False,
         args: Tuple[LiteralType, ...] = (),
         name: Optional[str] = None,
     ) -> RemoteBatch:
@@ -179,7 +183,7 @@ def _compile(
         )
 
         circuit, params = self.circuit, self.params
-        capabilities = self.backend.get_capabilities()
+        capabilities = self.backend.get_capabilities(use_experimental)
 
         tasks = OrderedDict()
 
@@ -215,6 +219,7 @@ def run_async(
         shots: int,
         args: Tuple[LiteralType, ...] = (),
         name: Optional[str] = None,
+        use_experimental: bool = False,
         shuffle: bool = False,
         **kwargs,
     ) -> RemoteBatch:
@@ -233,7 +238,7 @@ def run_async(
             RemoteBatch
 
         """
-        batch = self._compile(shots, args, name)
+        batch = self._compile(shots, use_experimental, args, name)
         batch._submit(shuffle, **kwargs)
         return batch
 
@@ -243,10 +248,11 @@ def run(
         shots: int,
         args: Tuple[LiteralType, ...] = (),
         name: Optional[str] = None,
+        use_experimental: bool = False,
         shuffle: bool = False,
         **kwargs,
     ) -> RemoteBatch:
-        batch = self.run_async(shots, args, name, shuffle, **kwargs)
+        batch = self.run_async(shots, args, name, use_experimental, shuffle, **kwargs)
         batch.pull()
         return batch
 
@@ -256,7 +262,8 @@ def __call__(
         *args: LiteralType,
         shots: int = 1,
         name: Optional[str] = None,
+        use_experimental: bool = False,
         shuffle: bool = False,
         **kwargs,
     ) -> RemoteBatch:
-        return self.run(shots, args, name, shuffle, **kwargs)
+        return self.run(shots, args, name, use_experimental, shuffle, **kwargs)

src/bloqade/submission/base.py

@@ -15,8 +15,8 @@ class SubmissionBackend(BaseModel):
     class Config:
         extra = Extra.forbid
 
-    def get_capabilities(self) -> QuEraCapabilities:
-        return get_capabilities()
+    def get_capabilities(self, use_experimental: bool = False) -> QuEraCapabilities:
+        return get_capabilities(use_experimental)
 
     def validate_task(
         self, task_ir: Union[BraketTaskSpecification, QuEraTaskSpecification]

src/bloqade/submission/braket.py

@@ -31,20 +31,23 @@ def device(self) -> AwsDevice:
 
         return self._device
 
-    def get_capabilities(self) -> QuEraCapabilities:
+    def get_capabilities(self, use_experimental: bool = False) -> QuEraCapabilities:
         from botocore.exceptions import BotoCoreError, ClientError
 
+        if use_experimental:
+            return super().get_capabilities(use_experimental)
+
         try:
             return to_quera_capabilities(self.device.properties.paradigm)
         except BotoCoreError:
             warnings.warn(
                 "Could not retrieve device capabilities from braket API. "
-                "Using local capabiltiiies file for Aquila."
+                "Using local capabilities file for Aquila."
             )
         except ClientError:
             warnings.warn(
                 "Could not retrieve device capabilities from braket API. "
-                "Using local capabiltiiies file for Aquila."
+                "Using local capabilities file for Aquila."
             )
 
         return super().get_capabilities()

src/bloqade/submission/quera.py

@@ -41,11 +41,11 @@ def queue_api(self):
 
         return self._queue_api
 
-    def get_capabilities(self) -> QuEraCapabilities:
+    def get_capabilities(self, use_experimental: bool = False) -> QuEraCapabilities:
         try:
             return QuEraCapabilities(**self.queue_api.get_capabilities())
         except BaseException:
-            return super().get_capabilities()
+            return super().get_capabilities(use_experimental)
 
     def submit_task(self, task_ir: QuEraTaskSpecification) -> str:
         return self.queue_api.submit_task(