Allow setting qpu mode in device

demo/qasm_single_task_simpler_serialization.py

@@ -36,7 +36,7 @@ def bell():
 
 # 2. Create the task using the device -- optionally set some metadata.
 # NOTE: context_name and program_language are set to qasm for testing.
-device = Device(context_name="gemini-qasm")
+device = Device(context_name="gemini-qasm", qpu_mode="qasm-10q")
 task = device.task(
     kernel=bell,
     num_shots=2,

src/bloqade/core/device/device.py

@@ -22,6 +22,9 @@ class Device(Generic[FutureType], AuthMixin):
     that can be dry-run or submitted asynchronously.
 
     Attributes:
+        qpu_mode (str | None): Explicit qlam QPU mode used by tasks created
+            from this device. When None, qlam-core resolves it from
+            configuration.
         future_cls (type[FutureType]): Future class used by tasks created from
             this device. Defaults to `Future`.
         kernel_serializer (KernelSerializer): Default serializer passed to
@@ -94,6 +97,7 @@ def task(
 
         return self.single_kernel_task_cls(
             context_name=self.context_name,
+            qpu_mode=self.qpu_mode,
             kernel=kernel,
             num_shots=num_shots,
             arguments=arguments,
@@ -139,6 +143,7 @@ def batch_task(
 
         return self.kernel_batch_task_cls(
             context_name=self.context_name,
+            qpu_mode=self.qpu_mode,
             kernels=kernels,
             arguments=arguments,
             num_shots=num_shots,
@@ -183,6 +188,7 @@ def parameter_scan(
 
         return self.parameter_scan_task_cls(
             context_name=self.context_name,
+            qpu_mode=self.qpu_mode,
             kernel=kernel,
             num_shots=num_shots,
             arguments=arguments,

src/bloqade/core/device/future.py

@@ -57,6 +57,8 @@ class Future(AuthMixin, Generic[ResultType]):
     and construct result views over that storage using ``result_cls``.
 
     Attributes:
+        qpu_mode (str | None): Explicit qlam QPU mode used for backend API
+            calls. When None, qlam-core resolves it from configuration.
         task_id (str): Backend task ID.
         storage (StorageBackend): Storage backend used for fetched shots and
             task metadata. Defaults to a fresh `DictStorage` (in-memory; not
@@ -95,7 +97,10 @@ def get_task(self) -> "Task":
             # NOTE: typing issue in qlam-core
             # every client is BaseRestApi, which doesn't have get, but it actually does
             task = self.call_with_auth_refresh(
-                lambda: client.get(id=self.task_id)  # type: ignore
+                lambda: client.get(  # type: ignore
+                    qpu_mode=self.qpu_mode,
+                    id=self.task_id,
+                )
             )
             logger.info(
                 f"Fetched task with id {self.task_id}. Current status: {task.task_status}"
@@ -120,7 +125,10 @@ def get_compilation(self, compilation_id: str | None = None):
 
         with CompilationsClient(self.app_context) as client:
             return self.call_with_auth_refresh(
-                lambda: client.get(id=compilation_id)  # type: ignore
+                lambda: client.get(  # type: ignore
+                    qpu_mode=self.qpu_mode,
+                    id=compilation_id,
+                )
             )
 
     def fetch(self) -> None:
@@ -176,7 +184,10 @@ def cancel(self):
             try:
                 # NOTE: typing issue because client is seen as BaseClient instead of TaskClient
                 return self.call_with_auth_refresh(
-                    lambda: client.cancel(id=self.task_id)  # type: ignore
+                    lambda: client.cancel(  # type: ignore
+                        qpu_mode=self.qpu_mode,
+                        id=self.task_id,
+                    )
                 )
             except Exception as e:
                 warn(
@@ -327,6 +338,7 @@ def from_storage(
         task_id: str | None = None,
         fetch_options: ApiFetchOptions = ApiFetchOptions(),
         context_name: str | None = None,
+        qpu_mode: str | None = None,
     ) -> Self:
         """Create a future from task metadata already present in storage.
 
@@ -343,6 +355,9 @@ def from_storage(
             context_name (str | None): Name of the qlam context to attach to
                 the returned future. When None, the class-level default on
                 `cls` is used. Defaults to None.
+            qpu_mode (str | None): Explicit qlam QPU mode to attach to the
+                returned future. When None, qlam-core resolves it from
+                configuration. Defaults to None.
 
         Returns:
             Self: A future attached to the selected task ID.
@@ -382,6 +397,7 @@ def from_storage(
             fetch_options=fetch_options,
             result_cls=cls.result_cls,
             context_name=context_name,
+            qpu_mode=qpu_mode,
         )
 
     @classmethod
@@ -392,6 +408,7 @@ def from_task_id(
         storage: StorageBackend | None = None,
         fetch_options: ApiFetchOptions = ApiFetchOptions(),
         context_name: str | None = None,
+        qpu_mode: str | None = None,
     ) -> Self:
         """Create a future from a backend task ID.
 
@@ -410,6 +427,9 @@ def from_task_id(
             context_name (str | None): Name of the qlam context used to fetch
                 the task and attached to the returned future. When None, the
                 class-level default on `cls` is used. Defaults to None.
+            qpu_mode (str | None): Explicit qlam QPU mode used to fetch the
+                task and attached to the returned future. When None, qlam-core
+                resolves it from configuration. Defaults to None.
 
         Returns:
             Self: A future attached to `task_id`.
@@ -422,17 +442,23 @@ def from_task_id(
             storage = DictStorage()
 
         context_name = cls._resolve_context_name(context_name)
-        auth = AuthMixin(context_name=context_name)
+        auth = AuthMixin(context_name=context_name, qpu_mode=qpu_mode)
         auth.authenticate()
         with TasksClient(auth.app_context) as client:
             task = auth.call_with_auth_refresh(
-                lambda: client.get(id=task_id)  # type: ignore
+                lambda: client.get(  # type: ignore
+                    qpu_mode=qpu_mode,
+                    id=task_id,
+                )
             )
 
         # fetch subtasks for metadata
         with DefinitionsClient(auth.app_context) as client:
             task_def = auth.call_with_auth_refresh(
-                lambda: client.get(id=task.definition_id)  # type: ignore
+                lambda: client.get(  # type: ignore
+                    qpu_mode=qpu_mode,
+                    id=task.definition_id,
+                )
             )
 
         storage.add_task_definition(
@@ -445,6 +471,7 @@ def from_task_id(
             fetch_options=fetch_options,
             result_cls=cls.result_cls,
             context_name=context_name,
+            qpu_mode=qpu_mode,
         )
 
     def _wait_for_completion(self, timeout: float | None = None) -> TaskStatus:
@@ -513,6 +540,7 @@ def _fetch_subtask_page(
 
         while full_shots_page:
             response = client.get(
+                qpu_mode=self.qpu_mode,
                 id=self.task_id,
                 page=subtask_page,
                 size=self.fetch_options.subtasks_per_fetch,

src/bloqade/core/device/mixins.py

@@ -9,17 +9,27 @@
 
 
 @dataclass(kw_only=True)
-class AuthMixin:
-    """Mixin that provides authentication helpers for qlam API clients.
-
-    Manages an `AppContext` scoped to a qlam context name and ensures the
-    client is authenticated before making API calls.
+class ConfigMixin:
+    """Mixin that provides qlam connection and API configuration.
 
     Attributes:
         context_name (str): Name of the qlam context to use.
+        qpu_mode (str | None): Explicit qlam QPU mode to use for API calls.
+            When None, qlam-core resolves it from configuration.
     """
 
     context_name: str
+    qpu_mode: str | None = None
+
+
+@dataclass(kw_only=True)
+class AuthMixin(ConfigMixin):
+    """Mixin that provides authentication helpers for qlam API clients.
+
+    Inherits qlam configuration from `ConfigMixin` and manages an
+    `AppContext` scoped to `context_name`. Ensures the client is
+    authenticated before making API calls.
+    """
 
     @property
     def app_context(self) -> AppContext:

src/bloqade/core/device/task.py

@@ -37,6 +37,8 @@ class TaskABC(Generic[FutureType], AuthMixin, ABC):
     `TaskDefinition` that can be dry-run or submitted to the backend.
 
     Attributes:
+        qpu_mode (str | None): Explicit qlam QPU mode used for task
+            submission. When None, qlam-core resolves it from configuration.
         program_language (str): Program language identifier stored on the
             task definition and used when serializing kernels.
         language_version (str): Program language version stored on the task
@@ -348,7 +350,10 @@ def submit_task_definition(
         task_request = TaskCreationRequest(root=task_definition)
         with TasksClient(self.app_context) as tasks_client:
             created_task = self.call_with_auth_refresh(
-                lambda: tasks_client.create(body=task_request)  # type: ignore
+                lambda: tasks_client.create(  # type: ignore
+                    qpu_mode=self.qpu_mode,
+                    body=task_request,
+                )
             )
 
         task_id = created_task.id
@@ -367,6 +372,7 @@ def submit_task_definition(
             fetch_options=fetch_options,
             storage=storage,
             context_name=self.context_name,
+            qpu_mode=self.qpu_mode,
         )
 
 

test/device/fixtures/remote.py

@@ -378,17 +378,23 @@ def __init__(
         self.create_return = create_return
         self.cancel_raises = cancel_raises
 
-    def get(self, id):  # noqa: A002
-        self._record("get", id=id)
+    def get(self, id, qpu_mode=None):  # noqa: A002
+        kwargs = {"id": id}
+        if qpu_mode is not None:
+            kwargs["qpu_mode"] = qpu_mode
+        self._record("get", **kwargs)
         ret = self.get_return
         if ret is None:
             raise AssertionError("FakeTasksClient.get called but no get_return set")
         if callable(ret):
             return ret(id)
         return ret
 
-    def create(self, body):
-        self._record("create", body=body)
+    def create(self, body, qpu_mode=None):
+        kwargs = {"body": body}
+        if qpu_mode is not None:
+            kwargs["qpu_mode"] = qpu_mode
+        self._record("create", **kwargs)
         ret = self.create_return
         if ret is None:
             raise AssertionError(
@@ -398,8 +404,11 @@ def create(self, body):
             return ret(body)
         return ret
 
-    def cancel(self, id):  # noqa: A002
-        self._record("cancel", id=id)
+    def cancel(self, id, qpu_mode=None):  # noqa: A002
+        kwargs = {"id": id}
+        if qpu_mode is not None:
+            kwargs["qpu_mode"] = qpu_mode
+        self._record("cancel", **kwargs)
         if self.cancel_raises is not None:
             raise self.cancel_raises
         return None
@@ -425,8 +434,11 @@ def __init__(
         self.app_context = app_context
         self.get_return = get_return
 
-    def get(self, id):  # noqa: A002
-        self._record("get", id=id)
+    def get(self, id, qpu_mode=None):  # noqa: A002
+        kwargs = {"id": id}
+        if qpu_mode is not None:
+            kwargs["qpu_mode"] = qpu_mode
+        self._record("get", **kwargs)
         if self.get_return is None:
             raise AssertionError(
                 "FakeDefinitionsClient.get called but no get_return set"
@@ -445,8 +457,11 @@ def __init__(
         self.app_context = app_context
         self.get_return = get_return
 
-    def get(self, id):  # noqa: A002
-        self._record("get", id=id)
+    def get(self, id, qpu_mode=None):  # noqa: A002
+        kwargs = {"id": id}
+        if qpu_mode is not None:
+            kwargs["qpu_mode"] = qpu_mode
+        self._record("get", **kwargs)
         if self.get_return is None:
             raise AssertionError(
                 "FakeCompilationsClient.get called but no get_return set"
@@ -502,6 +517,8 @@ def __init__(
         self.envelope_fn = envelope_fn
 
     def get(self, **kwargs):
+        if kwargs.get("qpu_mode") is None:
+            kwargs.pop("qpu_mode", None)
         self._record("get", **kwargs)
         if self.envelope_fn is not None:
             return self.envelope_fn(**kwargs)

test/device/test_device.py

@@ -149,3 +149,22 @@ def second():
         ).kernel_serializer
         is override_serializer
     )
+
+
+def test_device_passes_qpu_mode_to_all_task_shapes():
+    @basic_no_opt
+    def first():
+        return
+
+    @basic_no_opt
+    def second():
+        return
+
+    device = Device(context_name="ctx", qpu_mode="squin-256q")
+
+    assert device.task(first).qpu_mode == "squin-256q"
+    assert device.batch_task([first, second]).qpu_mode == "squin-256q"
+    assert (
+        device.parameter_scan(first, arguments=[{"theta": 0.5}]).qpu_mode
+        == "squin-256q"
+    )