Merge branch 'main' into integ

.github/workflows/code-freeze.yml

@@ -0,0 +1,39 @@
+name: Code Freeze
+
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+env:
+  FROZEN: ${{ vars.FROZEN }}
+  UNFROZEN_PREFIX: ${{ vars.UNFROZEN_PREFIX }}
+
+jobs:
+  check-pr-frozen-status:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Fetch PR data and check if merge allowed
+      if: env.FROZEN == 'true'
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      run: |
+        PR_DATA=$(curl -s \
+          -H "Authorization: Bearer $GITHUB_TOKEN" \
+          -H "Accept: application/vnd.github.v3+json" \
+          https://api.github.com/repos/${{ github.repository }}/pulls/${{ github.event.pull_request.number }})
+        BRANCH_NAME=$(echo $PR_DATA | jq .head.ref -r)
+        PR_TITLE=$(echo $PR_DATA | jq .title -r)
+
+        echo $BRANCH_NAME
+        echo $PR_TITLE
+
+        if [[ "$BRANCH_NAME" != $UNFROZEN_PREFIX* ]] && 
+           [[ "$PR_TITLE" != fix:* && "$PR_TITLE" != *"[critical]"* ]]; then
+          echo "Error: You can only merge from branches that start with '$UNFROZEN_PREFIX', or PRs titled with 'fix: ' and containing '[critical]'."
+          exit 1
+        fi

.github/workflows/publish-to-pypi.yml

@@ -26,6 +26,6 @@ jobs:
       - name: Build a binary wheel and a source tarball
         run: python setup.py sdist bdist_wheel
       - name: Publish distribution to PyPI
-        uses: pypa/gh-action-pypi-publish@b7f401de30cb6434a1e19f805ff006643653240e # release/v1
+        uses: pypa/gh-action-pypi-publish@2f6f737ca5f74c637829c0f5c3acd0e29ea5e8bf # release/v1
         with:
           password: ${{ secrets.pypi_token }}

CHANGELOG.md

@@ -1,5 +1,26 @@
 # Changelog
 
+## v1.63.0 (2023-12-05)
+
+### Features
+
+ * Allow reservation ARN in task and job creation
+
+### Bug Fixes and Other Changes
+
+ * Add Forte 1 device
+
+## v1.62.1 (2023-11-17)
+
+### Bug Fixes and Other Changes
+
+ * Fix broken link to example notebook
+ * update: default no longer returning RETIRED devices from get_devices
+
+### Documentation Changes
+
+ * Add matrix expressions to docstrings
+
 ## v1.62.0 (2023-11-09)
 
 ### Features

doc/examples-hybrid-jobs.rst

@@ -8,7 +8,7 @@ Learn more about hybrid jobs on Amazon Braket.
     :maxdepth: 2
 
 ************************************************************************************************************************************************************************************************
-`Creating your first Hybrid Job <https://github.com/amazon-braket/amazon-braket-examples/blob/main/examples/hybrid_jobs/0_Creating_your_first_Hybrid_Job/Creating_your_first_Hybrid_Job.ipynb>`_
+`Creating your first Hybrid Job <https://github.com/amazon-braket/amazon-braket-examples/blob/main/examples/hybrid_jobs/0_Creating_your_first_Hybrid_Job/0_Creating_your_first_Hybrid_Job.ipynb>`_
 ************************************************************************************************************************************************************************************************
 
 This tutorial shows how to run your first Amazon Braket Hybrid Job.

examples/hybrid_job.py

@@ -18,7 +18,14 @@
 from braket.jobs.metrics import log_metric
 
 
-@hybrid_job(device=Devices.Amazon.SV1, wait_until_complete=True)
+@hybrid_job(
+    device=Devices.Amazon.SV1,
+    wait_until_complete=True,
+    # If you want to run the job in a device reservation window,
+    # change the device to the one you've reserved,
+    # uncomment the following line and fill in your reservation ARN
+    # reservation_arn="<reservation ARN>"
+)
 def run_hybrid_job(num_tasks=1):
     # declare AwsDevice within the hybrid job
     device = AwsDevice(get_job_device_arn())

examples/reservation.py

@@ -0,0 +1,24 @@
+# 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 braket.aws import AwsDevice
+from braket.circuits import Circuit
+from braket.devices import Devices
+
+bell = Circuit().h(0).cnot(0, 1)
+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")
+print(task.result().measurement_counts)

src/braket/_sdk/_version.py

@@ -15,4 +15,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "1.62.1.dev0"
+__version__ = "1.63.1.dev0"

src/braket/aws/aws_device.py

@@ -121,6 +121,7 @@ def run(
         poll_interval_seconds: Optional[float] = None,
         inputs: Optional[dict[str, float]] = None,
         gate_definitions: Optional[dict[tuple[Gate, QubitSet], PulseSequence]] = None,
+        reservation_arn: str | None = None,
         *aws_quantum_task_args,
         **aws_quantum_task_kwargs,
     ) -> AwsQuantumTask:
@@ -150,6 +151,11 @@ def run(
                 The calibration is defined for a particular `Gate` on a particular `QubitSet`
                 and is represented by a `PulseSequence`.
                 Default: None.
+            reservation_arn (str | None): The reservation ARN provided by Braket Direct
+                to reserve exclusive usage for the device to run the quantum task on.
+                Note: If you are creating tasks in a job that itself was created reservation ARN,
+                those tasks do not need to be created with the reservation ARN.
+                Default: None.
 
         Returns:
             AwsQuantumTask: An AwsQuantumTask that tracks the execution on the device.
@@ -199,6 +205,7 @@ def run(
             poll_interval_seconds=poll_interval_seconds or self._poll_interval_seconds,
             inputs=inputs,
             gate_definitions=gate_definitions,
+            reservation_arn=reservation_arn,
             *aws_quantum_task_args,
             **aws_quantum_task_kwargs,
         )
@@ -233,6 +240,7 @@ def run_batch(
         poll_interval_seconds: float = AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL,
         inputs: Optional[Union[dict[str, float], list[dict[str, float]]]] = None,
         gate_definitions: Optional[dict[tuple[Gate, QubitSet], PulseSequence]] = None,
+        reservation_arn: Optional[str] = None,
         *aws_quantum_task_args,
         **aws_quantum_task_kwargs,
     ) -> AwsQuantumTaskBatch:
@@ -263,7 +271,11 @@ def run_batch(
             gate_definitions (Optional[dict[tuple[Gate, QubitSet], PulseSequence]]): A
                 `dict[tuple[Gate, QubitSet], PulseSequence]]` for a user defined gate calibration.
                 The calibration is defined for a particular `Gate` on a particular `QubitSet`
-                and is represented by a `PulseSequence`.
+                and is represented by a `PulseSequence`. Default: None.
+            reservation_arn (Optional[str]): The reservation ARN provided by Braket Direct
+                to reserve exclusive usage for the device to run the quantum task on.
+                Note: If you are creating tasks in a job that itself was created reservation ARN,
+                those tasks do not need to be created with the reservation ARN.
                 Default: None.
 
         Returns:
@@ -290,6 +302,7 @@ def run_batch(
             poll_interval_seconds=poll_interval_seconds or self._poll_interval_seconds,
             inputs=inputs,
             gate_definitions=gate_definitions,
+            reservation_arn=reservation_arn,
             *aws_quantum_task_args,
             **aws_quantum_task_kwargs,
         )
@@ -564,13 +577,15 @@ def get_devices(
             >>> AwsDevice.get_devices(types=['SIMULATOR'])
 
         Args:
-            arns (Optional[list[str]]): device ARN list, default is `None`
-            names (Optional[list[str]]): device name list, default is `None`
-            types (Optional[list[AwsDeviceType]]): device type list, default is `None`
+            arns (Optional[list[str]]): device ARN filter, default is `None`
+            names (Optional[list[str]]): device name filter, default is `None`
+            types (Optional[list[AwsDeviceType]]): device type filter, default is `None`
                 QPUs will be searched for all regions and simulators will only be
                 searched for the region of the current session.
-            statuses (Optional[list[str]]): device status list, default is `None`
-            provider_names (Optional[list[str]]): provider name list, default is `None`
+            statuses (Optional[list[str]]): device status filter, default is `None`. When `None`
+                is used, RETIRED devices will not be returned. To include RETIRED devices in
+                the results, use a filter that includes "RETIRED" for this parameter.
+            provider_names (Optional[list[str]]): provider name filter, default is `None`
             order_by (str): field to order result by, default is `name`.
                 Accepted values are ['arn', 'name', 'type', 'provider_name', 'status']
             aws_session (Optional[AwsSession]): An AWS session object.

src/braket/aws/aws_quantum_job.py

@@ -81,6 +81,7 @@ def create(
         aws_session: AwsSession | None = None,
         tags: dict[str, str] | None = None,
         logger: Logger = getLogger(__name__),
+        reservation_arn: str | None = None,
     ) -> AwsQuantumJob:
         """Creates a hybrid job by invoking the Braket CreateJob API.
 
@@ -175,6 +176,10 @@ def create(
                 while waiting for quantum task to be in a terminal state. Default is
                 `getLogger(__name__)`
 
+            reservation_arn (str | None): the reservation window arn provided by Braket
+                Direct to reserve exclusive usage for the device to run the hybrid job on.
+                Default: None.
+
         Returns:
             AwsQuantumJob: Hybrid job tracking the execution on Amazon Braket.
 
@@ -201,6 +206,7 @@ def create(
             checkpoint_config=checkpoint_config,
             aws_session=aws_session,
             tags=tags,
+            reservation_arn=reservation_arn,
         )
 
         job_arn = aws_session.create_job(**create_job_kwargs)

src/braket/aws/aws_quantum_task.py

@@ -105,6 +105,7 @@ def create(
         tags: dict[str, str] | None = None,
         inputs: dict[str, float] | None = None,
         gate_definitions: Optional[dict[tuple[Gate, QubitSet], PulseSequence]] | None = None,
+        reservation_arn: str | None = None,
         *args,
         **kwargs,
     ) -> AwsQuantumTask:
@@ -151,6 +152,12 @@ def create(
                 a `PulseSequence`.
                 Default: None.
 
+            reservation_arn (str | None): The reservation ARN provided by Braket Direct
+                to reserve exclusive usage for the device to run the quantum task on.
+                Note: If you are creating tasks in a job that itself was created reservation ARN,
+                those tasks do not need to be created with the reservation ARN.
+                Default: None.
+
         Returns:
             AwsQuantumTask: AwsQuantumTask tracking the quantum task execution on the device.
 
@@ -179,6 +186,18 @@ def create(
             create_task_kwargs.update({"tags": tags})
         inputs = inputs or {}
 
+        if reservation_arn:
+            create_task_kwargs.update(
+                {
+                    "associations": [
+                        {
+                            "arn": reservation_arn,
+                            "type": "RESERVATION_TIME_WINDOW_ARN",
+                        }
+                    ]
+                }
+            )
+
         if isinstance(task_specification, Circuit):
             param_names = {param.name for param in task_specification.parameters}
             unbounded_parameters = param_names - set(inputs.keys())
@@ -477,6 +496,12 @@ async def _wait_for_completion(
         self._result = None
         return None
 
+    def _has_reservation_arn_from_metadata(self, current_metadata: dict[str, Any]) -> bool:
+        for association in current_metadata.get("associations", []):
+            if association.get("type") == "RESERVATION_TIME_WINDOW_ARN":
+                return True
+        return False
+
     def _download_result(
         self,
     ) -> Union[
@@ -488,7 +513,12 @@ def _download_result(
             current_metadata["outputS3Directory"] + f"/{AwsQuantumTask.RESULTS_FILENAME}",
         )
         self._result = _format_result(BraketSchemaBase.parse_raw_schema(result_string))
-        task_event = {"arn": self.id, "status": self.state(), "execution_duration": None}
+        task_event = {
+            "arn": self.id,
+            "status": self.state(),
+            "execution_duration": None,
+            "has_reservation_arn": self._has_reservation_arn_from_metadata(current_metadata),
+        }
         try:
             task_event[
                 "execution_duration"

src/braket/aws/aws_quantum_task_batch.py

@@ -61,6 +61,7 @@ def __init__(
         poll_timeout_seconds: float = AwsQuantumTask.DEFAULT_RESULTS_POLL_TIMEOUT,
         poll_interval_seconds: float = AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL,
         inputs: Union[dict[str, float], list[dict[str, float]]] | None = None,
+        reservation_arn: str | None = None,
         *aws_quantum_task_args,
         **aws_quantum_task_kwargs,
     ):
@@ -91,6 +92,11 @@ def __init__(
             inputs (Union[dict[str, float], list[dict[str, float]]] | None): Inputs to be passed
                 along with the IR. If the IR supports inputs, the inputs will be updated
                 with this value. Default: {}.
+            reservation_arn (str | None): The reservation ARN provided by Braket Direct
+                to reserve exclusive usage for the device to run the quantum task on.
+                Note: If you are creating tasks in a job that itself was created reservation ARN,
+                those tasks do not need to be created with the reservation ARN.
+                Default: None.
         """
         self._tasks = AwsQuantumTaskBatch._execute(
             aws_session,
@@ -103,6 +109,7 @@ def __init__(
             poll_timeout_seconds,
             poll_interval_seconds,
             inputs,
+            reservation_arn,
             *aws_quantum_task_args,
             **aws_quantum_task_kwargs,
         )
@@ -120,6 +127,7 @@ def __init__(
         self._poll_timeout_seconds = poll_timeout_seconds
         self._poll_interval_seconds = poll_interval_seconds
         self._inputs = inputs
+        self._reservation_arn = reservation_arn
         self._aws_quantum_task_args = aws_quantum_task_args
         self._aws_quantum_task_kwargs = aws_quantum_task_kwargs
 
@@ -197,6 +205,7 @@ def _execute(
         poll_timeout_seconds: float = AwsQuantumTask.DEFAULT_RESULTS_POLL_TIMEOUT,
         poll_interval_seconds: float = AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL,
         inputs: Union[dict[str, float], list[dict[str, float]]] = None,
+        reservation_arn: str | None = None,
         *args,
         **kwargs,
     ) -> list[AwsQuantumTask]:
@@ -217,6 +226,7 @@ def _execute(
                         poll_timeout_seconds=poll_timeout_seconds,
                         poll_interval_seconds=poll_interval_seconds,
                         inputs=input_map,
+                        reservation_arn=reservation_arn,
                         *args,
                         **kwargs,
                     )
@@ -248,6 +258,7 @@ def _create_task(
         shots: int,
         poll_interval_seconds: float = AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL,
         inputs: dict[str, float] = None,
+        reservation_arn: str | None = None,
         *args,
         **kwargs,
     ) -> AwsQuantumTask:
@@ -259,6 +270,7 @@ def _create_task(
             shots,
             poll_interval_seconds=poll_interval_seconds,
             inputs=inputs,
+            reservation_arn=reservation_arn,
             *args,
             **kwargs,
         )
@@ -347,6 +359,7 @@ def retry_unsuccessful_tasks(self) -> bool:
             self._max_workers,
             self._poll_timeout_seconds,
             self._poll_interval_seconds,
+            self._reservation_arn,
             *self._aws_quantum_task_args,
             **self._aws_quantum_task_kwargs,
         )