Skip to content

Commit

Permalink
Add user agent for Braket interactions (#99)
Browse files Browse the repository at this point in the history
  • Loading branch information
kshitijc authored Jul 1, 2022
1 parent 7214eca commit ccee356
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 6 deletions.
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
packages=find_namespace_packages(where="src", exclude=("test",)),
package_dir={"": "src"},
install_requires=[
"amazon-braket-sdk",
"amazon-braket-sdk>=1.25.0",
"pennylane>=0.23.0",
],
entry_points={
Expand Down
2 changes: 2 additions & 0 deletions src/braket/pennylane_plugin/braket_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -319,6 +319,8 @@ def __init__(
**run_kwargs,
):
device = AwsDevice(device_arn, aws_session=aws_session)
user_agent = f"BraketPennylanePlugin/{__version__}"
device.aws_session.add_braket_user_agent(user_agent)
if DeviceActionType.JAQCD not in device.properties.action:
raise ValueError(f"Device {device.name} does not support quantum circuits")

Expand Down
30 changes: 25 additions & 5 deletions test/unit_tests/test_braket_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
from pennylane.tape import QuantumTape

import braket.pennylane_plugin.braket_device
from braket.pennylane_plugin import BraketAwsQubitDevice, BraketLocalQubitDevice
from braket.pennylane_plugin import BraketAwsQubitDevice, BraketLocalQubitDevice, __version__
from braket.pennylane_plugin.braket_device import BraketQubitDevice, Shots

SHOTS = 10000
Expand Down Expand Up @@ -753,6 +753,18 @@ def test_run_task_unimplemented():
dev.execute(circuit)


@patch("braket.pennylane_plugin.braket_device.AwsDevice")
def test_add_braket_user_agent_invoked(aws_device_mock):
aws_device_mock_instance = aws_device_mock.return_value
aws_device_mock_instance.properties.action = {DeviceActionType.JAQCD: ACTION_PROPERTIES}
aws_device_mock_instance.type = AwsDeviceType.SIMULATOR
BraketAwsQubitDevice(wires=2, device_arn="foo", shots=10)
expected_user_agent = f"BraketPennylanePlugin/{__version__}"
aws_device_mock_instance.aws_session.add_braket_user_agent.assert_called_with(
expected_user_agent
)


class DummyLocalQubitDevice(BraketQubitDevice):
short_name = "dummy"

Expand Down Expand Up @@ -827,10 +839,17 @@ def _noop(*args, **kwargs):


@patch.object(AwsDevice, "__init__", _noop)
@patch.object(AwsDevice, "aws_session", new_callable=mock.PropertyMock)
@patch.object(AwsDevice, "type", new_callable=mock.PropertyMock)
@patch.object(AwsDevice, "properties")
def _aws_device(
properties_mock, type_mock, wires, device_type=AwsDeviceType.QPU, shots=SHOTS, **kwargs
properties_mock,
type_mock,
session_mock,
wires,
device_type=AwsDeviceType.QPU,
shots=SHOTS,
**kwargs,
):
properties_mock.action = {DeviceActionType.JAQCD: ACTION_PROPERTIES}
properties_mock.return_value.action.return_value = {DeviceActionType.JAQCD: ACTION_PROPERTIES}
Expand All @@ -841,13 +860,14 @@ def _aws_device(
device_arn="baz",
aws_session=Mock(),
shots=shots,
**kwargs
**kwargs,
)


@patch.object(AwsDevice, "__init__", _noop)
@patch.object(AwsDevice, "aws_session", new_callable=mock.PropertyMock)
@patch.object(AwsDevice, "properties")
def _bad_aws_device(properties_mock, wires, **kwargs):
def _bad_aws_device(properties_mock, session_mock, wires, **kwargs):
properties_mock.action = {DeviceActionType.ANNEALING: ACTION_PROPERTIES}
properties_mock.type = AwsDeviceType.QPU
return BraketAwsQubitDevice(
Expand All @@ -856,5 +876,5 @@ def _bad_aws_device(properties_mock, wires, **kwargs):
device_arn=DEVICE_ARN,
aws_session=Mock(),
shots=SHOTS,
**kwargs
**kwargs,
)

0 comments on commit ccee356

Please sign in to comment.