iqm-station-control-client 11.2.1__py3-none-any.whl → 12.0.0__py3-none-any.whl

iqm/station_control/client/station_control.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Client implementation for station control service REST API."""
+"""Client implementation for Station Control service REST API."""
 
 from __future__ import annotations
 
@@ -40,6 +40,7 @@ from exa.common.errors.station_control_errors import (
     map_from_status_code_to_error,
 )
 from exa.common.qcm_data.qcm_data_client import QCMDataClient
+from iqm.station_control.client.authentication import TokenManager
 from iqm.station_control.client.list_models import (
     DutFieldDataList,
     DutList,
@@ -77,6 +78,7 @@ from iqm.station_control.interface.models import (
     ObservationSetDefinition,
     ObservationSetUpdate,
     ObservationUpdate,
+    ProgressCallback,
     QualityMetrics,
     RunData,
     RunDefinition,
@@ -87,12 +89,11 @@ from iqm.station_control.interface.models import (
     SequenceResultDefinition,
     SoftwareVersionSet,
     StaticQuantumArchitecture,
-    Statuses,
+    StrUUID,
     SweepData,
     SweepDefinition,
     SweepResults,
 )
-from iqm.station_control.interface.models.type_aliases import StrUUID
 from iqm.station_control.interface.pydantic_base import PydanticBase
 from iqm.station_control.interface.station_control import StationControlInterface
 
@@ -100,16 +101,19 @@ logger = logging.getLogger(__name__)
 TypePydanticBase = TypeVar("TypePydanticBase", bound=PydanticBase)
 
 
-class _StationControlClientBase(StationControlInterface):
-    """Shared functionality for StationControlClient and IqmServerClient.
+class StationControlClient(StationControlInterface):
+    """Client implementation for station control service REST API.
 
     Args:
-        root_url: Remote server URL.
-        get_token_callback: A callback function that returns a token
-            which will be passed in Authorization header in all requests.
+        root_url: Remote station control service URL.
+        token: Long-lived authentication token in plain text format.
+            If ``token`` is given no other user authentication parameters should be given.
+        tokens_file: Path to a tokens file used for authentication.
+            If ``tokens_file`` is given no other user authentication parameters should be given.
+        get_token_callback: Callback function that returns an Authorization header containing a bearer token.
+            If ``get_token_callback`` is given no other user authentication parameters should be given.
         client_signature: String that is added to the User-Agent header of requests
             sent to the server.
-        enable_opentelemetry: Iff True, enable Jaeger/OpenTelemetry tracing.
 
     """
 
@@ -117,144 +121,18 @@ class _StationControlClientBase(StationControlInterface):
         self,
         root_url: str,
         *,
+        token: str | None = None,
+        tokens_file: str | None = None,
         get_token_callback: Callable[[], str] | None = None,
         client_signature: str | None = None,
-        enable_opentelemetry: bool = False,
     ):
         self.root_url = root_url
-        self._get_token_callback = get_token_callback
-        self._signature = self._create_signature(client_signature)
-        self._enable_opentelemetry = enable_opentelemetry
-
-    @classmethod
-    def _create_signature(cls, client_signature: str | None) -> str:
-        signature = f"{platform.platform(terse=True)}"
-        signature += f", python {platform.python_version()}"
-        dist_pkg_name = "iqm-station-control-client"
-        signature += f", {cls.__name__} {dist_pkg_name} {version(dist_pkg_name)}"
-        if client_signature:
-            signature += f", {client_signature}"
-        return signature
-
-    def _send_request(
-        self,
-        http_method: Callable[..., requests.Response],
-        url_path: str,
-        *,
-        headers: dict[str, str] | None = None,
-        params: dict[str, Any] | None = None,
-        json_str: str | None = None,
-        octets: bytes | None = None,
-        timeout: int = 600,
-    ) -> requests.Response:
-        """Send an HTTP request.
-
-        Parameters ``json_str`` and ``octets`` are mutually exclusive.
-        The first non-None argument (in this order) will be used to construct the body of the request.
-
-        Args:
-            http_method: HTTP method to use for the request, any of requests.[post|get|put|head|delete|patch|options].
-            url_path: URL for the request.
-            headers: Additional HTTP headers for the request. Some may be overridden.
-            params: HTTP query parameters to store in the query string of the request URL.
-            json_str: JSON string to store in the body, may contain arbitrary Unicode characters.
-            octets: Pre-serialized binary data to store in the body.
-            timeout: Timeout for the request in seconds.
-
-        Returns:
-            Response to the request.
-
-        Raises:
-            StationControlError: Request was not successful.
-
-        """
-        # Will raise an error if respectively an error response code is returned.
-        # http_method should be any of requests.[post|get|put|head|delete|patch|options]
-
-        request_kwargs = self._build_request_kwargs(
-            headers=headers or {}, params=params or {}, json_str=json_str, octets=octets, timeout=timeout
-        )
-        url = f"{self.root_url}/{url_path}"
-        # TODO SW-1387: Use v1 API
-        # url = f"{self.root_url}/{self.version}/{url_path}"
-        response = http_method(url, **request_kwargs)
-        if not response.ok:
-            try:
-                response_json = response.json()
-                error_message = response_json.get("message")
-            except (json.JSONDecodeError, KeyError):
-                error_message = response.text
-
-            try:
-                error_class = map_from_status_code_to_error(response.status_code)  # type: ignore[arg-type]
-            except KeyError:
-                raise RuntimeError(f"Unexpected response status code {response.status_code}: {error_message}")
-
-            raise error_class(error_message)
-        return response
-
-    def _build_request_kwargs(
-        self,
-        *,
-        headers: dict[str, str],
-        params: dict[str, Any],
-        json_str: str | None = None,
-        octets: bytes | None = None,
-        timeout: int,
-    ) -> dict[str, Any]:
-        """Prepare the keyword arguments for an HTTP request."""
-        # add default headers
-        headers["User-Agent"] = self._signature
-
-        # json_str and octets are mutually exclusive
-        data: bytes | None = None
-        if json_str is not None:
-            # Must be able to handle JSON strings with arbitrary unicode characters, so we use an explicit
-            # encoding into bytes, and set the headers so the recipient can decode the request body correctly.
-            data = json_str.encode("utf-8")
-            headers["Content-Type"] = "application/json; charset=UTF-8"
-        elif octets is not None:
-            data = octets
-            headers["Content-Type"] = "application/protobuf"
-
-        if self._enable_opentelemetry:
-            parent_span_context = trace.set_span_in_context(trace.get_current_span())
-            propagate.inject(carrier=headers, context=parent_span_context)
-
-        # If token callback exists, use it to retrieve the token and add it to the headers
-        if self._get_token_callback:
-            headers["Authorization"] = self._get_token_callback()
-
-        kwargs = {
-            "headers": headers,
-            "params": params,
-            "data": data,
-            "timeout": timeout,
-        }
-        return _remove_empty_values(kwargs)
-
-
-class StationControlClient(_StationControlClientBase):
-    """Client implementation for station control service REST API.
-
-    Args:
-        root_url: Remote station control service URL.
-        get_token_callback: A callback function that returns a token (str)
-            which will be passed in Authorization header in all requests.
-        client_signature: String that is added to the User-Agent header of requests
-            sent to the server.
-
-    """
+        tm = TokenManager(token, tokens_file, get_token_callback, use_env_vars=False)
+        self._token_manager = tm
+        self._auth_header_callback = tm.get_auth_header_callback()
 
-    def __init__(
-        self, root_url: str, *, get_token_callback: Callable[[], str] | None = None, client_signature: str | None = None
-    ):
-        super().__init__(
-            root_url,
-            get_token_callback=get_token_callback,
-            client_signature=client_signature,  # type: ignore[arg-type]
-            enable_opentelemetry=os.environ.get("JAEGER_OPENTELEMETRY_COLLECTOR_ENDPOINT", None) is not None,
-        )
+        self._signature = self._create_signature(client_signature)
+        self._enable_opentelemetry = os.environ.get("JAEGER_OPENTELEMETRY_COLLECTOR_ENDPOINT", None) is not None
         # TODO SW-1387: Remove this when using v1 API, not needed
         self._check_api_versions()
         qcm_data_url = os.environ.get("CHIP_DESIGN_RECORD_FALLBACK_URL", None)
@@ -262,7 +140,7 @@ class StationControlClient(_StationControlClientBase):
 
     @property
     def version(self) -> str:
-        """Return the version of the station control API this client is using."""
+        """Version of the Station Control API this client is using."""
         return "v1"
 
     @cache
@@ -336,7 +214,7 @@ class StationControlClient(_StationControlClientBase):
     def run(
         self,
         run_definition: RunDefinition,
-        update_progress_callback: Callable[[Statuses], None] | None = None,
+        update_progress_callback: ProgressCallback | None = None,
         wait_job_completion: bool = True,
     ) -> bool:
         data = serialize_run_job_request(run_definition, queue_name="sweeps")
@@ -496,7 +374,7 @@ class StationControlClient(_StationControlClientBase):
     def abort_job(self, job_id: StrUUID) -> None:
         self._send_request(requests.post, f"jobs/{job_id}/abort")
 
-    def _wait_job_completion(self, job_id: str, update_progress_callback: Callable[[Statuses], None] | None) -> bool:
+    def _wait_job_completion(self, job_id: str, update_progress_callback: ProgressCallback | None) -> bool:
         logger.info("Waiting for job ID: %s", job_id)
         update_progress_callback = update_progress_callback or (lambda status: None)
         try:
@@ -510,7 +388,7 @@ class StationControlClient(_StationControlClientBase):
         return False
 
     def _poll_job_status_until_execution_start(
-        self, job_id: str, update_progress_callback: Callable[[Statuses], None]
+        self, job_id: str, update_progress_callback: ProgressCallback
     ) -> JobExecutorStatus:
         # Keep polling job status as long as it's PENDING, and update progress with `update_progress_callback`.
         max_seen_position = 0
@@ -532,7 +410,7 @@ class StationControlClient(_StationControlClientBase):
     def _poll_job_status_until_terminal(
         self,
         job_id: str,
-        update_progress_callback: Callable[[Statuses], None],
+        update_progress_callback: ProgressCallback,
     ) -> None:
         # Keep polling job status until it finishes, and update progress with `update_progress_callback`.
         while True:
@@ -566,7 +444,6 @@ class StationControlClient(_StationControlClientBase):
         # using the \uXXXX syntax, BaseModel.model_dump_json() keeps them in the produced JSON str.
         return model.model_dump_json()
 
-    # TODO SW-1387: Remove this when using v1 API, not needed
     def _check_api_versions(self):  # noqa: ANN202
         client_api_version = self._get_client_api_version()
         # Parse versions using standard packaging.version implementation.
@@ -575,12 +452,6 @@ class StationControlClient(_StationControlClientBase):
             self.get_about()["software_versions"]["iqm-station-control-client"].replace(" (local editable)", "+local")
         )
 
-        if client_api_version.major != server_api_version.major:
-            raise ValueError(
-                f"station-control-client version '{client_api_version}' is not compatible with the station control "
-                f"server, please use station-control-client version compatible with version '{server_api_version}'."
-            )
-
         if client_api_version.local or server_api_version.local:
             logger.warning(
                 "Client ('%s') and/or server ('%s') is using a local version of the station-control-client. "
@@ -588,6 +459,14 @@ class StationControlClient(_StationControlClientBase):
                 client_api_version,
                 server_api_version,
             )
+        elif client_api_version.major != server_api_version.major:
+            logger.warning(
+                "Client and server compatibility cannot be guaranteed. "
+                "The installed station-control-client package version '%s' has a different major to "
+                "station-control server version '%s'.",
+                client_api_version,
+                server_api_version,
+            )
         elif client_api_version.minor > server_api_version.minor:
             logger.warning(
                 "station-control-client version '%s' is newer minor version than '%s' used by the station control "
@@ -630,6 +509,113 @@ class StationControlClient(_StationControlClientBase):
             return model.root
         return model
 
+    @classmethod
+    def _create_signature(cls, client_signature: str | None) -> str:
+        signature = f"{platform.platform(terse=True)}"
+        signature += f", python {platform.python_version()}"
+        dist_pkg_name = "iqm-station-control-client"
+        signature += f", {cls.__name__} {dist_pkg_name} {version(dist_pkg_name)}"
+        if client_signature:
+            signature += f", {client_signature}"
+        return signature
+
+    def _send_request(
+        self,
+        http_method: Callable[..., requests.Response],
+        url_path: str,
+        *,
+        headers: dict[str, str] | None = None,
+        params: dict[str, Any] | None = None,
+        json_str: str | None = None,
+        octets: bytes | None = None,
+        timeout: int = 600,
+    ) -> requests.Response:
+        """Send an HTTP request.
+
+        Parameters ``json_str`` and ``octets`` are mutually exclusive.
+        The first non-None argument (in this order) will be used to construct the body of the request.
+
+        Args:
+            http_method: HTTP method to use for the request, any of requests.[post|get|put|head|delete|patch|options].
+            url_path: URL for the request.
+            headers: Additional HTTP headers for the request. Some may be overridden.
+            params: HTTP query parameters to store in the query string of the request URL.
+            json_str: JSON string to store in the body, may contain arbitrary Unicode characters.
+            octets: Pre-serialized binary data to store in the body.
+            timeout: Timeout for the request in seconds.
+
+        Returns:
+            Response to the request.
+
+        Raises:
+            StationControlError: Request was not successful.
+
+        """
+        # Will raise an error if respectively an error response code is returned.
+        # http_method should be any of requests.[post|get|put|head|delete|patch|options]
+
+        request_kwargs = self._build_request_kwargs(
+            headers=headers or {}, params=params or {}, json_str=json_str, octets=octets, timeout=timeout
+        )
+        url = f"{self.root_url}/{url_path}"
+        # TODO SW-1387: Use v1 API
+        # url = f"{self.root_url}/{self.version}/{url_path}"
+        response = http_method(url, **request_kwargs)
+        if not response.ok:
+            try:
+                response_json = response.json()
+                error_message = response_json.get("message")
+            except (json.JSONDecodeError, KeyError):
+                error_message = response.text
+
+            try:
+                error_class = map_from_status_code_to_error(response.status_code)  # type: ignore[arg-type]
+            except KeyError:
+                raise RuntimeError(f"Unexpected response status code {response.status_code}: {error_message}")
+
+            raise error_class(error_message)
+        return response
+
+    def _build_request_kwargs(
+        self,
+        *,
+        headers: dict[str, str],
+        params: dict[str, Any],
+        json_str: str | None = None,
+        octets: bytes | None = None,
+        timeout: int,
+    ) -> dict[str, Any]:
+        """Prepare the keyword arguments for an HTTP request."""
+        # add default headers
+        headers["User-Agent"] = self._signature
+
+        # json_str and octets are mutually exclusive
+        data: bytes | None = None
+        if json_str is not None:
+            # Must be able to handle JSON strings with arbitrary unicode characters, so we use an explicit
+            # encoding into bytes, and set the headers so the recipient can decode the request body correctly.
+            data = json_str.encode("utf-8")
+            headers["Content-Type"] = "application/json; charset=UTF-8"
+        elif octets is not None:
+            data = octets
+            headers["Content-Type"] = "application/protobuf"
+
+        if self._enable_opentelemetry:
+            parent_span_context = trace.set_span_in_context(trace.get_current_span())
+            propagate.inject(carrier=headers, context=parent_span_context)
+
+        # If auth header callback exists, use it to add the header
+        if self._auth_header_callback:
+            headers["Authorization"] = self._auth_header_callback()
+
+        kwargs = {
+            "headers": headers,
+            "params": params,
+            "data": data,
+            "timeout": timeout,
+        }
+        return _remove_empty_values(kwargs)
+
 
 def _remove_empty_values(kwargs: dict[str, Any]) -> dict[str, Any]:
     """Return a copy of the given dict without values that are None or {}."""

iqm/station_control/client/utils.py

@@ -13,22 +13,17 @@
 # limitations under the License.
 """Utility functions for IQM Station Control Client."""
 
-from collections.abc import Callable
-
-import requests
 from tqdm.auto import tqdm
 
-from iqm.station_control.client.iqm_server.iqm_server_client import IqmServerClient
-from iqm.station_control.client.station_control import StationControlClient
-from iqm.station_control.interface.models import Statuses
-from iqm.station_control.interface.station_control import StationControlInterface
+from iqm.station_control.interface.models import ProgressCallback
+from iqm.station_control.interface.models.jobs import _Progress
 
 
-def get_progress_bar_callback() -> Callable[[Statuses], None]:
+def get_progress_bar_callback() -> ProgressCallback:
     """Returns a callback function that creates or updates existing progressbars when called."""
     progress_bars = {}
 
-    def _create_and_update_progress_bars(statuses: Statuses) -> None:
+    def _create_and_update_progress_bars(statuses: list[_Progress]) -> None:
         for label, value, total in statuses:
             if label not in progress_bars:
                 progress_bars[label] = tqdm(total=total, desc=label, leave=True)
@@ -36,36 +31,3 @@ def get_progress_bar_callback() -> Callable[[Statuses], None]:
             progress_bars[label].refresh()
 
     return _create_and_update_progress_bars
-
-
-def init_station_control(
-    root_url: str, get_token_callback: Callable[[], str] | None = None, **kwargs
-) -> StationControlInterface:
-    """Initialize a new station control instance connected to the given remote.
-
-    Client implementation is selected automatically based on the remote station: if the remote station
-    is running the IQM Server software stack, then the IQM Server client implementation (with a limited
-    feature set) is chosen. If the remote station is running the SC software stack, then the Station
-    Control client implementation (with the full feature set) is chosen.
-
-    Args:
-        root_url: Remote station control service URL. For IQM Server remotes, this is the "Quantum Computer URL"
-            value from the web dashboard.
-        get_token_callback: A callback function that returns a token (str) which will be passed in Authorization
-            header in all requests.
-
-    """
-    try:
-        headers = {"Authorization": get_token_callback()} if get_token_callback else {}
-        response = requests.get(f"{root_url}/about", headers=headers)
-        response.raise_for_status()
-        about = response.json()
-        if isinstance(about, dict) and about.get("iqm_server") is True:
-            # If about information has iqm_server flag, it means that we're communicating
-            # with IQM server instead of direct Station Control service, hence we need to
-            # use the specialized client
-            return IqmServerClient(root_url, get_token_callback=get_token_callback, **kwargs)
-        # Using direct station control by default
-        return StationControlClient(root_url=root_url, get_token_callback=get_token_callback, **kwargs)
-    except Exception as e:
-        raise RuntimeError("Failed to initialize the client.") from e

iqm/station_control/interface/models/__init__.py

@@ -13,6 +13,22 @@
 # limitations under the License.
 """Station control interface models."""
 
+from iqm.station_control.interface.models.circuit import (
+    CircuitBatch,
+    CircuitMeasurementCounts,
+    CircuitMeasurementCountsBatch,
+    CircuitMeasurementResults,
+    CircuitMeasurementResultsBatch,
+    DDMode,
+    DDStrategy,
+    HeraldingMode,
+    MoveGateFrameTrackingMode,
+    MoveGateValidationMode,
+    PRXSequence,
+    QIRCode,
+    QubitMapping,
+    RunRequest,
+)
 from iqm.station_control.interface.models.dut import DutData, DutFieldData
 from iqm.station_control.interface.models.dynamic_quantum_architecture import (
     DynamicQuantumArchitecture,
@@ -22,8 +38,11 @@ from iqm.station_control.interface.models.dynamic_quantum_architecture import (
 )
 from iqm.station_control.interface.models.jobs import (
     JobData,
+    JobError,
     JobExecutorStatus,
     JobResult,
+    ProgressCallback,
+    Statuses,
     TimelineEntry,
 )
 from iqm.station_control.interface.models.observation import (
@@ -49,11 +68,11 @@ from iqm.station_control.interface.models.sequence import (
     SequenceResultDefinition,
 )
 from iqm.station_control.interface.models.static_quantum_architecture import StaticQuantumArchitecture
-from iqm.station_control.interface.models.sweep import SweepData, SweepDefinition
+from iqm.station_control.interface.models.sweep import SweepBase, SweepData, SweepDefinition
 from iqm.station_control.interface.models.type_aliases import (
     DutType,
     GetObservationsMode,
     SoftwareVersionSet,
-    Statuses,
+    StrUUID,
     SweepResults,
 )