Update: support qcs account headers on engagement creation

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
 

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

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

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,14 +799,26 @@ 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
 
     .. _QCS API Docs: https://docs.api.qcs.rigetti.com/#tag/endpoints
     """
 
     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)

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