From d9662227874f23d3b85622254e242dbba15569dc Mon Sep 17 00:00:00 2001 From: mbeach-aws <85963088+mbeach-aws@users.noreply.github.com> Date: Fri, 3 May 2024 13:37:46 -0400 Subject: [PATCH 01/11] feature: Direct Reservation context manager (#955) * feature: context manager for reservation arns Co-authored-by: Cody Wang --- examples/reservation.py | 17 +- src/braket/aws/__init__.py | 1 + src/braket/aws/aws_session.py | 27 +++ src/braket/aws/direct_reservations.py | 98 ++++++++++ test/integ_tests/test_reservation_arn.py | 24 +-- .../braket/aws/test_direct_reservations.py | 181 ++++++++++++++++++ 6 files changed, 333 insertions(+), 15 deletions(-) create mode 100644 src/braket/aws/direct_reservations.py create mode 100644 test/unit_tests/braket/aws/test_direct_reservations.py diff --git a/examples/reservation.py b/examples/reservation.py index 682f71f50..83be87ebd 100644 --- a/examples/reservation.py +++ b/examples/reservation.py @@ -11,7 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. -from braket.aws import AwsDevice +from braket.aws import AwsDevice, DirectReservation from braket.circuits import Circuit from braket.devices import Devices @@ -19,6 +19,17 @@ device = AwsDevice(Devices.IonQ.Aria1) # To run a task in a device reservation, change the device to the one you reserved -# and fill in your reservation ARN -task = device.run(bell, shots=100, reservation_arn="reservation ARN") +# and fill in your reservation ARN. +with DirectReservation(device, reservation_arn=""): + task = device.run(bell, shots=100) +print(task.result().measurement_counts) + +# Alternatively, you may start the reservation globally +reservation = DirectReservation(device, reservation_arn="").start() +task = device.run(bell, shots=100) +print(task.result().measurement_counts) +reservation.stop() # stop creating tasks in the reservation + +# Lastly, you may pass the reservation ARN directly to a quantum task +task = device.run(bell, shots=100, reservation_arn="") print(task.result().measurement_counts) diff --git a/src/braket/aws/__init__.py b/src/braket/aws/__init__.py index d0b3a3411..3be348f34 100644 --- a/src/braket/aws/__init__.py +++ b/src/braket/aws/__init__.py @@ -16,3 +16,4 @@ from braket.aws.aws_quantum_task import AwsQuantumTask # noqa: F401 from braket.aws.aws_quantum_task_batch import AwsQuantumTaskBatch # noqa: F401 from braket.aws.aws_session import AwsSession # noqa: F401 +from braket.aws.direct_reservations import DirectReservation # noqa: F401 diff --git a/src/braket/aws/aws_session.py b/src/braket/aws/aws_session.py index 16a021e7d..b4cdfcd31 100644 --- a/src/braket/aws/aws_session.py +++ b/src/braket/aws/aws_session.py @@ -17,6 +17,7 @@ import os import os.path import re +import warnings from functools import cache from pathlib import Path from typing import Any, NamedTuple, Optional @@ -235,6 +236,32 @@ def create_quantum_task(self, **boto3_kwargs) -> str: Returns: str: The ARN of the quantum task. """ + # Add reservation arn if available and device is correct. + context_device_arn = os.getenv("AMZN_BRAKET_RESERVATION_DEVICE_ARN") + context_reservation_arn = os.getenv("AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN") + + # if the task has a reservation_arn and also context does, raise a warning + # Raise warning if reservation ARN is found in both context and task parameters + task_has_reservation = any( + item.get("type") == "RESERVATION_TIME_WINDOW_ARN" + for item in boto3_kwargs.get("associations", []) + ) + if task_has_reservation and context_reservation_arn: + warnings.warn( + "A reservation ARN was passed to 'CreateQuantumTask', but it is being overridden " + "by a 'DirectReservation' context. If this was not intended, please review your " + "reservation ARN settings or the context in which 'CreateQuantumTask' is called." + ) + + # Ensure reservation only applies to specific device + if context_device_arn == boto3_kwargs["deviceArn"] and context_reservation_arn: + boto3_kwargs["associations"] = [ + { + "arn": context_reservation_arn, + "type": "RESERVATION_TIME_WINDOW_ARN", + } + ] + # Add job token to request, if available. job_token = os.getenv("AMZN_BRAKET_JOB_TOKEN") if job_token: diff --git a/src/braket/aws/direct_reservations.py b/src/braket/aws/direct_reservations.py new file mode 100644 index 000000000..4ffcb8fce --- /dev/null +++ b/src/braket/aws/direct_reservations.py @@ -0,0 +1,98 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, 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. + +from __future__ import annotations + +import os +import warnings +from contextlib import AbstractContextManager + +from braket.aws.aws_device import AwsDevice +from braket.devices import Device + + +class DirectReservation(AbstractContextManager): + """ + Context manager that modifies AwsQuantumTasks created within the context to use a reservation + ARN for all tasks targeting the specified device. Note: this context manager only allows for + one reservation at a time. + + Reservations are AWS account and device specific. Only the AWS account that created the + reservation can use your reservation ARN. Additionally, the reservation ARN is only valid on the + reserved device at the chosen start and end times. + + Args: + device (Device | str | None): The Braket device for which you have a reservation ARN, or + optionally the device ARN. + reservation_arn (str | None): The Braket Direct reservation ARN to be applied to all + quantum tasks run within the context. + + Examples: + As a context manager + >>> with DirectReservation(device_arn, reservation_arn=""): + ... task1 = device.run(circuit, shots) + ... task2 = device.run(circuit, shots) + + or start the reservation + >>> DirectReservation(device_arn, reservation_arn="").start() + ... task1 = device.run(circuit, shots) + ... task2 = device.run(circuit, shots) + + References: + + [1] https://docs.aws.amazon.com/braket/latest/developerguide/braket-reservations.html + """ + + _is_active = False # Class variable to track active reservation context + + def __init__(self, device: Device | str | None, reservation_arn: str | None): + if isinstance(device, AwsDevice): + self.device_arn = device.arn + elif isinstance(device, str): + self.device_arn = AwsDevice(device).arn # validate ARN early + elif isinstance(device, Device) or device is None: # LocalSimulator + warnings.warn( + "Using a local simulator with the reservation. For a reservation on a QPU, please " + "ensure the device matches the reserved Braket device." + ) + self.device_arn = "" # instead of None, use empty string + else: + raise TypeError("Device must be an AwsDevice or its ARN, or a local simulator device.") + + self.reservation_arn = reservation_arn + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self.stop() + + def start(self) -> None: + """Start the reservation context.""" + if DirectReservation._is_active: + raise RuntimeError("Another reservation is already active.") + + os.environ["AMZN_BRAKET_RESERVATION_DEVICE_ARN"] = self.device_arn + if self.reservation_arn: + os.environ["AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN"] = self.reservation_arn + DirectReservation._is_active = True + + def stop(self) -> None: + """Stop the reservation context.""" + if not DirectReservation._is_active: + warnings.warn("Reservation context is not active.") + return + os.environ.pop("AMZN_BRAKET_RESERVATION_DEVICE_ARN", None) + os.environ.pop("AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN", None) + DirectReservation._is_active = False diff --git a/test/integ_tests/test_reservation_arn.py b/test/integ_tests/test_reservation_arn.py index 98b87f075..f5efd7227 100644 --- a/test/integ_tests/test_reservation_arn.py +++ b/test/integ_tests/test_reservation_arn.py @@ -15,12 +15,12 @@ import pytest from botocore.exceptions import ClientError -from test_create_quantum_job import decorator_python_version -from braket.aws import AwsDevice +from braket.aws import AwsDevice, DirectReservation from braket.circuits import Circuit from braket.devices import Devices from braket.jobs import get_job_device_arn, hybrid_job +from braket.test.integ_tests.test_create_quantum_job import decorator_python_version @pytest.fixture @@ -36,11 +36,11 @@ def test_create_task_via_invalid_reservation_arn_on_qpu(reservation_arn): device = AwsDevice(Devices.IonQ.Harmony) with pytest.raises(ClientError, match="Reservation arn is invalid"): - device.run( - circuit, - shots=10, - reservation_arn=reservation_arn, - ) + device.run(circuit, shots=10, reservation_arn=reservation_arn) + + with pytest.raises(ClientError, match="Reservation arn is invalid"): + with DirectReservation(device, reservation_arn=reservation_arn): + device.run(circuit, shots=10) def test_create_task_via_reservation_arn_on_simulator(reservation_arn): @@ -48,11 +48,11 @@ def test_create_task_via_reservation_arn_on_simulator(reservation_arn): device = AwsDevice(Devices.Amazon.SV1) with pytest.raises(ClientError, match="Braket Direct is not supported for"): - device.run( - circuit, - shots=10, - reservation_arn=reservation_arn, - ) + device.run(circuit, shots=10, reservation_arn=reservation_arn) + + with pytest.raises(ClientError, match="Braket Direct is not supported for"): + with DirectReservation(device, reservation_arn=reservation_arn): + device.run(circuit, shots=10) @pytest.mark.xfail( diff --git a/test/unit_tests/braket/aws/test_direct_reservations.py b/test/unit_tests/braket/aws/test_direct_reservations.py new file mode 100644 index 000000000..332421e7a --- /dev/null +++ b/test/unit_tests/braket/aws/test_direct_reservations.py @@ -0,0 +1,181 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, 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. + +import os +from unittest.mock import MagicMock, patch + +import pytest + +from braket.aws import AwsDevice, AwsSession, DirectReservation +from braket.devices import LocalSimulator + +RESERVATION_ARN = "arn:aws:braket:us-east-1:123456789:reservation/uuid" +DEVICE_ARN = "arn:aws:braket:us-east-1:123456789:device/qpu/ionq/Forte-1" +VALUE_ERROR_MESSAGE = "Device must be an AwsDevice or its ARN, or a local simulator device." +RUNTIME_ERROR_MESSAGE = "Another reservation is already active." + + +@pytest.fixture +def aws_device(): + mock_device = MagicMock(spec=AwsDevice) + mock_device._arn = DEVICE_ARN + type(mock_device).arn = property(lambda x: DEVICE_ARN) + return mock_device + + +def test_direct_reservation_aws_device(aws_device): + with DirectReservation(aws_device, RESERVATION_ARN) as reservation: + assert reservation.device_arn == DEVICE_ARN + assert reservation.reservation_arn == RESERVATION_ARN + assert reservation._is_active + + +def test_direct_reservation_device_str(aws_device): + with patch( + "braket.aws.AwsDevice.__init__", + side_effect=lambda self, *args, **kwargs: setattr(self, "_arn", DEVICE_ARN), + autospec=True, + ): + with patch("braket.aws.AwsDevice", return_value=aws_device, autospec=True): + with DirectReservation(DEVICE_ARN, RESERVATION_ARN) as reservation: + assert reservation.device_arn == DEVICE_ARN + assert reservation.reservation_arn == RESERVATION_ARN + assert reservation._is_active + + +def test_direct_reservation_local_simulator(): + mock_device = MagicMock(spec=LocalSimulator) + with pytest.warns(UserWarning): + with DirectReservation(mock_device, RESERVATION_ARN) as reservation: + assert os.environ["AMZN_BRAKET_RESERVATION_DEVICE_ARN"] == "" + assert os.environ["AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN"] == RESERVATION_ARN + assert reservation._is_active is True + + +@pytest.mark.parametrize("device", [123, False, [aws_device], {"a": 1}]) +def test_direct_reservation_invalid_inputs(device): + with pytest.raises(TypeError): + DirectReservation(device, RESERVATION_ARN) + + +def test_direct_reservation_local_no_reservation(): + mock_device = MagicMock(spec=LocalSimulator) + mock_device.create_quantum_task = MagicMock() + kwargs = { + "program": {"ir": '{"instructions":[]}', "qubitCount": 4}, + "shots": 1, + } + with DirectReservation(mock_device, None): + mock_device.create_quantum_task(**kwargs) + mock_device.create_quantum_task.assert_called_once_with(**kwargs) + + +def test_context_management(aws_device): + with DirectReservation(aws_device, RESERVATION_ARN): + assert os.getenv("AMZN_BRAKET_RESERVATION_DEVICE_ARN") == DEVICE_ARN + assert os.getenv("AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN") == RESERVATION_ARN + assert not os.getenv("AMZN_BRAKET_RESERVATION_DEVICE_ARN") + assert not os.getenv("AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN") + + +def test_start_reservation_already_active(aws_device): + reservation = DirectReservation(aws_device, RESERVATION_ARN) + reservation.start() + with pytest.raises(RuntimeError, match=RUNTIME_ERROR_MESSAGE): + reservation.start() + reservation.stop() + + +def test_stop_reservation_not_active(aws_device): + reservation = DirectReservation(aws_device, RESERVATION_ARN) + with pytest.warns(UserWarning): + reservation.stop() + + +def test_multiple_start_stop_cycles(aws_device): + reservation = DirectReservation(aws_device, RESERVATION_ARN) + reservation.start() + reservation.stop() + reservation.start() + reservation.stop() + assert not os.getenv("AMZN_BRAKET_RESERVATION_DEVICE_ARN") + assert not os.getenv("AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN") + + +def test_two_direct_reservations(aws_device): + with pytest.raises(RuntimeError, match=RUNTIME_ERROR_MESSAGE): + with DirectReservation(aws_device, RESERVATION_ARN): + with DirectReservation(aws_device, "reservation_arn_example_2"): + pass + + +def test_create_quantum_task_with_correct_device_and_reservation(aws_device): + kwargs = {"deviceArn": DEVICE_ARN, "shots": 1} + with patch("boto3.client"): + mock_client = MagicMock() + aws_session = AwsSession(braket_client=mock_client) + with DirectReservation(aws_device, RESERVATION_ARN): + aws_session.create_quantum_task(**kwargs) + kwargs["associations"] = [ + { + "arn": RESERVATION_ARN, + "type": "RESERVATION_TIME_WINDOW_ARN", + } + ] + mock_client.create_quantum_task.assert_called_once_with(**kwargs) + + +def test_warning_for_overridden_reservation_arn(aws_device): + kwargs = { + "deviceArn": DEVICE_ARN, + "shots": 1, + "associations": [ + { + "arn": "task_reservation_arn", + "type": "RESERVATION_TIME_WINDOW_ARN", + } + ], + } + correct_kwargs = { + "deviceArn": DEVICE_ARN, + "shots": 1, + "associations": [ + { + "arn": RESERVATION_ARN, + "type": "RESERVATION_TIME_WINDOW_ARN", + } + ], + } + with patch("boto3.client"): + mock_client = MagicMock() + aws_session = AwsSession(braket_client=mock_client) + with pytest.warns( + UserWarning, + match="A reservation ARN was passed to 'CreateQuantumTask', but it is being overridden", + ): + with DirectReservation(aws_device, RESERVATION_ARN): + aws_session.create_quantum_task(**kwargs) + mock_client.create_quantum_task.assert_called_once_with(**correct_kwargs) + + +def test_warning_not_triggered_wrong_association_type(): + kwargs = { + "deviceArn": DEVICE_ARN, + "shots": 1, + "associations": [{"type": "OTHER_TYPE"}], + } + with patch("boto3.client"): + mock_client = MagicMock() + aws_session = AwsSession(braket_client=mock_client) + aws_session.create_quantum_task(**kwargs) + mock_client.create_quantum_task.assert_called_once_with(**kwargs) From 84c3bb7c7cecefbe0676c92c1748aa729bdfc6cf Mon Sep 17 00:00:00 2001 From: Ashlyn Hanson <65787294+ashlhans@users.noreply.github.com> Date: Fri, 3 May 2024 14:44:14 -0700 Subject: [PATCH 02/11] doc: correct the example in the measure docstring (#965) --- src/braket/circuits/circuit.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/braket/circuits/circuit.py b/src/braket/circuits/circuit.py index 94d53248c..f15e03647 100644 --- a/src/braket/circuits/circuit.py +++ b/src/braket/circuits/circuit.py @@ -737,7 +737,6 @@ def measure(self, target_qubits: QubitSetInput) -> Circuit: [Instruction('operator': H('qubit_count': 1), 'target': QubitSet([Qubit(0)]), Instruction('operator': CNot('qubit_count': 2), 'target': QubitSet([Qubit(0), Qubit(1)]), - Instruction('operator': H('qubit_count': 1), 'target': QubitSet([Qubit(2)]), Instruction('operator': Measure, 'target': QubitSet([Qubit(0)])] """ if not isinstance(target_qubits, Iterable): From 2730aa1405695d389e747815a20ef2ee37d39389 Mon Sep 17 00:00:00 2001 From: Abe Coull <85974725+math411@users.noreply.github.com> Date: Mon, 6 May 2024 14:23:45 -0700 Subject: [PATCH 03/11] infra: add opts for tox builds ran in parallel (#759) --- .github/workflows/check-format.yml | 2 +- README.md | 6 ++++++ tox.ini | 3 +-- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/.github/workflows/check-format.yml b/.github/workflows/check-format.yml index 8f6807b5e..a6106b2d7 100644 --- a/.github/workflows/check-format.yml +++ b/.github/workflows/check-format.yml @@ -26,4 +26,4 @@ jobs: pip install tox - name: Run code format checks run: | - tox -e linters_check + tox -e linters_check -p auto diff --git a/README.md b/README.md index b03bd1ef4..0c935853d 100644 --- a/README.md +++ b/README.md @@ -206,6 +206,12 @@ To run linters and doc generators and unit tests: tox ``` +or if your machine can handle multithreaded workloads, run them in parallel with: + +```bash +tox -p auto +``` + ### Integration Tests First, configure a profile to use your account to interact with AWS. To learn more, see [Configure AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html). diff --git a/tox.ini b/tox.ini index 98a9b30e3..d4fae64c5 100644 --- a/tox.ini +++ b/tox.ini @@ -1,7 +1,6 @@ [tox] envlist = clean,linters,docs,unit-tests - [testenv] parallel_show_output = true package = wheel @@ -111,7 +110,7 @@ deps = sphinx-rtd-theme sphinxcontrib-apidoc commands = - sphinx-build -E -T -b html doc build/documentation/html + sphinx-build -E -T -b html doc build/documentation/html -j auto [testenv:serve-docs] basepython = python3 From 7824bfbb4ed7215143d578f1754bc5881458cf19 Mon Sep 17 00:00:00 2001 From: Abe Coull <85974725+math411@users.noreply.github.com> Date: Mon, 6 May 2024 15:04:00 -0700 Subject: [PATCH 04/11] infra: allow worksteal for testing (#960) --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index d75c4f034..ab93f5955 100644 --- a/setup.cfg +++ b/setup.cfg @@ -5,7 +5,7 @@ test=pytest xfail_strict = true # https://pytest-xdist.readthedocs.io/en/latest/known-limitations.html addopts = - --verbose -n logical --durations=0 --durations-min=1 + --verbose -n logical --durations=0 --durations-min=1 --dist worksteal testpaths = test/unit_tests filterwarnings= # Issue #557 in `pytest-cov` (currently v4.x) has not moved for a while now, From 6d30734ace5d5b98884626aba627cd421bb4a647 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Mon, 6 May 2024 15:20:42 -0700 Subject: [PATCH 05/11] test: Extract `decorator_python_version` (#968) --- model.tar.gz | Bin 336 -> 0 bytes test/integ_tests/job_testing_utils.py | 25 ++++++++++++++++++++ test/integ_tests/test_create_quantum_job.py | 12 ++-------- test/integ_tests/test_reservation_arn.py | 2 +- 4 files changed, 28 insertions(+), 11 deletions(-) delete mode 100644 model.tar.gz create mode 100644 test/integ_tests/job_testing_utils.py diff --git a/model.tar.gz b/model.tar.gz deleted file mode 100644 index 93bf6a4a03f7d08314601e2907a704651eb0b07a..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 336 zcmV-W0k8faiwFP!000001MHQ}OT#c2#(Va!2svx^rcE0sc<_@AJcxpL8(AB)b4^B) z%4F<+H{HjjuuWKXsQq1x%3_`*Q35O_NgDPfvx=n;VBX>FXTDp6uL>oc|;iH ztQ*0R_tGt1vB_fzy6azFJY4nqPd6kr(*HrLP1T3Kf&b071ir@3{KvGOe~7{W?VZW5 zu+G2H+HI@b<^R(B&+yQQH|ZYJS6PVVLx9iF3@cGcKUmphq=$Bp2`9+JzZAK3G8=ep zA>m_$-z!zCY6Zn}FI2{Lo>s{h=3~(^)ykK>$jr~2DW$KH$_tfy0wi27yVa%;u4*+I ii(EN5b$EX0i)v|UY58M(0ssL2{{sNvw;Ch>3;+NyA)$Q$ diff --git a/test/integ_tests/job_testing_utils.py b/test/integ_tests/job_testing_utils.py new file mode 100644 index 000000000..4493df180 --- /dev/null +++ b/test/integ_tests/job_testing_utils.py @@ -0,0 +1,25 @@ +# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, 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. + +import re + +from braket.aws import AwsSession +from braket.jobs import Framework, retrieve_image + + +def decorator_python_version(): + aws_session = AwsSession() + image_uri = retrieve_image(Framework.BASE, aws_session.region) + tag = aws_session.get_full_image_tag(image_uri) + major_version, minor_version = re.search(r"-py(\d)(\d+)-", tag).groups() + return int(major_version), int(minor_version) diff --git a/test/integ_tests/test_create_quantum_job.py b/test/integ_tests/test_create_quantum_job.py index be0e49a85..ce88d122b 100644 --- a/test/integ_tests/test_create_quantum_job.py +++ b/test/integ_tests/test_create_quantum_job.py @@ -22,19 +22,11 @@ import job_test_script import pytest from job_test_module.job_test_submodule.job_test_submodule_file import submodule_helper +from job_testing_utils import decorator_python_version -from braket.aws import AwsSession from braket.aws.aws_quantum_job import AwsQuantumJob from braket.devices import Devices -from braket.jobs import Framework, get_input_data_dir, hybrid_job, retrieve_image, save_job_result - - -def decorator_python_version(): - aws_session = AwsSession() - image_uri = retrieve_image(Framework.BASE, aws_session.region) - tag = aws_session.get_full_image_tag(image_uri) - major_version, minor_version = re.search(r"-py(\d)(\d+)-", tag).groups() - return int(major_version), int(minor_version) +from braket.jobs import get_input_data_dir, hybrid_job, save_job_result def test_failed_quantum_job(aws_session, capsys, failed_quantum_job): diff --git a/test/integ_tests/test_reservation_arn.py b/test/integ_tests/test_reservation_arn.py index f5efd7227..64135f76e 100644 --- a/test/integ_tests/test_reservation_arn.py +++ b/test/integ_tests/test_reservation_arn.py @@ -15,12 +15,12 @@ import pytest from botocore.exceptions import ClientError +from job_testing_utils import decorator_python_version from braket.aws import AwsDevice, DirectReservation from braket.circuits import Circuit from braket.devices import Devices from braket.jobs import get_job_device_arn, hybrid_job -from braket.test.integ_tests.test_create_quantum_job import decorator_python_version @pytest.fixture From 376fb90d122ba88d2b312081705a8d1185903a95 Mon Sep 17 00:00:00 2001 From: ci Date: Mon, 6 May 2024 22:40:25 +0000 Subject: [PATCH 06/11] prepare release v1.79.0 --- CHANGELOG.md | 10 ++++++++++ src/braket/_sdk/_version.py | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4cb8ae8a6..49eba0f5b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## v1.79.0 (2024-05-06) + +### Features + + * Direct Reservation context manager + +### Documentation Changes + + * correct the example in the measure docstring + ## v1.78.0 (2024-04-18) ### Features diff --git a/src/braket/_sdk/_version.py b/src/braket/_sdk/_version.py index fec1b4f08..1d9dd72b1 100644 --- a/src/braket/_sdk/_version.py +++ b/src/braket/_sdk/_version.py @@ -15,4 +15,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "1.78.1.dev0" +__version__ = "1.79.0" From b791858ca6a5e036d727a9fb4e0f377b0b0b4bff Mon Sep 17 00:00:00 2001 From: ci Date: Mon, 6 May 2024 22:40:25 +0000 Subject: [PATCH 07/11] update development version to v1.79.1.dev0 --- src/braket/_sdk/_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/braket/_sdk/_version.py b/src/braket/_sdk/_version.py index 1d9dd72b1..1f46f945a 100644 --- a/src/braket/_sdk/_version.py +++ b/src/braket/_sdk/_version.py @@ -15,4 +15,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "1.79.0" +__version__ = "1.79.1.dev0" From 6c4282e7eafcff906c4cfc6804bf06209d82a4fc Mon Sep 17 00:00:00 2001 From: Abe Coull <85974725+math411@users.noreply.github.com> Date: Tue, 7 May 2024 18:09:55 -0700 Subject: [PATCH 08/11] fix: check the qubit set length against observables (#970) --- src/braket/circuits/result_type.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/braket/circuits/result_type.py b/src/braket/circuits/result_type.py index b66d4da67..3e9f0dfad 100644 --- a/src/braket/circuits/result_type.py +++ b/src/braket/circuits/result_type.py @@ -221,7 +221,7 @@ def __init__( "target length is equal to the observable term's qubits count." ) self._target = [QubitSet(term_target) for term_target in target] - for term_target, obs in zip(target, observable.summands): + for term_target, obs in zip(self._target, observable.summands): if obs.qubit_count != len(term_target): raise ValueError( "Sum observable's target shape must be a nested list where each term's " From dfd75b38aba98ba748ef7c99d8e241f527760c39 Mon Sep 17 00:00:00 2001 From: ci Date: Wed, 8 May 2024 22:01:49 +0000 Subject: [PATCH 09/11] prepare release v1.79.1 --- CHANGELOG.md | 6 ++++++ src/braket/_sdk/_version.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 49eba0f5b..e54b3f658 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## v1.79.1 (2024-05-08) + +### Bug Fixes and Other Changes + + * check the qubit set length against observables + ## v1.79.0 (2024-05-06) ### Features diff --git a/src/braket/_sdk/_version.py b/src/braket/_sdk/_version.py index 1f46f945a..c15ebeb09 100644 --- a/src/braket/_sdk/_version.py +++ b/src/braket/_sdk/_version.py @@ -15,4 +15,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "1.79.1.dev0" +__version__ = "1.79.1" From ec5edafd43dbb498bf5e0dc623f6b1ff74b4698c Mon Sep 17 00:00:00 2001 From: ci Date: Wed, 8 May 2024 22:01:49 +0000 Subject: [PATCH 10/11] update development version to v1.79.2.dev0 --- src/braket/_sdk/_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/braket/_sdk/_version.py b/src/braket/_sdk/_version.py index c15ebeb09..0215a4833 100644 --- a/src/braket/_sdk/_version.py +++ b/src/braket/_sdk/_version.py @@ -15,4 +15,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "1.79.1" +__version__ = "1.79.2.dev0" From 7ba54cedaa300f55dd6031a47f66d3fce7eec17f Mon Sep 17 00:00:00 2001 From: Ryan Shaffer <3620100+rmshaffer@users.noreply.github.com> Date: Tue, 21 May 2024 12:54:59 -0400 Subject: [PATCH 11/11] feature: Add support for SerializableProgram abstraction to Device interface (#976) --- src/braket/aws/aws_quantum_task.py | 29 +++++++++ src/braket/circuits/serialization.py | 21 +++++++ src/braket/devices/local_simulator.py | 60 +++++++++++++++---- .../braket/aws/test_aws_quantum_task.py | 28 +++++++++ .../braket/devices/test_local_simulator.py | 38 ++++++++++++ 5 files changed, 164 insertions(+), 12 deletions(-) diff --git a/src/braket/aws/aws_quantum_task.py b/src/braket/aws/aws_quantum_task.py index 17ae29252..a21c7782c 100644 --- a/src/braket/aws/aws_quantum_task.py +++ b/src/braket/aws/aws_quantum_task.py @@ -34,6 +34,7 @@ IRType, OpenQASMSerializationProperties, QubitReferenceType, + SerializableProgram, ) from braket.device_schema import GateModelParameters from braket.device_schema.dwave import ( @@ -623,6 +624,34 @@ def _( return AwsQuantumTask(task_arn, aws_session, *args, **kwargs) +@_create_internal.register +def _( + serializable_program: SerializableProgram, + aws_session: AwsSession, + create_task_kwargs: dict[str, Any], + device_arn: str, + device_parameters: Union[dict, BraketSchemaBase], + _disable_qubit_rewiring: bool, + inputs: dict[str, float], + gate_definitions: Optional[dict[tuple[Gate, QubitSet], PulseSequence]], + *args, + **kwargs, +) -> AwsQuantumTask: + openqasm_program = OpenQASMProgram(source=serializable_program.to_ir(ir_type=IRType.OPENQASM)) + return _create_internal( + openqasm_program, + aws_session, + create_task_kwargs, + device_arn, + device_parameters, + _disable_qubit_rewiring, + inputs, + gate_definitions, + *args, + **kwargs, + ) + + @_create_internal.register def _( blackbird_program: BlackbirdProgram, diff --git a/src/braket/circuits/serialization.py b/src/braket/circuits/serialization.py index afcb5d118..fdee7d144 100644 --- a/src/braket/circuits/serialization.py +++ b/src/braket/circuits/serialization.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. +from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum @@ -32,6 +33,26 @@ class QubitReferenceType(str, Enum): PHYSICAL = "PHYSICAL" +class SerializableProgram(ABC): + @abstractmethod + def to_ir( + self, + ir_type: IRType = IRType.OPENQASM, + ) -> str: + """Serializes the program into an intermediate representation. + + Args: + ir_type (IRType): The IRType to use for converting the program to its + IR representation. Defaults to IRType.OPENQASM. + + Raises: + ValueError: Raised if the supplied `ir_type` is not supported. + + Returns: + str: A representation of the program in the `ir_type` format. + """ + + @dataclass class OpenQASMSerializationProperties: """Properties for serializing a circuit to OpenQASM. diff --git a/src/braket/devices/local_simulator.py b/src/braket/devices/local_simulator.py index 1dec56d37..15ec904de 100644 --- a/src/braket/devices/local_simulator.py +++ b/src/braket/devices/local_simulator.py @@ -25,11 +25,11 @@ from braket.circuits import Circuit from braket.circuits.circuit_helpers import validate_circuit_and_shots from braket.circuits.noise_model import NoiseModel -from braket.circuits.serialization import IRType +from braket.circuits.serialization import IRType, SerializableProgram from braket.device_schema import DeviceActionType, DeviceCapabilities from braket.devices.device import Device from braket.ir.ahs import Program as AHSProgram -from braket.ir.openqasm import Program +from braket.ir.openqasm import Program as OpenQASMProgram from braket.simulator import BraketSimulator from braket.tasks import AnnealingQuantumTaskResult, GateModelQuantumTaskResult from braket.tasks.analog_hamiltonian_simulation_quantum_task_result import ( @@ -80,7 +80,9 @@ def __init__( def run( self, - task_specification: Union[Circuit, Problem, Program, AnalogHamiltonianSimulation], + task_specification: Union[ + Circuit, Problem, OpenQASMProgram, AnalogHamiltonianSimulation, SerializableProgram + ], shots: int = 0, inputs: Optional[dict[str, float]] = None, *args: Any, @@ -89,7 +91,7 @@ def run( """Runs the given task with the wrapped local simulator. Args: - task_specification (Union[Circuit, Problem, Program, AnalogHamiltonianSimulation]): + task_specification (Union[Circuit, Problem, OpenQASMProgram, AnalogHamiltonianSimulation, SerializableProgram]): # noqa E501 The quantum task specification. shots (int): The number of times to run the circuit or annealing problem. Default is 0, which means that the simulator will compute the exact @@ -122,8 +124,18 @@ def run( def run_batch( # noqa: C901 self, task_specifications: Union[ - Union[Circuit, Problem, Program, AnalogHamiltonianSimulation], - list[Union[Circuit, Problem, Program, AnalogHamiltonianSimulation]], + Union[ + Circuit, Problem, OpenQASMProgram, AnalogHamiltonianSimulation, SerializableProgram + ], + list[ + Union[ + Circuit, + Problem, + OpenQASMProgram, + AnalogHamiltonianSimulation, + SerializableProgram, + ] + ], ], shots: Optional[int] = 0, max_parallel: Optional[int] = None, @@ -134,7 +146,7 @@ def run_batch( # noqa: C901 """Executes a batch of quantum tasks in parallel Args: - task_specifications (Union[Union[Circuit, Problem, Program, AnalogHamiltonianSimulation], list[Union[Circuit, Problem, Program, AnalogHamiltonianSimulation]]]): + task_specifications (Union[Union[Circuit, Problem, OpenQASMProgram, AnalogHamiltonianSimulation, SerializableProgram], list[Union[Circuit, Problem, OpenQASMProgram, AnalogHamiltonianSimulation, SerializableProgram]]]): # noqa Single instance or list of quantum task specification. shots (Optional[int]): The number of times to run the quantum task. Default: 0. @@ -163,7 +175,7 @@ def run_batch( # noqa: C901 single_task = isinstance( task_specifications, - (Circuit, Program, Problem, AnalogHamiltonianSimulation), + (Circuit, OpenQASMProgram, Problem, AnalogHamiltonianSimulation), ) single_input = isinstance(inputs, dict) @@ -220,7 +232,9 @@ def registered_backends() -> set[str]: def _run_internal_wrap( self, - task_specification: Union[Circuit, Problem, Program, AnalogHamiltonianSimulation], + task_specification: Union[ + Circuit, Problem, OpenQASMProgram, AnalogHamiltonianSimulation, SerializableProgram + ], shots: Optional[int] = None, inputs: Optional[dict[str, float]] = None, *args, @@ -250,7 +264,12 @@ def _(self, backend_impl: BraketSimulator): def _run_internal( self, task_specification: Union[ - Circuit, Problem, Program, AnalogHamiltonianSimulation, AHSProgram + Circuit, + Problem, + OpenQASMProgram, + AnalogHamiltonianSimulation, + AHSProgram, + SerializableProgram, ], shots: Optional[int] = None, *args, @@ -296,7 +315,7 @@ def _(self, problem: Problem, shots: Optional[int] = None, *args, **kwargs): @_run_internal.register def _( self, - program: Program, + program: OpenQASMProgram, shots: Optional[int] = None, inputs: Optional[dict[str, float]] = None, *args, @@ -308,13 +327,30 @@ def _( if inputs: inputs_copy = program.inputs.copy() if program.inputs is not None else {} inputs_copy.update(inputs) - program = Program( + program = OpenQASMProgram( source=program.source, inputs=inputs_copy, ) + results = simulator.run(program, shots, *args, **kwargs) + + if isinstance(results, GateModelQuantumTaskResult): + return results + return GateModelQuantumTaskResult.from_object(results) + @_run_internal.register + def _( + self, + program: SerializableProgram, + shots: Optional[int] = None, + inputs: Optional[dict[str, float]] = None, + *args, + **kwargs, + ): + program = OpenQASMProgram(source=program.to_ir(ir_type=IRType.OPENQASM)) + return self._run_internal(program, shots, inputs, *args, **kwargs) + @_run_internal.register def _( self, diff --git a/test/unit_tests/braket/aws/test_aws_quantum_task.py b/test/unit_tests/braket/aws/test_aws_quantum_task.py index 16a72da7a..28032d943 100644 --- a/test/unit_tests/braket/aws/test_aws_quantum_task.py +++ b/test/unit_tests/braket/aws/test_aws_quantum_task.py @@ -33,6 +33,7 @@ IRType, OpenQASMSerializationProperties, QubitReferenceType, + SerializableProgram, ) from braket.device_schema import GateModelParameters, error_mitigation from braket.device_schema.dwave import ( @@ -123,6 +124,19 @@ def openqasm_program(): return OpenQASMProgram(source="OPENQASM 3.0; h $0;") +class DummySerializableProgram(SerializableProgram): + def __init__(self, source: str): + self.source = source + + def to_ir(self, ir_type: IRType = IRType.OPENQASM) -> str: + return self.source + + +@pytest.fixture +def serializable_program(): + return DummySerializableProgram(source="OPENQASM 3.0; h $0;") + + @pytest.fixture def blackbird_program(): return BlackbirdProgram(source="Vac | q[0]") @@ -614,6 +628,20 @@ def test_create_openqasm_program_em_serialized(aws_session, arn, openqasm_progra ) +def test_create_serializable_program(aws_session, arn, serializable_program): + aws_session.create_quantum_task.return_value = arn + shots = 21 + AwsQuantumTask.create(aws_session, SIMULATOR_ARN, serializable_program, S3_TARGET, shots) + + _assert_create_quantum_task_called_with( + aws_session, + SIMULATOR_ARN, + OpenQASMProgram(source=serializable_program.to_ir()).json(), + S3_TARGET, + shots, + ) + + def test_create_blackbird_program(aws_session, arn, blackbird_program): aws_session.create_quantum_task.return_value = arn shots = 21 diff --git a/test/unit_tests/braket/devices/test_local_simulator.py b/test/unit_tests/braket/devices/test_local_simulator.py index a7e8bfe17..216d161c7 100644 --- a/test/unit_tests/braket/devices/test_local_simulator.py +++ b/test/unit_tests/braket/devices/test_local_simulator.py @@ -27,6 +27,7 @@ from braket.annealing import Problem, ProblemType from braket.circuits import Circuit, FreeParameter, Gate, Noise from braket.circuits.noise_model import GateCriteria, NoiseModel, NoiseModelInstruction +from braket.circuits.serialization import IRType, SerializableProgram from braket.device_schema import DeviceActionType, DeviceCapabilities from braket.device_schema.openqasm_device_action_properties import OpenQASMDeviceActionProperties from braket.devices import LocalSimulator, local_simulator @@ -250,6 +251,24 @@ def properties(self) -> DeviceCapabilities: return device_properties +class DummySerializableProgram(SerializableProgram): + def __init__(self, source: str): + self.source = source + + def to_ir(self, ir_type: IRType = IRType.OPENQASM) -> str: + return self.source + + +class DummySerializableProgramSimulator(DummyProgramSimulator): + def run( + self, + program: SerializableProgram, + shots: int = 0, + batch_size: int = 1, + ) -> GateModelQuantumTaskResult: + return GateModelQuantumTaskResult.from_object(GATE_MODEL_RESULT) + + class DummyProgramDensityMatrixSimulator(BraketSimulator): def run( self, program: ir.openqasm.Program, shots: Optional[int], *args, **kwargs @@ -556,6 +575,25 @@ def test_run_program_model(): assert task.result() == GateModelQuantumTaskResult.from_object(GATE_MODEL_RESULT) +def test_run_serializable_program_model(): + dummy = DummySerializableProgramSimulator() + sim = LocalSimulator(dummy) + task = sim.run( + DummySerializableProgram( + source=""" +qubit[2] q; +bit[2] c; + +h q[0]; +cnot q[0], q[1]; + +c = measure q; +""" + ) + ) + assert task.result() == GateModelQuantumTaskResult.from_object(GATE_MODEL_RESULT) + + @pytest.mark.xfail(raises=ValueError) def test_run_gate_model_value_error(): dummy = DummyCircuitSimulator()