diff --git a/CHANGELOG.md b/CHANGELOG.md index 3d14c756b..1a239dda6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,8 @@ Changelog - `QAMExecutionResult` now includes `execution_duration_microseconds`, providing the amount of time a job held exclusive hardware access. (@randall-fulton, #1436) + +- `get_qc` accepts `account_id` and `account_type` keyword arguments, which `EngagementManager` can use to specify `X-QCS-ACCOUNT-{ID/TYPE}` headers on engagement requests. (@erichulburd, #1438) ### Bugfixes diff --git a/poetry.lock b/poetry.lock index 1e1625a74..0af665e9f 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1005,7 +1005,7 @@ py = {version = "*", markers = "implementation_name == \"pypy\""} [[package]] name = "qcs-api-client" -version = "0.20.10" +version = "0.20.12" description = "A client library for accessing the Rigetti QCS API" category = "main" optional = false @@ -1411,7 +1411,7 @@ latex = ["ipython"] [metadata] lock-version = "1.1" python-versions = "^3.7" -content-hash = "a2b41292a997e87f5cc8c599e8a824feb1d7f75c98cec276e04903282007a89e" +content-hash = "d2af98c6620da01edbf16ab8815b833152357794f62122671ca937e3a5ecede3" [metadata.files] alabaster = [ @@ -2057,8 +2057,8 @@ pyzmq = [ {file = "pyzmq-22.1.0.tar.gz", hash = "sha256:7040d6dd85ea65703904d023d7f57fab793d7ffee9ba9e14f3b897f34ff2415d"}, ] qcs-api-client = [ - {file = "qcs-api-client-0.20.10.tar.gz", hash = "sha256:4859884f43a3a29a90171c0470fb8a8e3524a50d6986ff5e12c4b877594ee837"}, - {file = "qcs_api_client-0.20.10-py3-none-any.whl", hash = "sha256:8ba34444f623cb684b9ee6717c95490fbf20e2209856964c02c3e5578a4dc16c"}, + {file = "qcs-api-client-0.20.12.tar.gz", hash = "sha256:e7815d0d3e819c95aac741716f664172e0e8d84caad8e29c1f6c2bf02cf497b4"}, + {file = "qcs_api_client-0.20.12-py3-none-any.whl", hash = "sha256:075b13bd97ed624d7f702c2f6a43dd7f1caf947bf73db5031e947377a63eb39c"}, ] recommonmark = [ {file = "recommonmark-0.7.1-py2.py3-none-any.whl", hash = "sha256:1b1db69af0231efce3fa21b94ff627ea33dee7079a01dd0a7f8482c3da148b3f"}, diff --git a/pyproject.toml b/pyproject.toml index fce69ff55..1f30998a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ lark = "^0.11.1" rpcq = "^3.10.0" networkx = "^2.5" importlib-metadata = { version = "^3.7.3", python = "<3.8" } -qcs-api-client = ">=0.8.1,<0.21.0" +qcs-api-client = ">=0.20.12,<0.21.0" retry = "^0.9.2" # latex extra diff --git a/pyquil/api/_engagement_manager.py b/pyquil/api/_engagement_manager.py index 5ff22523b..98384addb 100644 --- a/pyquil/api/_engagement_manager.py +++ b/pyquil/api/_engagement_manager.py @@ -19,7 +19,7 @@ from dateutil.parser import parse as parsedate from dateutil.tz import tzutc -from qcs_api_client.client import QCSClientConfiguration +from qcs_api_client.client import QCSClientConfiguration, QCSAccountType from qcs_api_client.models import EngagementWithCredentials, CreateEngagementRequest from qcs_api_client.operations.sync import create_engagement from qcs_api_client.types import UNSET @@ -43,15 +43,29 @@ class EngagementManager: _lock: threading.Lock """Lock used to ensure that only one engagement request is in flight at once.""" - def __init__(self, *, client_configuration: QCSClientConfiguration) -> None: + def __init__( + self, + *, + client_configuration: QCSClientConfiguration, + account_id: Optional[str] = None, + account_type: Optional[QCSAccountType] = None, + ) -> None: """ Instantiate a new engagement manager. :param client_configuration: Client configuration, used for refreshing engagements. + :param account_id: Optional account id. In practice, this should be left blank unless specifying a QCS + group account name, which will be used for the purposes of billing and metrics. This will + override the QCS account id set on your QCS profile. + :param account_type: Optional account type. In practice, this should be left blank unless specifying a QCS + group account type, which will be used for the purposes of billing and metrics. This will + override the QCS account type set on your QCS profile. """ self._client_configuration = client_configuration self._cached_engagements: Dict[EngagementCacheKey, EngagementWithCredentials] = {} self._lock = threading.Lock() + self._account_id = account_id + self._account_type = account_type def get_engagement( self, *, quantum_processor_id: str, request_timeout: float = 10.0, endpoint_id: Optional[str] = None @@ -73,12 +87,21 @@ def get_engagement( with self._lock: if not self._engagement_valid(self._cached_engagements.get(key)): with qcs_client( - client_configuration=self._client_configuration, request_timeout=request_timeout + client_configuration=self._client_configuration, + request_timeout=request_timeout, ) as client: # type: httpx.Client request = CreateEngagementRequest( - quantum_processor_id=quantum_processor_id, endpoint_id=endpoint_id or UNSET + quantum_processor_id=quantum_processor_id, + endpoint_id=endpoint_id or UNSET, ) - self._cached_engagements[key] = create_engagement(client=client, json_body=request).parsed + headers = {} + if self._account_id is not None: + headers["X-QCS-ACCOUNT-ID"] = self._account_id + if self._account_type is not None: + headers["X-QCS-ACCOUNT-TYPE"] = self._account_type.value + self._cached_engagements[key] = create_engagement( + client=client, json_body=request, httpx_request_kwargs={"headers": headers} + ).parsed return self._cached_engagements[key] @staticmethod diff --git a/pyquil/api/_quantum_computer.py b/pyquil/api/_quantum_computer.py index c9e31b588..a97261a6e 100644 --- a/pyquil/api/_quantum_computer.py +++ b/pyquil/api/_quantum_computer.py @@ -36,7 +36,7 @@ import httpx import networkx as nx import numpy as np -from qcs_api_client.client import QCSClientConfiguration +from qcs_api_client.client import QCSClientConfiguration, QCSAccountType from qcs_api_client.models import ListQuantumProcessorsResponse from qcs_api_client.operations.sync import list_quantum_processors from rpcq.messages import ParameterAref @@ -727,6 +727,8 @@ def get_qc( client_configuration: Optional[QCSClientConfiguration] = None, endpoint_id: Optional[str] = None, engagement_manager: Optional[EngagementManager] = None, + account_id: Optional[str] = None, + account_type: Optional[QCSAccountType] = None, ) -> QuantumComputer: """ Get a quantum computer. @@ -797,6 +799,14 @@ def get_qc( :param client_configuration: Optional client configuration. If none is provided, a default one will be loaded. :param endpoint_id: Optional quantum processor endpoint ID, as used in the `QCS API Docs`_. :param engagement_manager: Optional engagement manager. If none is provided, a default one will be created. + :param account_id: Optional account id. If engagement manager is not set, ``EngagementManager`` will be + initialized with this ``account_id``. In practice, this should be left blank unless specifying a QCS + group account name, which will be used for the purposes of billing and metrics. This will + override the QCS account id set on your QCS profile. + :param account_type: Optional account type. If engagement manager is not set, ``EngagementManager`` will be + initialized with this ``account_type``. In practice, this should be left blank unless specifying a QCS + group account type, which will be used for the purposes of billing and metrics. This will + override the QCS account type set on your QCS profile. :return: A pre-configured QuantumComputer @@ -804,7 +814,11 @@ def get_qc( """ client_configuration = client_configuration or QCSClientConfiguration.load() - engagement_manager = engagement_manager or EngagementManager(client_configuration=client_configuration) + engagement_manager = engagement_manager or EngagementManager( + client_configuration=client_configuration, + account_id=account_id, + account_type=account_type, + ) # 1. Parse name, check for redundant options, canonicalize names. prefix, qvm_type, noisy = _parse_name(name, as_qvm, noisy) diff --git a/test/unit/test_quantum_computer.py b/test/unit/test_quantum_computer.py index fd8207f51..1cee5254a 100644 --- a/test/unit/test_quantum_computer.py +++ b/test/unit/test_quantum_computer.py @@ -1,6 +1,7 @@ import itertools import random from test.unit.utils import DummyCompiler +from typing import cast import networkx as nx import numpy as np @@ -23,6 +24,7 @@ _symmetrization, get_qc, ) +from pyquil.api._qpu import QPU from pyquil.api._qvm import QVM from pyquil.experiment import Experiment, ExperimentSetting from pyquil.experiment._main import _pauli_to_product_state @@ -32,6 +34,7 @@ from pyquil.pyqvm import PyQVM from pyquil.quantum_processor import NxQuantumProcessor from pyquil.quilbase import Declare, MemoryReference +from qcs_api_client.client import QCSAccountType from qcs_api_client.models.instruction_set_architecture import InstructionSetArchitecture from rpcq.messages import ParameterAref @@ -849,3 +852,91 @@ def test_get_qc_endpoint_id(client_configuration: QCSClientConfiguration, qcs_as qc = get_qc("test", endpoint_id="test-endpoint") assert qc.qam._qpu_client._endpoint_id == "test-endpoint" + + +@respx.mock +def test_get_qc_with_explicit_account(client_configuration: QCSClientConfiguration, qcs_aspen8_isa: InstructionSetArchitecture): + """ + Assert that get_qc will pass QCS account via kwargs to ``EngagementManager`` which will in turn set the appropriate + QCS Account headers. + """ + respx.get( + url=f"{client_configuration.profile.api_url}/v1/quantumProcessors/test/instructionSetArchitecture", + ).respond(json=qcs_aspen8_isa.to_dict()) + + qc = get_qc("test", endpoint_id="test-endpoint", account_id="group0", account_type=QCSAccountType.group) + + assert isinstance(qc, QuantumComputer) + quantum_computer = cast(QuantumComputer, qc) + assert isinstance(quantum_computer.qam, QPU) + qpu = cast(QPU, quantum_computer.qam) + engagement_manager = qpu._qpu_client._engagement_manager + + assert "group0" == engagement_manager._account_id + assert QCSAccountType.group == engagement_manager._account_type + + respx.post( + url=f"{client_configuration.profile.api_url}/v1/engagements", + headers__contains={ + 'X-QCS-ACCOUNT-ID': 'group0', + 'X-QCS-ACCOUNT-TYPE': QCSAccountType.group.value, + }, + ).respond(json={ + 'address': 'address', + 'endpointId': 'endpointId', + 'quantumProcessorId': 'quantumProcessorId', + 'userId': 'userId', + 'expiresAt': '01-01-2200T00:00:00Z', + 'credentials': { + 'clientPublic': 'faux', + 'clientSecret': 'faux', + 'serverPublic': 'faux', + } + }) + + engagement = engagement_manager.get_engagement(quantum_processor_id='test') + assert 'faux' == engagement.credentials.client_public + + +@respx.mock +def test_get_qc_with_account_profile(client_configuration: QCSClientConfiguration, qcs_aspen8_isa: InstructionSetArchitecture): + """ + Assert that get_qc will pass QCS account via QCS API Profile to ``EngagementManager`` which will in turn set the appropriate + QCS Account headers. + """ + respx.get( + url=f"{client_configuration.profile.api_url}/v1/quantumProcessors/test/instructionSetArchitecture", + ).respond(json=qcs_aspen8_isa.to_dict()) + + client_configuration.profile.account_id = 'group0' + client_configuration.profile.account_type = QCSAccountType.group + + qc = get_qc("test", client_configuration=client_configuration, account_type=QCSAccountType.group) + + assert isinstance(qc, QuantumComputer) + quantum_computer = cast(QuantumComputer, qc) + assert isinstance(quantum_computer.qam, QPU) + qpu = cast(QPU, quantum_computer.qam) + engagement_manager = qpu._qpu_client._engagement_manager + + respx.post( + url=f"{client_configuration.profile.api_url}/v1/engagements", + headers__contains={ + 'X-QCS-ACCOUNT-ID': 'group0', + 'X-QCS-ACCOUNT-TYPE': QCSAccountType.group.value, + }, + ).respond(json={ + 'address': 'address', + 'endpointId': 'endpointId', + 'quantumProcessorId': 'quantumProcessorId', + 'userId': 'userId', + 'expiresAt': '01-01-2200T00:00:00Z', + 'credentials': { + 'clientPublic': 'faux', + 'clientSecret': 'faux', + 'serverPublic': 'faux', + } + }) + + engagement = engagement_manager.get_engagement(quantum_processor_id='test') + assert 'faux' == engagement.credentials.client_public