feature: Direct Reservation context manager (#955)

* feature: context manager for reservation arns
Co-authored-by: Cody Wang <speller26@gmail.com>

examples/reservation.py

@@ -11,14 +11,25 @@
 # 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
 
 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")
+# and fill in your reservation ARN.
+with DirectReservation(device, reservation_arn="<my_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="<my_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="<my_reservation_arn>")
 print(task.result().measurement_counts)

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

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:

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="<my_reservation_arn>"):
+        ...     task1 = device.run(circuit, shots)
+        ...     task2 = device.run(circuit, shots)
+
+        or start the reservation
+        >>> DirectReservation(device_arn, reservation_arn="<my_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

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,23 +36,23 @@ 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):
     circuit = Circuit().h(0)
     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(