open-standards/wai/wai-quantum-py/python/wai_quantum/service.py

"""A client for the WAI quantum job service, in the standard library alone.

The public instance is at ``https://wai.transaction.science/quantum-jobs/v1``;
no account and no key. Every result the service returns carries a receipt
sealed with the service's key, and :meth:`Service.verify` checks it here,
offline, with the engine in this package: the network is needed to run a
job, never to trust its answer.

Routes (``/v1/openapi.json`` describes them): ``GET /v1`` the service,
``GET /v1/methods``, ``POST /v1/jobs`` submit, ``GET /v1/jobs/{id}`` a job's
status and result, ``DELETE /v1/jobs/{id}`` cancel, ``POST /v1/run`` submit
and wait.
"""

from __future__ import annotations

import json
import os
import time
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, Dict, Mapping, Optional, Tuple, Union

from . import _engine
from ._api import Program, job_request, verify
from ._job import Job, Verification
from .errors import ServiceError

__all__ = ["Service", "DEFAULT_SERVICE", "USER_AGENT"]

#: The public instance. ``$WAI_QUANTUM_SERVICE`` overrides it.
DEFAULT_SERVICE = "https://wai.transaction.science/quantum-jobs"

#: How the client names itself. The service's edge refuses requests that
#: name no client or a generic library's default; this names the package
#: and its version.
USER_AGENT = f"wai-quantum-python/{_engine.__version__}"

_FINISHED = ("done", "failed", "cancelled")


class Service:
    """The job service at ``base`` (``…/quantum-jobs`` or ``…/quantum-jobs/v1``;
    by default ``$WAI_QUANTUM_SERVICE`` or the public instance).

    ``timeout`` bounds each HTTP request, in seconds; waiting for a job is
    bounded separately (``wait=`` on :meth:`result` and :meth:`run`).
    """

    def __init__(self, base: Optional[str] = None, *, timeout: float = 60.0, user_agent: Optional[str] = None) -> None:
        base = (base or os.environ.get("WAI_QUANTUM_SERVICE") or DEFAULT_SERVICE).rstrip("/")
        if base.endswith("/v1"):
            base = base[: -len("/v1")]
        self.base = base
        self.timeout = timeout
        self.user_agent = user_agent or USER_AGENT

    def __repr__(self) -> str:
        return f"Service({self.base!r})"

    # -- HTTP ---------------------------------------------------------------

    def _call(self, method: str, path: str, body: Any = None) -> Tuple[int, Any]:
        url = f"{self.base}/v1{path}"
        data = None if body is None else json.dumps(body, allow_nan=False).encode()
        headers = {"User-Agent": self.user_agent, "Accept": "application/json"}
        if data is not None:
            headers["Content-Type"] = "application/json"
        req = urllib.request.Request(url, data=data, method=method, headers=headers)
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as r:
                status, raw = r.status, r.read()
        except urllib.error.HTTPError as e:
            status, raw = e.code, e.read()
        except (urllib.error.URLError, OSError) as e:
            reason = getattr(e, "reason", e)
            raise ServiceError(f"{url}: {reason}") from e
        try:
            answer = json.loads(raw) if raw else None
        except ValueError:
            raise ServiceError(f"{url}: the answer is not JSON", status, raw[:200].decode("utf-8", "replace")) from None
        if status >= 400:
            why = answer.get("error") if isinstance(answer, dict) else None
            raise ServiceError(why or f"{method} {url} failed", status, answer)
        return status, answer

    # -- the service ----------------------------------------------------------

    def describe(self) -> Dict[str, Any]:
        """The service: its methods, its limits and the key it seals with."""
        return self._call("GET", "")[1]

    def health(self) -> bool:
        try:
            return self._call("GET", "/health")[1].get("ok") is True
        except ServiceError:
            return False

    def methods(self) -> Any:
        return self._call("GET", "/methods")[1]

    # -- jobs -------------------------------------------------------------

    def submit(self, program: Union[Program, Mapping[str, Any]], **job: Any) -> Dict[str, Any]:
        """Submit a job and return at once: its record, with its ``id`` and
        ``status`` (``queued`` or ``running``). The keywords are
        :func:`wai_quantum.job_request`'s."""
        return self._call("POST", "/jobs", job_request(program, **job))[1]

    def status(self, job_id: str) -> Dict[str, Any]:
        """A job's record as it stands: its status, its plan, the request as
        submitted, and once it is done its result with the receipt."""
        return self._call("GET", f"/jobs/{urllib.parse.quote(job_id, safe='')}")[1]

    def cancel(self, job_id: str) -> Dict[str, Any]:
        """Cancel a queued or running job; its record afterwards. A job that
        has already finished is refused (``ServiceError``, status 409)."""
        return self._call("DELETE", f"/jobs/{urllib.parse.quote(job_id, safe='')}")[1]

    def result(self, job_id: str, *, wait: float = 600.0, poll: float = 1.0) -> Job:
        """The job's result, waiting up to ``wait`` seconds for it to finish.
        A job that failed or was cancelled raises ``ServiceError`` with the
        service's reason; one still running after ``wait`` raises too."""
        return self._finish(self.status(job_id), wait, poll)

    def run(self, program: Union[Program, Mapping[str, Any]], *, wait: float = 600.0, poll: float = 1.0, **job: Any) -> Job:
        """Run a job on the service and wait for its result (``POST /v1/run``,
        then the job's record while it is still running)."""
        return self._finish(self._call("POST", "/run", job_request(program, **job))[1], wait, poll)

    def _finish(self, record: Dict[str, Any], wait: float, poll: float) -> Job:
        deadline = time.monotonic() + wait
        while record.get("status") not in _FINISHED:
            job_id = record.get("id")
            if not job_id:
                raise ServiceError("the service answered with no job id", None, record)
            if time.monotonic() >= deadline:
                raise ServiceError(f"job {job_id} is still {record.get('status')} after {wait:g} s", None, record)
            time.sleep(poll)
            record = self.status(job_id)
        if record["status"] != "done":
            raise ServiceError(f"job {record.get('id')} {record['status']}: {record.get('error') or 'no reason given'}", None, record)
        return Job(record)

    # -- receipts ---------------------------------------------------------

    def verify(self, job: Any) -> Verification:
        """Whether a job's receipt seals its request and its result, checked
        here, offline (no request is made). Compare its ``key`` with the
        service's (``describe()["receipt_key"]``) to know who sealed it."""
        return verify(job)